squirreling 0.6.0 → 0.7.0

package/README.md

@@ -19,6 +19,8 @@ Squirreling is a streaming async SQL engine for JavaScript. It is designed to pr
 - Lets you move query execution closer to your users
 - Supports standard SQL queries
 - Async streaming for large datasets
+- Native javascript Promises, AsyncGenerators, AbortSignals
+- Async user-defined functions (UDFs)
 - Constant memory usage for simple queries with LIMIT
 - Robust error handling and validation designed for LLM tool use
 - In-memory data option for simple use cases
@@ -27,12 +29,12 @@ Squirreling is a streaming async SQL engine for JavaScript. It is designed to pr
 
 ## Usage
 
-Squirreling returns an async generator, allowing you to process rows one at a time without loading everything into memory.
+Squirreling returns an AsyncGenerator of AsyncRows, allowing you to process rows one at a time without loading everything into memory. AsyncRows are made up of AsyncCells, allowing for late materialization of values.
 
 ```typescript
 import { executeSql } from 'squirreling'
 
-// In-memory table
+// Input table (in-memory for this example)
 const users = [
   { id: 1, name: 'Alice', active: true },
   { id: 2, name: 'Bob', active: false },
@@ -40,35 +42,37 @@ const users = [
   // ...more rows
 ]
 
+// Squirreling return types
 interface AsyncRow {
   columns: string[]
   cells: Record<string, AsyncCell>
 }
 type AsyncCell = () => Promise<SqlPrimitive>
 
-// Returns an async iterable of rows with async cells
+// Returns an AsyncIterable of rows with async cell loading
 const asyncRows: AsyncIterable<AsyncRow> = executeSql({
   tables: { users },
-  query: 'SELECT count(*) as cnt FROM users WHERE active = TRUE LIMIT 10',
+  query: 'SELECT * FROM users',
 })
 
 // Process rows as they arrive (streaming)
-for await (const { cnt } of asyncRows) {
-  console.log('Count', await cnt())
+for await (const { id, name } of asyncRows) {
+  console.log(`User id=${await id()}, name=${await name()}`)
 }
 ```
 
-There is an exported helper function `collect` to gather all rows into an array if needed:
+Squirreling exports a helper function `collect` to gather all rows into an array:
 
 ```javascript
 import { collect, executeSql } from 'squirreling'
 
 // Collect all rows and cells into a materialized array
-const allUsers: Record<string, SqlPrimitive>[] = await collect(executeSql({
+const rows: Record<string, SqlPrimitive>[] = await collect(executeSql({
   tables: { users },
-  query: 'SELECT * FROM users',
+  query: 'SELECT active, count(*) as cnt FROM users GROUP BY active',
 }))
-console.log(allUsers)
+console.log(`Collected rows:`, rows)
+// Collected rows: [ { active: true, cnt: 2 }, { active: false, cnt: 1 } ]
 ```
 
 ## Supported SQL Features

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "squirreling",
-  "version": "0.6.0",
+  "version": "0.7.0",
   "description": "Squirreling SQL Engine",
   "author": "Hyperparam",
   "homepage": "https://hyperparam.app",
@@ -37,11 +37,11 @@
     "test": "vitest run"
   },
   "devDependencies": {
-    "@types/node": "24.10.2",
-    "@vitest/coverage-v8": "4.0.15",
-    "eslint": "9.39.1",
+    "@types/node": "24.10.4",
+    "@vitest/coverage-v8": "4.0.16",
+    "eslint": "9.39.2",
     "eslint-plugin-jsdoc": "61.5.0",
     "typescript": "5.9.3",
-    "vitest": "4.0.15"
+    "vitest": "4.0.16"
   }
 }

package/src/backend/dataSource.js

@@ -1,8 +1,7 @@
 /**
- * @import { AsyncCell, AsyncCells, AsyncDataSource, AsyncRow, SqlPrimitive } from '../types.js'
+ * @import { AsyncCell, AsyncCells, AsyncDataSource, AsyncRow, ScanOptions, SqlPrimitive } from '../types.js'
  */
 
