corp-transfer/scripts/migrate-corpid-receiver.js

"use strict";

const express = require("express");
const crypto = require("crypto");
const { MongoClient } = require("mongodb");
const { EJSON } = require("bson");

const CONFIG = {
  // 目标环境 MongoDB（内网）
  targetMongoUri: "mongodb://root:ykt%26GK87892@10.16.146.81:27017/admin",
  // receiver 监听端口（与内网 nginx proxy_pass 对应）
  port: 8081,
  // 迁移鉴权 token（请改成高强度随机串）
  token: "ca764511bdace06c64f22bb97dd19911bcc4f0c9a07b2211cbbc1d43e9271657",
};

const args = parseArgs(process.argv.slice(2));
const PORT = Number(args.port || CONFIG.port);
const TARGET_MONGO_URI = args.targetMongoUri || CONFIG.targetMongoUri;
const TOKEN = args.token || CONFIG.token;

if (!TARGET_MONGO_URI) {
  console.error("[receiver] Missing TARGET_MONGO_URI in environment");
  process.exit(1);
}

if (!TOKEN) {
  console.error("[receiver] Missing migration token");
  process.exit(1);
}

function parseArgs(argv) {
  const out = {};
  for (let i = 0; i < argv.length; i += 1) {
    const a = argv[i];
    if (a === "--port") out.port = argv[++i];
    else if (a === "--targetMongoUri") out.targetMongoUri = argv[++i];
    else if (a === "--token") out.token = argv[++i];
  }
  return out;
}

const app = express();
app.use(express.json({ limit: "30mb" }));

let client;
const appliedIndexes = new Set();

function safeDecode(value) {
  return EJSON.deserialize(value, { relaxed: false });
}

function auth(req, res, next) {
  const incoming = req.headers["x-migration-token"];
  if (!incoming || incoming !== TOKEN) {
    return res.status(401).json({ success: false, message: "unauthorized" });
  }
  next();
}

function getJobId(corpId) {
  return crypto.createHash("sha1").update(String(corpId)).digest("hex");
}

async function ensureClient() {
  if (client) return client;
  client = new MongoClient(TARGET_MONGO_URI, {
    maxPoolSize: 30,
    minPoolSize: 2,
    connectTimeoutMS: 30000,
    socketTimeoutMS: 60000,
    serverSelectionTimeoutMS: 30000,
  });
  await client.connect();
  return client;
}

app.get("/health", (_req, res) => {
  res.json({ success: true, service: "migration-receiver", port: PORT });
});

app.post("/migrate/init", auth, async (req, res) => {
  try {
    const { corpId } = req.body || {};
    if (!corpId) {
      return res.status(400).json({ success: false, message: "corpId required" });
    }

    await ensureClient();

    return res.json({ success: true, jobId: getJobId(corpId) });
  } catch (err) {
    return res.status(500).json({ success: false, message: err.message });
  }
});

app.post("/migrate/indexes", auth, async (req, res) => {
  try {
    const { dbName, collectionName, indexes } = req.body || {};
    if (!dbName || !collectionName || !Array.isArray(indexes)) {
      return res.status(400).json({ success: false, message: "invalid payload" });
    }

    const cli = await ensureClient();
    const col = cli.db(dbName).collection(collectionName);

    let created = 0;
    for (const rawIdx of indexes) {
      const idx = safeDecode(rawIdx);
      if (!idx || idx.name === "_id_" || !idx.key) continue;

      const cacheKey = `${dbName}.${collectionName}.${idx.name}`;
      if (appliedIndexes.has(cacheKey)) continue;

      const options = { ...idx };
      delete options.key;
      delete options.v;
      delete options.ns;

      await col.createIndex(idx.key, options);
      appliedIndexes.add(cacheKey);
      created += 1;
    }

    return res.json({ success: true, created });
  } catch (err) {
    return res.status(500).json({ success: false, message: err.message });
  }
});

app.post("/migrate/chunk", auth, async (req, res) => {
  try {
    const { dbName, collectionName, docs } = req.body || {};
    if (!dbName || !collectionName || !Array.isArray(docs)) {
      return res.status(400).json({ success: false, message: "invalid payload" });
    }

    const cli = await ensureClient();
    const col = cli.db(dbName).collection(collectionName);

    if (!docs.length) {
      return res.json({ success: true, upserted: 0, modified: 0 });
    }

    const operations = docs.map((raw) => {
      const doc = safeDecode(raw);
      if (!doc || typeof doc !== "object") {
        throw new Error("invalid document in chunk");
      }
      if (!doc._id) {
        throw new Error("_id is required for each document");
      }
      return {
        replaceOne: {
          filter: { _id: doc._id },
          replacement: doc,
          upsert: true,
        },
      };
    });

    const result = await col.bulkWrite(operations, { ordered: false });

    return res.json({
      success: true,
      upserted: result.upsertedCount || 0,
      modified: result.modifiedCount || 0,
      matched: result.matchedCount || 0,
    });
  } catch (err) {
    return res.status(500).json({ success: false, message: err.message });
  }
});

app.post("/migrate/finish", auth, async (req, res) => {
  const { corpId } = req.body || {};
  return res.json({ success: true, corpId, jobId: getJobId(corpId || "") });
});

app.listen(PORT, () => {
  console.log(`[receiver] Listening on 0.0.0.0:${PORT}`);
});

process.on("SIGINT", async () => {
  try {
    if (client) await client.close();
  } finally {
    process.exit(0);
  }
});

process.on("SIGTERM", async () => {
  try {
    if (client) await client.close();
  } finally {
    process.exit(0);
  }
});