squirreling 0.12.4 → 0.12.6

package/src/parse/parse.js

@@ -1,4 +1,6 @@
+import { derivedAlias } from '../expression/alias.js'
 import { expectNoAggregate, findAggregate } from '../validation/aggregates.js'
+import { isTableFunction, validateFunctionArgs } from '../validation/functions.js'
 import { RESERVED_AFTER_COLUMN, RESERVED_AFTER_TABLE } from '../validation/keywords.js'
 import { ParseError } from '../validation/parseErrors.js'
 import { parseExpression } from './expression.js'
@@ -7,7 +9,7 @@ import { consume, current, expect, match, parseError, peekToken } from './state.
 import { tokenizeSql } from './tokenize.js'
 
 /**
- * @import { CTEDefinition, ExprNode, FromSubquery, FromTable, OrderByItem, ParseSqlOptions, ParserState, SelectColumn, SelectStatement, SetOperationStatement, SetOperator, Statement } from '../types.js'
+ * @import { CTEDefinition, ExprNode, FromFunction, FromSubquery, FromTable, OrderByItem, ParseSqlOptions, ParserState, SelectColumn, SelectStatement, SetOperationStatement, SetOperator, Statement } from '../types.js'
  */
 
 /**
@@ -183,8 +185,8 @@ function parseSelect(state) {
     expect(state, 'keyword', 'FROM')
   }
 
-  // Check if it's a subquery or table name
-  /** @type {FromTable | FromSubquery} */
+  // Check if it's a subquery, table function, or table name
+  /** @type {FromTable | FromSubquery | FromFunction} */
   let from
   const fromTok = current(state)
   if (fromTok.type === 'paren' && fromTok.value === '(') {
@@ -200,6 +202,9 @@ function parseSelect(state) {
       positionStart: fromTok.positionStart,
       positionEnd: state.lastPos,
     }
+  } else if (isTableFunctionStart(state)) {
+    // Table function: SELECT * FROM UNNEST(expr) [AS alias[(col_alias)]]
+    from = parseFromFunction(state)
   } else {
     // Simple table name: SELECT * FROM users
     expect(state, 'identifier')
@@ -237,7 +242,7 @@ function parseSelect(state) {
   if (match(state, 'keyword', 'GROUP')) {
     expect(state, 'keyword', 'BY')
     while (true) {
-      const expr = parseExpression(state)
+      const expr = resolvePositionalRef(parseExpression(state), columns, 'GROUP BY')
       expectNoAggregate(expr, 'GROUP BY')
       groupBy.push(expr)
       if (!match(state, 'comma')) break
@@ -255,7 +260,7 @@ function parseSelect(state) {
   if (match(state, 'keyword', 'ORDER')) {
     expect(state, 'keyword', 'BY')
     while (true) {
-      const expr = parseExpression(state)
+      const expr = resolvePositionalRef(parseExpression(state), columns, 'ORDER BY', hasAggregate)
       if (!hasAggregate) {
         expectNoAggregate(expr, 'ORDER BY')
       }
@@ -333,6 +338,53 @@ function parseSelect(state) {
   }
 }
 
+/**
+ * Resolves a positive integer literal in GROUP BY or ORDER BY as a positional
+ * reference to the Nth SELECT column. Non-integer and non-literal expressions
+ * pass through unchanged.
+ *
+ * In post-aggregation contexts (ORDER BY when the query aggregates), the
+ * reference resolves to an identifier on the output column name, so it
+ * matches columns produced by the aggregate projection. In pre-projection
+ * contexts (GROUP BY, non-aggregated ORDER BY), the target column's full
+ * expression is substituted, since the output name won't exist yet.
+ *
+ * @param {ExprNode} expr
+ * @param {SelectColumn[]} columns
+ * @param {string} clauseName
+ * @param {boolean} [postAgg]
+ * @returns {ExprNode}
+ */
+function resolvePositionalRef(expr, columns, clauseName, postAgg) {
+  if (expr.type !== 'literal') return expr
+  const n = expr.value
+  if (typeof n !== 'number' || !Number.isInteger(n) || n <= 0) return expr
+  if (n > columns.length) {
+    throw new ParseError({
+      message: `${clauseName} position ${n} is out of range (expected 1..${columns.length}) at position ${expr.positionStart}`,
+      positionStart: expr.positionStart,
+      positionEnd: expr.positionEnd,
+    })
+  }
+  const col = columns[n - 1]
+  if (col.type === 'star') {
+    throw new ParseError({
+      message: `${clauseName} position ${n} refers to * which is not supported at position ${expr.positionStart}`,
+      positionStart: expr.positionStart,
+      positionEnd: expr.positionEnd,
+    })
+  }
+  if (postAgg) {
+    return {
+      type: 'identifier',
+      name: col.alias ?? derivedAlias(col.expr),
+      positionStart: expr.positionStart,
+      positionEnd: expr.positionEnd,
+    }
+  }
+  return col.expr
+}
+
 /**
  * @param {ParserState} state
  * @returns {SelectColumn[]}
@@ -376,6 +428,73 @@ function parseSelectList(state) {
   return cols
 }
 
+/**
+ * Peeks whether the current token starts a table-valued function call
+ * like `UNNEST(...)`.
+ *
+ * @param {ParserState} state
+ * @returns {boolean}
+ */
+export function isTableFunctionStart(state) {
+  const tok = current(state)
+  if (tok.type !== 'identifier') return false
+  if (!isTableFunction(tok.value.toUpperCase())) return false
+  const next = peekToken(state, 1)
+  return next.type === 'paren' && next.value === '('
+}
+
+/**
+ * Parses a table function source: UNNEST(args...) [AS alias[(col_alias)]]
+ *
+ * @param {ParserState} state
+ * @returns {FromFunction}
+ */
+export function parseFromFunction(state) {
+  const funcTok = consume(state)
+  const funcName = funcTok.value.toUpperCase()
+  const { positionStart } = funcTok
+
+  expect(state, 'paren', '(')
+  /** @type {ExprNode[]} */
+  const args = []
+  if (!match(state, 'paren', ')')) {
+    while (true) {
+      args.push(parseExpression(state))
+      if (!match(state, 'comma')) break
+    }
+    expect(state, 'paren', ')')
+  }
+
+  validateFunctionArgs(funcName, args.length, positionStart, state.lastPos, state.functions)
+
+  const alias = parseTableAlias(state)
+  /** @type {string | undefined} */
+  let columnAlias
+  if (alias && match(state, 'paren', '(')) {
+    const colStart = state.lastPos
+    const colTok = expect(state, 'identifier')
+    columnAlias = colTok.value
+    if (match(state, 'comma')) {
+      throw new ParseError({
+        message: `${funcName} produces a single column; only one column alias is allowed`,
+        positionStart: colStart,
+        positionEnd: state.lastPos,
+      })
+    }
+    expect(state, 'paren', ')')
+  }
+
+  return {
+    type: 'function',
+    funcName,
+    args,
+    alias,
+    columnAlias,
+    positionStart,
+    positionEnd: state.lastPos,
+  }
+}
+
 /**
  * Parses an optional table alias (e.g., "FROM users u" or "FROM users AS u")
  * @param {ParserState} state

package/src/parse/primary.js

@@ -1,4 +1,4 @@
-import { isCastType, isExtractField, isIntervalUnit, isKnownFunction, niladicFuncs } from '../validation/functions.js'
+import { isCastType, isExtractField, isIntervalUnit, isKnownFunction, isTableFunction, niladicFuncs } from '../validation/functions.js'
 import { InvalidLiteralError, ParseError, SyntaxError, UnknownFunctionError } from '../validation/parseErrors.js'
 import { RESERVED_KEYWORDS } from '../validation/keywords.js'
 import { parseExpression } from './expression.js'
@@ -7,7 +7,7 @@ import { parseStatement } from './parse.js'
 import { consume, current, expect, match, parseError, peekToken } from './state.js'
 
 /**
- * @import { ExprNode, IntervalNode, ParserState, WhenClause } from '../types.js'
+ * @import { ExprNode, IntervalNode, ParserState, SqlPrimitive, WhenClause } from '../types.js'
  */
 
 /**
@@ -40,6 +40,34 @@ export function parsePrimary(state) {
     return expr
   }
 
+  // Array literal: [elem, elem, ...] — elements must be literals
+  if (match(state, 'bracket', '[')) {
+    /** @type {SqlPrimitive[]} */
+    const values = []
+    if (!match(state, 'bracket', ']')) {
+      while (true) {
+        const elemStart = current(state).positionStart
+        const elem = parseExpression(state)
+        if (elem.type !== 'literal') {
+          throw new ParseError({
+            message: 'Array literal elements must be constant literals',
+            positionStart: elemStart,
+            positionEnd: state.lastPos,
+          })
+        }
+        values.push(elem.value)
+        if (!match(state, 'comma')) break
+      }
+      expect(state, 'bracket', ']')
+    }
+    return {
+      type: 'literal',
+      value: values,
+      positionStart,
+      positionEnd: state.lastPos,
+    }
+  }
+
   if (tok.type === 'identifier') {
     const next = peekToken(state, 1)
     const funcNameUpper = tok.value.toUpperCase()
@@ -109,6 +137,14 @@ export function parsePrimary(state) {
         })
       }
 
+      if (isTableFunction(funcNameUpper)) {
+        throw new ParseError({
+          message: `${funcNameUpper} is a table function and can only be used in FROM clauses at position ${positionStart}`,
+          positionStart,
+          positionEnd: tok.positionEnd,
+        })
+      }
+
       return parseFunctionCall(state, positionStart)
     }
 

package/src/parse/tokenize.js

@@ -222,6 +222,17 @@ export function tokenizeSql(query) {
       continue
     }
 
+    if (ch === '[' || ch === ']') {
+      i++
+      tokens.push({
+        type: 'bracket',
+        value: ch,
+        positionStart,
+        positionEnd: i,
+      })
+      continue
+    }
+
     if (ch === ';') {
       i++
       tokens.push({

package/src/parse/types.d.ts

@@ -17,6 +17,7 @@ export type TokenType =
   | 'comma'
   | 'dot'
   | 'paren'
+  | 'bracket'
   | 'semicolon'
   | 'eof'
 

package/src/plan/columns.js

@@ -1,15 +1,30 @@
 import { derivedAlias } from '../expression/alias.js'
 
 /**
- * @import { AsyncDataSource, ExprNode, FromSubquery, FromTable, IdentifierNode, SelectStatement, Statement } from '../types.js'
+ * @import { AsyncDataSource, ExprNode, FromFunction, FromSubquery, FromTable, IdentifierNode, SelectStatement, Statement } from '../types.js'
  */
 
 /**
- * @param {FromTable | FromSubquery} from
+ * @param {FromTable | FromSubquery | FromFunction} from
  * @returns {string}
  */
 export function fromAlias(from) {
-  return from.alias ?? (from.type === 'table' ? from.table : 'table')
+  if (from.alias) return from.alias
+  if (from.type === 'table') return from.table
+  if (from.type === 'function') return from.funcName.toLowerCase()
+  // Unaliased subquery: no natural name. Callers that need a stable alias
+  // should require one at parse time.
+  return 'table'
+}
+
+/**
+ * Returns the single output column name for a FROM table function.
+ *
+ * @param {FromFunction} from
+ * @returns {string}
+ */
+export function tableFunctionColumnName(from) {
+  return from.columnAlias ?? from.funcName.toLowerCase()
 }
 
 /**
@@ -54,8 +69,24 @@ export function extractColumns({ select, parentColumns }) {
   // directly. For non-star queries, parent names may be aliases and are
   // handled below by filtering derived columns and collecting from expressions.
   const hasStar = select.columns.some(col => col.type === 'star' && !col.table)
+  // Exclude parent names that match a derived alias in this SELECT — those are
+  // produced by projection (e.g. `SELECT *, a+b AS c`), not by the source.
+  /** @type {Set<string>} */
+  const derivedAliases = new Set()
+  for (const col of select.columns) {
+    if (col.type === 'derived') {
+      derivedAliases.add(col.alias ?? derivedAlias(col.expr))
+    }
+  }
   /** @type {IdentifierNode[]} */
-  const identifiers = hasStar && parentColumns ? [...parentColumns] : []
+  const identifiers = hasStar && parentColumns
+    ? parentColumns.filter(id => !derivedAliases.has(id.name))
+    : []
+  // Identifiers collected from lateral table function arguments, grouped with
+  // the aliases visible to that argument. Earlier lateral outputs are visible;
+  // the current function alias and later joins are not.
+  /** @type {{ identifiers: IdentifierNode[], visibleAliases: string[] }[]} */
+  const lateralArgGroups = []
 
   // Collect ORDER BY identifiers, excluding SELECT aliases (their underlying
   // columns are already collected from select.columns expressions above)
@@ -88,8 +119,19 @@ export function extractColumns({ select, parentColumns }) {
     collectColumnsFromExpr(expr, identifiers, selectAliases)
   }
   collectColumnsFromExpr(select.having, identifiers, selectAliases)
+  const visibleLateralAliases = [fromAlias(select.from)]
   for (const join of select.joins) {
     collectColumnsFromExpr(join.on, identifiers)
+    const joinAlias = join.alias ?? join.table
+    if (join.fromFunction) {
+      /** @type {IdentifierNode[]} */
+      const lateralArgIdentifiers = []
+      for (const arg of join.fromFunction.args) {
+        collectColumnsFromExpr(arg, lateralArgIdentifiers)
+      }
+      lateralArgGroups.push({ identifiers: lateralArgIdentifiers, visibleAliases: [...visibleLateralAliases] })
+    }
+    visibleLateralAliases.push(joinAlias)
   }
 
   // Partition identifiers by table prefix
@@ -111,6 +153,26 @@ export function extractColumns({ select, parentColumns }) {
     }
   }
 
+  // Partition identifiers from lateral UNNEST args using only the left-side
+  // aliases that are in scope for that specific join.
+  for (const { identifiers, visibleAliases } of lateralArgGroups) {
+    for (const { prefix, name } of identifiers) {
+      if (prefix) {
+        const set = perTable.get(prefix)
+        if (set) set.add(name)
+      } else {
+        if (visibleAliases.length > 1) {
+          for (const alias of visibleAliases) {
+            perTable.set(alias, undefined)
+          }
+        } else if (visibleAliases.length === 1) {
+          const set = perTable.get(visibleAliases[0])
+          if (set) set.add(name)
+        }
+      }
+    }
+  }
+
   // Build result map: convert Sets to arrays, undefined for all-columns tables
   for (const alias of aliases) {
     const set = perTable.get(alias)
@@ -207,7 +269,14 @@ function collectColumnsFromStatement(stmt, columns) {
   if (stmt.from?.type === 'subquery') {
     collectColumnsFromStatement(stmt.from.query, columns)
   }
-  for (const join of stmt.joins) collectColumnsFromExpr(join.on, columns)
+  for (const join of stmt.joins) {
+    collectColumnsFromExpr(join.on, columns)
+    if (join.fromFunction) {
+      for (const arg of join.fromFunction.args) {
+        collectColumnsFromExpr(arg, columns)
+      }
+    }
+  }
   for (const expr of stmt.groupBy) collectColumnsFromExpr(expr, columns)
   collectColumnsFromExpr(stmt.having, columns)
   for (const item of stmt.orderBy) collectColumnsFromExpr(item.expr, columns)
@@ -255,11 +324,33 @@ export function inferStatementColumns({ stmt, cteColumns, tables }) {
  * @param {Record<string, AsyncDataSource>} [options.tables]
  * @returns {string[]}
  */
-function inferSelectSourceColumns({ select, cteColumns, tables }) {
+export function inferSelectSourceColumns({ select, cteColumns, tables }) {
   if (select.from.type === 'subquery') {
     return inferStatementColumns({ stmt: select.from.query, cteColumns, tables })
   }
 
+  if (select.from.type === 'function') {
+    // Table functions currently produce a single column
+    if (!select.joins.length) {
+      return [tableFunctionColumnName(select.from)]
+    }
+    /** @type {string[]} */
+    const result = []
+    const alias = fromAlias(select.from)
+    result.push(`${alias}.${tableFunctionColumnName(select.from)}`)
+    for (const join of select.joins) {
+      const joinAlias = join.alias ?? join.table
+      if (join.fromFunction) {
+        result.push(`${joinAlias}.${tableFunctionColumnName(join.fromFunction)}`)
+      } else {
+        for (const col of lookupTableColumns(join.table, cteColumns, tables)) {
+          result.push(`${joinAlias}.${col}`)
+        }
+      }
+    }
+    return result
+  }
+
   if (!select.joins.length) {
     return lookupTableColumns(select.from.table, cteColumns, tables)
   }
@@ -267,14 +358,18 @@ function inferSelectSourceColumns({ select, cteColumns, tables }) {
   // Collect all sources, then prefix each table's columns
   /** @type {string[]} */
   const result = []
-  const fromAlias = select.from.alias ?? select.from.table
+  const sourceAlias = select.from.alias ?? select.from.table
   for (const col of lookupTableColumns(select.from.table, cteColumns, tables)) {
-    result.push(`${fromAlias}.${col}`)
+    result.push(`${sourceAlias}.${col}`)
   }
   for (const join of select.joins) {
     const joinAlias = join.alias ?? join.table
-    for (const col of lookupTableColumns(join.table, cteColumns, tables)) {
-      result.push(`${joinAlias}.${col}`)
+    if (join.fromFunction) {
+      result.push(`${joinAlias}.${tableFunctionColumnName(join.fromFunction)}`)
+    } else {
+      for (const col of lookupTableColumns(join.table, cteColumns, tables)) {
+        result.push(`${joinAlias}.${col}`)
+      }
     }
   }
   return result