-
 /**
  * Wraps an async generator of plain objects into an AsyncDataSource
  *
@@ -11,8 +10,11 @@
  */
 export function generatorSource(gen) {
   return {
-    async *scan() {
-      yield* gen
+    async *scan({ signal }) {
+      for await (const row of gen) {
+        if (signal?.aborted) break
+        yield row
+      }
     },
   }
 }
@@ -40,8 +42,9 @@ function asyncRow(obj) {
  */
 export function memorySource(data) {
   return {
-    async *scan() {
+    async *scan({ signal }) {
       for (const item of data) {
+        if (signal?.aborted) break
         yield asyncRow(item)
       }
     },
@@ -58,11 +61,14 @@ export function cachedDataSource(source) {
   const cache = new Map()
   return {
     /**
+     * @param {ScanOptions} options
      * @yields {AsyncRow}
      */
-    async *scan() {
+    async *scan(options) {
+      const { signal } = options
       let index = 0
-      for await (const row of source.scan()) {
+      for await (const row of source.scan(options)) {
+        if (signal?.aborted) break
         const rowIndex = index
         /** @type {AsyncCells} */
         const cells = {}

package/src/execute/columns.js

@@ -1,7 +1,46 @@
+import { isAggregateFunc } from '../validation.js'
+
 /**
  * @import { ExprNode, SelectStatement, SelectColumn } from '../types.js'
  */
 
+/**
+ * Checks if an expression contains any aggregate function calls
+ *
+ * @param {ExprNode | undefined} expr
+ * @returns {boolean}
+ */
+export function containsAggregate(expr) {
+  if (!expr) return false
+  if (expr.type === 'function' && isAggregateFunc(expr.name.toUpperCase())) {
+    return true
+  }
+  if (expr.type === 'binary') {
+    return containsAggregate(expr.left) || containsAggregate(expr.right)
+  }
+  if (expr.type === 'unary') {
+    return containsAggregate(expr.argument)
+  }
+  if (expr.type === 'cast') {
+    return containsAggregate(expr.expr)
+  }
+  if (expr.type === 'case') {
+    if (expr.caseExpr && containsAggregate(expr.caseExpr)) return true
+    for (const when of expr.whenClauses) {
+      if (containsAggregate(when.condition) || containsAggregate(when.result)) return true
+    }
+    if (containsAggregate(expr.elseResult)) return true
+  }
+  if (expr.type === 'in valuelist') {
+    if (containsAggregate(expr.expr)) return true
+    for (const val of expr.values) {
+      if (containsAggregate(val)) return true
+    }
+  }
+  // Note: Don't recurse into subqueries - they have their own aggregate scope
+  return false
+}
+
 /**
  * Extracts column names needed from a SELECT statement.
  *
@@ -50,11 +89,6 @@ export function extractColumns(select) {
 function collectColumnsFromSelectColumn(col, columns) {
   if (col.kind === 'derived') {
     collectColumnsFromExpr(col.expr, columns)
-  } else if (col.kind === 'aggregate') {
-    if (col.arg.kind === 'expression') {
-      collectColumnsFromExpr(col.arg.expr, columns)
-    }
-    // 'star' aggregate (COUNT(*)) doesn't reference specific columns
   }
   // 'star' columns handled separately (returns undefined for all columns)
 }

package/src/execute/execute.js

@@ -2,15 +2,14 @@ import { missingClauseError } from '../parseErrors.js'
 import { tableNotFoundError, unsupportedOperationError } from '../executionErrors.js'
 import { generatorSource, memorySource } from '../backend/dataSource.js'
 import { parseSql } from '../parse/parse.js'
-import { defaultAggregateAlias, evaluateAggregate } from './aggregates.js'
-import { extractColumns } from './columns.js'
+import { containsAggregate, extractColumns } from './columns.js'
 import { evaluateExpr } from './expression.js'
 import { evaluateHavingExpr } from './having.js'
 import { executeJoins } from './join.js'
 import { compareForTerm, defaultDerivedAlias, stringify } from './utils.js'
 
 /**
- * @import { AsyncCells, AsyncDataSource, AsyncRow, ExecuteSqlOptions, OrderByItem, QueryHints, SelectStatement, SqlPrimitive } from '../types.js'
+ * @import { AsyncCells, AsyncDataSource, AsyncRow, ExecuteSqlOptions, OrderByItem, QueryHints, SelectStatement, SqlPrimitive, UserDefinedFunction } from '../types.js'
  */
 
 /**
@@ -19,8 +18,8 @@ import { compareForTerm, defaultDerivedAlias, stringify } from './utils.js'
  * @param {ExecuteSqlOptions} options - the execution options
  * @yields {AsyncRow} async generator yielding result rows
  */
-export async function* executeSql({ tables, query }) {
-  const select = typeof query === 'string' ? parseSql(query) : query
+export async function* executeSql({ tables, query, functions, signal }) {
+  const select = typeof query === 'string' ? parseSql({ query, functions }) : query
 
   // Check for unsupported operations
   if (!select.from) {
@@ -41,41 +40,44 @@ export async function* executeSql({ tables, query }) {
     }
   }
 
-  yield* executeSelect(select, normalizedTables)
+  yield* executeSelect({ select, tables: normalizedTables, functions, signal })
 }
 
 /**
  * Executes a SELECT query against the provided tables
  *
- * @param {SelectStatement} select
- * @param {Record<string, AsyncDataSource>} tables
+ * @param {Object} options
+ * @param {SelectStatement} options.select
+ * @param {Record<string, AsyncDataSource>} options.tables
+ * @param {Record<string, UserDefinedFunction>} [options.functions]
+ * @param {AbortSignal} [options.signal]
  * @yields {AsyncRow}
  */
-export async function* executeSelect(select, tables) {
+export async function* executeSelect({ select, tables, functions, signal }) {
   /** @type {AsyncDataSource} */
   let dataSource
   /** @type {string} */
-  let fromTableName
+  let leftTable
 
   if (select.from.kind === 'table') {
     // Use alias for column prefixing, but look up the actual table name
-    fromTableName = select.from.alias ?? select.from.table
+    leftTable = select.from.alias ?? select.from.table
     dataSource = tables[select.from.table]
     if (dataSource === undefined) {
       throw tableNotFoundError({ tableName: select.from.table })
     }
   } else {
     // Nested subquery - recursively resolve
-    fromTableName = select.from.alias
-    dataSource = generatorSource(executeSelect(select.from.query, tables))
+    leftTable = select.from.alias
+    dataSource = generatorSource(executeSelect({ select: select.from.query, tables, functions, signal }))
   }
 
   // Execute JOINs if present
   if (select.joins.length) {
-    dataSource = await executeJoins(dataSource, select.joins, fromTableName, tables)
+    dataSource = await executeJoins({ leftSource: dataSource, joins: select.joins, leftTable, tables, functions })
   }
 
-  yield* evaluateSelectAst(select, dataSource, tables)
+  yield* evaluateSelectAst({ select, dataSource, tables, functions, signal })
 }
 
 /**
@@ -116,17 +118,20 @@ async function applyDistinct(rows, distinct) {
   }
   return result
 }
+
 /**
  * Applies ORDER BY sorting to rows using multi-pass lazy evaluation.
  * Secondary ORDER BY columns are only evaluated for rows that tie on
  * previous columns, reducing expensive cell evaluations.
  *
- * @param {AsyncRow[]} rows - the input rows
- * @param {OrderByItem[]} orderBy - the sort specifications
- * @param {Record<string, AsyncDataSource>} tables
+ * @param {Object} options
+ * @param {AsyncRow[]} options.rows - the input rows
+ * @param {OrderByItem[]} options.orderBy - the sort specifications
+ * @param {Record<string, AsyncDataSource>} options.tables
+ * @param {Record<string, UserDefinedFunction>} [options.functions]
  * @returns {Promise<AsyncRow[]>} the sorted rows
  */
-async function sortRows(rows, orderBy, tables) {
+async function sortRows({ rows, orderBy, tables, functions }) {
   if (!orderBy.length) return rows
 
   // Cache for evaluated values: evaluatedValues[rowIdx][colIdx]
@@ -157,6 +162,7 @@ async function sortRows(rows, orderBy, tables) {
             node: term.expr,
             row: rows[idx],
             tables,
+            functions,
           })
         }
       }
@@ -204,24 +210,27 @@ async function sortRows(rows, orderBy, tables) {
 /**
  * Evaluates a select with a resolved FROM data source
  *
- * @param {SelectStatement} select
- * @param {AsyncDataSource} dataSource
- * @param {Record<string, AsyncDataSource>} tables
+ * @param {Object} options
+ * @param {SelectStatement} options.select
+ * @param {AsyncDataSource} options.dataSource
+ * @param {Record<string, AsyncDataSource>} options.tables
+ * @param {Record<string, UserDefinedFunction>} [options.functions]
+ * @param {AbortSignal} [options.signal]
  * @yields {AsyncRow}
  */
-async function* evaluateSelectAst(select, dataSource, tables) {
+async function* evaluateSelectAst({ select, dataSource, tables, functions, signal }) {
   // SQL priority: from, where, group by, having, select, order by, offset, limit
 
-  const hasAggregate = select.columns.some(col => col.kind === 'aggregate')
+  const hasAggregate = select.columns.some(col => col.kind === 'derived' && containsAggregate(col.expr))
   const useGrouping = hasAggregate || select.groupBy.length > 0
   const needsBuffering = useGrouping || select.orderBy.length > 0
 
   if (needsBuffering) {
     // BUFFERING PATH: Collect all rows, process, then yield
-    yield* evaluateBuffered(select, dataSource, tables, hasAggregate, useGrouping)
+    yield* evaluateBuffered({ select, dataSource, tables, functions, hasAggregate, useGrouping, signal })
   } else {
     // STREAMING PATH: Yield rows one by one
-    yield* evaluateStreaming(select, dataSource, tables)
+    yield* evaluateStreaming({ select, dataSource, tables, functions, signal })
   }
 }
 
@@ -229,12 +238,15 @@ async function* evaluateSelectAst(select, dataSource, tables) {
  * Streaming evaluation for simple queries (no ORDER BY or GROUP BY)
  * Supports DISTINCT by tracking seen row keys without buffering full rows
  *
- * @param {SelectStatement} select
- * @param {AsyncDataSource} dataSource
- * @param {Record<string, AsyncDataSource>} tables
+ * @param {Object} options
+ * @param {SelectStatement} options.select
+ * @param {AsyncDataSource} options.dataSource
+ * @param {Record<string, AsyncDataSource>} options.tables
+ * @param {Record<string, UserDefinedFunction>} [options.functions]
+ * @param {AbortSignal} [options.signal]
  * @yields {AsyncRow}
  */
-async function* evaluateStreaming(select, dataSource, tables) {
+async function* evaluateStreaming({ select, dataSource, tables, functions, signal }) {
   let rowsYielded = 0
   let rowsSkipped = 0
   let rowIndex = 0
@@ -255,11 +267,11 @@ async function* evaluateStreaming(select, dataSource, tables) {
     offset: select.offset,
   }
 
-  for await (const row of dataSource.scan(hints)) {
+  for await (const row of dataSource.scan({ hints, signal })) {
     rowIndex++
     // WHERE filter
     if (select.where) {
-      const pass = await evaluateExpr({ node: select.where, row, tables, rowIndex })
+      const pass = await evaluateExpr({ node: select.where, row, tables, functions, rowIndex })
       if (!pass) continue
     }
 
@@ -284,11 +296,7 @@ async function* evaluateStreaming(select, dataSource, tables) {
       } else if (col.kind === 'derived') {
         const alias = col.alias ?? defaultDerivedAlias(col.expr)
         columns.push(alias)
-        cells[alias] = () => evaluateExpr({ node: col.expr, row, tables, rowIndex: currentRowIndex })
-      } else if (col.kind === 'aggregate') {
-        throw new Error(
-          'Aggregate functions require GROUP BY or will act on the whole dataset; add GROUP BY or remove aggregates'
-        )
+        cells[alias] = () => evaluateExpr({ node: col.expr, row, tables, functions, rowIndex: currentRowIndex })
       }
     }
 
@@ -315,14 +323,17 @@ async function* evaluateStreaming(select, dataSource, tables) {
 /**
  * Buffered evaluation for complex queries (with ORDER BY or GROUP BY)
  *
- * @param {SelectStatement} select
- * @param {AsyncDataSource} dataSource
- * @param {Record<string, AsyncDataSource>} tables
- * @param {boolean} hasAggregate
- * @param {boolean} useGrouping
+ * @param {Object} options
+ * @param {SelectStatement} options.select
+ * @param {AsyncDataSource} options.dataSource
+ * @param {Record<string, AsyncDataSource>} options.tables
+ * @param {Record<string, UserDefinedFunction>} [options.functions]
+ * @param {boolean} options.hasAggregate
+ * @param {boolean} options.useGrouping
+ * @param {AbortSignal} [options.signal]
  * @yields {AsyncRow}
  */
-async function* evaluateBuffered(select, dataSource, tables, hasAggregate, useGrouping) {
+async function* evaluateBuffered({ select, dataSource, tables, functions, hasAggregate, useGrouping, signal }) {
   // Build hints for data source optimization
   // Note: limit/offset not passed here since buffering needs all rows for sorting/grouping
   /** @type {QueryHints} */
@@ -334,7 +345,7 @@ async function* evaluateBuffered(select, dataSource, tables, hasAggregate, useGr
   // Step 1: Collect all rows from data source
   /** @type {AsyncRow[]} */
   const working = []
-  for await (const row of dataSource.scan(hints)) {
+  for await (const row of dataSource.scan({ hints, signal })) {
     working.push(row)
   }
 
@@ -346,7 +357,7 @@ async function* evaluateBuffered(select, dataSource, tables, hasAggregate, useGr
     const row = working[i]
     const rowIndex = i + 1 // 1-based
     if (select.where) {
-      const passes = await evaluateExpr({ node: select.where, row, tables, rowIndex })
+      const passes = await evaluateExpr({ node: select.where, row, tables, functions, rowIndex })
 
       if (!passes) {
         continue
@@ -371,7 +382,7 @@ async function* evaluateBuffered(select, dataSource, tables, hasAggregate, useGr
         /** @type {string[]} */
         const keyParts = []
         for (const expr of select.groupBy) {
-          const v = await evaluateExpr({ node: expr, row, tables })
+          const v = await evaluateExpr({ node: expr, row, tables, functions })
           keyParts.push(stringify(v))
         }
         const key = keyParts.join('|')
@@ -414,18 +425,9 @@ async function* evaluateBuffered(select, dataSource, tables, hasAggregate, useGr
         if (col.kind === 'derived') {
           const alias = col.alias ?? defaultDerivedAlias(col.expr)
           columns.push(alias)
-          if (group.length > 0) {
-            cells[alias] = () => evaluateExpr({ node: col.expr, row: group[0], tables })
-          } else {
-            delete cells[alias]
-          }
-          continue
-        }
-
-        if (col.kind === 'aggregate') {
-          const alias = col.alias ?? defaultAggregateAlias(col)
-          columns.push(alias)
-          cells[alias] = () => evaluateAggregate({ col, rows: group, tables })
+          // Pass group to evaluateExpr so it can handle aggregate functions within expressions
+          // For empty groups, still provide an empty row context for aggregates to return appropriate values
+          cells[alias] = () => evaluateExpr({ node: col.expr, row: group[0] ?? { columns: [], cells: {} }, tables, functions, rows: group })
           continue
         }
       }
@@ -433,7 +435,7 @@ async function* evaluateBuffered(select, dataSource, tables, hasAggregate, useGr
 
       // Apply HAVING filter before adding to projected results
       if (select.having) {
-        if (!await evaluateHavingExpr(select.having, asyncRow, group, tables)) {
+        if (!await evaluateHavingExpr({ expr: select.having, row: asyncRow, group, tables, functions })) {
           continue
         }
       }
@@ -443,7 +445,7 @@ async function* evaluateBuffered(select, dataSource, tables, hasAggregate, useGr
   } else {
     // No grouping, simple projection
     // Sort before projection so ORDER BY can access columns not in SELECT
-    const sorted = await sortRows(filtered, select.orderBy, tables)
+    const sorted = await sortRows({ rows: filtered, orderBy: select.orderBy, tables, functions })
 
     // OPTIMIZATION: For non-DISTINCT queries, apply OFFSET/LIMIT before projection
     // to avoid reading expensive cells for rows that won't be in the final result
@@ -467,7 +469,7 @@ async function* evaluateBuffered(select, dataSource, tables, hasAggregate, useGr
         } else if (col.kind === 'derived') {
           const alias = col.alias ?? defaultDerivedAlias(col.expr)
           columns.push(alias)
-          cells[alias] = () => evaluateExpr({ node: col.expr, row, tables })
+          cells[alias] = () => evaluateExpr({ node: col.expr, row, tables, functions })
         }
       }
       projected.push({ columns, cells })
@@ -479,7 +481,7 @@ async function* evaluateBuffered(select, dataSource, tables, hasAggregate, useGr
 
   // Step 5: ORDER BY (final sort for grouped queries)
   if (useGrouping) {
-    projected = await sortRows(projected, select.orderBy, tables)
+    projected = await sortRows({ rows: projected, orderBy: select.orderBy, tables, functions })
   }
 
   // Step 6: OFFSET and LIMIT