Add Database Explorer query policy tests

.github/extensions/database-explorer/extension.mjs

@@ -3,6 +3,7 @@ import { fileURLToPath } from "node:url";
 import { dirname, join, resolve } from "node:path";
 import { createClient } from "@libsql/client";
 import { CanvasError, createCanvas, joinSession } from "@github/copilot-sdk/extension";
+import { executeReadOnlyQuery } from "./query-policy.mjs";
 
 const servers = new Map();
 const MAX_DISPLAY_ROWS = 250;
@@ -28,86 +29,6 @@ function readRequestBody(request) {
     });
 }
 
-function stripCommentsAndValidateSingleStatement(sql) {
-    let output = "";
-    let quote = "";
-    let statementEnded = false;
-
-    for (let index = 0; index < sql.length; index += 1) {
-        const character = sql[index];
-        const next = sql[index + 1];
-
-        if (quote) {
-            output += character;
-            if (character === quote) {
-                if (next === quote) {
-                    output += next;
-                    index += 1;
-                } else {
-                    quote = "";
-                }
-            }
-            continue;
-        }
-
-        if (character === "'" || character === '"' || character === "`") {
-            quote = character;
-            output += character;
-            continue;
-        }
-
-        if (character === "-" && next === "-") {
-            index = sql.indexOf("\n", index + 2);
-            if (index === -1) {
-                break;
-            }
-            output += " ";
-            continue;
-        }
-
-        if (character === "/" && next === "*") {
-            const end = sql.indexOf("*/", index + 2);
-            if (end === -1) {
-                throw new Error("The SQL comment is not closed.");
-            }
-            index = end + 1;
-            output += " ";
-            continue;
-        }
-
-        if (character === ";") {
-            statementEnded = true;
-            continue;
-        }
-
-        if (statementEnded && !/\s/.test(character)) {
-            throw new Error("Run one SQL statement at a time.");
-        }
-
-        output += character;
-    }
-
-    if (quote) {
-        throw new Error("The SQL string is not closed.");
-    }
-
-    return output.trim();
-}
-
-function validateReadOnlyQuery(query) {
-    if (typeof query !== "string" || !query.trim()) {
-        throw new Error("Enter a SQL query.");
-    }
-
-    const statement = stripCommentsAndValidateSingleStatement(query);
-    const keyword = statement.match(/^([a-zA-Z]+)/)?.[1]?.toUpperCase();
-    if (keyword !== "SELECT" && keyword !== "WITH") {
-        throw new Error("Only read-only SELECT and WITH queries are allowed.");
-    }
-
-    return statement;
-}
-
 function normalizeValue(value) {
     return typeof value === "bigint" ? value.toString() : value;
 }
