@atproto/ozone 0.0.2

package/src/db/index.ts

@@ -0,0 +1,197 @@
+import assert from 'assert'
+import {
+  Kysely,
+  PostgresDialect,
+  Migrator,
+  KyselyPlugin,
+  PluginTransformQueryArgs,
+  PluginTransformResultArgs,
+  RootOperationNode,
+  QueryResult,
+  UnknownRow,
+} from 'kysely'
+import TypedEmitter from 'typed-emitter'
+import { Pool as PgPool, types as pgTypes } from 'pg'
+import DatabaseSchema, { DatabaseSchemaType } from './schema'
+import { PgOptions } from './types'
+import { dbLogger } from '../logger'
+import { EventEmitter } from 'stream'
+import * as migrations from './migrations'
+import { CtxMigrationProvider } from './migrations/provider'
+
+export class Database {
+  pool: PgPool
+  db: DatabaseSchema
+  migrator: Migrator
+  txEvt = new EventEmitter() as TxnEmitter
+  destroyed = false
+  isPrimary = false
+
+  constructor(
+    public opts: PgOptions,
+    instances?: { db: DatabaseSchema; pool: PgPool },
+  ) {
+    // if instances are provided, use those
+    if (instances) {
+      this.db = instances.db
+      this.pool = instances.pool
+      return
+    }
+
+    // else create a pool & connect
+    const { schema, url } = opts
+    const pool =
+      opts.pool ??
+      new PgPool({
+        connectionString: url,
+        max: opts.poolSize,
+        maxUses: opts.poolMaxUses,
+        idleTimeoutMillis: opts.poolIdleTimeoutMs,
+      })
+
+    // Select count(*) and other pg bigints as js integer
+    pgTypes.setTypeParser(pgTypes.builtins.INT8, (n) => parseInt(n, 10))
+
+    // Setup schema usage, primarily for test parallelism (each test suite runs in its own pg schema)
+    if (schema && !/^[a-z_]+$/i.test(schema)) {
+      throw new Error(`Postgres schema must only contain [A-Za-z_]: ${schema}`)
+    }
+
+    pool.on('error', onPoolError)
+    pool.on('connect', (client) => {
+      client.on('error', onClientError)
+      // Used for trigram indexes, e.g. on actor search
+      client.query('SET pg_trgm.word_similarity_threshold TO .4;')
+      if (schema) {
+        // Shared objects such as extensions will go in the public schema
+        client.query(`SET search_path TO "${schema}",public;`)
+      }
+    })
+
+    this.pool = pool
+    this.db = new Kysely<DatabaseSchemaType>({
+      dialect: new PostgresDialect({ pool }),
+    })
+    this.migrator = new Migrator({
+      db: this.db,
+      migrationTableSchema: opts.schema,
+      provider: new CtxMigrationProvider(migrations, 'pg'),
+    })
+  }
+
+  get schema(): string | undefined {
+    return this.opts.schema
+  }
+
+  get isTransaction() {
+    return this.db.isTransaction
+  }
+
+  assertTransaction() {
+    assert(this.isTransaction, 'Transaction required')
+  }
+
+  assertNotTransaction() {
+    assert(!this.isTransaction, 'Cannot be in a transaction')
+  }
+
+  async transaction<T>(fn: (db: Database) => Promise<T>): Promise<T> {
+    const leakyTxPlugin = new LeakyTxPlugin()
+    const { dbTxn, txRes } = await this.db
+      .withPlugin(leakyTxPlugin)
+      .transaction()
+      .execute(async (txn) => {
+        const dbTxn = new Database(this.opts, {
+          db: txn,
+          pool: this.pool,
+        })
+        const txRes = await fn(dbTxn)
+          .catch(async (err) => {
+            leakyTxPlugin.endTx()
+            // ensure that all in-flight queries are flushed & the connection is open
+            await dbTxn.db.getExecutor().provideConnection(noopAsync)
+            throw err
+          })
+          .finally(() => leakyTxPlugin.endTx())
+        return { dbTxn, txRes }
+      })
+    dbTxn?.txEvt.emit('commit')
+    return txRes
+  }
+
+  onCommit(fn: () => void) {
+    this.assertTransaction()
+    this.txEvt.once('commit', fn)
+  }
+
+  async close(): Promise<void> {
+    if (this.destroyed) return
+    await this.db.destroy()
+    this.destroyed = true
+  }
+
+  async migrateToOrThrow(migration: string) {
+    if (this.schema) {
+      await this.db.schema.createSchema(this.schema).ifNotExists().execute()
+    }
+    const { error, results } = await this.migrator.migrateTo(migration)
+    if (error) {
+      throw error
+    }
+    if (!results) {
+      throw new Error('An unknown failure occurred while migrating')
+    }
+    return results
+  }
+
+  async migrateToLatestOrThrow() {
+    if (this.schema) {
+      await this.db.schema.createSchema(this.schema).ifNotExists().execute()
+    }
+    const { error, results } = await this.migrator.migrateToLatest()
+    if (error) {
+      throw error
+    }
+    if (!results) {
+      throw new Error('An unknown failure occurred while migrating')
+    }
+    return results
+  }
+}
+
+export default Database
+
+const onPoolError = (err: Error) => dbLogger.error({ err }, 'db pool error')
+const onClientError = (err: Error) => dbLogger.error({ err }, 'db client error')
+
+// utils
+// -------
+
+class LeakyTxPlugin implements KyselyPlugin {
+  private txOver: boolean
+
+  endTx() {
+    this.txOver = true
+  }
+
+  transformQuery(args: PluginTransformQueryArgs): RootOperationNode {
+    if (this.txOver) {
+      throw new Error('tx already failed')
+    }
+    return args.node
+  }
+
+  async transformResult(
+    args: PluginTransformResultArgs,
+  ): Promise<QueryResult<UnknownRow>> {
+    return args.result
+  }
+}
+
+type TxnEmitter = TypedEmitter<TxnEvents>
+
+type TxnEvents = {
+  commit: () => void
+}
+
+const noopAsync = async () => {}

