feat: streaming responses from rust server

.github/dependabot.yml

@@ -16,3 +16,7 @@ updates:
     directory: "/packages/duckdb-server"
     schedule:
       interval: monthly
+  - package-ecosystem: cargo
+    directory: "/packages/duckdb-server-rust"
+    schedule:
+      interval: monthly

.github/workflows/test.yml

@@ -34,7 +34,7 @@ jobs:
       - run: npm run test
 
   python:
-    name: Test on Python
+    name: Test in Python
 
     runs-on: ubuntu-latest
 
@@ -64,3 +64,36 @@ jobs:
           hatch build
           hatch fmt --check
           hatch run test:cov
+
+  rust:
+    name: Test in Rust
+
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v4
+      - name: Install Rust
+        uses: dtolnay/rust-toolchain@stable
+        with:
+          components: clippy, rust-src, rustfmt
+
+      - name: Check formatting
+        run: |
+          cd packages/duckdb-server-rust
+          cargo fmt -- --check
+      - name: Check
+        run: |
+          cd packages/duckdb-server-rust
+          cargo check
+      - name: Clippy
+        run: |
+          cd packages/duckdb-server-rust
+          cargo clippy --all -- -D warnings
+      - name: Build
+        run: |
+          cd packages/duckdb-server-rust
+          cargo build --verbose
+      - name: Test
+        run: |
+          cd packages/duckdb-server-rust
+          cargo test --verbose

dev/bundle.html

@@ -17,6 +17,7 @@
       <select id="connectors">
         <option value="socket" selected>Socket</option>
         <option value="rest">REST</option>
+        <option value="rest_https">REST (HTTPS)</option>
       </select>
     </span>
   </header>
@@ -47,28 +48,31 @@
 
   async function create() {
     try {
-      await coordinator().createBundle('test', queries);
-      output.innerText = "Created bundle"
+      await coordinator.createBundle('test', queries);
+      output.innerText = "Created bundle";
     } catch(err) {
-      output.innerText = `Error: ${err}`
+      output.innerText = `Error: ${err}`;
+      console.log(err.stack);
     }
   }
 
   async function load() {
     try {
-      await coordinator().loadBundle('test');
-      output.innerText = "Loaded bundle"
+      await coordinator.loadBundle('test');
+      output.innerText = "Loaded bundle";
     } catch(err) {
-      output.innerText = `Error: ${err}`
+      output.innerText = `Error: ${err}`;
+      console.log(err.stack);
     }
   }
 
   async function query() {
     try {
-      const result = await coordinator().query(queries[1], {cache: false});
-      output.innerText = `Result = ${result}`
+      const result = await coordinator.query(queries[1], {cache: false});
+      output.innerText = `Result = ${result}`;
     } catch(err) {
-      output.innerText = `Error: ${err}`
+      output.innerText = `Error: ${err}`;
+      console.log(err.stack);
     }
   }
 </script>

dev/index.html

@@ -65,6 +65,7 @@
         <option value="wasm" selected>WASM</option>
         <option value="socket">Socket</option>
         <option value="rest">REST</option>
+        <option value="rest_https">REST (HTTPS)</option>
       </select>
     </span>
 

dev/setup.js

