104 lines
3.3 KiB
JavaScript
104 lines
3.3 KiB
JavaScript
import { readFile } from "node:fs/promises";
|
||
import mysql from "mysql2/promise";
|
||
|
||
const allowedTables = [
|
||
"partners",
|
||
"tasks",
|
||
"contents",
|
||
"accounts",
|
||
"claims",
|
||
"delegation_bundles",
|
||
"distributions",
|
||
"collection_runs",
|
||
"users",
|
||
"auth_sessions",
|
||
"mcp_export_tokens",
|
||
];
|
||
|
||
function databaseUrl() {
|
||
if (process.env.DATABASE_URL) return process.env.DATABASE_URL;
|
||
const host = process.env.MYSQL_HOST ?? "127.0.0.1";
|
||
const port = process.env.MYSQL_PORT ?? "3306";
|
||
const user = encodeURIComponent(process.env.MYSQL_USER ?? "koc");
|
||
const password = encodeURIComponent(process.env.MYSQL_PASSWORD ?? "");
|
||
const database = process.env.MYSQL_DATABASE ?? "koc_loop";
|
||
return `mysql://${user}:${password}@${host}:${port}/${database}`;
|
||
}
|
||
|
||
function normalizeImportedValue(value) {
|
||
if (
|
||
typeof value === "string" &&
|
||
/^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(?:\.\d{1,3})?Z$/.test(value)
|
||
) {
|
||
return value.replace("T", " ").replace(/Z$/, "");
|
||
}
|
||
return value;
|
||
}
|
||
|
||
const inputPath = process.argv[2];
|
||
if (!inputPath) {
|
||
throw new Error("用法：npm run db:import-json -- /path/to/d1-export.json");
|
||
}
|
||
const parsed = JSON.parse(await readFile(inputPath, "utf8"));
|
||
const tables = parsed.tables && typeof parsed.tables === "object"
|
||
? parsed.tables
|
||
: parsed;
|
||
const pool = mysql.createPool({
|
||
uri: databaseUrl(),
|
||
connectionLimit: 2,
|
||
charset: "utf8mb4",
|
||
timezone: "Z",
|
||
dateStrings: true,
|
||
});
|
||
|
||
try {
|
||
for (const table of allowedTables) {
|
||
const rows = Array.isArray(tables[table]) ? tables[table] : [];
|
||
if (rows.length === 0) continue;
|
||
const [columnRows] = await pool.query(
|
||
`SELECT column_name, extra FROM information_schema.columns
|
||
WHERE table_schema = DATABASE() AND table_name = ? ORDER BY ordinal_position`,
|
||
[table],
|
||
);
|
||
const allowedColumns = new Set(
|
||
columnRows
|
||
.filter((column) => !String(column.extra ?? "").includes("GENERATED"))
|
||
.map((column) => String(column.column_name)),
|
||
);
|
||
const columns = [...new Set(rows.flatMap((row) => Object.keys(row)))]
|
||
.filter((column) => allowedColumns.has(column));
|
||
if (columns.length === 0) continue;
|
||
const quotedColumns = columns.map((column) => `\`${column}\``).join(", ");
|
||
const updates = columns
|
||
.filter((column) => column !== "id" && column !== "token_hash")
|
||
.map((column) => `\`${column}\` = VALUES(\`${column}\`)`)
|
||
.join(", ");
|
||
const connection = await pool.getConnection();
|
||
try {
|
||
await connection.beginTransaction();
|
||
for (let offset = 0; offset < rows.length; offset += 200) {
|
||
const chunk = rows.slice(offset, offset + 200);
|
||
const placeholders = chunk
|
||
.map(() => `(${columns.map(() => "?").join(", ")})`)
|
||
.join(", ");
|
||
const sql = `INSERT INTO \`${table}\` (${quotedColumns}) VALUES ${placeholders}${
|
||
updates ? ` ON DUPLICATE KEY UPDATE ${updates}` : ""
|
||
}`;
|
||
const values = chunk.flatMap((row) =>
|
||
columns.map((column) => normalizeImportedValue(row[column] ?? null)),
|
||
);
|
||
await connection.query(sql, values);
|
||
}
|
||
await connection.commit();
|
||
console.info(`imported ${table}: ${rows.length}`);
|
||
} catch (error) {
|
||
await connection.rollback();
|
||
throw error;
|
||
} finally {
|
||
connection.release();
|
||
}
|
||
}
|
||
} finally {
|
||
await pool.end();
|
||
}
|