package/src/db/migrations/20231219T205730722Z-init.ts

@@ -0,0 +1,164 @@
+import { Kysely } from 'kysely'
+
+export async function up(db: Kysely<unknown>): Promise<void> {
+  // Moderation event
+  await db.schema
+    .createTable('moderation_event')
+    .addColumn('id', 'serial', (col) => col.primaryKey())
+    .addColumn('action', 'varchar', (col) => col.notNull())
+    .addColumn('subjectType', 'varchar', (col) => col.notNull())
+    .addColumn('subjectDid', 'varchar', (col) => col.notNull())
+    .addColumn('subjectUri', 'varchar')
+    .addColumn('subjectCid', 'varchar')
+    .addColumn('comment', 'text')
+    .addColumn('meta', 'jsonb')
+    .addColumn('createdAt', 'varchar', (col) => col.notNull())
+    .addColumn('createdBy', 'varchar', (col) => col.notNull())
+    .addColumn('reversedAt', 'varchar')
+    .addColumn('reversedBy', 'varchar')
+    .addColumn('durationInHours', 'integer')
+    .addColumn('expiresAt', 'varchar')
+    .addColumn('reversedReason', 'text')
+    .addColumn('createLabelVals', 'varchar')
+    .addColumn('negateLabelVals', 'varchar')
+    .addColumn('legacyRefId', 'integer')
+    .execute()
+
+  // Moderation subject status
+  await db.schema
+    .createTable('moderation_subject_status')
+    .addColumn('id', 'serial', (col) => col.primaryKey())
+
+    // Identifiers
+    .addColumn('did', 'varchar', (col) => col.notNull())
+    // Default to '' so that we can apply unique constraints on did and recordPath columns
+    .addColumn('recordPath', 'varchar', (col) => col.notNull().defaultTo(''))
+    .addColumn('blobCids', 'jsonb')
+    .addColumn('recordCid', 'varchar')
+
+    // human review team state
+    .addColumn('reviewState', 'varchar', (col) => col.notNull())
+    .addColumn('comment', 'varchar')
+    .addColumn('muteUntil', 'varchar')
+    .addColumn('lastReviewedAt', 'varchar')
+    .addColumn('lastReviewedBy', 'varchar')
+
+    // report state
+    .addColumn('lastReportedAt', 'varchar')
+    .addColumn('lastAppealedAt', 'varchar')
+
+    // visibility/intervention state
+    .addColumn('takendown', 'boolean', (col) => col.defaultTo(false).notNull())
+    .addColumn('suspendUntil', 'varchar')
+    .addColumn('appealed', 'boolean')
+
+    // timestamps
+    .addColumn('createdAt', 'varchar', (col) => col.notNull())
+    .addColumn('updatedAt', 'varchar', (col) => col.notNull())
+    .addUniqueConstraint('moderation_status_unique_idx', ['did', 'recordPath'])
+    .execute()
+
+  await db.schema
+    .createIndex('moderation_subject_status_blob_cids_idx')
+    .on('moderation_subject_status')
+    .using('gin')
+    .column('blobCids')
+    .execute()
+
+  // Label
+  await db.schema
+    .createTable('label')
+    .addColumn('src', 'varchar', (col) => col.notNull())
+    .addColumn('uri', 'varchar', (col) => col.notNull())
+    .addColumn('cid', 'varchar', (col) => col.notNull())
+    .addColumn('val', 'varchar', (col) => col.notNull())
+    .addColumn('neg', 'boolean', (col) => col.notNull())
+    .addColumn('cts', 'varchar', (col) => col.notNull())
+    .addPrimaryKeyConstraint('label_pkey', ['src', 'uri', 'cid', 'val'])
+    .execute()
+  await db.schema
+    .createIndex('label_uri_index')
+    .on('label')
+    .column('uri')
+    .execute()
+
+  // Push Events
+  await db.schema
+    .createTable('repo_push_event')
+    .addColumn('id', 'serial', (col) => col.primaryKey())
+    .addColumn('eventType', 'varchar', (col) => col.notNull())
+    .addColumn('subjectDid', 'varchar', (col) => col.notNull())
+    .addColumn('takedownRef', 'varchar')
+    .addColumn('confirmedAt', 'timestamptz')
+    .addColumn('lastAttempted', 'timestamptz')
+    .addColumn('attempts', 'integer', (col) => col.notNull().defaultTo(0))
+    .addUniqueConstraint('repo_push_event_unique_evt', [
+      'subjectDid',
+      'eventType',
+    ])
+    .execute()
+  await db.schema
+    .createIndex('repo_push_confirmation_idx')
+    .on('repo_push_event')
+    .columns(['confirmedAt', 'attempts'])
+    .execute()
+
+  await db.schema
+    .createTable('record_push_event')
+    .addColumn('id', 'serial', (col) => col.primaryKey())
+    .addColumn('eventType', 'varchar', (col) => col.notNull())
+    .addColumn('subjectDid', 'varchar', (col) => col.notNull())
+    .addColumn('subjectUri', 'varchar', (col) => col.notNull())
+    .addColumn('subjectCid', 'varchar')
+    .addColumn('takedownRef', 'varchar')
+    .addColumn('confirmedAt', 'timestamptz')
+    .addColumn('lastAttempted', 'timestamptz')
+    .addColumn('attempts', 'integer', (col) => col.notNull().defaultTo(0))
+    .addUniqueConstraint('record_push_event_unique_evt', [
+      'subjectUri',
+      'eventType',
+    ])
+    .execute()
+  await db.schema
+    .createIndex('record_push_event_did_type_idx')
+    .on('record_push_event')
+    .columns(['subjectDid', 'eventType'])
+    .execute()
+  await db.schema
+    .createIndex('record_push_confirmation_idx')
+    .on('record_push_event')
+    .columns(['confirmedAt', 'attempts'])
+    .execute()
+
+  await db.schema
+    .createTable('blob_push_event')
+    .addColumn('id', 'serial', (col) => col.primaryKey())
+    .addColumn('eventType', 'varchar', (col) => col.notNull())
+    .addColumn('subjectDid', 'varchar', (col) => col.notNull())
+    .addColumn('subjectBlobCid', 'varchar', (col) => col.notNull())
+    .addColumn('subjectUri', 'varchar')
+    .addColumn('takedownRef', 'varchar')
+    .addColumn('confirmedAt', 'timestamptz')
+    .addColumn('lastAttempted', 'timestamptz')
+    .addColumn('attempts', 'integer', (col) => col.notNull().defaultTo(0))
+    .addUniqueConstraint('blob_push_event_unique_evt', [
+      'subjectDid',
+      'subjectBlobCid',
+      'eventType',
+    ])
+    .execute()
+  await db.schema
+    .createIndex('blob_push_confirmation_idx')
+    .on('blob_push_event')
+    .columns(['confirmedAt', 'attempts'])
+    .execute()
+}
+
+export async function down(db: Kysely<unknown>): Promise<void> {
+  await db.schema.dropTable('moderation_event').execute()
+  await db.schema.dropTable('moderation_subject_status').execute()
+  await db.schema.dropTable('label').execute()
+  await db.schema.dropTable('repo_push_event').execute()
+  await db.schema.dropTable('record_push_event').execute()
+  await db.schema.dropTable('blob_push_event').execute()
+}

