squirreling 0.6.1 → 0.7.0

package/src/execute/join.js

@@ -4,20 +4,22 @@ import { evaluateExpr } from './expression.js'
 import { stringify } from './utils.js'
 
 /**
- * @import { AsyncRow, AsyncDataSource, JoinClause, ExprNode, AsyncCells } from '../types.js'
+ * @import { AsyncRow, AsyncDataSource, JoinClause, ExprNode, AsyncCells, UserDefinedFunction } from '../types.js'
  */
 
 /**
  * Executes JOIN operations against a base data source
  *
- * @param {AsyncDataSource} leftSource - the left side of the join (FROM table)
- * @param {JoinClause[]} joins - array of join clauses to execute
- * @param {string} leftTableName - name of the left table (for column prefixing)
- * @param {Record<string, AsyncDataSource>} tables - all available tables
+ * @param {Object} options
+ * @param {AsyncDataSource} options.leftSource - the left side of the join (FROM table)
+ * @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 {Record<string, UserDefinedFunction>} [options.functions]
  * @returns {Promise<AsyncDataSource>} data source yielding joined rows
  */
-export async function executeJoins(leftSource, joins, leftTableName, tables) {
-  let currentLeftTable = leftTableName
+export async function executeJoins({ leftSource, joins, leftTable, tables, functions }) {
+  let currentLeftTable = leftTable
 
   // Single join optimization: stream left rows without buffering
   if (joins.length === 1) {
@@ -35,7 +37,7 @@ export async function executeJoins(leftSource, joins, leftTableName, tables) {
     }
 
     // Use alias for column prefixing if present
-    const rightTableName = join.alias ?? join.table
+    const rightTable = join.alias ?? join.table
 
     // Return streaming data source - left rows stream through without buffering
     return {
@@ -46,8 +48,9 @@ export async function executeJoins(leftSource, joins, leftTableName, tables) {
           rightRows,
           join,
           leftTable: currentLeftTable,
-          rightTable: rightTableName,
+          rightTable,
           tables,
+          functions,
           signal,
         })
       },
@@ -76,7 +79,7 @@ export async function executeJoins(leftSource, joins, leftTableName, tables) {
     }
 
     // Use alias for column prefixing if present
-    const rightTableName = join.alias ?? join.table
+    const rightTable = join.alias ?? join.table
 
     // Collect intermediate results into array for next join
     /** @type {AsyncRow[]} */
@@ -86,8 +89,9 @@ export async function executeJoins(leftSource, joins, leftTableName, tables) {
       rightRows,
       join,
       leftTable: currentLeftTable,
-      rightTable: rightTableName,
+      rightTable,
       tables,
+      functions,
     })
     for await (const row of joined) {
       newLeftRows.push(row)
@@ -95,14 +99,14 @@ export async function executeJoins(leftSource, joins, leftTableName, tables) {
     leftRows = newLeftRows
 
     // After join, the "left" table for the next join includes all joined tables
-    currentLeftTable = `${currentLeftTable}_${rightTableName}`
+    currentLeftTable = `${currentLeftTable}_${rightTable}`
   }
 
   // Final join: stream the results
-  const lastJoin = joins[joins.length - 1]
-  const rightSource = tables[lastJoin.table]
+  const join = joins[joins.length - 1]
+  const rightSource = tables[join.table]
   if (rightSource === undefined) {
-    throw tableNotFoundError({ tableName: lastJoin.table })
+    throw tableNotFoundError({ tableName: join.table })
   }
 
   /** @type {AsyncRow[]} */
@@ -112,7 +116,7 @@ export async function executeJoins(leftSource, joins, leftTableName, tables) {
   }
 
   // Use alias for column prefixing if present
-  const lastRightTableName = lastJoin.alias ?? lastJoin.table
+  const rightTable = join.alias ?? join.table
 
   return {
     async *scan(options) {
@@ -120,10 +124,11 @@ export async function executeJoins(leftSource, joins, leftTableName, tables) {
       yield* hashJoin({
         leftRows,
         rightRows,
-        join: lastJoin,
+        join,
         leftTable: currentLeftTable,
-        rightTable: lastRightTableName,
+        rightTable,
         tables,
+        functions,
         signal,
       })
     },
@@ -236,10 +241,11 @@ function mergeRows(leftRow, rightRow, leftTable, rightTable) {
  * @param {string} params.leftTable - name of left table (for column prefixing)
  * @param {string} params.rightTable - name of right table (for column prefixing, may be alias)
  * @param {Record<string, AsyncDataSource>} params.tables - all tables for expression evaluation
+ * @param {Record<string, UserDefinedFunction>} [params.functions]
  * @param {AbortSignal} [params.signal] - abort signal for cancellation
  * @yields {AsyncRow} joined rows
  */
-async function* hashJoin({ leftRows, rightRows, join, leftTable, rightTable, tables, signal }) {
+async function* hashJoin({ leftRows, rightRows, join, leftTable, rightTable, tables, functions, signal }) {
   const { joinType, on: onCondition } = join
 
   if (!onCondition) {
@@ -269,7 +275,7 @@ async function* hashJoin({ leftRows, rightRows, join, leftTable, rightTable, tab
     // BUILD PHASE: Index right rows by join key
     // Skip null keys - SQL semantics: NULL != NULL
     for (const rightRow of rightRows) {
-      const keyValue = await evaluateExpr({ node: keys.rightKey, row: rightRow, tables })
+      const keyValue = await evaluateExpr({ node: keys.rightKey, row: rightRow, tables, functions })
       if (keyValue == null) continue // NULL keys never match
       const keyStr = stringify(keyValue)
       let bucket = hashMap.get(keyStr)
@@ -294,7 +300,7 @@ async function* hashJoin({ leftRows, rightRows, join, leftTable, rightTable, tab
         )
       }
 
-      const keyValue = await evaluateExpr({ node: keys.leftKey, row: leftRow, tables })
+      const keyValue = await evaluateExpr({ node: keys.leftKey, row: leftRow, tables, functions })
       const keyStr = stringify(keyValue)
 
       const matchingRightRows = hashMap.get(keyStr)
@@ -340,7 +346,7 @@ async function* hashJoin({ leftRows, rightRows, join, leftTable, rightTable, tab
 
       for (const rightRow of rightRows) {
         const tempMerged = mergeRows(leftRow, rightRow, leftTable, rightTable)
-        const matches = await evaluateExpr({ node: onCondition, row: tempMerged, tables })
+        const matches = await evaluateExpr({ node: onCondition, row: tempMerged, tables, functions })
 
         if (matches) {
           hasMatch = true

package/src/execute/math.js

@@ -1,7 +1,6 @@
 /**
  * @import { MathFunc, SqlPrimitive } from '../types.js'
  */
-import { argCountError } from '../validationErrors.js'
 
 /**
  * Evaluate a math function
@@ -9,71 +8,28 @@ import { argCountError } from '../validationErrors.js'
  * @param {Object} options
  * @param {MathFunc} options.funcName - Uppercase function name
  * @param {SqlPrimitive[]} options.args - Function arguments
- * @param {number} options.positionStart - Start position in query
- * @param {number} options.positionEnd - End position in query
- * @param {number} [options.rowNumber] - 1-based row number for error reporting
  * @returns {SqlPrimitive} Result
  */
-export function evaluateMathFunc({ funcName, args, positionStart, positionEnd, rowNumber }) {
+export function evaluateMathFunc({ funcName, args }) {
   if (funcName === 'FLOOR') {
-    if (args.length !== 1) {
-      throw argCountError({
-        funcName: 'FLOOR',
-        expected: 1,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     const val = args[0]
     if (val == null) return null
     return Math.floor(Number(val))
   }
 
   if (funcName === 'CEIL' || funcName === 'CEILING') {
-    if (args.length !== 1) {
-      throw argCountError({
-        funcName,
-        expected: 1,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     const val = args[0]
     if (val == null) return null
     return Math.ceil(Number(val))
   }
 
   if (funcName === 'ABS') {
-    if (args.length !== 1) {
-      throw argCountError({
-        funcName: 'ABS',
-        expected: 1,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     const val = args[0]
     if (val == null) return null
     return Math.abs(Number(val))
   }
 
   if (funcName === 'MOD') {
-    if (args.length !== 2) {
-      throw argCountError({
-        funcName: 'MOD',
-        expected: 2,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     const dividend = args[0]
     const divisor = args[1]
     if (dividend == null || divisor == null) return null
@@ -81,64 +37,24 @@ export function evaluateMathFunc({ funcName, args, positionStart, positionEnd, r
   }
 
   if (funcName === 'EXP') {
-    if (args.length !== 1) {
-      throw argCountError({
-        funcName: 'EXP',
-        expected: 1,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     const val = args[0]
     if (val == null) return null
     return Math.exp(Number(val))
   }
 
   if (funcName === 'LN') {
-    if (args.length !== 1) {
-      throw argCountError({
-        funcName: 'LN',
-        expected: 1,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     const val = args[0]
     if (val == null) return null
     return Math.log(Number(val))
   }
 
   if (funcName === 'LOG10') {
-    if (args.length !== 1) {
-      throw argCountError({
-        funcName: 'LOG10',
-        expected: 1,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     const val = args[0]
     if (val == null) return null
     return Math.log10(Number(val))
   }
 
   if (funcName === 'POWER') {
-    if (args.length !== 2) {
-      throw argCountError({
-        funcName: 'POWER',
-        expected: 2,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     const base = args[0]
     const exponent = args[1]
     if (base == null || exponent == null) return null
@@ -146,144 +62,61 @@ export function evaluateMathFunc({ funcName, args, positionStart, positionEnd, r
   }
 
   if (funcName === 'SQRT') {
-    if (args.length !== 1) {
-      throw argCountError({
-        funcName: 'SQRT',
-        expected: 1,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     const val = args[0]
     if (val == null) return null
     return Math.sqrt(Number(val))
   }
 
   if (funcName === 'SIN') {
-    if (args.length !== 1) {
-      throw argCountError({
-        funcName: 'SIN',
-        expected: 1,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     const val = args[0]
     if (val == null) return null
     return Math.sin(Number(val))
   }
 
   if (funcName === 'COS') {
-    if (args.length !== 1) {
-      throw argCountError({
-        funcName: 'COS',
-        expected: 1,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     const val = args[0]
     if (val == null) return null
     return Math.cos(Number(val))
   }
 
   if (funcName === 'TAN') {
-    if (args.length !== 1) {
-      throw argCountError({
-        funcName: 'TAN',
-        expected: 1,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     const val = args[0]
     if (val == null) return null
     return Math.tan(Number(val))
   }
 
   if (funcName === 'COT') {
-    if (args.length !== 1) {
-      throw argCountError({
-        funcName: 'COT',
-        expected: 1,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     const val = args[0]
     if (val == null) return null
     return 1 / Math.tan(Number(val))
   }
 
   if (funcName === 'ASIN') {
-    if (args.length !== 1) {
-      throw argCountError({
-        funcName: 'ASIN',
-        expected: 1,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     const val = args[0]
     if (val == null) return null
     return Math.asin(Number(val))
   }
 
   if (funcName === 'ACOS') {
-    if (args.length !== 1) {
-      throw argCountError({
-        funcName: 'ACOS',
-        expected: 1,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     const val = args[0]
     if (val == null) return null
     return Math.acos(Number(val))
   }
 
   if (funcName === 'ATAN') {
-    if (args.length !== 1) {
-      throw argCountError({
-        funcName: 'ATAN',
-        expected: 1,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
+    if (args.length === 1) {
+      const val = args[0]
+      if (val == null) return null
+      return Math.atan(Number(val))
+    } else {
+      const y = args[0]
+      const x = args[1]
+      if (y == null || x == null) return null
+      return Math.atan2(Number(y), Number(x))
     }
-    const val = args[0]
-    if (val == null) return null
-    return Math.atan(Number(val))
   }
 
   if (funcName === 'ATAN2') {
-    if (args.length !== 2) {
-      throw argCountError({
-        funcName: 'ATAN2',
-        expected: 2,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     const y = args[0]
     const x = args[1]
     if (y == null || x == null) return null
@@ -291,50 +124,18 @@ export function evaluateMathFunc({ funcName, args, positionStart, positionEnd, r
   }
 
   if (funcName === 'DEGREES') {
-    if (args.length !== 1) {
-      throw argCountError({
-        funcName: 'DEGREES',
-        expected: 1,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     const val = args[0]
     if (val == null) return null
     return Number(val) * 180 / Math.PI
   }
 
   if (funcName === 'RADIANS') {
-    if (args.length !== 1) {
-      throw argCountError({
-        funcName: 'RADIANS',
-        expected: 1,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     const val = args[0]
     if (val == null) return null
     return Number(val) * Math.PI / 180
   }
 
   if (funcName === 'PI') {
-    if (args.length !== 0) {
-      throw argCountError({
-        funcName: 'PI',
-        expected: 0,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     return Math.PI
   }
-
-  return null
 }

package/src/index.d.ts

@@ -14,10 +14,11 @@ export function executeSql(options: ExecuteSqlOptions): AsyncGenerator<AsyncRow>
 /**
  * Parses a SQL query string into an abstract syntax tree
  *
- * @param query - SQL query string to parse
+ * @param options
+ * @param options.query - SQL query string to parse
  * @returns parsed SQL select statement
  */
-export function parseSql(query: string): SelectStatement
+export function parseSql(options: { query: string }): SelectStatement
 
 /**
  * Collects all results from an async generator into an array

package/src/parse/expression.js

@@ -1,10 +1,11 @@
 import {
+  argCountParseError,
   invalidLiteralError,
   missingClauseError,
   syntaxError,
   unknownFunctionError,
 } from '../parseErrors.js'
-import { isAggregateFunc, isIntervalUnit, isMathFunc, isStringFunc } from '../validation.js'
+import { isIntervalUnit, isKnownFunction, validateFunctionArgCount } from '../validation.js'
 import { parseComparison } from './comparison.js'
 import { parseSelectInternal } from './parse.js'
 import { consume, current, expect, expectIdentifier, lastPosition, match, peekToken } from './state.js'
@@ -122,10 +123,15 @@ export function parsePrimary(state) {
     // function call
     if (next.type === 'paren' && next.value === '(') {
       const funcName = tok.value
-
-      // validate function names
-      if (!isStringFunc(funcName) && !isAggregateFunc(funcName) && !isMathFunc(funcName)) {
-        throw unknownFunctionError({ funcName, positionStart: tok.positionStart, positionEnd: tok.positionEnd })
+      const funcNameUpper = funcName.toUpperCase()
+
+      // Validate function existence early for better error messages
+      if (!isKnownFunction(funcNameUpper, state.functions)) {
+        throw unknownFunctionError({
+          funcName,
+          positionStart,
+          positionEnd: tok.positionEnd,
+        })
       }
 
       consume(state) // function name
@@ -165,12 +171,14 @@ export function parsePrimary(state) {
 
       expect(state, 'paren', ')')
 
-      // Aggregate functions require at least one argument
-      if (isAggregateFunc(funcName) && args.length === 0) {
-        throw syntaxError({
-          expected: 'expression',
-          received: '")"',
-          positionStart: positionStart + funcName.length + 1, // position after opening paren
+      // Validate argument count at parse time
+      const validation = validateFunctionArgCount(funcNameUpper, args.length)
+      if (!validation.valid) {
+        throw argCountParseError({
+          funcName,
+          expected: validation.expected,
+          received: args.length,
+          positionStart,
           positionEnd: lastPosition(state),
         })
       }

package/src/parse/parse.js

@@ -5,17 +5,17 @@ import { consume, current, expect, expectIdentifier, match, parseError, peekToke
 import { parseJoins } from './joins.js'
 
 /**
- * @import { ExprNode, FromSubquery, FromTable, OrderByItem, ParserState, SelectStatement, SelectColumn } from '../types.js'
+ * @import { ExprNode, FromSubquery, FromTable, OrderByItem, ParserState, SelectStatement, SelectColumn, UserDefinedFunction } from '../types.js'
  */
 
 /**
- * @param {string} query
+ * @param {{ query: string, functions?: Record<string, UserDefinedFunction> }} options
  * @returns {SelectStatement}
  */
-export function parseSql(query) {
+export function parseSql({ query, functions }) {
   const tokens = tokenize(query)
   /** @type {ParserState} */
-  const state = { tokens, pos: 0 }
+  const state = { tokens, pos: 0, functions }
   const select = parseSelectInternal(state)
 
   const tok = current(state)

package/src/parseErrors.js

@@ -2,6 +2,8 @@
 // PARSE ERRORS - Issues during SQL tokenization and parsing
 // ============================================================================
 
+import { FUNCTION_SIGNATURES } from './validationErrors.js'
+
 /**
  * Structured parse error with position range.
  */
@@ -103,6 +105,33 @@ export function unknownFunctionError({ funcName, positionStart, positionEnd, val
   })
 }
 
+/**
+ * Error for wrong number of function arguments at parse time.
+ *
+ * @param {Object} options
+ * @param {string} options.funcName - The function name
+ * @param {number | string} options.expected - Expected count (number or range like "2 to 3")
+ * @param {number} options.received - Actual argument count
+ * @param {number} options.positionStart - Start position in query
+ * @param {number} options.positionEnd - End position in query
+ * @returns {ParseError}
+ */
+export function argCountParseError({ funcName, expected, received, positionStart, positionEnd }) {
+  const signature = FUNCTION_SIGNATURES[funcName] ?? ''
+  let expectedStr = `${expected} arguments`
+  if (expected === 0) expectedStr = 'no arguments'
+  if (expected === 1) expectedStr = '1 argument'
+  if (typeof expected === 'string' && expected.endsWith(' 1')) {
+    expectedStr = `${expected} argument`
+  }
+
+  return new ParseError({
+    message: `${funcName}(${signature}) function requires ${expectedStr}, got ${received}`,
+    positionStart,
+    positionEnd,
+  })
+}
+
 /**
  * Error for missing required clause or structure.
  *

package/src/types.d.ts

@@ -1,7 +1,11 @@
+// User-defined function type
+export type UserDefinedFunction = (...args: SqlPrimitive[]) => SqlPrimitive | Promise<SqlPrimitive>
+
 // executeSql(options)
 export interface ExecuteSqlOptions {
   tables: Record<string, Row | AsyncDataSource>
   query: string | SelectStatement
+  functions?: Record<string, UserDefinedFunction>
   signal?: AbortSignal
 }
 
@@ -254,6 +258,7 @@ export interface ParserState {
   tokens: Token[]
   pos: number
   lastPos?: number
+  functions?: Record<string, UserDefinedFunction>
 }
 
 // Tokenizer types