squirreling 0.7.2 → 0.7.3

package/README.md

@@ -78,6 +78,7 @@ console.log(`Collected rows:`, rows)
 ## Supported SQL Features
 
 - `SELECT` statements with `WHERE`, `ORDER BY`, `LIMIT`, `OFFSET`
+- `WITH` clause for Common Table Expressions (CTEs)
 - Subqueries in `SELECT`, `FROM`, and `WHERE` clauses
 - `JOIN` operations: `INNER JOIN`, `LEFT JOIN`, `RIGHT JOIN`, `FULL JOIN`, `POSITIONAL JOIN`
 - `GROUP BY` and `HAVING` clauses

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "squirreling",
-  "version": "0.7.2",
+  "version": "0.7.3",
   "description": "Squirreling SQL Engine",
   "author": "Hyperparam",
   "homepage": "https://hyperparam.app",

package/src/execute/execute.js

@@ -1,15 +1,16 @@
 import { missingClauseError } from '../parseErrors.js'
-import { tableNotFoundError, unsupportedOperationError } from '../executionErrors.js'
+import { unsupportedOperationError } from '../executionErrors.js'
 import { generatorSource, memorySource } from '../backend/dataSource.js'
 import { parseSql } from '../parse/parse.js'
 import { containsAggregate, extractColumns } from './columns.js'
 import { evaluateExpr } from './expression.js'
 import { evaluateHavingExpr } from './having.js'
 import { executeJoins } from './join.js'
+import { resolveTableSource } from './tableSource.js'
 import { compareForTerm, defaultDerivedAlias, stringify } from './utils.js'
 
 /**
- * @import { AsyncCells, AsyncDataSource, AsyncRow, ExecuteSqlOptions, OrderByItem, QueryHints, SelectStatement, SqlPrimitive, UserDefinedFunction } from '../types.js'
+ * @import { AsyncCells, AsyncDataSource, AsyncRow, CTEDefinition, ExecuteSqlOptions, OrderByItem, QueryHints, SelectStatement, SqlPrimitive, UserDefinedFunction, WithClause } from '../types.js'
  */
 
 /**
@@ -40,7 +41,7 @@ export async function* executeSql({ tables, query, functions, signal }) {
     }
   }
 
-  yield* executeSelect({ select, tables: normalizedTables, functions, signal })
+  yield* executeSelect({ select, tables: normalizedTables, withClause: select.with, functions, signal })
 }
 
 /**
@@ -49,32 +50,45 @@ export async function* executeSql({ tables, query, functions, signal }) {
  * @param {Object} options
  * @param {SelectStatement} options.select
  * @param {Record<string, AsyncDataSource>} options.tables
+ * @param {WithClause} [options.withClause] - WITH clause containing CTE definitions
  * @param {Record<string, UserDefinedFunction>} [options.functions]
  * @param {AbortSignal} [options.signal]
  * @yields {AsyncRow}
  */