package/src/db/migrations/index.ts

@@ -0,0 +1,5 @@
+// NOTE this file can be edited by hand, but it is also appended to by the migration:create command.
+// It's important that every migration is exported from here with the proper name. We'd simplify
+// this with kysely's FileMigrationProvider, but it doesn't play nicely with the build process.
+
+export * as _20231219T205730722Z from './20231219T205730722Z-init'

package/src/db/migrations/provider.ts

@@ -0,0 +1,25 @@
+import { Kysely, Migration, MigrationProvider } from 'kysely'
+
+// Passes a context argument to migrations. We use this to thread the dialect into migrations
+
+export class CtxMigrationProvider<T> implements MigrationProvider {
+  constructor(
+    private migrations: Record<string, CtxMigration<T>>,
+    private ctx: T,
+  ) {}
+  async getMigrations(): Promise<Record<string, Migration>> {
+    const ctxMigrations: Record<string, Migration> = {}
+    Object.entries(this.migrations).forEach(([name, migration]) => {
+      ctxMigrations[name] = {
+        up: async (db) => await migration.up(db, this.ctx),
+        down: async (db) => await migration.down?.(db, this.ctx),
+      }
+    })
+    return ctxMigrations
+  }
+}
+
+export interface CtxMigration<T> {
+  up(db: Kysely<unknown>, ctx: T): Promise<void>
+  down?(db: Kysely<unknown>, ctx: T): Promise<void>
+}

