chore(extention-driver-duckdb): use promise all to pareralize duckdb queries, and use connection.all to speed up

Author: onlyjackfrost

packages/extension-driver-duckdb/src/lib/duckdbDataSource.ts

@@ -14,7 +14,7 @@ import {
   VulcanExtensionId,
 } from '@vulcan-sql/core';
 import * as path from 'path';
-import { buildSQL } from './sqlBuilder';
+import { buildSQL, chunkSize } from './sqlBuilder';
 import { DuckDBExtensionLoader } from './duckdbExtensionLoader';
 
 const getType = (value: any) => {
@@ -99,49 +99,74 @@ export class DuckDBDataSource extends DataSource<any, DuckDBOptions> {
     }
     const { db, configurationParameters, ...options } =
       this.dbMapping.get(profileName)!;
-    const builtSQL = buildSQL(sql, operations);
+    const [builtSQL, streamSQL] = buildSQL(sql, operations);
+
     // create new connection for each query
     const connection = db.connect();
     await this.loadExtensions(connection, configurationParameters);
-    const statement = connection.prepare(builtSQL);
     const parameters = Array.from(bindParams.values());
     this.logRequest(builtSQL, parameters, options);
+    if (streamSQL) this.logRequest(streamSQL, parameters, options);
 
-    const result = await statement.stream(...parameters);
-    const firstChunk = await result.nextChunk();
+    const [result, asyncIterable] = await Promise.all([
+      new Promise<duckdb.TableData>((resolve, reject) => {
+        const c = db.connect();
+        c.all(
+          builtSQL,
+          ...parameters,
+          (err: duckdb.DuckDbError | null, res: duckdb.TableData) => {
+            if (err) {
+              reject(err);
+            }
+            resolve(res);
+          }
+        );
+      }),
+      new Promise<duckdb.QueryResult | undefined>((resolve, reject) => {
+        if (!streamSQL) resolve(undefined);
+        try {
+          const c = db.connect();
+          const result = c.stream(streamSQL, ...parameters);
+          resolve(result);
+        } catch (err: any) {
+          reject(err);
+        }
+      }),
+    ]);
+    const asyncIterableStream = new Readable({
+      objectMode: true,
+      read: function () {
+        for (const row of result) {
+          this.push(row);
+        }
+        this.push(null);
+      },
+    });
+    if (result.length >= chunkSize) {
+      asyncIterableStream._read = async function () {
+        if (asyncIterable) {
+          for await (const row of asyncIterable) {
+            this.push(row);
+          }
+          this.push(null);
+        }
+      };
+      if (result) {
+        for (const row of result) {
+          asyncIterableStream.push(row);
+        }
+      }
+    }
     return {
       getColumns: () => {
-        if (!firstChunk || firstChunk.length === 0) return [];
-        return Object.keys(firstChunk[0]).map((columnName) => ({
+        if (!result || result.length === 0) return [];
+        return Object.keys(result[0]).map((columnName) => ({
           name: columnName,
-          type: getType(firstChunk[0][columnName as any]),
+          type: getType(result[0][columnName as any]),
         }));
       },
       getData: () => {
-        const stream = new Readable({
-          objectMode: true,
-          read() {
-            result.nextChunk().then((chunk) => {
-              if (!chunk) {
-                this.push(null);
-                return;
-              }
-              for (const row of chunk) {
-                this.push(row);
-              }
-            });
-          },
-        });
-        // Send the first chunk
-        if (firstChunk) {
-          for (const row of firstChunk) {
-            stream.push(row);
-          }
-        } else {
-          // If there is no data, close the stream.
-          stream.push(null);
-        }
-        return stream;
+        return asyncIterableStream;
       },
     };
   }

packages/extension-driver-duckdb/src/lib/sqlBuilder.ts

@@ -26,15 +26,23 @@ export const isNoOP = (
   return true;
 };
 
+const duckdbExecuteChunkSize =
+  process.env['DUCKDB_EXECUTE_CHUNK_SIZE'] || '2000';
+
+export const chunkSize = Number(duckdbExecuteChunkSize);
+
 export const buildSQL = (
   sql: string,
   operations: Partial<Parameterized<SQLClauseOperation>>
-): string => {
-  if (isNoOP(operations)) return sql;
+): string[] => {
+  if (isNoOP(operations) && !/^select/.test(sql.toLowerCase()))
+    return [sql, ''];
   let builtSQL = '';
   builtSQL += `SELECT * FROM (${removeEndingSemiColon(sql)})`;
   builtSQL = addLimit(builtSQL, operations.limit);
   builtSQL = addOffset(builtSQL, operations.offset);
-  builtSQL += ';';
-  return builtSQL;
+
+  const chunkSql = `SELECT * FROM (${builtSQL}) LIMIT ${chunkSize};`;
+  const streamSql = `SELECT * FROM (${builtSQL}) OFFSET ${chunkSize};`;
+  return [chunkSql, streamSql];
 };

packages/extension-driver-duckdb/tests/duckdbDataSource.spec.ts

@@ -206,7 +206,7 @@ it('Should throw error from upstream', async () => {
       operations: {} as any,
       profileName: 'mocked-profile',
     })
-  ).rejects.toThrow(/^Parser Error: syntax error at or near "wrong"/);
+  ).rejects.toThrow(/^Parser Error: syntax error at or near/);
 });
 
 it('Should return empty data and column with zero result', async () => {
@@ -275,7 +275,7 @@ it('Should print queries without binding when log-queries = true', async () => {
     profileName: 'mocked-profile',
   });
   // Assert
-  expect(logs.slice(-1)[0][0]).toBe(`select $1::INTEGER as test`);
+  expect(/select \$1::INTEGER as test/.test(logs.slice(-1)[0][0])).toBe(true);
 });
 
 it('Should print queries with binding when log-queries = true and log-parameters = true', async () => {
@@ -316,7 +316,7 @@ it('Should print queries with binding when log-queries = true and log-parameters
     profileName: 'mocked-profile',
   });
   // Assert
-  expect(logs.slice(-1)[0][0]).toBe(`select $1::INTEGER as test`);
+  expect(/select \$1::INTEGER as test/.test(logs.slice(-1)[0][0])).toBe(true);
   expect(logs.slice(-1)[0][1]).toEqual([1234]);
 });
 

packages/extension-driver-duckdb/tests/sqlBuilder.spec.ts

@@ -67,7 +67,10 @@ it('BuildSQL function should build sql with operations', async () => {
   // Act
   const result = builder.buildSQL(statement, { limit: '$1', offset: '$2' });
   // Arrange
-  expect(result).toBe(
-    'SELECT * FROM (SELECT * FROM users) LIMIT $1 OFFSET $2;'
+  expect(result[0]).toBe(
+    `SELECT * FROM (SELECT * FROM (SELECT * FROM users) LIMIT $1 OFFSET $2) LIMIT ${builder.chunkSize};`
+  );
+  expect(result[1]).toBe(
+    `SELECT * FROM (SELECT * FROM (SELECT * FROM users) LIMIT $1 OFFSET $2) OFFSET ${builder.chunkSize};`
   );
 });

types/duckdb.d.ts

@@ -40,6 +40,7 @@ declare module 'duckdb' {
       sql: string,
       ...params: [...any, Callback<Statement>] | []
     ): Statement;
+    stream(sql: any, ...args: any[]): QueryResult;
   }
 
   export class Database {