From 2a289936d61725b488ba62cd43665940526d4b19 Mon Sep 17 00:00:00 2001
From: Meng Lin <menglinmaker@gmail.com>
Date: Sun, 1 Mar 2026 08:16:03 +1100
Subject: [PATCH] fix: include selected columns for query interpreter

---
 .../tanstack-db-sql-interpreter/src/ast/groupClause.ts    | 2 +-
 .../tanstack-db-sql-interpreter/src/ast/orderClause.ts    | 6 +++---
 .../src/ast/selectStatementNode.ts                        | 2 +-
 .../tanstack-db-sql-interpreter/src/ast/shared/column.ts  | 8 ++++++--
 .../src/ast/shared/expression/expression.ts               | 2 +-
 .../tanstack-db-sql-interpreter/src/util/collection.ts    | 2 ++
 6 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/packages/tanstack-db-sql-interpreter/src/ast/groupClause.ts b/packages/tanstack-db-sql-interpreter/src/ast/groupClause.ts
index 2c5f047..f0d2247 100644
--- a/packages/tanstack-db-sql-interpreter/src/ast/groupClause.ts
+++ b/packages/tanstack-db-sql-interpreter/src/ast/groupClause.ts
@@ -11,7 +11,7 @@ export const groupNode = (node: Node.GROUP, collections: Collections) => {
       case 'BY__':
         break
       case 'COLUMN': {
-        columns.push(columnNode(n, collections))
+        columns.push(columnNode(n, collections, true))
         break
       }
       default:
diff --git a/packages/tanstack-db-sql-interpreter/src/ast/orderClause.ts b/packages/tanstack-db-sql-interpreter/src/ast/orderClause.ts
index 8c6a3f8..877df65 100644
--- a/packages/tanstack-db-sql-interpreter/src/ast/orderClause.ts
+++ b/packages/tanstack-db-sql-interpreter/src/ast/orderClause.ts
@@ -33,17 +33,17 @@ const orderExpressionNode = (
   switch (n.name) {
     case 'ORDER_COLUMN':
       return {
-        ...columnNode(n.children[0], collections),
+        ...columnNode(n.children[0], collections, true),
         type: 'asc',
       } as const
     case 'ORDER_ASC':
       return {
-        ...columnNode(n.children[0], collections),
+        ...columnNode(n.children[0], collections, true),
         type: 'asc',
       } as const
     case 'ORDER_DESC':
       return {
-        ...columnNode(n.children[0], collections),
+        ...columnNode(n.children[0], collections, true),
         type: 'desc',
       } as const
     default:
diff --git a/packages/tanstack-db-sql-interpreter/src/ast/selectStatementNode.ts b/packages/tanstack-db-sql-interpreter/src/ast/selectStatementNode.ts
index c7c2d5b..7fe58e1 100644
--- a/packages/tanstack-db-sql-interpreter/src/ast/selectStatementNode.ts
+++ b/packages/tanstack-db-sql-interpreter/src/ast/selectStatementNode.ts
@@ -8,8 +8,8 @@ import { groupNode } from './groupClause.ts'
 import { limitNode } from './limitClause.ts'
 import { orderNode } from './orderClause.ts'
 import { selectNode } from './selectClause.ts'
-import { expressionNode } from './shared/expression/expression.ts'
 import { applyExpression } from './shared/expression/apply.ts'
+import { expressionNode } from './shared/expression/expression.ts'
 import { stringifyExpression } from './shared/expression/stringify.ts'
 
 export const selectStatementNode = (
diff --git a/packages/tanstack-db-sql-interpreter/src/ast/shared/column.ts b/packages/tanstack-db-sql-interpreter/src/ast/shared/column.ts
index 4141905..ed09b4c 100644
--- a/packages/tanstack-db-sql-interpreter/src/ast/shared/column.ts
+++ b/packages/tanstack-db-sql-interpreter/src/ast/shared/column.ts
@@ -6,12 +6,16 @@ import {
 import { defaultSwitchNodeError } from '../../util/error.ts'
 import type { Collections } from '../../util/types.ts'
 
-export const columnNode = (node: Node.COLUMN, collections: Collections) => {
+export const columnNode = (
+  node: Node.COLUMN,
+  collections: Collections,
+  includeSelected: boolean = false,
+) => {
   const n = node.children[0]
   switch (n.name) {
     case 'COLUMN_NAME__': {
       const column = n.value
-      return findColumnFromTables(collections, column)
+      return findColumnFromTables(collections, column, includeSelected)
     }
     case 'TABLE_COLUMN_NAME': {
       const table = n.children[0].value
diff --git a/packages/tanstack-db-sql-interpreter/src/ast/shared/expression/expression.ts b/packages/tanstack-db-sql-interpreter/src/ast/shared/expression/expression.ts
index 4c558a0..b4a296c 100644
--- a/packages/tanstack-db-sql-interpreter/src/ast/shared/expression/expression.ts
+++ b/packages/tanstack-db-sql-interpreter/src/ast/shared/expression/expression.ts
@@ -85,7 +85,7 @@ export const expressionNode = (
     case 'COLUMN': {
       return {
         type: 'column',
-        column: columnNode(n, collections),
+        column: columnNode(n, collections, true),
       }
     }
 
diff --git a/packages/tanstack-db-sql-interpreter/src/util/collection.ts b/packages/tanstack-db-sql-interpreter/src/util/collection.ts
index b980977..3b25ee2 100644
--- a/packages/tanstack-db-sql-interpreter/src/util/collection.ts
+++ b/packages/tanstack-db-sql-interpreter/src/util/collection.ts
@@ -45,12 +45,14 @@ export const collectionProperties = (
 export const findColumnFromTables = (
   collections: Collections,
   column: string,
+  includeSelected: boolean = false,
 ) => {
   const tableNames = Object.keys(collections)
   for (const table of tableNames) {
     const properties = collectionProperties(collections, table)
     if (properties.includes(column)) return { table, column }
   }
+  if (includeSelected) return { table: '$selected', column }
   throw new LiveQuerySqlError(
     `Cannot find column '${column}' in tables: '${tableNames.join(`', '`)}'`,
   )