package/src/db/pagination.ts

@@ -0,0 +1,216 @@
+import { sql, DynamicModule } from 'kysely'
+import { InvalidRequestError } from '@atproto/xrpc-server'
+import { AnyQb, DbRef } from './types'
+
+export type Cursor = { primary: string; secondary: string }
+export type LabeledResult = {
+  primary: string | number
+  secondary: string | number
+}
+
+/**
+ * The GenericKeyset is an abstract class that sets-up the interface and partial implementation
+ * of a keyset-paginated cursor with two parts. There are three types involved:
+ *  - Result: a raw result (i.e. a row from the db) containing data that will make-up a cursor.
+ *    - E.g. { createdAt: '2022-01-01T12:00:00Z', cid: 'bafyx' }
+ *  - LabeledResult: a Result processed such that the "primary" and "secondary" parts of the cursor are labeled.
+ *    - E.g. { primary: '2022-01-01T12:00:00Z', secondary: 'bafyx' }
+ *  - Cursor: the two string parts that make-up the packed/string cursor.
+ *    - E.g. packed cursor '1641038400000::bafyx' in parts { primary: '1641038400000', secondary: 'bafyx' }
+ *
+ * These types relate as such. Implementers define the relations marked with a *:
+ *   Result -*-> LabeledResult <-*-> Cursor <--> packed/string cursor
+ *                     ↳ SQL Condition
+ */
+export abstract class GenericKeyset<R, LR extends LabeledResult> {
+  constructor(public primary: DbRef, public secondary: DbRef) {}
+  abstract labelResult(result: R): LR
+  abstract labeledResultToCursor(labeled: LR): Cursor
+  abstract cursorToLabeledResult(cursor: Cursor): LR
+  packFromResult(results: R | R[]): string | undefined {
+    const result = Array.isArray(results) ? results.at(-1) : results
+    if (!result) return
+    return this.pack(this.labelResult(result))
+  }
+  pack(labeled?: LR): string | undefined {
+    if (!labeled) return
+    const cursor = this.labeledResultToCursor(labeled)
+    return this.packCursor(cursor)
+  }
+  unpack(cursorStr?: string): LR | undefined {
+    const cursor = this.unpackCursor(cursorStr)
+    if (!cursor) return
+    return this.cursorToLabeledResult(cursor)
+  }
+  packCursor(cursor?: Cursor): string | undefined {
+    if (!cursor) return
+    return `${cursor.primary}::${cursor.secondary}`
+  }
+  unpackCursor(cursorStr?: string): Cursor | undefined {
+    if (!cursorStr) return
+    const result = cursorStr.split('::')
+    const [primary, secondary, ...others] = result
+    if (!primary || !secondary || others.length > 0) {
+      throw new InvalidRequestError('Malformed cursor')
+    }
+    return {
+      primary,
+      secondary,
+    }
+  }
+  getSql(labeled?: LR, direction?: 'asc' | 'desc', tryIndex?: boolean) {
+    if (labeled === undefined) return
+    if (tryIndex) {
+      // The tryIndex param will likely disappear and become the default implementation: here for now for gradual rollout query-by-query.
+      if (direction === 'asc') {
+        return sql`((${this.primary}, ${this.secondary}) > (${labeled.primary}, ${labeled.secondary}))`
+      } else {
+        return sql`((${this.primary}, ${this.secondary}) < (${labeled.primary}, ${labeled.secondary}))`
+      }
+    } else {
+      // @NOTE this implementation can struggle to use an index on (primary, secondary) for pagination due to the "or" usage.
+      if (direction === 'asc') {
+        return sql`((${this.primary} > ${labeled.primary}) or (${this.primary} = ${labeled.primary} and ${this.secondary} > ${labeled.secondary}))`
+      } else {
+        return sql`((${this.primary} < ${labeled.primary}) or (${this.primary} = ${labeled.primary} and ${this.secondary} < ${labeled.secondary}))`
+      }
+    }
+  }
+}
+
+type StatusKeysetParam = {
+  lastReviewedAt: string | null
+  lastReportedAt: string | null
+  id: number
+}
+
+export class StatusKeyset extends GenericKeyset<StatusKeysetParam, Cursor> {
+  labelResult(result: StatusKeysetParam): Cursor
+  labelResult(result: StatusKeysetParam) {
+    const primaryField = (
+      this.primary as ReturnType<DynamicModule['ref']>
+    ).dynamicReference.includes('lastReviewedAt')
+      ? 'lastReviewedAt'
+      : 'lastReportedAt'
+
+    return {
+      primary: result[primaryField]
+        ? new Date(`${result[primaryField]}`).getTime().toString()
+        : '',
+      secondary: result.id.toString(),
+    }
+  }
+  labeledResultToCursor(labeled: Cursor) {
+    return {
+      primary: labeled.primary,
+      secondary: labeled.secondary,
+    }
+  }
+  cursorToLabeledResult(cursor: Cursor) {
+    return {
+      primary: cursor.primary
+        ? new Date(parseInt(cursor.primary, 10)).toISOString()
+        : '',
+      secondary: cursor.secondary,
+    }
+  }
+  unpackCursor(cursorStr?: string): Cursor | undefined {
+    if (!cursorStr) return
+    const result = cursorStr.split('::')
+    const [primary, secondary, ...others] = result
+    if (!secondary || others.length > 0) {
+      throw new InvalidRequestError('Malformed cursor')
+    }
+    return {
+      primary,
+      secondary,
+    }
+  }
+  // This is specifically built to handle nullable columns as primary sorting column
+  getSql(labeled?: Cursor, direction?: 'asc' | 'desc') {
+    if (labeled === undefined) return
+    if (direction === 'asc') {
+      return !labeled.primary
+        ? sql`(${this.primary} IS NULL AND ${this.secondary} > ${labeled.secondary})`
+        : sql`((${this.primary}, ${this.secondary}) > (${labeled.primary}, ${labeled.secondary}) OR (${this.primary} is null))`
+    } else {
+      return !labeled.primary
+        ? sql`(${this.primary} IS NULL AND ${this.secondary} < ${labeled.secondary})`
+        : sql`((${this.primary}, ${this.secondary}) < (${labeled.primary}, ${labeled.secondary}) OR (${this.primary} is null))`
+    }
+  }
+}
+
+type TimeIdKeysetParam = {
+  id: number
+  createdAt: string
+}
+type TimeIdResult = TimeIdKeysetParam
+
+export class TimeIdKeyset extends GenericKeyset<TimeIdKeysetParam, Cursor> {
+  labelResult(result: TimeIdResult): Cursor
+  labelResult(result: TimeIdResult) {
+    return { primary: result.createdAt, secondary: result.id.toString() }
+  }
+  labeledResultToCursor(labeled: Cursor) {
+    return {
+      primary: new Date(labeled.primary).getTime().toString(),
+      secondary: labeled.secondary,
+    }
+  }
+  cursorToLabeledResult(cursor: Cursor) {
+    const primaryDate = new Date(parseInt(cursor.primary, 10))
+    if (isNaN(primaryDate.getTime())) {
+      throw new InvalidRequestError('Malformed cursor')
+    }
+    return {
+      primary: primaryDate.toISOString(),
+      secondary: cursor.secondary,
+    }
+  }
+}
+
+export const paginate = <
+  QB extends AnyQb,
+  K extends GenericKeyset<unknown, any>,
+>(
+  qb: QB,
+  opts: {
+    limit?: number
+    cursor?: string
+    direction?: 'asc' | 'desc'
+    keyset: K
+    tryIndex?: boolean
+    // By default, pg does nullsFirst
+    nullsLast?: boolean
+  },
+): QB => {
+  const {
+    limit,
+    cursor,
+    keyset,
+    direction = 'desc',
+    tryIndex,
+    nullsLast,
+  } = opts
+  const keysetSql = keyset.getSql(keyset.unpack(cursor), direction, tryIndex)
+  return qb
+    .if(!!limit, (q) => q.limit(limit as number))
+    .if(!nullsLast, (q) =>
+      q.orderBy(keyset.primary, direction).orderBy(keyset.secondary, direction),
+    )
+    .if(!!nullsLast, (q) =>
+      q
+        .orderBy(
+          direction === 'asc'
+            ? sql`${keyset.primary} asc nulls last`
+            : sql`${keyset.primary} desc nulls last`,
+        )
+        .orderBy(
+          direction === 'asc'
+            ? sql`${keyset.secondary} asc nulls last`
+            : sql`${keyset.secondary} desc nulls last`,
+        ),
+    )
+    .if(!!keysetSql, (qb) => (keysetSql ? qb.where(keysetSql) : qb)) as QB
+}

