squirreling 0.8.0 → 0.9.0

package/src/parse/parse.js

@@ -3,62 +3,12 @@ import { parseJoins } from './joins.js'
 import { consume, current, expect, expectIdentifier, match, parseError, peekToken } from './state.js'
 import { tokenizeSql } from './tokenize.js'
 import { duplicateCTEError } from '../parseErrors.js'
-import { RESERVED_AFTER_COLUMN, RESERVED_AFTER_TABLE, expectNoAggregate, isKnownFunction } from '../validation.js'
+import { RESERVED_AFTER_COLUMN, RESERVED_AFTER_TABLE, expectNoAggregate, findAggregate } from '../validation.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}
@@ -66,7 +16,7 @@ function parseWithClause(state) {
 export function parseSql({ query, functions }) {
   const tokens = tokenizeSql(query)
   /** @type {ParserState} */
-  const state = { tokens, pos: 0, functions }
+  const state = { tokens, pos: 0, lastPos: 0, functions }
 
   // Check for WITH clause
   /** @type {WithClause | undefined} */
@@ -130,29 +80,59 @@ function parseSelectList(state) {
   return cols
 }
 
-// Keywords that can start a valid expression in SELECT
-const EXPRESSION_START_KEYWORDS = new Set([
-  'CASE', 'TRUE', 'FALSE', 'NULL', 'EXISTS', 'NOT', 'INTERVAL',
-])
-
 /**
+ * Parses a WITH clause containing one or more CTEs
+ *
  * @param {ParserState} state
- * @returns {SelectColumn}
+ * @returns {WithClause}
  */
-function parseSelectItem(state) {
-  const tok = current(state)
+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 statement
+    expect(state, 'keyword', 'AS')
+    expect(state, 'paren', '(')
+
+    // Parse the CTE's SELECT statement
+    const query = parseSelectInternal(state)
+
+    expect(state, 'paren', ')')
 
-  // Check if keyword followed by ( is a known function (e.g., LEFT, RIGHT)
-  const isKeywordFunction = tok.type === 'keyword' &&
-    peekToken(state, 1).type === 'paren' &&
-    peekToken(state, 1).value === '(' &&
-    isKnownFunction(tok.value, state.functions)
+    ctes.push({ name, query })
 
-  if (tok.type === 'keyword' && !EXPRESSION_START_KEYWORDS.has(tok.value) && !isKeywordFunction || tok.type === 'eof') {
-    throw parseError(state, 'column name or expression')
+    // Check for comma (more CTEs) or end of WITH clause
+    if (!match(state, 'comma')) {
+      break
+    }
   }
 
-  // Delegate to expression parser (handles all expressions including aggregates)
+  return { ctes }
+}
+
+/**
+ * @param {ParserState} state
+ * @returns {SelectColumn}
+ */
+function parseSelectItem(state) {
   const expr = parseExpression(state)
   const alias = parseAs(state)
   return { kind: 'derived', expr, alias }
@@ -283,10 +263,17 @@ export function parseSelectInternal(state) {
     having = parseExpression(state)
   }
 
+  const hasAggregate = groupBy.length > 0 || columns.some(col =>
+    col.kind === 'derived' && findAggregate(col.expr)
+  )
+
   if (match(state, 'keyword', 'ORDER')) {
     expect(state, 'keyword', 'BY')
     while (true) {
       const expr = parseExpression(state)
+      if (!hasAggregate) {
+        expectNoAggregate(expr, 'ORDER BY')
+      }
       /** @type {'ASC' | 'DESC'} */
       let direction = 'ASC'
       if (match(state, 'keyword', 'ASC')) {

package/src/parse/state.js

@@ -39,15 +39,6 @@ export function consume(state) {
   return tok
 }
 
-/**
- * Gets the position after the last consumed token.
- * @param {ParserState} state
- * @returns {number}
- */
-export function lastPosition(state) {
-  return state.lastPos ?? 0
-}
-
 /**
  * @param {ParserState} state
  * @param {TokenType} type

package/src/parse/tokenize.js

@@ -229,7 +229,7 @@ export function tokenizeSql(sql) {
       let text = ''
       while (i <= length) {
         if (i === length) {
-          throw unterminatedError('string', pos, length)
+          throw unterminatedError({ type: 'string', positionStart: pos, positionEnd: length })
         }
         const c = nextChar()
         if (c === quote) {
@@ -258,7 +258,7 @@ export function tokenizeSql(sql) {
       let text = ''
       while (i <= length) {
         if (i === length) {
-          throw unterminatedError('identifier', pos, length)
+          throw unterminatedError({ type: 'identifier', positionStart: pos, positionEnd: length })
         }
         const c = nextChar()
         if (c === quote) {

package/src/parse/types.d.ts

@@ -3,7 +3,7 @@ import { UserDefinedFunction } from '../types.js'
 export interface ParserState {
   tokens: Token[]
   pos: number
-  lastPos?: number
+  lastPos: number
   functions?: Record<string, UserDefinedFunction>
 }
 

package/src/parseErrors.js

@@ -41,12 +41,13 @@ export function syntaxError({ expected, received, positionStart, positionEnd, af
 /**
  * Error for unterminated literals (strings, identifiers).
  *
- * @param {'string' | 'identifier'} type - Type of unterminated literal
- * @param {number} positionStart - Starting position
- * @param {number} positionEnd - End position
+ * @param {Object} options
+ * @param {'string' | 'identifier'} options.type - Type of unterminated literal
+ * @param {number} options.positionStart - Starting position
+ * @param {number} options.positionEnd - End position
  * @returns {ParseError}
  */
-export function unterminatedError(type, positionStart, positionEnd) {
+export function unterminatedError({ type, positionStart, positionEnd }) {
   const name = type === 'string' ? 'string literal' : 'identifier'
   return new ParseError({ message: `Unterminated ${name} starting at position ${positionStart}`, positionStart, positionEnd })
 }

package/src/plan/columns.js

@@ -0,0 +1,149 @@
+/**
+ * @import { ExprNode, SelectStatement } from '../types.js'
+ */
+
+/**
+ * Extracts per-table column names needed from a SELECT statement with joins.
+ * Returns a Map from table alias to column names, or undefined if all columns needed.
+ *
+ * @param {SelectStatement} select
+ * @returns {Map<string, string[] | undefined>}
+ */
+export function extractColumns(select) {
+  // Build alias list from FROM + JOINs
+  const fromAlias = select.from.kind === 'table'
+    ? select.from.alias ?? select.from.table
+    : select.from.alias
+  /** @type {string[]} */
+  const aliases = [fromAlias]
+  for (const join of select.joins) {
+    aliases.push(join.alias ?? join.table)
+  }
+
+  // If any unqualified SELECT * exists, all tables need all columns
+  if (select.columns.some(col => col.kind === 'star' && !col.table)) {
+    /** @type {Map<string, string[] | undefined>} */
+    const result = new Map()
+    for (const alias of aliases) {
+      result.set(alias, undefined)
+    }
+    return result
+  }
+
+  // Track which tables need all columns (SELECT table.*)
+  /** @type {Set<string>} */
+  const allColumnsNeeded = new Set()
+  for (const col of select.columns) {
+    if (col.kind === 'star' && col.table) {
+      allColumnsNeeded.add(col.table)
+    }
+  }
+
+  // Collect all identifiers from all clauses
+  /** @type {Set<string>} */
+  const identifiers = new Set()
+  for (const col of select.columns) {
+    if (col.kind === 'derived') {
+      collectColumnsFromExpr(col.expr, identifiers)
+    }
+  }
+  collectColumnsFromExpr(select.where, identifiers)
+  for (const item of select.orderBy) {
+    collectColumnsFromExpr(item.expr, identifiers)
+  }
+  for (const expr of select.groupBy) {
+    collectColumnsFromExpr(expr, identifiers)
+  }
+  collectColumnsFromExpr(select.having, identifiers)
+  for (const join of select.joins) {
+    collectColumnsFromExpr(join.on, identifiers)
+  }
+
+  // Initialize per-table sets (skip tables needing all columns)
+  /** @type {Map<string, Set<string>>} */
+  const perTable = new Map()
+  for (const alias of aliases) {
+    if (!allColumnsNeeded.has(alias)) {
+      perTable.set(alias, new Set())
+    }
+  }
+
+  // Partition identifiers by table prefix
+  for (const name of identifiers) {
+    const dotIndex = name.indexOf('.')
+    if (dotIndex >= 0) {
+      // Qualified: add to matching table only
+      const tablePrefix = name.substring(0, dotIndex)
+      const columnName = name.substring(dotIndex + 1)
+      const set = perTable.get(tablePrefix)
+      if (set) {
+        set.add(columnName)
+      }
+    } else {
+      // Unqualified: add to all tables (ambiguous)
+      for (const [, set] of perTable) {
+        set.add(name)
+      }
+    }
+  }
+
+  // Build result map: convert Sets to arrays, undefined for all-columns tables
+  /** @type {Map<string, string[] | undefined>} */
+  const result = new Map()
+  for (const alias of aliases) {
+    if (allColumnsNeeded.has(alias)) {
+      result.set(alias, undefined)
+    } else {
+      const set = perTable.get(alias)
+      result.set(alias, set ? [...set] : undefined)
+    }
+  }
+  return result
+}
+
+/**
+ * Recursively collects column names (identifiers) from an expression
+ *
+ * @param {ExprNode | undefined} expr
+ * @param {Set<string>} columns
+ */
+function collectColumnsFromExpr(expr, columns) {
+  if (!expr) return
+  if (expr.type === 'identifier' && expr.name !== '*') {
+    columns.add(expr.name)
+  } else if (expr.type === 'literal') {
+    // No columns
+  } else if (expr.type === 'binary') {
+    collectColumnsFromExpr(expr.left, columns)
+    collectColumnsFromExpr(expr.right, columns)
+  } else if (expr.type === 'unary') {
+    collectColumnsFromExpr(expr.argument, columns)
+  } else if (expr.type === 'function') {
+    for (const arg of expr.args) {
+      collectColumnsFromExpr(arg, columns)
+    }
+  } else if (expr.type === 'cast') {
+    collectColumnsFromExpr(expr.expr, columns)
+  } else if (expr.type === 'in valuelist') {
+    collectColumnsFromExpr(expr.expr, columns)
+    for (const val of expr.values) {
+      collectColumnsFromExpr(val, columns)
+    }
+  } else if (expr.type === 'in') {
+    collectColumnsFromExpr(expr.expr, columns)
+    // Subquery columns are from a different scope, don't collect
+  } else if (expr.type === 'exists' || expr.type === 'not exists') {
+    // Subquery columns are from a different scope, don't collect
+  } else if (expr.type === 'case') {
+    if (expr.caseExpr) {
+      collectColumnsFromExpr(expr.caseExpr, columns)
+    }
+    for (const when of expr.whenClauses) {
+      collectColumnsFromExpr(when.condition, columns)
+      collectColumnsFromExpr(when.result, columns)
+    }
+    if (expr.elseResult) {
+      collectColumnsFromExpr(expr.elseResult, columns)
+    }
+  }
+}

package/src/plan/plan.js

@@ -1,40 +1,44 @@
-import { extractColumns } from '../execute/columns.js'
+import { parseSql } from '../parse/parse.js'
 import { findAggregate } from '../validation.js'
+import { extractColumns } from './columns.js'
 
 /**
- * @import { ExprNode, JoinClause, ScanOptions, SelectStatement } from '../types.js'
- * @import { QueryPlan, ScanNode } from './types.d.ts'
+ * @import { ExprNode, JoinClause, PlanSqlOptions, ScanOptions, SelectStatement } from '../types.js'
+ * @import { QueryPlan } from './types.d.ts'
  */
 
 /**
  * Builds a query plan from a SELECT statement AST.
  * Resolves CTEs at plan time so no planning occurs during execution.
  *
- * @param {SelectStatement} select - the SELECT statement AST
+ * @param {PlanSqlOptions} options
  * @returns {QueryPlan} the root of the query plan tree
  */
-export function queryPlan(select) {
+export function planSql({ query, functions }) {
+  const select = typeof query === 'string' ? parseSql({ query, functions }) : query
+
   // Build CTE plans in order (each CTE can reference preceding CTEs)
   /** @type {Map<string, QueryPlan>} */
   const ctePlans = new Map()
   if (select.with) {
     for (const cte of select.with.ctes) {
-      const ctePlan = buildSelectPlan(cte.query, ctePlans)
+      const ctePlan = planSelect({ select: cte.query, ctePlans })
       ctePlans.set(cte.name.toLowerCase(), ctePlan)
     }
   }
 
-  return buildSelectPlan(select, ctePlans)
+  return planSelect({ select, ctePlans })
 }
 
 /**
- * Builds a plan for a SELECT statement with CTE resolution.
+ * Builds a plan for a SELECT statement with CTEs pre-resolved.
  *
- * @param {SelectStatement} select - the SELECT statement AST
- * @param {Map<string, QueryPlan>} ctePlans
- * @returns {QueryPlan} the root of the query plan tree
+ * @param {object} options
+ * @param {SelectStatement} options.select
+ * @param {Map<string, QueryPlan>} options.ctePlans
+ * @returns {QueryPlan}
  */
-function buildSelectPlan(select, ctePlans) {
+function planSelect({ select, ctePlans }) {
   // Check for aggregation
   const hasAggregate = select.columns.some(col =>
     col.kind === 'derived' && findAggregate(col.expr)
@@ -42,18 +46,24 @@ function buildSelectPlan(select, ctePlans) {
   const useGrouping = hasAggregate || select.groupBy.length > 0
   const needsBuffering = useGrouping || select.orderBy.length > 0
 
-  // Start with the data source (FROM clause)
+  // Source alias for FROM clause
+  const sourceAlias = select.from.kind === 'table'
+    ? select.from.alias ?? select.from.table
+    : select.from.alias
+
+  // Determine per-table column hints for pushdown
   /** @type {ScanOptions} */
-  const hints = { columns: extractColumns(select) }
+  const hints = {}
+  const perTableColumns = extractColumns(select)
+  hints.columns = perTableColumns.get(sourceAlias)
+
+  // Start with the data source (FROM clause)
   /** @type {QueryPlan} */
-  let plan = buildFromPlan(select, ctePlans, hints)
+  let plan = planFrom({ select, ctePlans, hints })
 
   // Add JOINs
   if (select.joins.length) {
-    const sourceAlias = select.from.kind === 'table'
-      ? select.from.alias ?? select.from.table
-      : select.from.alias
-    plan = buildJoinPlan(plan, select.joins, sourceAlias, ctePlans)
+    plan = planJoin({ left: plan, joins: select.joins, leftTable: sourceAlias, ctePlans, perTableColumns })
   }
 
   // Delegate WHERE and LIMIT/OFFSET to scan when plan is a direct table scan
@@ -96,7 +106,7 @@ function buildSelectPlan(select, ctePlans) {
     // Non-aggregation path
 
     // ORDER BY (before projection so it can access all columns)
-    // Pass aliases so ORDER BY can reference SELECT column aliases
+    // Resolve SELECT aliases in ORDER BY expressions at plan time
     if (select.orderBy.length) {
       /** @type {Map<string, ExprNode>} */
       const aliases = new Map()
@@ -105,7 +115,10 @@ function buildSelectPlan(select, ctePlans) {
           aliases.set(col.alias, col.expr)
         }
       }
-      plan = { type: 'Sort', orderBy: select.orderBy, aliases: aliases.size > 0 ? aliases : undefined, child: plan }
+      const orderBy = aliases.size > 0
+        ? select.orderBy.map(term => ({ ...term, expr: resolveAliases(term.expr, aliases) }))
+        : select.orderBy
+      plan = { type: 'Sort', orderBy, child: plan }
     }
 
     // DISTINCT needs to come after projection but before LIMIT
@@ -126,39 +139,37 @@ function buildSelectPlan(select, ctePlans) {
 }
 
 /**
- * Builds a plan for the FROM clause
- *
- * @param {SelectStatement} select
- * @param {Map<string, QueryPlan>} ctePlans
- * @param {ScanOptions} hints - scan options to pass to data source
+ * @param {object} options
+ * @param {SelectStatement} options.select
+ * @param {Map<string, QueryPlan>} options.ctePlans
+ * @param {ScanOptions} options.hints
  * @returns {QueryPlan}
  */
-function buildFromPlan(select, ctePlans, hints) {
+function planFrom({ select, ctePlans, hints }) {
   if (select.from.kind === 'table') {
     const ctePlan = ctePlans.get(select.from.table.toLowerCase())
     if (ctePlan) {
       return ctePlan
     }
-    return {
-      type: 'Scan',
-      table: select.from.table,
-      hints,
-    }
+    return { type: 'Scan', table: select.from.table, hints }
   } else {
-    return queryPlan(select.from.query)
+    if (select.from.query.with) {
+      throw new Error('WITH clause is not supported inside subqueries')
+    }
+    return planSelect({ select: select.from.query, ctePlans })
   }
 }
 
 /**
- * Builds join plan nodes for all joins
- *
- * @param {QueryPlan} left - the left side of the join (FROM or previous joins)
- * @param {JoinClause[]} joins - array of join clauses
- * @param {string} leftTable - name/alias of the left table
- * @param {Map<string, QueryPlan>} ctePlans
+ * @param {object} options
+ * @param {QueryPlan} options.left - the left side of the join (FROM or previous joins)
+ * @param {JoinClause[]} options.joins - array of join clauses
+ * @param {string} options.leftTable - name/alias of the left table
+ * @param {Map<string, QueryPlan>} options.ctePlans
+ * @param {Map<string, string[] | undefined>} options.perTableColumns
  * @returns {QueryPlan}
  */
-function buildJoinPlan(left, joins, leftTable, ctePlans) {
+function planJoin({ left, joins, leftTable, ctePlans, perTableColumns }) {
   let plan = left
   let currentLeftTable = leftTable
 
@@ -166,13 +177,18 @@ function buildJoinPlan(left, joins, leftTable, ctePlans) {
     const rightTable = join.alias ?? join.table
 
     const ctePlan = ctePlans.get(join.table.toLowerCase())
+    /** @type {ScanOptions} */
+    const rightHints = {}
+    if (!ctePlan) {
+      rightHints.columns = perTableColumns.get(rightTable)
+    }
     /** @type {QueryPlan} */
-    const rightScan = ctePlan ?? { type: 'Scan', table: join.table } // TODO: pass hints
+    const rightScan = ctePlan ?? { type: 'Scan', table: join.table, hints: rightHints }
 
     if (join.joinType === 'POSITIONAL') {
       plan = { type: 'PositionalJoin', leftAlias: currentLeftTable, rightAlias: rightTable, left: plan, right: rightScan }
     } else {
-      const keys = join.on && extractSimpleJoinKeys(join.on, currentLeftTable, rightTable)
+      const keys = join.on && extractSimpleJoinKeys({ condition: join.on, leftTable: currentLeftTable, rightTable })
       if (keys) {
         plan = {
           type: 'HashJoin',
@@ -204,16 +220,70 @@ function buildJoinPlan(left, joins, leftTable, ctePlans) {
   return plan
 }
 
+/**
+ * Recursively replaces identifier nodes that match SELECT aliases
+ * with their aliased expressions.
+ *
+ * @param {ExprNode} node
+ * @param {Map<string, ExprNode>} aliases
+ * @returns {ExprNode}
+ */
+function resolveAliases(node, aliases) {
+  if (node.type === 'identifier') {
+    const resolved = aliases.get(node.name)
+    if (resolved) return resolved
+    return node
+  }
+  if (node.type === 'unary') {
+    const argument = resolveAliases(node.argument, aliases)
+    return argument === node.argument ? node : { ...node, argument }
+  }
+  if (node.type === 'binary') {
+    const left = resolveAliases(node.left, aliases)
+    const right = resolveAliases(node.right, aliases)
+    return left === node.left && right === node.right ? node : { ...node, left, right }
+  }
+  if (node.type === 'function') {
+    const args = node.args.map(arg => resolveAliases(arg, aliases))
+    const changed = args.some((arg, i) => arg !== node.args[i])
+    return changed ? { ...node, args } : node
+  }
+  if (node.type === 'cast') {
+    const expr = resolveAliases(node.expr, aliases)
+    return expr === node.expr ? node : { ...node, expr }
+  }
+  if (node.type === 'in valuelist') {
+    const expr = resolveAliases(node.expr, aliases)
+    const values = node.values.map(v => resolveAliases(v, aliases))
+    const changed = expr !== node.expr || values.some((v, i) => v !== node.values[i])
+    return changed ? { ...node, expr, values } : node
+  }
+  if (node.type === 'case') {
+    const caseExpr = node.caseExpr ? resolveAliases(node.caseExpr, aliases) : node.caseExpr
+    const whenClauses = node.whenClauses.map(w => {
+      const condition = resolveAliases(w.condition, aliases)
+      const result = resolveAliases(w.result, aliases)
+      return condition === w.condition && result === w.result ? w : { ...w, condition, result }
+    })
+    const elseResult = node.elseResult ? resolveAliases(node.elseResult, aliases) : node.elseResult
+    const changed = caseExpr !== node.caseExpr || elseResult !== node.elseResult || whenClauses.some((w, i) => w !== node.whenClauses[i])
+    return changed ? { ...node, caseExpr, whenClauses, elseResult } : node
+  }
+  // literal, interval, subquery, in, exists: no identifiers to resolve
+  return node
+}
+
 /**
  * Extracts left and right key expressions from a simple equality join condition.
  * Returns undefined if the condition is not a simple equality between identifiers.
  *
- * @param {ExprNode} condition
- * @param {string} leftTable
- * @param {string} rightTable
+ * @param {object} options
+ * @param {ExprNode} options.condition
+ * @param {string} options.leftTable
+ * @param {string} options.rightTable
  * @returns {{ leftKey: ExprNode, rightKey: ExprNode } | undefined}
  */
-function extractSimpleJoinKeys(condition, leftTable, rightTable) {
+function extractSimpleJoinKeys({ condition, leftTable, rightTable }) {
   if (condition.type !== 'binary' || condition.op !== '=') {
     return undefined
   }