auron/src/services/db.ts

Interact with your ATProto labeler from your terminal

import { Kysely, sql, SqliteDialect } from "kysely";
import SQLite from "better-sqlite3";

import { DatabaseSchema } from "../schemas";
import { SubjectRow } from "../schemas/subject";
import {
  initializeRecordTable,
  initializeRepoTable,
  initializeSubjectTable,
} from "../schemas/initializers";
import { RepoRow } from "../schemas/repo";
import { RecordRow } from "../schemas/record";
import { transformFullStatusFromDatabase } from "./subject/transformers";

// Initialize Kysely with SQLite
const db = new Kysely<DatabaseSchema>({
  dialect: new SqliteDialect({
    database: new SQLite(process.env.LOCAL_DB_PATH || "aurondb.sqlite"),
  }),
});

// Create the subjects table if it doesn't exist
export const database = {
  async initialize() {
    await initializeSubjectTable(db);
    await initializeRepoTable(db);
    await initializeRecordTable(db);
  },

  async insertSubjects(subjectStatuses: SubjectRow[]) {
    await this.initialize();
    const ids = subjectStatuses.map(({ id }) => id);
    // This is questionable, probably could do update on conflict or something but quick n dirty for now
    await db.deleteFrom("subjects").where("id", "in", ids).execute();
    const { ref } = db.dynamic;

    return await db
      .insertInto("subjects")
      .values(subjectStatuses)
      .onConflict((oc) => {
        return oc.doUpdateSet({
          reviewState: sql`${ref(`excluded.reviewState`)}`,
          lastReviewedAt: sql`${ref(`excluded.lastReviewedAt`)}`,
          lastReportedAt: sql`${ref(`excluded.lastReportedAt`)}`,
          updatedAt: sql`${ref("excluded.updatedAt")}`,
          lastReviewedBy: sql`${ref("excluded.lastReviewedBy")}`,
          takendown: sql`${ref("excluded.takendown")}`,
          tags: sql`${ref("excluded.tags")}`,
          lastAppealedAt: sql`${ref("excluded.lastAppealedAt")}`,
          suspendUntil: sql`${ref("excluded.suspendUntil")}`,
          muteUntil: sql`${ref("excluded.muteUntil")}`,
          comment: sql`${ref("excluded.comment")}`,
        });
      })
      .execute();
  },

  async listSubjects({
    subjectType,
    cursor,
    limit,
  }: {
    subjectType?: string;
    limit?: number;
    cursor?: string;
  }) {
    await this.initialize();
    let builder = db
      .selectFrom("subjects")
      .leftJoin("repos", "subjects.did", "repos.did")
      .leftJoin("records", (join) =>
        join.onRef(
          "records.uri",
          "=",
          sql`'at://' || subjects.did || '/' || subjects.recordPath`
        )
      )
      .selectAll(["subjects", "repos", "records"]);

    if (subjectType === "account") {
      builder = builder.where("recordPath", "=", "");
    } else if (subjectType === "record") {
      builder = builder.where("recordPath", "!=", "");
    }

    if (limit) {
      builder = builder.limit(limit);
    }

    if (cursor) {
      builder = builder.where("lastReportedAt", "<", cursor);
    }

    const results = await builder.orderBy("lastReportedAt", "desc").execute();
    // @ts-ignore
    return results.map(transformFullStatusFromDatabase);
  },

  async clearSubjects() {
    await this.initialize();
    await db.deleteFrom("subjects").execute();
  },

  async getMissingRepoDids() {
    await this.initialize();
    const results = await db
      .selectFrom("subjects")
      .select("did")
      .distinct()
      .where("did", "not in", (qb) => qb.selectFrom("repos").select(["did"]))
      .execute();

    return results.map(({ did }) => did);
  },

  async getMissingRecordUris() {
    await this.initialize();
    const results = await db
      .selectFrom("subjects")
      .leftJoin("records", (join) =>
        join.onRef(
          "records.uri",
          "=",
          sql`'at://' || subjects.did || '/' || subjects.recordPath`
        )
      )
      .select(["subjects.did", "subjects.recordPath"])
      .where("records.uri", "is", null)
      .where("subjects.recordPath", "!=", "")
      .execute();

    return results.map(({ did, recordPath }) => `at://${did}/${recordPath}`);
  },

  async saveRepos(repos: RepoRow[]) {
    await this.initialize();
    await db.insertInto("repos").values(repos).execute();
  },

  async saveRecords(records: RecordRow[]) {
    await this.initialize();
    await db.insertInto("records").values(records).execute();
  },

  async clear() {
    await db.schema.dropTable("subjects").execute();
    await db.schema.dropTable("repos").execute();
  },
};