xinmi/scripts/migrate_native_graph.mjs

#!/usr/bin/env node
import { spawn, spawnSync } from 'node:child_process';
import { createHash } from 'node:crypto';
import { createReadStream } from 'node:fs';
import { once } from 'node:events';
import { createInterface } from 'node:readline';
import { fileURLToPath } from 'node:url';

const snapshot = process.argv[2];
if (!snapshot || !process.argv.includes('--apply')) throw new Error('Usage: node scripts/migrate_native_graph.mjs SNAPSHOT --apply');
const database = process.env.GRAPH_MIGRATION_DATABASE || 'yawan_native_acceptance_20260915';
if (!/^yawan_native_acceptance_[0-9]+$/.test(database)) throw new Error('Only isolated acceptance databases are allowed');
const tenant = Number(process.env.GRAPH_MIGRATION_TENANT || 10059);
if (!Number.isSafeInteger(tenant) || tenant < 1) throw new Error('Invalid tenant');
const docker = ['--context', 'colima-yawan', 'exec', '-i', 'WeKnora-postgres', 'psql', '-X', '-v', 'ON_ERROR_STOP=1', '-U', 'postgres', '-d', database];
const hash = createHash('sha256');
for await (const part of createReadStream(snapshot)) hash.update(part);
const digest = hash.digest('hex');
const sqlite = (sql) => {
  const result = spawnSync('sqlite3', ['-readonly', '-json', snapshot, sql], {encoding:'utf8', maxBuffer:8*1024*1024});
  if (result.status !== 0) throw new Error(result.stderr);
  return JSON.parse(result.stdout || '[]');
};
const quote = value => `'${String(value).replaceAll("'", "''")}'`;
const sourceKBs = sqlite('SELECT id,name FROM knowledgebase ORDER BY id');
const proc = spawn('docker', docker, {stdio:['pipe','pipe','pipe']});
proc.stdout.pipe(process.stdout); proc.stderr.pipe(process.stderr);
const completion = once(proc, 'close');
const write = async data => {if (!proc.stdin.write(data)) await once(proc.stdin,'drain');};
await write('BEGIN; SET statement_timeout = 0;\n');
await write(await (await import('node:fs/promises')).readFile(fileURLToPath(new URL('../migrations/versioned/000200_native_graph.up.sql', import.meta.url)), 'utf8'));
await write(`\nCREATE TEMP TABLE graph_kb_map(legacy_id bigint PRIMARY KEY,native_id varchar(36) NOT NULL);\n`);
for (const kb of sourceKBs) {
  await write(`DO $$ BEGIN IF (SELECT count(*) FROM knowledge_bases WHERE tenant_id=${tenant} AND name=${quote(kb.name)} AND deleted_at IS NULL)<>1 THEN RAISE EXCEPTION 'Ambiguous or missing knowledge base ${kb.id}'; END IF; END $$;\nINSERT INTO graph_kb_map SELECT ${kb.id},id FROM knowledge_bases WHERE tenant_id=${tenant} AND name=${quote(kb.name)} AND deleted_at IS NULL;\n`);
}
const tables = ['knowledgegraphversion','knowledgeentity','knowledgerelation','knowledgegraphevidence','knowledgegraphchunkstate','knowledgegraphsourceunit','knowledgegraphclaim','knowledgeentityalias','knowledgegraphcommunity','knowledgegraphcommunityreport'];
const expected = {};
for (const table of tables) {
  const columns = sqlite(`PRAGMA table_info(${table})`).map(c=>c.name);
  expected[table] = sqlite(`SELECT count(*) AS n FROM ${table}`)[0].n;
  const sql = `SELECT json_object(${columns.flatMap(c=>[quote(c), '"'+c+'"']).join(',')}) FROM ${table} ORDER BY id;`;
  await write(`CREATE TEMP TABLE graph_import_${table}(payload jsonb);\nCOPY graph_import_${table}(payload) FROM STDIN WITH (FORMAT csv);\n`);
  const source = spawn('sqlite3',['-readonly','-noheader',snapshot,sql],{stdio:['ignore','pipe','pipe']});
  const sourceCompletion = once(source,'close');
  source.stderr.pipe(process.stderr);
  for await (const line of createInterface({input:source.stdout,crlfDelay:Infinity})) await write('"'+line.replaceAll('"','""')+'"\n');
  const [sourceCode]=await sourceCompletion;
  if(sourceCode!==0){proc.kill();throw new Error(`SQLite export failed: ${table}`);}
  await write(`\\.\nDO $$ BEGIN IF EXISTS(SELECT 1 FROM graph_import_${table} s JOIN graph_kb_map k ON k.legacy_id=(s.payload->>'kb_id')::bigint JOIN native_graph_records r ON r.knowledge_base_id=k.native_id AND r.tenant_id=${tenant} AND r.source_table=${quote(table)} AND r.legacy_id=(s.payload->>'id')::bigint WHERE r.payload<>s.payload) THEN RAISE EXCEPTION 'Source record differs from preserved graph archive'; END IF; END $$;\nINSERT INTO native_graph_records(tenant_id,knowledge_base_id,source_table,legacy_id,graph_version_id,payload,source_snapshot_sha256) SELECT ${tenant},k.native_id,${quote(table)},(s.payload->>'id')::bigint,COALESCE((s.payload->>'graph_version_id')::bigint,(s.payload->>'id')::bigint),s.payload,${quote(digest)} FROM graph_import_${table} s JOIN graph_kb_map k ON k.legacy_id=(s.payload->>'kb_id')::bigint ON CONFLICT DO NOTHING;\nDO $$ BEGIN IF (SELECT count(*) FROM native_graph_records WHERE tenant_id=${tenant} AND knowledge_base_id IN (SELECT native_id FROM graph_kb_map) AND source_table=${quote(table)})<>${expected[table]} THEN RAISE EXCEPTION 'Graph count mismatch ${table}'; END IF; END $$;\nDROP TABLE graph_import_${table};\n`);
}
await write(`
DO $$ BEGIN IF EXISTS(SELECT 1 FROM native_graph_records WHERE tenant_id=${tenant} AND knowledge_base_id IN (SELECT native_id FROM graph_kb_map) AND source_table='knowledgegraphevidence' GROUP BY knowledge_base_id,payload->>'source_chunk_id' HAVING count(DISTINCT payload->>'source_content_hash')>1) THEN RAISE EXCEPTION 'Conflicting source content hashes'; END IF; END $$;
INSERT INTO native_graph_sources(tenant_id,knowledge_base_id,source_chunk_id,source_file,source_content_hash,excerpt)
SELECT DISTINCT ON (knowledge_base_id,payload->>'source_chunk_id') tenant_id,knowledge_base_id,payload->>'source_chunk_id',payload->>'source_file',payload->>'source_content_hash',payload->>'excerpt'
FROM native_graph_records WHERE tenant_id=${tenant} AND knowledge_base_id IN (SELECT native_id FROM graph_kb_map) AND source_table='knowledgegraphevidence' ORDER BY knowledge_base_id,payload->>'source_chunk_id',length(payload->>'excerpt') DESC,legacy_id ON CONFLICT DO NOTHING;
CREATE TEMP TABLE graph_native_chunks AS SELECT c.id,c.knowledge_id,c.knowledge_base_id,k.file_name,encode(sha256(convert_to(c.content,'UTF8')),'hex') AS content_hash FROM chunks c JOIN knowledges k ON k.id=c.knowledge_id AND k.tenant_id=c.tenant_id AND k.knowledge_base_id=c.knowledge_base_id AND k.deleted_at IS NULL WHERE c.tenant_id=${tenant} AND c.knowledge_base_id IN (SELECT native_id FROM graph_kb_map) AND c.deleted_at IS NULL AND c.is_enabled=true AND c.chunk_type='text';
CREATE INDEX ON graph_native_chunks(knowledge_base_id,content_hash);
ANALYZE graph_native_chunks;
UPDATE native_graph_sources s SET native_chunk_id=m.id,native_knowledge_id=m.knowledge_id,native_content_sha256=m.content_hash,mapping_state='exact_content_hash' FROM (SELECT * FROM (SELECT *,count(*) OVER (PARTITION BY knowledge_base_id,file_name,content_hash) AS matches FROM graph_native_chunks) candidates WHERE matches=1) m WHERE s.tenant_id=${tenant} AND s.knowledge_base_id=m.knowledge_base_id AND regexp_replace(s.source_file,'\\.jsonl$','.md')=m.file_name AND s.source_content_hash=m.content_hash AND s.native_chunk_id IS NULL;
COMMIT;
VACUUM (ANALYZE) native_graph_records, native_graph_sources;
SELECT source_table,count(*) AS preserved FROM native_graph_records WHERE tenant_id=${tenant} GROUP BY source_table ORDER BY source_table;
SELECT mapping_state,count(*) AS source_chunks FROM native_graph_sources WHERE tenant_id=${tenant} GROUP BY mapping_state;
`);
proc.stdin.end();
const [code] = await completion;
if(code!==0)throw new Error(`PostgreSQL import failed (${code}); transaction rolled back`);
console.log(JSON.stringify({snapshot_sha256:digest,tenant,database,expected}));