pdslabeler / knownPDSStorage.ts

import { LabelerServer } from "@skyware/labeler";

export type KnownPDS = {
    uri: string,
    lastCrawled: Date,
    totalRepos: number,
}

export class KnownPDSStorage {
    private db: typeof LabelerServer.prototype.db;

    /**
     * Promise that resolves when database initialization is complete.
     * This should be awaited before any database operations.
     */
    private readonly dbInitLock?: Promise<void>;

    constructor(db: typeof this.db) {
        this.db = db;

        this.dbInitLock = this.initializeDatabase();
    }

    async initializeDatabase() {
        await this.db.execute(`
            CREATE TABLE IF NOT EXISTS known_pds (
                uri TEXT NOT NULL,
                last_crawled DATETIME NOT NULL,
                total_repos INTEGER NOT NULL,
                PRIMARY KEY (uri)
            );
        `);
        await this.db.execute(`
            CREATE TABLE IF NOT EXISTS pds_crawling_failures (
                uri TEXT NOT NULL,
                last_attempt DATETIME NOT NULL,
                PRIMARY KEY (uri)
            );
        `);
    }

    private rowToKnownPDS(row: any): KnownPDS {
        const v: KnownPDS = {
            uri: row.uri as string,
            lastCrawled: new Date(row.last_crawled as string),
            totalRepos: row.total_repos as number,
        };
        return v
    }

    async countKnownPDSs(): Promise<number> {
        await this.dbInitLock;

        const result = await this.db.execute({
            sql: `
                SELECT
                    COUNT(*) AS c
                FROM known_pds
			    `,
        });

        if (!result.rows || !result.rows.length) {
            return 0;
        }

        return result.rows[0]!.c as number;
    }

    async* getKnownPDSs(lastCrawledBefore: Date = new Date(), lastFailureBefore: Date = new Date()) {
        await this.dbInitLock;

        let cursor = "";
        while (true) {
            const result = await this.db.execute({
                sql: `
                    SELECT
                        k.uri, k.last_crawled, k.total_repos
                    FROM known_pds k
                    WHERE
                        k.uri > ? AND
                        k.last_crawled < ? AND
                        k.uri NOT IN (
                            SELECT f.uri
                            FROM pds_crawling_failures f
                            WHERE f.last_attempt >= ?
                        )
                    ORDER BY uri
                    LIMIT 100
			    `,
                args: [cursor, lastCrawledBefore.toISOString(), lastFailureBefore.toISOString()],
            });

            if (!result.rows.length) {
                return;
            }

            for (const row of result.rows) {
                const v = this.rowToKnownPDS(row);
                yield v;
                cursor = v.uri;
            }
        }
    }

    async* getBiggestKnownPDSs(limit: number) {
        await this.dbInitLock;

        const result = await this.db.execute({
            sql: `
                SELECT
                    uri, last_crawled, total_repos
                FROM known_pds
                ORDER BY total_repos DESC, uri
                LIMIT ?
			`,
            args: [limit],
        });

        for (const row of result.rows) {
            const v = this.rowToKnownPDS(row);
            yield v;
        }
    }

    async getKnownPDS(uri: string): Promise<KnownPDS | undefined> {
        await this.dbInitLock;

        const result = await this.db.execute({
            sql: `
                SELECT
                    uri, last_crawled, total_repos
                FROM known_pds
                WHERE uri = ?
			`,
            args: [uri],
        });

        if (!result.rows.length) {
            return undefined;
        }

        return this.rowToKnownPDS(result.rows[0]);
    }

    async upsertKnownPDS(knownPDS: KnownPDS) {
        await this.dbInitLock;

        const transaction = await this.db.transaction("write");

        try {
            const result = await transaction.execute({
                sql: `
                    INSERT INTO known_pds (uri, last_crawled, total_repos)
                    VALUES (?, ?, ?)
                    ON CONFLICT (uri) DO UPDATE SET
                        last_crawled = excluded.last_crawled,
                        total_repos = excluded.total_repos
			    `,
                args: [knownPDS.uri, knownPDS.lastCrawled.toISOString(), knownPDS.totalRepos],
            });

            if (!result.rowsAffected) throw new Error("Failed to upsert known PDS");

            await transaction.execute({
                sql: `
                    DELETE FROM pds_crawling_failures
                    WHERE uri = ?
                `,
                args: [knownPDS.uri],
            });
            await transaction.commit();
        } finally {
            transaction.close();
        }
    }

    async removeKnownPDS(knownPDSURI: string) {
        await this.dbInitLock;

        const transaction = await this.db.transaction("write");

        try {
            await transaction.execute({
                sql: `
                    DELETE FROM known_pds
                    WHERE uri = ?
			    `,
                args: [knownPDSURI],
            });

            await transaction.execute({
                sql: `
                    DELETE FROM pds_crawling_failures
                    WHERE uri = ?
                `,
                args: [knownPDSURI],
            });
            await transaction.commit();
        } finally {
            transaction.close();
        }
    }

    async countPDSCrawlingFailures(): Promise<number> {
        await this.dbInitLock;

        const result = await this.db.execute({
            sql: `
                SELECT
                    COUNT(*) AS c
                FROM pds_crawling_failures
			    `,
        });

        if (!result.rows || !result.rows.length) {
            return 0;
        }

        return result.rows[0]!.c as number;
    }

    async markPDSCrawlFailure(uri: string, at: Date) {
        const result = await this.db.execute({
            sql: `
                INSERT INTO pds_crawling_failures (uri, last_attempt)
                VALUES (?, ?)
                ON CONFLICT (uri) DO UPDATE SET
                    last_attempt = excluded.last_attempt
            `,
            args: [uri, at.toISOString()],
        });

        if (!result.rowsAffected) throw new Error("Failed to upsert known PDS");
    }

    async hasPDSCrawlFailed(uri: string, since: Date): Promise<boolean> {
        await this.dbInitLock;

        const result = await this.db.execute({
            sql: `
                SELECT 1
                FROM pds_crawling_failures
                WHERE uri = ? AND last_attempt > ?
			`,
            args: [uri, since.toISOString()],
        });

        return result.rows.length > 0;
    }
}