Add splitSQLStatements to db-commons

ebardelli · ebardelli · commit e23de237a7ef · 2025-11-12T12:13:25.000-08:00
diff --git a/packages/datasources/duckdb/index.cjs b/packages/datasources/duckdb/index.cjs
@@ -3,12 +3,14 @@ const {
 	TypeFidelity,
 	asyncIterableToBatchedAsyncGenerator,
 	cleanQuery,
-	exhaustStream
+	exhaustStream,
+	splitSQLStatements
 } = require('@evidence-dev/db-commons');
 const { DuckDBInstance } = require('@duckdb/node-api');
 const path = require('path');
 const fs = require('fs/promises');
 
+
 /**
  * Converts BigInt values to Numbers in an object.
  * @param {Record<string, unknown>} obj - The input object with potential BigInt values.
@@ -150,49 +152,78 @@ const runQuery = async (queryString, database, batchSize = 100000) => {
 		}
 	}
 
-	// Prepare a reader that can stream chunks
-	const count_query = `WITH root as (${cleanQuery(queryString)}) SELECT COUNT(*) FROM root`;
+	// Split the incoming SQL into statements and treat all but the last as prefix
+	// statements which should be executed before metadata queries or the main
+	// streaming query. This allows SET/USE/CREATE statements to affect the
+	// session.
+	const statements = splitSQLStatements(queryString);
+	const prefixStatements = statements.slice(0, -1);
+	const mainStatement = statements.length > 0 ? statements[statements.length - 1] : '';
+
+	const count_query = mainStatement
+		? `WITH root as (${cleanQuery(mainStatement)}) SELECT COUNT(*) FROM root`
+		: null;
 	let expected_row_count = null;
 	try {
-		const countReader = await conn.runAndReadAll(count_query);
-		await countReader.readAll();
-		const countRow = countReader.getRowObjectsJS()?.[0];
-		if (countRow) {
-			// take the first value from the row (column name may vary)
-			expected_row_count = Object.values(countRow)[0];
+		// Execute prefix statements first so they apply to the session
+		for (const ps of prefixStatements) {
+			try {
+				await conn.run(ps).catch(() => null);
+			} catch (err) {
+				// ignore errors from prefix statements to mimic existing behavior
+			}
+		}
+
+		if (count_query) {
+			const countReader = await conn.runAndReadAll(count_query);
+			await countReader.readAll();
+			const countRow = countReader.getRowObjectsJS()?.[0];
+			if (countRow) {
+				// take the first value from the row (column name may vary)
+				expected_row_count = Object.values(countRow)[0];
+			}
 		}
 	} catch (e) {
 		expected_row_count = null;
 	}
 
-	const column_query = `DESCRIBE ${cleanQuery(queryString)}`;
+	const column_query = mainStatement ? `DESCRIBE ${cleanQuery(mainStatement)}` : null;
 	let column_types = null;
 	try {
-		const colReader = await conn.runAndReadAll(column_query);
-		await colReader.readAll();
-		const describeRows = colReader.getRowObjectsJS();
-		column_types = duckdbDescribeToEvidenceType(describeRows);
+		if (column_query) {
+			const colReader = await conn.runAndReadAll(column_query);
+			await colReader.readAll();
+			const describeRows = colReader.getRowObjectsJS();
+			column_types = duckdbDescribeToEvidenceType(describeRows);
+		}
 	} catch (e) {
 		column_types = null;
 	}
 
 	// Create an async generator that yields batches using the DuckDBResultReader
-	const reader = await conn.streamAndRead(queryString);
+	// Stream only the main statement (last statement). If there is no main
+	// statement, create an empty generator.
+	let reader = null;
+	if (mainStatement) {
+		reader = await conn.streamAndRead(mainStatement);
+	}
 
 	const rowsAsyncIterable = (async function* () {
 		try {
 			let prev = 0;
-			// keep reading until done
-			while (!reader.done) {
-				const target = prev + batchSize;
-				await reader.readUntil(target);
-				const allRows = reader.getRowObjectsJS();
-				const newRows = allRows.slice(prev).map(standardizeRow);
-				// yield rows one-by-one (asyncIterableToBatchedAsyncGenerator expects single-row yields)
-				for (const r of newRows) {
-					yield r;
+			if (reader) {
+				// keep reading until done
+				while (!reader.done) {
+					const target = prev + batchSize;
+					await reader.readUntil(target);
+					const allRows = reader.getRowObjectsJS();
+					const newRows = allRows.slice(prev).map(standardizeRow);
+					// yield rows one-by-one (asyncIterableToBatchedAsyncGenerator expects single-row yields)
+					for (const r of newRows) {
+						yield r;
+					}
+					prev = reader.currentRowCount;
 				}
-				prev = reader.currentRowCount;
 			}
 		} finally {
 			// disconnect the connection and close the instance synchronously
diff --git a/packages/datasources/duckdb/test/test.js b/packages/datasources/duckdb/test/test.js
@@ -4,6 +4,7 @@ import path from 'path';
 import { fileURLToPath } from 'url';
 import runQuery from '../index.cjs';
 import { batchedAsyncGeneratorToArray, TypeFidelity } from '@evidence-dev/db-commons';
+import fs from 'fs/promises';
 import 'dotenv/config';
 
 test('basic select from needful_things.duckdb', async () => {
@@ -198,6 +199,20 @@ test('query runs', async () => {
 	}
 });
 
+test('expectedRowCount present for UNION with NULL-first-row', async () => {
+	// Simple reproduction: first row has NULL, second row has a value
+ 	const query = `select NULL as a UNION ALL select 1 as a`;
+	const { rows, expectedRowCount } = await runQuery(query, undefined, 2);
+	// consume rows to ensure the stream runs
+	const arr = [];
+	for await (const batch of rows()) {
+		arr.push(...batch);
+	}
+	assert.equal(arr.length, 2);
+	assert.type(expectedRowCount, 'number');
+	assert.equal(expectedRowCount, 2);
+});
+
 test('query batches results properly', async () => {
 	try {
 		const { rows, expectedRowCount } = await runQuery(
@@ -220,4 +235,131 @@ test('query batches results properly', async () => {
 	}
 });
 
+test('handles leading SET statements before main query', async () => {
+	try {
+		const query = `
+		SET VARIABLE VAR1 = DATE '2023-10-04';
+		SET VARIABLE VAR2 = '23';
+		select 1 union all select 2 union all select 3;
+		`;
+		const { rows, expectedRowCount } = await runQuery(query, undefined, 2);
+		const arr = [];
+		for await (const batch of rows()) {
+			arr.push(batch);
+		}
+		// should batch into [2,1]
+		assert.equal(arr[0].length, 2);
+		assert.equal(arr[1].length, 1);
+		assert.equal(expectedRowCount, 3);
+	} catch (e) {
+		throw Error(e);
+	}
+});
+
+test('semicolon inside single-quoted string should not split statements', async () => {
+	try {
+		const query = "SET V='this;is;a;string'; select 1 union all select 2;";
+		const { rows, expectedRowCount } = await runQuery(query, undefined, 2);
+		const arr = [];
+		for await (const batch of rows()) {
+			arr.push(batch);
+		}
+		assert.equal(expectedRowCount, 2);
+		assert.equal(arr[0].length, 2);
+	} catch (e) {
+		throw Error(e);
+	}
+});
+
+test('semicolon inside block comment should not split statements', async () => {
+	try {
+		const query = "/* comment; still comment; */ select 1 union all select 2;";
+		const { rows, expectedRowCount } = await runQuery(query, undefined, 2);
+		const arr = [];
+		for await (const batch of rows()) {
+			arr.push(batch);
+		}
+		assert.equal(expectedRowCount, 2);
+		assert.equal(arr[0].length, 2);
+	} catch (e) {
+		throw Error(e);
+	}
+});
+
+test('handles leading SET statements before main query', async () => {
+	try {
+		const query = `
+		SET VARIABLE VAR1 = DATE '2023-10-04';
+		SET VARIABLE VAR2 = '23';
+		select 1 union all select 2 union all select 3;
+		`;
+		const { rows, expectedRowCount } = await runQuery(query, undefined, 2);
+		const arr = [];
+		for await (const batch of rows()) {
+			arr.push(batch);
+		}
+		// should batch into [2,1]
+		assert.equal(arr[0].length, 2);
+		assert.equal(arr[1].length, 1);
+		assert.equal(expectedRowCount, 3);
+	} catch (e) {
+		throw Error(e);
+	}
+});
+
+test('semicolon inside single-quoted string should not split statements', async () => {
+	try {
+		const query = "SET V='this;is;a;string'; select 1 union all select 2;";
+		const { rows, expectedRowCount } = await runQuery(query, undefined, 2);
+		const arr = [];
+		for await (const batch of rows()) {
+			arr.push(batch);
+		}
+		assert.equal(expectedRowCount, 2);
+		assert.equal(arr[0].length, 2);
+	} catch (e) {
+		throw Error(e);
+	}
+});
+
+test('semicolon inside block comment should not split statements', async () => {
+	try {
+		const query = '/* comment; still comment; */ select 1 union all select 2;';
+		const { rows, expectedRowCount } = await runQuery(query, undefined, 2);
+		const arr = [];
+		for await (const batch of rows()) {
+			arr.push(batch);
+		}
+		assert.equal(expectedRowCount, 2);
+		assert.equal(arr[0].length, 2);
+	} catch (e) {
+		throw Error(e);
+	}
+});
+
+test('USE statement before query should apply to metadata queries', async () => {
+	try {
+		// Create a temporary on-disk DuckDB file to test that USE affects
+		// subsequent queries in the same session. We create a schema and table
+		// in prefix statements, then USE the schema and select from the table.
+		// Use an in-memory database so we can CREATE schema/table in the test
+		const dbFile = ':memory:';
+		const query = `
+		CREATE SCHEMA s;
+		CREATE TABLE s.t AS SELECT 1 AS x UNION ALL SELECT 2 AS x;
+		USE s;
+		SELECT x FROM t;
+		`;
+		const { rows, expectedRowCount } = await runQuery(query, { filename: dbFile }, 2);
+		const arr = [];
+		for await (const batch of rows()) {
+			arr.push(batch);
+		}
+		assert.equal(expectedRowCount, 2);
+		assert.equal(arr[0].length, 2);
+	} catch (e) {
+		throw Error(e);
+	}
+});
+
 test.run();
