pdslabeler / activeLabelsStorage.ts

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

export class ActiveLabelsStorage {
    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 active_labels (
                uri TEXT NOT NULL,
                val TEXT NOT NULL,
                PRIMARY KEY (uri, val)
            );
        `);
        await this.db.execute(`
            CREATE INDEX IF NOT EXISTS idx_active_labels_uri ON active_labels (uri);
        `);
    }

    async computeLabelsDiff(uri: string, desiredOutcome: string[]): Promise<{
        labelsToAdd: string[],
        labelsToNegate: string[],
    }> {
        await this.dbInitLock;

        const diff: {
            labelsToAdd: string[],
            labelsToNegate: string[],
        } = {
            labelsToAdd: [],
            labelsToNegate: [],
        };

        const desiredOutcomeSet = new Set(desiredOutcome);

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

        const activeLabels = result.rows.map((row) => row.val as string);
        const activeLabelSet = new Set(activeLabels);

        for (const label of activeLabels) {
            if (!desiredOutcomeSet.has(label)) {
                diff.labelsToNegate.push(label);
            }
        }

        for (const label of desiredOutcome) {
            if (!activeLabelSet.has(label)) {
                diff.labelsToAdd.push(label);
            }
        }

        return diff;
    }

    async registerLabelActivation(uri: string, val: string) {
        await this.dbInitLock;

        const result = await this.db.execute({
            sql: `
				INSERT INTO active_labels (uri, val)
                VALUES (?, ?)
			`,
            args: [uri, val],
        });

        if (!result.rowsAffected) throw new Error("Failed to register label activation");
    }

    async registerLabelDeactivation(uri: string, val: string) {
        await this.dbInitLock;

        const result = await this.db.execute({
            sql: `
				DELETE FROM active_labels
                WHERE uri = ? AND val = ?
			`,
            args: [uri, val],
        });

        if (!result.rowsAffected) throw new Error("Failed to register label deactivation");
    }

    async countActiveLabelSubjects() {
        await this.dbInitLock;

        const result = await this.db.execute({
            sql: `
                    SELECT COUNT (DISTINCT uri) AS count
                    FROM active_labels
                `,
            args: [],
        });

        return result.rows[0]?.count as number ?? 0;
    }

    async* getActiveLabelSubjects() {
        await this.dbInitLock;

        let cursor = "";
        while (true) {
            const result = await this.db.execute({
                sql: `
                    SELECT DISTINCT uri
                    FROM active_labels
                    WHERE uri > ?
                    ORDER BY uri
                    LIMIT 100
                `,
                args: [cursor],
            });

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

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