-export async function* executeSelect({ select, tables, functions, signal }) {
+export async function* executeSelect({ select, tables, withClause, functions, signal }) {
   /** @type {AsyncDataSource} */
   let dataSource
   /** @type {string} */
   let leftTable
 
   if (select.from.kind === 'table') {
-    // Use alias for column prefixing, but look up the actual table name
-    leftTable = select.from.alias ?? select.from.table
-    dataSource = tables[select.from.table]
-    if (dataSource === undefined) {
-      throw tableNotFoundError({ tableName: select.from.table })
-    }
+    const tableName = select.from.table
+    leftTable = select.from.alias ?? tableName
+    dataSource = resolveTableSource(tableName, tables, withClause, executeSelect, functions, signal)
   } else {
     // Nested subquery - recursively resolve
     leftTable = select.from.alias
-    dataSource = generatorSource(executeSelect({ select: select.from.query, tables, functions, signal }))
+    dataSource = generatorSource(executeSelect({
+      select: select.from.query,
+      tables,
+      withClause,
+      functions,
+      signal,
+    }))
   }
 
   // Execute JOINs if present
   if (select.joins.length) {
-    dataSource = await executeJoins({ leftSource: dataSource, joins: select.joins, leftTable, tables, functions })
+    dataSource = await executeJoins({
+      leftSource: dataSource,
+      joins: select.joins,
+      leftTable,
+      tables,
+      withClause,
+      functions,
+      executeSelectFn: executeSelect,
+      signal,
+    })
   }
 
   yield* evaluateSelectAst({ select, dataSource, tables, functions, signal })

package/src/execute/join.js

@@ -1,10 +1,10 @@
 import { missingClauseError } from '../parseErrors.js'
-import { tableNotFoundError } from '../executionErrors.js'
 import { evaluateExpr } from './expression.js'
+import { resolveTableSource } from './tableSource.js'
 import { stringify } from './utils.js'
 
 /**
- * @import { AsyncRow, AsyncDataSource, JoinClause, ExprNode, AsyncCells, UserDefinedFunction } from '../types.js'
+ * @import { AsyncRow, AsyncDataSource, JoinClause, ExprNode, AsyncCells, UserDefinedFunction, WithClause } from '../types.js'
  */
 
 /**
@@ -15,19 +15,19 @@ import { stringify } from './utils.js'
  * @param {JoinClause[]} options.joins - array of join clauses to execute
  * @param {string} options.leftTable - name of the left table (for column prefixing)
  * @param {Record<string, AsyncDataSource>} options.tables - all available tables
+ * @param {WithClause} [options.withClause] - WITH clause containing CTE definitions
  * @param {Record<string, UserDefinedFunction>} [options.functions]
+ * @param {Function} [options.executeSelectFn] - function to execute SELECT for CTEs (passed to avoid circular dep)
+ * @param {AbortSignal} [options.signal]
  * @returns {Promise<AsyncDataSource>} data source yielding joined rows
  */
-export async function executeJoins({ leftSource, joins, leftTable, tables, functions }) {
+export async function executeJoins({ leftSource, joins, leftTable, tables, withClause, functions, executeSelectFn, signal }) {
   let currentLeftTable = leftTable
 
   // Single join optimization: stream left rows without buffering
   if (joins.length === 1) {
     const join = joins[0]
-    const rightSource = tables[join.table]
-    if (rightSource === undefined) {
-      throw tableNotFoundError({ tableName: join.table })
-    }
+    const rightSource = resolveTableSource(join.table, tables, withClause, executeSelectFn, functions, signal)
 
     // Buffer right rows for hash index (required for hash join)
     /** @type {AsyncRow[]} */
@@ -77,10 +77,7 @@ export async function executeJoins({ leftSource, joins, leftTable, tables, funct
   // Process all but the last join, buffering intermediate results
   for (let i = 0; i < joins.length - 1; i++) {
     const join = joins[i]
-    const rightSource = tables[join.table]
-    if (rightSource === undefined) {
-      throw tableNotFoundError({ tableName: join.table })
-    }
+    const rightSource = resolveTableSource(join.table, tables, withClause, executeSelectFn, functions, signal)
 
     /** @type {AsyncRow[]} */
     const rightRows = []
@@ -121,10 +118,7 @@ export async function executeJoins({ leftSource, joins, leftTable, tables, funct
 
   // Final join: stream the results
   const join = joins[joins.length - 1]
-  const rightSource = tables[join.table]
-  if (rightSource === undefined) {
-    throw tableNotFoundError({ tableName: join.table })
-  }
+  const rightSource = resolveTableSource(join.table, tables, withClause, executeSelectFn, functions, signal)
 
   /** @type {AsyncRow[]} */
   const rightRows = []

package/src/execute/tableSource.js

@@ -0,0 +1,63 @@
+import { tableNotFoundError } from '../executionErrors.js'
+import { generatorSource } from '../backend/dataSource.js'
+
+/**
+ * @import { AsyncDataSource, CTEDefinition, UserDefinedFunction, WithClause } from '../types.js'
+ */
+
+/**
+ * Gets CTEs defined before the target CTE (excluding the target itself).
+ * Enforces SQL scoping rules: each CTE can only reference CTEs defined before it.
+ *
+ * @param {CTEDefinition[]} allCtes - all CTE definitions in order
+ * @param {string} targetCteName - the CTE name (case-insensitive)
+ * @returns {WithClause} CTEs available to the target
+ */
+export function getCtesDefinedBefore(allCtes, targetCteName) {
+  const available = []
+  for (const cte of allCtes) {
+    if (cte.name.toLowerCase() === targetCteName) break
+    available.push(cte)
+  }
+  return { ctes: available }
+}
+
+/**
+ * Resolves a table name to an AsyncDataSource, checking CTEs first
+ *
+ * @param {string} tableName - the table name to resolve
+ * @param {Record<string, AsyncDataSource>} tables - regular tables
+ * @param {import('../types.js').WithClause} [withClause] - WITH clause containing CTE definitions
+ * @param {Function} [executeSelectFn] - function to execute SELECT for CTEs
+ * @param {Record<string, UserDefinedFunction>} [functions]
+ * @param {AbortSignal} [signal]
+ * @returns {AsyncDataSource}
+ */
+export function resolveTableSource(tableName, tables, withClause, executeSelectFn, functions, signal) {
+  // Check CTEs first (case-insensitive) - only build map when CTE is actually found
+  if (withClause && executeSelectFn) {
+    const lowerName = tableName.toLowerCase()
+    const cte = withClause.ctes.find(c => c.name.toLowerCase() === lowerName)
+
+    if (cte) {
+      // CTE reference: wrap in generatorSource, re-execute each time (streaming)
+      // Pass only CTEs defined before this one to prevent self-reference
+      const availableCtes = getCtesDefinedBefore(withClause.ctes, lowerName)
+
+      return generatorSource(executeSelectFn({
+        select: cte.query,
+        tables,
+        withClause: availableCtes,
+        functions,
+        signal,
+      }))
+    }
+  }
+
+  // Regular table lookup
+  const tableSource = tables[tableName]
+  if (tableSource === undefined) {
+    throw tableNotFoundError({ tableName })
+  }
+  return tableSource
+}

package/src/parse/parse.js

@@ -1,13 +1,64 @@
-import { tokenizeSql } from './tokenize.js'
 import { parseExpression } from './expression.js'
-import { RESERVED_AFTER_COLUMN, RESERVED_AFTER_TABLE, isKnownFunction } from '../validation.js'
+import { tokenizeSql } from './tokenize.js'
 import { consume, current, expect, expectIdentifier, match, parseError, peekToken } from './state.js'
 import { parseJoins } from './joins.js'
+import { duplicateCTEError } from '../parseErrors.js'
+import { RESERVED_AFTER_COLUMN, RESERVED_AFTER_TABLE, isKnownFunction } from '../validation.js'
 
 /**
- * @import { ExprNode, FromSubquery, FromTable, OrderByItem, ParseSqlOptions, ParserState, SelectStatement, SelectColumn } from '../types.js'
+ * @import { CTEDefinition, ExprNode, FromSubquery, FromTable, OrderByItem, ParseSqlOptions, ParserState, SelectStatement, SelectColumn, WithClause } from '../types.js'
  */
 
+/**
+ * Parses a WITH clause containing one or more CTEs
+ * @param {ParserState} state
+ * @returns {WithClause}
+ */
+function parseWithClause(state) {
+  /** @type {CTEDefinition[]} */
+  const ctes = []
+  /** @type {Set<string>} */
+  const seenNames = new Set()
+
+  while (true) {
+    // Parse CTE name
+    const nameTok = expectIdentifier(state)
+    const name = nameTok.value
+    const nameLower = name.toLowerCase()
+
+    // Check for duplicate CTE names
+    if (seenNames.has(nameLower)) {
+      throw duplicateCTEError({
+        cteName: name,
+        positionStart: nameTok.positionStart,
+        positionEnd: nameTok.positionEnd,
+      })
+    }
+    seenNames.add(nameLower)
+
+    // Expect AS keyword
+    expect(state, 'keyword', 'AS')
+
+    // Expect opening parenthesis
+    expect(state, 'paren', '(')
+
+    // Parse the CTE's SELECT statement
+    const query = parseSelectInternal(state)
+
+    // Expect closing parenthesis
+    expect(state, 'paren', ')')
+
+    ctes.push({ name, query })
+
+    // Check for comma (more CTEs) or end of WITH clause
+    if (!match(state, 'comma')) {
+      break
+    }
+  }
+
+  return { ctes }
+}
+
 /**
  * @param {ParseSqlOptions} options
  * @returns {SelectStatement}
@@ -16,8 +67,21 @@ export function parseSql({ query, functions }) {
   const tokens = tokenizeSql(query)
   /** @type {ParserState} */
   const state = { tokens, pos: 0, functions }
+
+  // Check for WITH clause
+  /** @type {WithClause | undefined} */
+  let withClause
+  if (match(state, 'keyword', 'WITH')) {
+    withClause = parseWithClause(state)
+  }
+
   const select = parseSelectInternal(state)
 
+  // Attach WITH clause to the select statement
+  if (withClause) {
+    select.with = withClause
+  }
+
   const tok = current(state)
   if (tok.type !== 'eof') {
     throw parseError(state, 'end of query')

package/src/parse/tokenize.js

@@ -9,6 +9,7 @@ import {
  */
 
 const KEYWORDS = new Set([
+  'WITH',
   'SELECT',
   'FROM',
   'WHERE',

package/src/parseErrors.js

@@ -79,7 +79,7 @@ export function invalidLiteralError({ type, value, positionStart, positionEnd, v
 export function unexpectedCharError({ char, positionStart, expectsSelect = false }) {
   const positionEnd = positionStart + 1
   if (expectsSelect) {
-    return new ParseError({ message: `Expected SELECT but found "${char}" at position ${positionStart}. Queries must start with SELECT.`, positionStart, positionEnd })
+    return new ParseError({ message: `Expected SELECT but found "${char}" at position ${positionStart}. Queries must start with SELECT or WITH.`, positionStart, positionEnd })
   }
   return new ParseError({ message: `Unexpected character "${char}" at position ${positionStart}`, positionStart, positionEnd })
 }
@@ -145,3 +145,20 @@ export function argCountParseError({ funcName, expected, received, positionStart
 export function missingClauseError({ missing, context, positionStart, positionEnd }) {
   return new ParseError({ message: `${context} requires ${missing}`, positionStart: positionStart ?? 0, positionEnd: positionEnd ?? 0 })
 }
+
+/**
+ * Error for duplicate CTE names in WITH clause.
+ *
+ * @param {Object} options
+ * @param {string} options.cteName - The duplicate CTE name
+ * @param {number} options.positionStart - Start position in query
+ * @param {number} options.positionEnd - End position in query
+ * @returns {ParseError}
+ */
+export function duplicateCTEError({ cteName, positionStart, positionEnd }) {
+  return new ParseError({
+    message: `CTE "${cteName}" is defined more than once at position ${positionStart}`,
+    positionStart,
+    positionEnd,
+  })
+}

package/src/types.d.ts

@@ -67,7 +67,17 @@ export interface UserDefinedFunction {
   }
 }
 
+export interface CTEDefinition {
+  name: string
+  query: SelectStatement
+}
+
+export interface WithClause {
+  ctes: CTEDefinition[]
+}
+
 export interface SelectStatement {
+  with?: WithClause
   distinct: boolean
   columns: SelectColumn[]
   from: FromTable | FromSubquery