diff --git a/packages/lib/db-commons/index.cjs b/packages/lib/db-commons/index.cjs
@@ -197,6 +197,88 @@ const cleanQuery = (query) => {
 	return cleanedString + '\n';
 };
 
+/**
+ * Split a SQL script into individual statements, respecting single/double/backtick
+ * quotes and block/line comments so semicolons inside those constructs don't split.
+ * Returns an array of statements (strings) without trailing semicolons and trimmed.
+ * @param {string} sql
+ * @returns {string[]}
+ */
+const splitSQLStatements = function (sql) {
+	const statements = [];
+	let cur = '';
+	let inSingle = false;
+	let inDouble = false;
+	let inBacktick = false;
+	let inLineComment = false;
+	let inBlockComment = false;
+	for (let i = 0; i < sql.length; i++) {
+		const ch = sql[i];
+		const next = sql[i + 1];
+
+		if (inLineComment) {
+			cur += ch;
+			if (ch === '\n') {
+				inLineComment = false;
+			}
+			continue;
+		}
+		if (inBlockComment) {
+			cur += ch;
+			if (ch === '*' && next === '/') {
+				cur += next;
+				i++;
+				inBlockComment = false;
+			}
+			continue;
+		}
+
+		// start of line comment
+		if (!inSingle && !inDouble && !inBacktick && ch === '-' && next === '-') {
+			cur += ch;
+			inLineComment = true;
+			continue;
+		}
+		// start of block comment
+		if (!inSingle && !inDouble && !inBacktick && ch === '/' && next === '*') {
+			cur += ch;
+			inBlockComment = true;
+			continue;
+		}
+
+		// quotes
+		if (!inDouble && !inBacktick && ch === "'") {
+			inSingle = !inSingle;
+			cur += ch;
+			continue;
+		}
+		if (!inSingle && !inBacktick && ch === '"') {
+			inDouble = !inDouble;
+			cur += ch;
+			continue;
+		}
+		if (!inSingle && !inDouble && ch === '`') {
+			inBacktick = !inBacktick;
+			cur += ch;
+			continue;
+		}
+
+		// semicolon splits statements only when not inside quotes or comments
+		if (ch === ';' && !inSingle && !inDouble && !inBacktick && !inLineComment && !inBlockComment) {
+			const s = cur.trim();
+			if (s.length > 0) statements.push(s);
+			cur = '';
+			continue;
+		}
+
+		cur += ch;
+	}
+
+	const last = cur.trim();
+	if (last.length > 0) statements.push(last);
+	return statements;
+};
+
 /**
  * @param {QueryResult} stream
  * @returns {Promise<void>}
@@ -217,5 +299,6 @@ exports.asyncIterableToBatchedAsyncGenerator = asyncIterableToBatchedAsyncGenera
 exports.batchedAsyncGeneratorToArray = batchedAsyncGeneratorToArray;
 exports.cleanQuery = cleanQuery;
 exports.exhaustStream = exhaustStream;
+exports.splitSQLStatements = splitSQLStatements;
 
 exports.getEnv = require('./src/getEnv.cjs').getEnv;
