ig-costing / scripts /migrate-legacy.mjs
imkrish's picture
deploy e26c13ff
296511e verified
Raw History Blame Contribute Delete
6.77 kB
/**
* Copy the legacy banaocosting cluster into the sharded clusters.
*
* node scripts/migrate-legacy.mjs llm_events --dry-run
* node scripts/migrate-legacy.mjs llm_events
* node scripts/migrate-legacy.mjs lambda_events
*
* Design notes, because this run takes hours and will be interrupted:
*
* - READ-ONLY on the source. Nothing is deleted from banaocosting, so a failed
* or half-finished run costs nothing but time.
* - RESUMABLE. Documents are walked in _id order and the last copied _id is
* checkpointed to disk, so re-running continues rather than restarting.
* - IDEMPOTENT. A unique index on event_id is created on each target first, so
* re-inserting an already-copied document fails with a duplicate-key error
* that is counted and ignored. Interrupting mid-batch is therefore safe.
* - SHARD-AWARE. Writes go to the first shard with room and roll to the next on
* a quota error, matching the app's fill-and-spill behaviour.
* - THROTTLED. Atlas free clusters cap at 100 operations/second and respond to
* bursts by throttling the connection, so batches are paced rather than fired
* as fast as the network allows.
*/
import { MongoClient } from 'mongodb'
import { readFileSync, writeFileSync, existsSync } from 'node:fs'
// ---- config -----------------------------------------------------------------
const COLLECTION = process.argv[2]
const DRY = process.argv.includes('--dry-run')
const BATCH = Number(process.env.MIGRATE_BATCH ?? 500)
const OPS_PER_SEC = Number(process.env.MIGRATE_OPS_PER_SEC ?? 80) // under Atlas's 100
const QUOTA_BYTES = 512 * 1024 * 1024
const ROLLOVER_AT = 0.9
if (!['llm_events', 'lambda_events'].includes(COLLECTION)) {
console.error('usage: node scripts/migrate-legacy.mjs <llm_events|lambda_events> [--dry-run]')
process.exit(1)
}
const GROUP = COLLECTION === 'llm_events' ? 'CostingDBURL' : 'LambdaDBURL'
const CHECKPOINT = new URL(`../.migrate-${COLLECTION}.json`, import.meta.url)
// ---- env --------------------------------------------------------------------
const env = {}
for (const line of readFileSync(new URL('../.env', import.meta.url), 'utf8').split('\n')) {
const t = line.trim()
if (!t || t.startsWith('#')) continue
const i = t.indexOf('=')
if (i > 0) env[t.slice(0, i).trim()] = t.slice(i + 1).trim()
}
const DB_NAME = env.MONGO_DB_NAME ?? 'costingDB'
const LEGACY = process.env.LEGACY_DB_URL ?? env.LEGACY_DB_URL
if (!LEGACY) {
console.error('LEGACY_DB_URL is not set (add it to backend/.env or pass it in the environment)')
process.exit(1)
}
const targetUris = Object.keys(env)
.map((k) => ({ k, n: k.match(new RegExp(`^${GROUP}(\\d+)$`, 'i'))?.[1] }))
.filter((x) => x.n !== undefined)
.sort((a, b) => Number(a.n) - Number(b.n))
.map((x) => ({ name: x.k, uri: env[x.k] }))
if (!targetUris.length) {
console.error(`no ${GROUP}* entries found in backend/.env`)
process.exit(1)
}
const sleep = (ms) => new Promise((r) => setTimeout(r, ms))
const isQuota = (e) => /over your space quota|quota exceeded|you are over/i.test(String(e))
// ---- connect ----------------------------------------------------------------
console.log(`migrating ${COLLECTION} -> ${targetUris.map((t) => t.name).join(', ')}${DRY ? ' [DRY RUN]' : ''}`)
const source = new MongoClient(LEGACY, { maxPoolSize: 4 })
await source.connect()
const src = source.db(DB_NAME).collection(COLLECTION)
const targets = []
for (const t of targetUris) {
const client = new MongoClient(t.uri, { maxPoolSize: 4 })
await client.connect()
targets.push({ ...t, client, col: client.db(DB_NAME).collection(COLLECTION) })
}
async function usedBytes(t) {
const s = await t.client.db(DB_NAME).command({ dbStats: 1, scale: 1 })
return (s.storageSize ?? 0) + (s.indexSize ?? 0)
}
// Unique event_id is what makes re-runs safe: duplicates are rejected by the
// index rather than silently doubling the data.
if (!DRY) {
for (const t of targets) {
await t.col.createIndex({ event_id: 1 }, { unique: true })
await t.col.createIndex({ created_at: -1 })
}
console.log('ensured event_id + created_at indexes on every target')
}
// ---- resume -----------------------------------------------------------------
let checkpoint = null
if (existsSync(CHECKPOINT)) {
checkpoint = JSON.parse(readFileSync(CHECKPOINT, 'utf8'))
console.log(`resuming after _id ${checkpoint.lastId} (${checkpoint.copied} already copied)`)
}
const total = await src.estimatedDocumentCount()
const query = checkpoint ? { _id: { $gt: new (await import('mongodb')).ObjectId(checkpoint.lastId) } } : {}
const cursor = src.find(query).sort({ _id: 1 }).batchSize(BATCH)
let copied = checkpoint?.copied ?? 0
let dupes = 0
let ti = 0
let batch = []
const started = Date.now()
async function pickTarget() {
while (ti < targets.length) {
const used = await usedBytes(targets[ti])
if (used < QUOTA_BYTES * ROLLOVER_AT) return targets[ti]
console.log(` ${targets[ti].name} is full (${(used / 1e6).toFixed(0)}MB) — rolling over`)
ti++
}
throw new Error('ALL TARGET SHARDS FULL — add more DB URLs to backend/.env and re-run; progress is checkpointed')
}
async function flush() {
if (!batch.length) return
if (DRY) { copied += batch.length; batch = []; return }
const target = await pickTarget()
try {
await target.col.insertMany(batch, { ordered: false })
copied += batch.length
} catch (err) {
if (isQuota(err)) {
ti++ // this shard filled mid-batch; retry the same batch on the next one
return flush()
}
// Duplicate keys mean those documents were already migrated - expected on
// resume, and the whole reason the run is safe to interrupt.
const dupCount = err?.writeErrors?.filter((w) => w.err?.code === 11000).length ?? 0
const other = (err?.writeErrors?.length ?? 0) - dupCount
if (other > 0) throw err
dupes += dupCount
copied += batch.length - dupCount
}
writeFileSync(CHECKPOINT, JSON.stringify({ lastId: String(batch.at(-1)._id), copied }))
batch = []
const pct = ((copied / total) * 100).toFixed(1)
const rate = copied / ((Date.now() - started) / 1000)
const eta = rate > 0 ? Math.round((total - copied) / rate / 60) : 0
process.stdout.write(`\r ${copied.toLocaleString()}/${total.toLocaleString()} (${pct}%) ${rate.toFixed(0)}/s eta ${eta}m dupes ${dupes} `)
await sleep((batch.length || BATCH) / OPS_PER_SEC * 1000)
}
for await (const doc of cursor) {
batch.push(doc)
if (batch.length >= BATCH) await flush()
}
await flush()
console.log(`\ndone: ${copied.toLocaleString()} copied, ${dupes.toLocaleString()} already present`)
console.log('source cluster was not modified')
await source.close()
await Promise.all(targets.map((t) => t.client.close()))