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;
  }
  };
 }