dbx/packages/node-core/tests/mongo-query.test.ts

import assert from "node:assert/strict";
import test from "node:test";
import {
  executeQuery,
  inferMongoColumns,
  mongoAggregateWriteStage,
  mongoDocumentsToQueryResult,
  parseMongoAggregateCommand,
  parseMongoCountDocumentsCommand,
  parseMongoFindCommand,
  parseMongoWriteCommand,
} from "../src/database.js";

test("parseMongoFindCommand accepts shell-style find commands", () => {
  assert.deepEqual(parseMongoFindCommand('db.getCollection("operation_logs").find({"level":"info"}).sort({"ts":-1}).skip(5).limit(10)'), {
    collection: "operation_logs",
    filter: '{"level":"info"}',
    skip: 5,
    limit: 10,
    sort: '{"ts":-1}',
  });
});

test("parseMongoFindCommand accepts Compass-style unquoted keys and ObjectId", () => {
  const command = parseMongoFindCommand("db.products.find({_id: ObjectId('6a045a92d2971e44243771a1')}).limit(1)");
  assert.ok(command);
  assert.equal(command.collection, "products");
  assert.equal(command.limit, 1);
  assert.deepEqual(JSON.parse(command.filter), { _id: { $oid: "6a045a92d2971e44243771a1" } });
});

test("parseMongoWriteCommand accepts unquoted update operator keys", () => {
  assert.deepEqual(parseMongoWriteCommand("db.projects.updateOne({_id: ObjectId('507f1f77bcf86cd799439011')}, {$set: {name: 'next'}})"), {
    kind: "update",
    collection: "projects",
    filter: '{"_id": {"$oid":"507f1f77bcf86cd799439011"}}',
    update: '{"$set": {"name": "next"}}',
    many: false,
  });
});

test("parseMongoCountDocumentsCommand accepts shell-style count commands", () => {
  assert.deepEqual(parseMongoCountDocumentsCommand('db.projects.countDocuments({"active":true})'), {
    collection: "projects",
    filter: '{"active":true}',
  });
});

test("parseMongoAggregateCommand accepts aggregate pipelines", () => {
  assert.deepEqual(parseMongoAggregateCommand('db.projects.aggregate([{"$match":{"active":true}},{"$group":{"_id":"$owner","total":{"$sum":1}}}])'), {
    collection: "projects",
    pipeline: '[{"$match":{"active":true}},{"$group":{"_id":"$owner","total":{"$sum":1}}}]',
  });
});

test("mongoAggregateWriteStage detects write stages", () => {
  assert.equal(mongoAggregateWriteStage('[{"$match":{"active":true}}]'), null);
  assert.equal(mongoAggregateWriteStage('[{"$match":{}},{"$out":"projects_dump"}]'), "$out");
  assert.equal(mongoAggregateWriteStage('[{"$merge":{"into":"projects_dump"}}]'), "$merge");
});

test("mongodb executeQuery blocks aggregate write stages until dangerous SQL is enabled", async () => {
  const oldAllowWrites = process.env.DBX_MCP_ALLOW_WRITES;
  const oldAllowDangerous = process.env.DBX_MCP_ALLOW_DANGEROUS_SQL;
  delete process.env.DBX_MCP_ALLOW_WRITES;
  delete process.env.DBX_MCP_ALLOW_DANGEROUS_SQL;
  const config = {
    id: "mongo",
    name: "mongo",
    db_type: "mongodb",
    host: "127.0.0.1",
    port: 27017,
    username: "",
    password: "",
    database: "app",
    ssh_enabled: false,
    ssl: false,
  } as const;

  await assert.rejects(
    executeQuery(config, 'db.projects.aggregate([{"$merge":{"into":"projects_dump"}}])'),
    /DBX_MCP_ALLOW_DANGEROUS_SQL=1/,
  );

  if (oldAllowWrites === undefined) delete process.env.DBX_MCP_ALLOW_WRITES;
  else process.env.DBX_MCP_ALLOW_WRITES = oldAllowWrites;
  if (oldAllowDangerous === undefined) delete process.env.DBX_MCP_ALLOW_DANGEROUS_SQL;
  else process.env.DBX_MCP_ALLOW_DANGEROUS_SQL = oldAllowDangerous;
});

test("parseMongoWriteCommand accepts supported write commands", () => {
  assert.deepEqual(parseMongoWriteCommand('db.projects.insertOne({"name":"demo"})'), {
    kind: "insert",
    collection: "projects",
    docsJson: '{"name":"demo"}',
  });
  assert.deepEqual(parseMongoWriteCommand('db.projects.updateOne({"_id":"1"},{"$set":{"name":"next"}})'), {
    kind: "update",
    collection: "projects",
    filter: '{"_id":"1"}',
    update: '{"$set":{"name":"next"}}',
    many: false,
  });
  assert.deepEqual(parseMongoWriteCommand('db.projects.deleteMany({"stale":true})'), {
    kind: "delete",
    collection: "projects",
    filter: '{"stale":true}',
    many: true,
  });
});

test("mongodb executeQuery blocks writes when writes are explicitly disabled", async () => {
  const oldAllowWrites = process.env.DBX_MCP_ALLOW_WRITES;
  process.env.DBX_MCP_ALLOW_WRITES = "0";
  await assert.rejects(
    executeQuery(
      {
        id: "mongo",
        name: "mongo",
        db_type: "mongodb",
        host: "127.0.0.1",
        port: 27017,
        username: "",
        password: "",
        database: "app",
        ssh_enabled: false,
        ssl: false,
      },
      'db.projects.insertOne({"name":"demo"})',
    ),
    /read-only/i,
  );
  if (oldAllowWrites === undefined) delete process.env.DBX_MCP_ALLOW_WRITES;
  else process.env.DBX_MCP_ALLOW_WRITES = oldAllowWrites;
});

test("mongoDocumentsToQueryResult turns documents into rows", () => {
  assert.deepEqual(mongoDocumentsToQueryResult([{ _id: "1", nested: { ok: true } }, { _id: "2", name: "demo" }], 2), {
    columns: ["_id", "nested", "name"],
    rows: [
      { _id: "1", nested: '{"ok":true}', name: undefined },
      { _id: "2", nested: undefined, name: "demo" },
    ],
    row_count: 2,
  });
});

test("inferMongoColumns marks _id as primary and reports observed types", () => {
  assert.deepEqual(inferMongoColumns([{ _id: "1", active: true }, { _id: "2", active: null }]), [
    {
      name: "_id",
      data_type: "string",
      is_nullable: false,
      column_default: null,
      is_primary_key: true,
      comment: null,
    },
    {
      name: "active",
      data_type: "boolean | null",
      is_nullable: true,
      column_default: null,
      is_primary_key: false,
      comment: null,
    },
  ]);
});