@@ -16,17 +16,20 @@ export function clear() {
 
 let wasm;
 
-export async function setDatabaseConnector(type, options) {
+export async function setDatabaseConnector(type) {
   let connector;
   switch (type) {
     case 'socket':
-      connector = socketConnector(options);
+      connector = socketConnector();
       break;
     case 'rest':
-      connector = restConnector(options);
+      connector = restConnector();
+      break;
+    case 'rest_https':
+      connector = restConnector('https://localhost:3000/');
       break;
     case 'wasm':
-      connector = wasm || (wasm = wasmConnector(options));
+      connector = wasm || (wasm = wasmConnector());
       break;
     default:
       throw new Error(`Unrecognized connector type: ${type}`);

package.json

@@ -19,6 +19,7 @@
     "lint": "lerna run lint",
     "test": "lerna run test",
     "server": "cd packages/duckdb-server && hatch run serve",
+    "server:rust": "cd packages/duckdb-server-rust && cargo run",
     "server:node": "nodemon packages/duckdb/bin/run-server.js",
     "dev": "vite",
     "release": "npm run test && npm run lint && lerna publish && npm run release:python",

packages/core/src/Coordinator.js

@@ -44,6 +44,7 @@ export class Coordinator {
     consolidate = true,
     indexes = {}
   } = {}) {
+    /** @type {QueryManager} */
     this.manager = manager;
     this.manager.cache(cache);
     this.manager.consolidate(consolidate);
@@ -152,11 +153,25 @@ export class Coordinator {
     return this.query(query, { ...options, cache: true, priority: Priority.Low });
   }
 
+  /**
+   * Create a bundle of queries that can be loaded into the cache.
+   *
+   * @param {string} name The name of the bundle.
+   * @param {[string | {sql: string},  {alias: string}]} queries The queries to save into the bundle.
+   * @param {number} priority Request priority.
+   * @returns 
+   */
   createBundle(name, queries, priority = Priority.Low) {
-    const options = { name, queries };
+    const options = { name, queries: queries.map(q => typeof q == 'string' ? {sql: q} : q) };
     return this.manager.request({ type: 'create-bundle', options }, priority);
   }
 
+  /**
+   * Load a bundle into the cache.
+   * @param {string} name The name of the bundle.
+   * @param {number} priority Request priority.
+   * @returns 
+   */
   loadBundle(name, priority = Priority.High) {
     const options = { name };
     return this.manager.request({ type: 'load-bundle', options }, priority);

packages/core/src/QueryManager.js

@@ -1,13 +1,13 @@
 import { consolidator } from './QueryConsolidator.js';
 import { lruCache, voidCache } from './util/cache.js';
-import { priorityQueue } from './util/priority-queue.js';
+import { PriorityQueue } from './util/priority-queue.js';
 import { QueryResult } from './util/query-result.js';
 
 export const Priority = { High: 0, Normal: 1, Low: 2 };
 
 export class QueryManager {
   constructor() {
-    this.queue = priorityQueue(3);
+    this.queue = new PriorityQueue(3);
     this.db = null;
     this.clientCache = null;
     this._logger = null;

packages/core/src/connectors/rest.js

@@ -5,7 +5,7 @@ export function restConnector(uri = 'http://localhost:3000/') {
     /**
      * Query the DuckDB server.
      * @param {object} query
-     * @param {'exec' | 'arrow' | 'json'} [query.type] The query type: 'exec', 'arrow', or 'json'.
+     * @param {'exec' | 'arrow' | 'json' | 'create-bundle' | 'load-bundle'} [query.type] The query type.
      * @param {string} query.sql A SQL query string.
      * @returns the query result
      */
@@ -19,9 +19,9 @@ export function restConnector(uri = 'http://localhost:3000/') {
         body: JSON.stringify(query)
       });
 
-      return query.type === 'exec' ? req
+      return query.type === 'json' ? (await req).json()
         : query.type === 'arrow' ? tableFromIPC(req)
-        : (await req).json();
+        : req;
     }
   };
 }

packages/core/src/connectors/socket.js

@@ -84,7 +84,7 @@ export function socketConnector(uri = 'ws://localhost:3000/') {
     /**
      * Query the DuckDB server.
      * @param {object} query
-     * @param {'exec' | 'arrow' | 'json'} [query.type] The query type: 'exec', 'arrow', or 'json'.
+     * @param {'exec' | 'arrow' | 'json' | 'create-bundle' | 'load-bundle'} [query.type] The query type.
      * @param {string} query.sql A SQL query string.
      * @returns the query result
      */

packages/core/src/connectors/wasm.js

@@ -45,17 +45,17 @@ export function wasmConnector(options = {}) {
     /**
      * Query the DuckDB-WASM instance.
      * @param {object} query
-     * @param {'exec' | 'arrow' | 'json'} [query.type] The query type: 'exec', 'arrow', or 'json'.
+     * @param {'exec' | 'arrow' | 'json' | 'create-bundle' | 'load-bundle'} [query.type] The query type.
      * @param {string} query.sql A SQL query string.
      * @returns the query result
      */
     query: async query => {
       const { type, sql } = query;
       const con = await getConnection();
       const result = await con.query(sql);
-      return type === 'exec' ? undefined
+      return type === 'json' ? result.toArray()
         : type === 'arrow' ? result
-        : result.toArray();
+        : undefined;
     }
   };
 }