@useatlas/create 0.0.1

package/templates/nextjs-standalone/src/lib/tools/__tests__/sql-duckdb.test.ts

@@ -0,0 +1,233 @@
+/**
+ * SQL validation tests specific to DuckDB queries.
+ *
+ * Verifies that:
+ * - Standard SELECT queries pass validation in DuckDB mode
+ * - DuckDB-specific forbidden operations (PRAGMA, ATTACH, etc.) are blocked
+ * - Common DuckDB query patterns work with the PostgreSQL parser mode
+ */
+import { describe, expect, it, beforeEach, afterEach, mock } from "bun:test";
+
+// Mock semantic layer
+mock.module("@atlas/api/lib/semantic", () => ({
+  getWhitelistedTables: () =>
+    new Set(["sales", "customers", "products", "orders"]),
+  _resetWhitelists: () => {},
+}));
+
+// Mock the DB connection
+const mockDBConnection = {
+  query: async () => ({ columns: [], rows: [] }),
+  close: async () => {},
+};
+
+const mockDetectDBType = () => {
+  const url = process.env.ATLAS_DATASOURCE_URL ?? "";
+  if (url.startsWith("duckdb://")) return "duckdb";
+  if (url.startsWith("postgresql://")) return "postgres";
+  throw new Error(`Unsupported: ${url}`);
+};
+
+mock.module("@atlas/api/lib/db/connection", () => ({
+  getDB: () => mockDBConnection,
+  connections: {
+    get: () => mockDBConnection,
+    getDefault: () => mockDBConnection,
+    getDBType: () => mockDetectDBType(),
+    getValidator: () => undefined,
+    list: () => ["default"],
+  },
+  detectDBType: mockDetectDBType,
+}));
+
+const { validateSQL } = await import("@atlas/api/lib/tools/sql");
+
+const origEnv = { ...process.env };
+
+describe("validateSQL — DuckDB mode", () => {
+  beforeEach(() => {
+    process.env.ATLAS_DATASOURCE_URL = "duckdb://:memory:";
+  });
+
+  afterEach(() => {
+    if (origEnv.ATLAS_DATASOURCE_URL === undefined) {
+      delete process.env.ATLAS_DATASOURCE_URL;
+    } else {
+      process.env.ATLAS_DATASOURCE_URL = origEnv.ATLAS_DATASOURCE_URL;
+    }
+  });
+
+  // --- Valid queries ---
+
+  it("allows simple SELECT", () => {
+    const result = validateSQL("SELECT * FROM sales");
+    expect(result.valid).toBe(true);
+  });
+
+  it("allows SELECT with aggregate functions", () => {
+    const result = validateSQL(
+      "SELECT COUNT(*), SUM(amount) FROM sales GROUP BY product_id"
+    );
+    expect(result.valid).toBe(true);
+  });
+
+  it("allows SELECT with CTE", () => {
+    const result = validateSQL(
+      "WITH totals AS (SELECT customer_id, SUM(amount) AS total FROM sales GROUP BY customer_id) SELECT * FROM totals"
+    );
+    expect(result.valid).toBe(true);
+  });
+
+  it("allows SELECT with window function", () => {
+    const result = validateSQL(
+      "SELECT *, ROW_NUMBER() OVER (PARTITION BY customer_id ORDER BY amount DESC) AS rn FROM sales"
+    );
+    expect(result.valid).toBe(true);
+  });
+
+  it("allows SELECT with JOIN", () => {
+    const result = validateSQL(
+      "SELECT s.*, c.name FROM sales s JOIN customers c ON s.customer_id = c.id"
+    );
+    expect(result.valid).toBe(true);
+  });
+
+  it("allows CAST expressions", () => {
+    const result = validateSQL(
+      "SELECT CAST(amount AS VARCHAR) FROM sales"
+    );
+    expect(result.valid).toBe(true);
+  });
+
+  it("allows COALESCE and CASE", () => {
+    const result = validateSQL(
+      "SELECT COALESCE(name, 'unknown'), CASE WHEN amount > 100 THEN 'high' ELSE 'low' END FROM sales"
+    );
+    expect(result.valid).toBe(true);
+  });
+
+  // --- Blocked operations ---
+
+  it("blocks PRAGMA", () => {
+    const result = validateSQL("PRAGMA database_list");
+    expect(result.valid).toBe(false);
+    expect(result.error).toContain("Forbidden");
+  });
+
+  it("blocks ATTACH", () => {
+    const result = validateSQL("ATTACH '/tmp/other.duckdb' AS other");
+    expect(result.valid).toBe(false);
+    expect(result.error).toContain("Forbidden");
+  });
+
+  it("blocks DETACH", () => {
+    const result = validateSQL("DETACH other");
+    expect(result.valid).toBe(false);
+    expect(result.error).toContain("Forbidden");
+  });
+
+  it("blocks INSTALL", () => {
+    const result = validateSQL("INSTALL httpfs");
+    expect(result.valid).toBe(false);
+    expect(result.error).toContain("Forbidden");
+  });
+
+  it("blocks EXPORT", () => {
+    const result = validateSQL("EXPORT DATABASE '/tmp/backup'");
+    expect(result.valid).toBe(false);
+    expect(result.error).toContain("Forbidden");
+  });
+
+  it("blocks IMPORT", () => {
+    const result = validateSQL("IMPORT DATABASE '/tmp/backup'");
+    expect(result.valid).toBe(false);
+    expect(result.error).toContain("Forbidden");
+  });
+
+  it("blocks CHECKPOINT", () => {
+    const result = validateSQL("CHECKPOINT");
+    expect(result.valid).toBe(false);
+    expect(result.error).toContain("Forbidden");
+  });
+
+  it("blocks INSERT", () => {
+    const result = validateSQL("INSERT INTO sales VALUES (1, 100)");
+    expect(result.valid).toBe(false);
+    expect(result.error).toContain("Forbidden");
+  });
+
+  it("blocks CREATE TABLE", () => {
+    const result = validateSQL("CREATE TABLE foo (id INT)");
+    expect(result.valid).toBe(false);
+    expect(result.error).toContain("Forbidden");
+  });
+
+  it("blocks DROP TABLE", () => {
+    const result = validateSQL("DROP TABLE sales");
+    expect(result.valid).toBe(false);
+    expect(result.error).toContain("Forbidden");
+  });
+
+  it("blocks DESCRIBE", () => {
+    const result = validateSQL("DESCRIBE sales");
+    expect(result.valid).toBe(false);
+    expect(result.error).toContain("Forbidden");
+  });
+
+  it("blocks SHOW", () => {
+    const result = validateSQL("SHOW TABLES");
+    expect(result.valid).toBe(false);
+    expect(result.error).toContain("Forbidden");
+  });
+
+  // --- Table whitelist ---
+
+  it("rejects queries on non-whitelisted tables", () => {
+    const result = validateSQL("SELECT * FROM secret_data");
+    expect(result.valid).toBe(false);
+    expect(result.error).toContain("not in the allowed list");
+  });
+
+  it("allows queries on whitelisted tables", () => {
+    const result = validateSQL("SELECT * FROM customers");
+    expect(result.valid).toBe(true);
+  });
+
+  // --- File-reading function blocks ---
+
+  it("blocks read_csv_auto", () => {
+    const result = validateSQL("SELECT * FROM read_csv_auto('/etc/passwd')");
+    expect(result.valid).toBe(false);
+    expect(result.error).toContain("Forbidden");
+  });
+
+  it("blocks read_parquet", () => {
+    const result = validateSQL("SELECT * FROM read_parquet('/data/secret.parquet')");
+    expect(result.valid).toBe(false);
+    expect(result.error).toContain("Forbidden");
+  });
+
+  it("blocks read_json_auto", () => {
+    const result = validateSQL("SELECT * FROM read_json_auto('/tmp/data.json')");
+    expect(result.valid).toBe(false);
+    expect(result.error).toContain("Forbidden");
+  });
+
+  it("blocks SET", () => {
+    const result = validateSQL("SET memory_limit='100GB'");
+    expect(result.valid).toBe(false);
+    expect(result.error).toContain("Forbidden");
+  });
+
+  it("blocks COPY ... TO (via base patterns)", () => {
+    const result = validateSQL("COPY sales TO '/tmp/exfil.csv'");
+    expect(result.valid).toBe(false);
+    expect(result.error).toContain("Forbidden");
+  });
+
+  it("blocks LOAD (via base patterns)", () => {
+    const result = validateSQL("LOAD httpfs");
+    expect(result.valid).toBe(false);
+    expect(result.error).toContain("Forbidden");
+  });
+});

