From 5480515ca95fdcdcfa0d50e9033286b59a84dcd5 Mon Sep 17 00:00:00 2001 From: dholms Date: Wed, 8 Nov 2023 17:41:39 -0600 Subject: [PATCH] refactor + better load-dids --- packages/pds/src/migrate-script/load-dids.ts | 25 +++++++++------ .../pds/src/migrate-script/migrate-all.ts | 28 +++-------------- .../pds/src/migrate-script/repair-blobs.ts | 28 ++--------------- .../pds/src/migrate-script/repair-prefs.ts | 27 ++-------------- packages/pds/src/migrate-script/util.ts | 31 +++++++++++++++++-- 5 files changed, 56 insertions(+), 83 deletions(-) diff --git a/packages/pds/src/migrate-script/load-dids.ts b/packages/pds/src/migrate-script/load-dids.ts index 807c8da4eaf..56e561fb772 100644 --- a/packages/pds/src/migrate-script/load-dids.ts +++ b/packages/pds/src/migrate-script/load-dids.ts @@ -1,19 +1,26 @@ -import fs from 'fs/promises' -import { getDb } from './db' import { chunkArray } from '@atproto/common' +import { setupEnv } from './util' const run = async () => { - const db = getDb() - const didsFile = await fs.readFile('dids.txt') - const dids = didsFile - .toString() - .split('\n') - .map((did) => ({ - did: did.trim(), + const amount = parseInt(process.argv[2]) + 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, failed: 0 as const, })) .filter((row) => row.did.length > 2) + await Promise.all( chunkArray(dids, 50).map((chunk) => db diff --git a/packages/pds/src/migrate-script/migrate-all.ts b/packages/pds/src/migrate-script/migrate-all.ts index b34cd33db9a..428ebe7b7eb 100644 --- a/packages/pds/src/migrate-script/migrate-all.ts +++ b/packages/pds/src/migrate-script/migrate-all.ts @@ -1,40 +1,25 @@ import assert from 'node:assert' -import dotenv from 'dotenv' import axios from 'axios' -import AtpAgent from '@atproto/api' import * as plcLib from '@did-plc/lib' import SqlRepoStorage from '../sql-repo-storage' import { createDeferrable } from '@atproto/common' -import { envToCfg, envToSecrets, readEnv } from '../config' import AppContext from '../context' -import { FailedTakedown, MigrateDb, Status, TransferPhase, getDb } from './db' +import { FailedTakedown, MigrateDb, Status, TransferPhase } from './db' import PQueue from 'p-queue' import { AdminHeaders, PdsInfo, - makeAdminHeaders, + getPds, repairBlob, repairFailedPrefs, retryOnce, + setupEnv, transferPreferences, } from './util' -dotenv.config() export const runScript = async () => { console.log('starting') - const db = getDb() - const env = readEnv() - const cfg = envToCfg(env) - const secrets = envToSecrets(env) - const ctx = await AppContext.fromConfig(cfg, secrets) - const adminHeaders = makeAdminHeaders(secrets) - const pdsRes = await ctx.db.db.selectFrom('pds').selectAll().execute() - const pdsInfos = pdsRes.map((row) => ({ - id: row.id, - did: row.did, - url: `https://${row.host}`, - agent: new AtpAgent({ service: `https://${row.host}` }), - })) + const { db, ctx, adminHeaders, pdsInfos } = await setupEnv() const todo = await db .selectFrom('status') .where('status.phase', '<', 7) @@ -61,10 +46,7 @@ export const runScript = async () => { status.pdsId = pdsInfos[pdsCounter % pdsInfos.length].id pdsCounter++ } - const pdsInfo = pdsInfos.find((info) => info.id === status.pdsId) - if (!pdsInfo) { - throw new Error(`could not find pds with id: ${status.pdsId}`) - } + const pdsInfo = getPds(pdsInfos, status.pdsId) migrateQueue.add(async () => { try { await retryOnce(() => diff --git a/packages/pds/src/migrate-script/repair-blobs.ts b/packages/pds/src/migrate-script/repair-blobs.ts index 2a836d3cc36..7574a05a6e5 100644 --- a/packages/pds/src/migrate-script/repair-blobs.ts +++ b/packages/pds/src/migrate-script/repair-blobs.ts @@ -1,26 +1,7 @@ -import dotenv from 'dotenv' -import AtpAgent from '@atproto/api' -import { envToCfg, envToSecrets, readEnv } from '../config' -import AppContext from '../context' -import { getDb } from './db' -import { makeAdminHeaders, repairBlob } from './util' - -dotenv.config() +import { getPds, repairBlob, setupEnv } from './util' export const runScript = async () => { - const db = getDb() - const env = readEnv() - const cfg = envToCfg(env) - const secrets = envToSecrets(env) - const ctx = await AppContext.fromConfig(cfg, secrets) - const adminHeaders = makeAdminHeaders(secrets) - const pdsRes = await ctx.db.db.selectFrom('pds').selectAll().execute() - const pdsInfos = pdsRes.map((row) => ({ - id: row.id, - did: row.did, - url: `https://${row.host}`, - agent: new AtpAgent({ service: `https://${row.host}` }), - })) + const { db, ctx, adminHeaders, pdsInfos } = await setupEnv() const failed = await db .selectFrom('failed_blob') .innerJoin('status', 'status.did', 'failed_blob.did') @@ -28,10 +9,7 @@ export const runScript = async () => { .execute() let count = 0 for (const blob of failed) { - const pdsInfo = pdsInfos.find((info) => info.id === blob.pdsId) - if (!pdsInfo) { - throw new Error(`could not find pds with id: ${blob.pdsId}`) - } + const pdsInfo = getPds(pdsInfos, blob.pdsId ?? -1) try { await repairBlob(ctx, db, pdsInfo, blob.did, blob.cid, adminHeaders) } catch (err) { diff --git a/packages/pds/src/migrate-script/repair-prefs.ts b/packages/pds/src/migrate-script/repair-prefs.ts index 4ef53bbc7da..78c86b81ab9 100644 --- a/packages/pds/src/migrate-script/repair-prefs.ts +++ b/packages/pds/src/migrate-script/repair-prefs.ts @@ -1,25 +1,7 @@ -import dotenv from 'dotenv' -import AtpAgent from '@atproto/api' -import { envToCfg, envToSecrets, readEnv } from '../config' -import AppContext from '../context' -import { getDb } from './db' -import { repairPrefs } from './util' - -dotenv.config() +import { getPds, repairPrefs, setupEnv } from './util' export const runScript = async () => { - const db = getDb() - const env = readEnv() - const cfg = envToCfg(env) - const secrets = envToSecrets(env) - const ctx = await AppContext.fromConfig(cfg, secrets) - const pdsRes = await ctx.db.db.selectFrom('pds').selectAll().execute() - const pdsInfos = pdsRes.map((row) => ({ - id: row.id, - did: row.did, - url: `https://${row.host}`, - agent: new AtpAgent({ service: `https://${row.host}` }), - })) + const { db, ctx, pdsInfos } = await setupEnv() const failed = await db .selectFrom('failed_pref') .innerJoin('status', 'status.did', 'failed_pref.did') @@ -27,10 +9,7 @@ export const runScript = async () => { .execute() let count = 0 for (const pref of failed) { - const pdsInfo = pdsInfos.find((info) => info.id === pref.pdsId) - if (!pdsInfo) { - throw new Error(`could not find pds with id: ${pref.pdsId}`) - } + const pdsInfo = getPds(pdsInfos, pref.pdsId ?? -1) try { await repairPrefs(ctx, db, pdsInfo, pref.did) } catch (err) { diff --git a/packages/pds/src/migrate-script/util.ts b/packages/pds/src/migrate-script/util.ts index b13af486cbf..2287c46ada4 100644 --- a/packages/pds/src/migrate-script/util.ts +++ b/packages/pds/src/migrate-script/util.ts @@ -1,10 +1,11 @@ +import dotenv from 'dotenv' import axios from 'axios' import * as ui8 from 'uint8arrays' import AtpAgent from '@atproto/api' import AppContext from '../context' -import { MigrateDb } from './db' +import { MigrateDb, getDb } from './db' import { CID } from 'multiformats/cid' -import { ServerSecrets } from '../config' +import { ServerSecrets, envToCfg, envToSecrets, readEnv } from '../config' export type PdsInfo = { id: number @@ -17,6 +18,32 @@ export type AdminHeaders = { authorization: string } +export const setupEnv = async () => { + dotenv.config() + const db = getDb() + const env = readEnv() + const cfg = envToCfg(env) + const secrets = envToSecrets(env) + const ctx = await AppContext.fromConfig(cfg, secrets) + const adminHeaders = makeAdminHeaders(secrets) + const pdsRes = await ctx.db.db.selectFrom('pds').selectAll().execute() + const pdsInfos = pdsRes.map((row) => ({ + id: row.id, + did: row.did, + url: `https://${row.host}`, + agent: new AtpAgent({ service: `https://${row.host}` }), + })) + return { db, ctx, adminHeaders, pdsInfos } +} + +export const getPds = (infos: PdsInfo[], id: number | null): PdsInfo => { + const pdsInfo = infos.find((info) => info.id === id) + if (!pdsInfo) { + throw new Error(`could not find pds with id: ${id}`) + } + return pdsInfo +} + export const makeAdminHeaders = (secrets: ServerSecrets): AdminHeaders => { const adminToken = ui8.toString( ui8.fromString(`admin:${secrets.adminPassword}`, 'utf8'),