Pds v2 migrate script

.github/workflows/build-and-push-pds-ghcr.yaml

@@ -3,6 +3,7 @@ on:
   push:
     branches:
       - main
+      - pds-v2-migrate-script
 env:
   REGISTRY: ghcr.io
   USERNAME: ${{ github.actor }}

packages/pds/package.json

@@ -48,6 +48,7 @@
     "compression": "^1.7.4",
     "cors": "^2.8.5",
     "disposable-email": "^0.2.3",
+    "dotenv": "^16.0.3",
     "express": "^4.17.2",
     "express-async-errors": "^3.1.1",
     "file-type": "^16.5.4",

packages/pds/src/index.ts

@@ -31,6 +31,7 @@ export { Database } from './db'
 export { DiskBlobStore, MemoryBlobStore } from './storage'
 export { AppContext } from './context'
 export { httpLogger } from './logger'
+export { runScript } from './migrate-script'
 
 export class PDS {
   public ctx: AppContext

packages/pds/src/migrate-script/check-failures.ts

@@ -0,0 +1,63 @@
+import { countAll } from '../db/util'
+import { MigrateDb, getDb } from './db'
+
+const run = async () => {
+  const db = getDb()
+  const results = await Promise.all([
+    totalCount(db),
+    failureCount(db),
+    failedBlobs(db),
+    failedPrefs(db),
+    failedTakedowns(db),
+  ])
+  console.log(`
+Total migrations: ${results[0]}
+Failed migrations: ${results[1]}
+Failed blobs: ${results[2]}
+Failed prefs: ${results[3]}
+Failed takedowns: ${results[4]}
+`)
+}
+
+const totalCount = async (db: MigrateDb) => {
+  const res = await db
+    .selectFrom('status')
+    .select(countAll.as('count'))
+    .executeTakeFirst()
+  return res?.count
+}
+
+const failureCount = async (db: MigrateDb) => {
+  const res = await db
+    .selectFrom('status')
+    .select(countAll.as('count'))
+    .where('failed', '=', 1)
+    .executeTakeFirst()
+  return res?.count
+}
+
+const failedBlobs = async (db: MigrateDb) => {
+  const res = await db
+    .selectFrom('failed_blob')
+    .select(countAll.as('count'))
+    .executeTakeFirst()
+  return res?.count
+}
+
+const failedPrefs = async (db: MigrateDb) => {
+  const res = await db
+    .selectFrom('failed_pref')
+    .select(countAll.as('count'))
+    .executeTakeFirst()
+  return res?.count
+}
+
+const failedTakedowns = async (db: MigrateDb) => {
+  const res = await db
+    .selectFrom('failed_takedown')
+    .select(countAll.as('count'))
+    .executeTakeFirst()
+  return res?.count
+}
+
+run()

packages/pds/src/migrate-script/db.ts

@@ -0,0 +1,139 @@
+import { Kysely, MigrationProvider, SqliteDialect } from 'kysely'
+import SqliteDB from 'better-sqlite3'
+
+const LOCATION = 'migrate.db'
+
+export const getDb = (): MigrateDb => {
+  const sqliteDb = new SqliteDB(LOCATION)
+  sqliteDb.pragma('journal_mode = WAL')
+  sqliteDb.pragma('busy_timeout = 5000')
+  return new Kysely<Schema>({
+    dialect: new SqliteDialect({
+      database: sqliteDb,
+    }),
+  })
+}
+
+export const dbMigrationProvider: MigrationProvider = {
+  async getMigrations() {
+    return {
+      '1': {
+        async up(db: Kysely<unknown>) {
+          await db.schema
+            .createTable('status')
+            .addColumn('did', 'varchar', (col) => col.primaryKey())
+            .addColumn('pdsId', 'integer')
+            .addColumn('signingKey', 'varchar')
+            .addColumn('phase', 'integer', (col) => col.notNull().defaultTo(0))
+            .addColumn('importedRev', 'varchar')
+            .addColumn('failed', 'integer', (col) => col.notNull().defaultTo(0))
+            .execute()
+          await db.schema
+            .createTable('failed_pref')
+            .addColumn('did', 'varchar', (col) => col.primaryKey())
+            .execute()
+          await db.schema
+            .createTable('failed_blob')
+            .addColumn('did', 'varchar', (col) => col.notNull())
+            .addColumn('cid', 'varchar', (col) => col.notNull())
+            .addPrimaryKeyConstraint('failed_blob_pkey', ['did', 'cid'])
+            .execute()
+          await db.schema
+            .createTable('failed_takedown')
+            .addColumn('did', 'varchar', (col) => col.notNull())
+            .addColumn('recordUri', 'varchar')
+            .addColumn('recordCid', 'varchar')
+            .addColumn('blobCid', 'varchar')
+            .execute()
+        },
+        async down() {},
+      },
+      '2': {
+        async up(db: Kysely<unknown>) {
+          await db.schema
+            .alterTable('status')
+            .addColumn('err', 'varchar')
+            .execute()
+          await db.schema
+            .alterTable('failed_pref')
+            .addColumn('err', 'varchar')
+            .execute()
+          await db.schema
+            .alterTable('failed_blob')
+            .addColumn('err', 'varchar')
+            .execute()
+          await db.schema
+            .alterTable('failed_takedown')
+            .addColumn('err', 'varchar')
+            .execute()
+        },
+        async down() {},
+      },
+      '3': {
+        async up(db: Kysely<unknown>) {
+          await db.schema
+            .createTable('failed_import')
+            .addColumn('did', 'varchar', (col) => col.primaryKey())
+            .addColumn('err', 'varchar')
+            .execute()
+        },
+        async down() {},
+      },
+    }
+  },
+}
+
+export type MigrateDb = Kysely<Schema>
+
+type Schema = {
+  status: Status
+  failed_pref: FailedPreference
+  failed_blob: FailedBlob
+  failed_takedown: FailedTakedown
+  failed_import: FailedImport
+}
+
+export enum TransferPhase {
+  notStarted = 0,
+  reservedKey = 1,
+  initImport = 2,
+  transferredPds = 3,
+  transferredEntryway = 4,
+  preferences = 5,
+  takedowns = 6,
+  completed = 7,
+}
+
+export type Status = {
+  did: string
+  pdsId: number | null
+  signingKey: string | null
+  phase: TransferPhase
+  importedRev: string | null
+  failed: 0 | 1
+  err: string | null
+}
+
+export type FailedPreference = {
+  did: string
+  err: string | null
+}
+
+export type FailedBlob = {
+  did: string
+  cid: string
+  err: string | null
+}
+
+export type FailedTakedown = {
+  did: string
+  blobCid?: string
+  recordUri?: string
+  recordCid?: string
+  err: string | null
+}
+
+export type FailedImport = {
+  did: string
+  err: string | null
+}

packages/pds/src/migrate-script/load-dids.ts

@@ -0,0 +1,35 @@
+import { chunkArray } from '@atproto/common'
+import { setupEnv } from './util'
+
+const run = async () => {
+  const amount = parseInt(process.argv[2])
+  const pdsId = parseInt(process.argv[3])
+  console.log(`loading next ${amount} dids`)
+  const { db, ctx } = await setupEnv()
+
+  const didsRes = await ctx.db.db
+    .selectFrom('user_account')
+    .select('did')
+    .where('pdsId', 'is', null)
+    .orderBy('did', 'asc')
+    .limit(amount)
+    .execute()
+  const dids = didsRes.map((row) => ({
+    did: row.did,
+    phase: 0,
+    pdsId,
+    failed: 0 as const,
+  }))
+
+  await Promise.all(
+    chunkArray(dids, 50).map((chunk) =>
+      db
+        .insertInto('status')
+        .values(chunk)
+        .onConflict((oc) => oc.doNothing())
+        .execute(),
+    ),
+  )
+}
+
+run()

packages/pds/src/migrate-script/load-failed-blobs.ts

@@ -0,0 +1,31 @@
+import fs from 'fs/promises'
+import { chunkArray } from '@atproto/common'
+import { getDb } from './db'
+
+const run = async () => {
+  const file = await fs.readFile('missing-blobs.txt')
+  const rows = file
+    .toString()
+    .split('\n')
+    .filter((row) => row.length > 5)
+    .map((row) => {
+      const [did, cid] = row.split(' ')
+      return {
+        did: did.trim(),
+        cid: cid.trim(),
+      }
+    })
+  const db = getDb()
+
+  await Promise.all(
+    chunkArray(rows, 500).map((chunk) =>
+      db
+        .insertInto('failed_blob')
+        .values(chunk)
+        .onConflict((oc) => oc.doNothing())
+        .execute(),
+    ),
+  )
+}
+
+run()