package/templates/nextjs-standalone/src/lib/tools/__tests__/sql-ratelimit.test.ts

@@ -0,0 +1,225 @@
+/**
+ * Tests: executeSQL blocks queries when rate-limited and
+ * concurrency decrement survives query failures.
+ *
+ * Separate file from sql-audit.test.ts because mock.module() is
+ * process-global and irreversible — we need different mock behavior
+ * for acquireSourceSlot.
+ */
+
+import { describe, expect, it, beforeEach, afterEach, mock, type Mock } from "bun:test";
+import { _resetPool, type InternalPool } from "@atlas/api/lib/db/internal";
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+type AnyResult = any;
+
+// --- Mocks ---
+
+mock.module("@atlas/api/lib/semantic", () => ({
+  getWhitelistedTables: () => new Set(["companies", "people"]),
+  _resetWhitelists: () => {},
+}));
+
+let queryFn: Mock<(sql: string, timeout: number) => Promise<{ columns: string[]; rows: Record<string, unknown>[] }>>;
+
+const mockConn = {
+  query: (...args: [string, number]) => queryFn(...args),
+  close: async () => {},
+};
+
+mock.module("@atlas/api/lib/db/connection", () => ({
+  getDB: () => mockConn,
+  connections: {
+    get: () => mockConn,
+    getDefault: () => mockConn,
+    getDBType: () => "postgres",
+    getTargetHost: () => "localhost",
+    getValidator: () => undefined,
+    list: () => ["default"],
+  },
+  detectDBType: () => "postgres",
+}));
+
+mock.module("@atlas/api/lib/tracing", () => ({
+  withSpan: async (
+    _name: string,
+    _attrs: Record<string, unknown>,
+    fn: () => Promise<unknown>,
+  ) => fn(),
+}));
+
+// Mutable rate-limit mock — tests override these.
+// acquireSourceSlot atomically checks QPM + concurrency and increments both,
+// so there is no separate incrementSourceConcurrency call in sql.ts.
+let slotResult: { acquired: boolean; reason?: string; retryAfterMs?: number } = { acquired: true };
+const decrementCalls: string[] = [];
+
+mock.module("@atlas/api/lib/db/source-rate-limit", () => ({
+  acquireSourceSlot: () => slotResult,
+  decrementSourceConcurrency: (id: string) => { decrementCalls.push(id); },
+}));
+
+// Import after mocks
+const { executeSQL } = await import("@atlas/api/lib/tools/sql");
+
+// --- Internal DB mock pool to capture audit inserts ---
+
+let auditInserts: Array<{ sql: string; params?: unknown[] }> = [];
+
+const mockPool: InternalPool = {
+  query: async (sql: string, params?: unknown[]) => {
+    auditInserts.push({ sql, params });
+    return { rows: [] };
+  },
+  end: async () => {},
+  on: () => {},
+};
+
+function extractAuditParams(params: unknown[]) {
+  return {
+    userId: params[0],
+    userLabel: params[1],
+    authMode: params[2],
+    sql: params[3] as string,
+    durationMs: params[4] as number,
+    rowCount: params[5] as number | null,
+    success: params[6] as boolean,
+    error: params[7] as string | null,
+    sourceId: params[8] as string | null,
+    sourceType: params[9] as string | null,
+    targetHost: params[10] as string | null,
+  };
+}
+
+// --- Test helpers ---
+
+const exec = (sql: string) =>
+  executeSQL.execute!(
+    { sql, explanation: "test" },
+    { toolCallId: "test", messages: [], abortSignal: undefined as never },
+  ) as Promise<AnyResult>;
+
+const getAuditInserts = () =>
+  auditInserts.filter((q) => q.sql.includes("INSERT INTO audit_log"));
+
+// --- Tests ---
+
+describe("executeSQL rate-limit rejection", () => {
+  const origDbUrl = process.env.DATABASE_URL;
+  const origDatasource = process.env.ATLAS_DATASOURCE_URL;
+
+  beforeEach(() => {
+    auditInserts = [];
+    decrementCalls.length = 0;
+    process.env.ATLAS_DATASOURCE_URL = "postgresql://test:test@localhost:5432/test";
+    process.env.DATABASE_URL = "postgresql://test:test@localhost:5432/atlas";
+    _resetPool(mockPool);
+    queryFn = mock(() =>
+      Promise.resolve({
+        columns: ["id", "name"],
+        rows: [{ id: 1, name: "Acme" }],
+      }),
+    );
+    // Default: rate-limited
+    slotResult = {
+      acquired: false,
+      reason: "QPM limit reached (3/min)",
+      retryAfterMs: 5000,
+    };
+  });
+
+  afterEach(() => {
+    slotResult = { acquired: true };
+    if (origDbUrl) process.env.DATABASE_URL = origDbUrl;
+    else delete process.env.DATABASE_URL;
+    if (origDatasource) process.env.ATLAS_DATASOURCE_URL = origDatasource;
+    else delete process.env.ATLAS_DATASOURCE_URL;
+    _resetPool(null);
+  });
+
+  it("returns success: false with the rate limit reason", async () => {
+    const result = await exec("SELECT id, name FROM companies");
+    expect(result.success).toBe(false);
+    expect(result.error).toContain("QPM limit reached");
+  });
+
+  it("does not execute the query against the database", async () => {
+    await exec("SELECT id, name FROM companies");
+    expect(queryFn).not.toHaveBeenCalled();
+  });
+
+  it("passes retryAfterMs through in the response", async () => {
+    const result = await exec("SELECT id, name FROM companies");
+    expect(result.retryAfterMs).toBe(5000);
+  });
+
+  it("creates an audit log entry with 'Rate limited' in the error field", async () => {
+    await exec("SELECT id, name FROM companies");
+    const inserts = getAuditInserts();
+    expect(inserts).toHaveLength(1);
+    const audit = extractAuditParams(inserts[0].params!);
+    expect(audit.success).toBe(false);
+    expect(audit.error).toContain("Rate limited");
+    expect(audit.error).toContain("QPM limit reached");
+  });
+
+  it("does not decrement concurrency when rate-limited (slot not acquired)", async () => {
+    await exec("SELECT id, name FROM companies");
+    expect(decrementCalls).toHaveLength(0);
+  });
+});
+
+describe("executeSQL concurrency decrement survives query failures", () => {
+  const origDbUrl = process.env.DATABASE_URL;
+  const origDatasource = process.env.ATLAS_DATASOURCE_URL;
+
+  beforeEach(() => {
+    auditInserts = [];
+    decrementCalls.length = 0;
+    process.env.ATLAS_DATASOURCE_URL = "postgresql://test:test@localhost:5432/test";
+    process.env.DATABASE_URL = "postgresql://test:test@localhost:5432/atlas";
+    _resetPool(mockPool);
+    // Allow rate limiting (acquireSourceSlot atomically increments concurrency)
+    slotResult = { acquired: true };
+  });
+
+  afterEach(() => {
+    if (origDbUrl) process.env.DATABASE_URL = origDbUrl;
+    else delete process.env.DATABASE_URL;
+    if (origDatasource) process.env.ATLAS_DATASOURCE_URL = origDatasource;
+    else delete process.env.ATLAS_DATASOURCE_URL;
+    _resetPool(null);
+  });
+
+  it("decrements concurrency even when query throws", async () => {
+    queryFn = mock(() => Promise.reject(new Error("connection refused")));
+
+    const result = await exec("SELECT id FROM companies");
+
+    expect(result.success).toBe(false);
+    // acquireSourceSlot increments atomically; decrementSourceConcurrency
+    // is called in the finally block regardless of success/failure.
+    expect(decrementCalls).toHaveLength(1);
+    expect(decrementCalls[0]).toBe("default");
+  });
+
+  it("decrements concurrency on successful query", async () => {
+    queryFn = mock(() =>
+      Promise.resolve({ columns: ["id"], rows: [{ id: 1 }] }),
+    );
+
+    const result = await exec("SELECT id FROM companies");
+
+    expect(result.success).toBe(true);
+    expect(decrementCalls).toHaveLength(1);
+  });
+
+  it("concurrency decrement called exactly once after failure", async () => {
+    queryFn = mock(() => Promise.reject(new Error("timeout")));
+
+    await exec("SELECT id FROM companies");
+
+    // The finally block should decrement exactly once
+    expect(decrementCalls).toHaveLength(1);
+  });
+});