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 {
  initializeEventTable,
  initializeProfileTable,
  initializeRecordTable,
  initializeRepoTable,
  initializeSubjectTable,
} from "../schemas/initializers";
import { RepoRow } from "../schemas/repo";
import { RecordRow } from "../schemas/record";
import { transformFullStatusFromDatabase } from "./subject/transformers";
import { ProfileRow } from "../schemas/profile";
import { EventRow } from "../schemas/event";

// 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 initializeEventTable(db);
    await initializeRepoTable(db);
    await initializeRecordTable(db);
    await initializeProfileTable(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 insertEvents(events: EventRow[]) {
    await this.initialize();
    const ids = events.map(({ id }) => id);
    // This is questionable, probably could do update on conflict or something but quick n dirty for now
    await db.deleteFrom("events").where("id", "in", ids).execute();
    return await db.insertInto("events").values(events).execute();
  },

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

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

    if (takendown !== undefined) {
      builder = builder.where("takendown", "=", takendown ? 1 : 0);
    }

    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 getMissingProfileDids() {
    await this.initialize();
    const syncPeriod = new Date(Date.now() - 1000 * 60 * 60 * 24).toISOString();
    const results = await db
      .selectFrom("subjects")
      .select("did")
      .distinct()
      .where("did", "not in", (qb) => {
        return qb
          .selectFrom("profiles")
          .select("did")
          .where("syncedAt", ">", syncPeriod);
      })
      .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 saveProfiles(profiles: ProfileRow[]) {
    await this.initialize();
    const dids = profiles.map(({ did }) => did);

    // Remove all existing profiles first
    if (dids.length) {
      await db.deleteFrom("profiles").where("did", "in", dids).execute();
    }

    await db.insertInto("profiles").values(profiles).execute();
  },

  async getEvents(opts: {
    cursor: number;
    types: string[];
    order?: "asc" | "desc";
    createdAfter?: string;
    createdBefore?: string;
  }) {
    let qb = db.selectFrom("events");

    if (opts.cursor) {
      qb = qb.where("id", opts.order === "asc" ? ">" : "<", opts.cursor);
    }

    if (opts.types.length) {
      qb = qb.where("action", "in", opts.types);
    }

    if (opts.createdAfter) {
      qb = qb.where("createdAt", ">", opts.createdAfter);
    }
    if (opts.createdBefore) {
      qb = qb.where("createdAt", "<", opts.createdBefore);
    }

    return qb
      .orderBy("id", opts.order || "asc")
      .select(["subjectDid", "subjectUri", "action", "createdAt", "createdBy"])
      .execute();
  },

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