Spaces:
Paused
Paused
Download scripts/migrate-legacy.mjs from imkrish/ig-costing: direct link, hf CLI and curl.
- Browser
- Download file 6.77 kB
-
https://huggingface.co/spaces/imkrish/ig-costing/resolve/main/scripts/migrate-legacy.mjs
- Command line
-
hf download hf://spaces/imkrish/ig-costing/scripts/migrate-legacy.mjs
-
curl -L -o migrate-legacy.mjs https://huggingface.co/spaces/imkrish/ig-costing/resolve/main/scripts/migrate-legacy.mjs
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())) | |