package/src/db/schema/blob_push_event.ts

@@ -0,0 +1,21 @@
+import { Generated } from 'kysely'
+
+export const eventTableName = 'blob_push_event'
+
+export type BlobPushEventType = 'pds_takedown' | 'appview_takedown'
+
+export interface BlobPushEvent {
+  id: Generated<number>
+  eventType: BlobPushEventType
+  subjectDid: string
+  subjectBlobCid: string
+  subjectUri: string | null
+  takedownRef: string | null
+  confirmedAt: Date | null
+  lastAttempted: Date | null
+  attempts: Generated<number>
+}
+
+export type PartialDB = {
+  [eventTableName]: BlobPushEvent
+}

package/src/db/schema/index.ts

@@ -0,0 +1,18 @@
+import { Kysely } from 'kysely'
+import * as modEvent from './moderation_event'
+import * as modSubjectStatus from './moderation_subject_status'
+import * as repoPushEvent from './repo_push_event'
+import * as recordPushEvent from './record_push_event'
+import * as blobPushEvent from './blob_push_event'
+import * as label from './label'
+
+export type DatabaseSchemaType = modEvent.PartialDB &
+  modSubjectStatus.PartialDB &
+  label.PartialDB &
+  repoPushEvent.PartialDB &
+  recordPushEvent.PartialDB &
+  blobPushEvent.PartialDB
+
+export type DatabaseSchema = Kysely<DatabaseSchemaType>
+
+export default DatabaseSchema

package/src/db/schema/label.ts

@@ -0,0 +1,12 @@
+export const tableName = 'label'
+
+export interface Label {
+  src: string
+  uri: string
+  cid: string
+  val: string
+  neg: boolean
+  cts: string
+}
+
+export type PartialDB = { [tableName]: Label }