@@ -294,7 +215,7 @@ async function startServer(url) {
 
             if (request.url === "/api/query" && request.method === "POST") {
                 const body = JSON.parse(await readRequestBody(request));
-                const data = await executeQuery(url, validateReadOnlyQuery(body.query));
+                const data = await executeReadOnlyQuery(body.query, (query) => executeQuery(url, query));
                 response.writeHead(200, { "Content-Type": "application/json; charset=utf-8" });
                 response.end(JSON.stringify(data));
                 return;
@@ -336,7 +257,7 @@ await joinSession({
                     },
                     handler: async (ctx) => {
                         try {
-                            return executeQuery(databaseUrl(), validateReadOnlyQuery(ctx.input?.query));
+                            return executeReadOnlyQuery(ctx.input?.query, (query) => executeQuery(databaseUrl(), query));
                         } catch (error) {
                             throw new CanvasError("database_query_invalid", error instanceof Error ? error.message : "The query is invalid.");
                         }

.github/extensions/database-explorer/query-policy.mjs

@@ -0,0 +1,83 @@
+function stripCommentsAndValidateSingleStatement(sql) {
+    let output = "";
+    let quote = "";
+    let statementEnded = false;
+
+    for (let index = 0; index < sql.length; index += 1) {
+        const character = sql[index];
+        const next = sql[index + 1];
+
+        if (quote) {
+            output += character;
+            if (character === quote) {
+                if (next === quote) {
+                    output += next;
+                    index += 1;
+                } else {
+                    quote = "";
+                }
+            }
+            continue;
+        }
+
+        if (character === "'" || character === '"' || character === "`") {
+            quote = character;
+            output += character;
+            continue;
+        }
+
+        if (character === "-" && next === "-") {
+            index = sql.indexOf("\n", index + 2);
+            if (index === -1) {
+                break;
+            }
+            output += " ";
+            continue;
+        }
+
+        if (character === "/" && next === "*") {
+            const end = sql.indexOf("*/", index + 2);
+            if (end === -1) {
+                throw new Error("The SQL comment is not closed.");
+            }
+            index = end + 1;
+            output += " ";
+            continue;
+        }
+
+        if (character === ";") {
+            statementEnded = true;
+            continue;
+        }
+
+        if (statementEnded && !/\s/.test(character)) {
+            throw new Error("Run one SQL statement at a time.");
+        }
+
+        output += character;
+    }
+
+    if (quote) {
+        throw new Error("The SQL string is not closed.");
+    }
+
+    return output.trim();
+}
+
+export function validateReadOnlyQuery(query) {
+    if (typeof query !== "string" || !query.trim()) {
+        throw new Error("Enter a SQL query.");
+    }
+
+    const statement = stripCommentsAndValidateSingleStatement(query);
+    const keyword = statement.match(/^([a-zA-Z]+)/)?.[1]?.toUpperCase();
+    if (keyword !== "SELECT" && keyword !== "WITH") {
+        throw new Error("Only read-only SELECT and WITH queries are allowed.");
+    }
+
+    return statement;
+}
+
+export function executeReadOnlyQuery(query, execute) {
+    return execute(validateReadOnlyQuery(query));
+}

.github/extensions/database-explorer/query-policy.test.mjs

@@ -0,0 +1,47 @@
+import { describe, expect, it, vi } from "vitest";
+import { executeReadOnlyQuery, validateReadOnlyQuery } from "./query-policy.mjs";
+
+describe("Database Explorer query policy", () => {
+    it.each([
+        ["plain SELECT", "SELECT id FROM games", "SELECT id FROM games"],
+        ["quoted semicolon", "SELECT ';' AS value;", "SELECT ';' AS value"],
+        ["comments", "/* note */ SELECT id FROM games -- trailing comment", "SELECT id FROM games"],
+        [
+            "read-only CTE",
+            "WITH selected AS (SELECT id FROM games) SELECT id FROM selected",
+            "WITH selected AS (SELECT id FROM games) SELECT id FROM selected",
+        ],
+    ])("normalizes a supported %s query", (_name, input, expected) => {
+        expect(validateReadOnlyQuery(input)).toBe(expected);
+    });
+
+    it.each([
+        ["non-string input", null],
+        ["empty input", ""],
+        ["whitespace input", " \t "],
+        ["direct mutation", "DELETE FROM games"],
+        ["multiple statements", "SELECT id FROM games; DELETE FROM games"],
+        ["unterminated string", "SELECT 'unfinished"],
+        ["unterminated comment", "SELECT 1 /* unfinished"],
+    ])("rejects %s", (_name, input) => {
+        expect(() => validateReadOnlyQuery(input)).toThrow();
+    });
+
+    it("passes a normalized supported query to the executor", () => {
+        const execute = vi.fn((query) => query);
+        const rows = executeReadOnlyQuery(
+            "SELECT id FROM games;",
+            execute,
+        );
+
+        expect(rows).toBe("SELECT id FROM games");
+        expect(execute).toHaveBeenCalledExactlyOnceWith("SELECT id FROM games");
+    });
+
+    it("does not invoke the executor for a rejected query", () => {
+        const execute = vi.fn();
+
+        expect(() => executeReadOnlyQuery("DELETE FROM games", execute)).toThrow();
+        expect(execute).not.toHaveBeenCalled();
+    });
+});

vitest.config.ts

@@ -2,7 +2,11 @@ import { defineConfig } from 'vitest/config';
 
 export default defineConfig({
     test: {
-        include: ['db/**/*.test.ts', 'src/**/*.test.ts'],
+        include: [
+            'db/**/*.test.ts',
+            'src/**/*.test.ts',
+            '.github/extensions/database-explorer/**/*.test.mjs',
+        ],
         environment: 'node',
     },
 });