feat(table): add Columns to Return to the legacy Table block's Query Rows (#9887)

* feat(table): add Columns to Return to the legacy Table block's Query Rows

* fix(table): treat a column selection with no resolvable columns as all columns

* Revert "fix(table): treat a column selection with no resolvable columns as all columns"

This reverts commit 1b59b12.

Author: waleedlatif1

apps/sim/app/api/table/[tableId]/rows/route.ts

@@ -124,6 +124,7 @@ export const GET = defineInternalJsonRoute({
       limit: query.limit,
       offset: query.offset,
       after: query.after,
+      columns: query.columns,
       includeTotal: query.includeTotal,
       includeRunState: query.limit !== undefined && query.limit <= TABLE_LIMITS.MAX_QUERY_LIMIT,
       allowExpandedLimit: true,

apps/sim/blocks/blocks/table.test.ts

@@ -23,4 +23,13 @@ describe('table query_rows transformer', () => {
       /Invalid number for Limit/
     )
   })
+
+  it('passes selected output columns to the query tool and treats an empty selection as all', () => {
+    expect(
+      params({ operation: 'query_rows', tableId: 'table-1', outputColumns: ['Email'] }).columns
+    ).toEqual(['Email'])
+    expect(
+      params({ operation: 'query_rows', tableId: 'table-1', outputColumns: [] }).columns
+    ).toBeUndefined()
+  })
 })

apps/sim/blocks/blocks/table.ts

@@ -3,7 +3,12 @@ import { toError } from '@sim/utils/errors'
 import { TABLE_LIMITS } from '@/lib/table/constants'
 import { filterRulesToFilter, sortRulesToSort } from '@/lib/table/query-builder/converters'
 import type { BlockConfig } from '@/blocks/types'
-import { parseOptionalNumberInput } from '@/blocks/utils'
+import {
+  parseOptionalNumberInput,
+  resolveTableColumnSelection,
+  TABLE_OUTPUT_COLUMNS_INPUT,
+  TABLE_OUTPUT_COLUMNS_SUBBLOCK,
+} from '@/blocks/utils'
 import { TABLE_ID_PARAM } from '@/tools/table/params'
 import { getTrigger } from '@/triggers'
 
@@ -52,6 +57,9 @@ interface TableBlockParams {
   rows?: string | unknown
   filter?: string | unknown
   sort?: string | unknown
+  outputColumns?: unknown
+  /** The tool's own `columns` param, supplied by a model when the block runs as an Agent tool. */
+  columns?: unknown
   limit?: string
   offset?: string
   builderMode?: string
@@ -178,6 +186,7 @@ const paramTransformers: Record<string, (params: TableBlockParams) => ParsedPara
 
     return {
       tableId: params.tableId,
+      columns: resolveTableColumnSelection(params),
       filter,
       sort,
       limit: parseOptionalNumberInput(params.limit, 'Limit', {
@@ -327,6 +336,8 @@ export const TableBlock: BlockConfig = {
       required: true,
     },
 
+    TABLE_OUTPUT_COLUMNS_SUBBLOCK,
+
     // Row ID for get/update/delete
     {
       id: 'rowId',
@@ -739,6 +750,7 @@ Return ONLY the sort JSON:`,
     data: { type: 'json', description: 'Row data for insert/update' },
     rows: { type: 'array', description: 'Array of row data for batch insert' },
     rowId: { type: 'string', description: 'Row identifier for ID-based operations' },
+    outputColumns: TABLE_OUTPUT_COLUMNS_INPUT,
     bulkFilterMode: {
       type: 'string',
       description: 'Filter input mode for bulk operations (builder or json)',

apps/sim/blocks/blocks/table_v2.test.ts

@@ -31,7 +31,7 @@ describe('table_v2 query_rows transformer', () => {
   it('fails fast instead of silently returning every column for a malformed selection', () => {
     expect(() =>
       params({ operation: 'query_rows', tableId: 't', outputColumns: 'col_name,col_email' })
-    ).toThrow(/Invalid JSON in Columns to Return/)
+    ).toThrow(/Invalid JSON for Columns to Return/)
     expect(() =>
       params({ operation: 'query_rows', tableId: 't', outputColumns: { col_name: true } })
     ).toThrow(/Invalid Columns to Return/)

apps/sim/blocks/blocks/table_v2.ts

@@ -12,6 +12,11 @@ import type {
   TablePredicateInput,
 } from '@/lib/table/types'
 import type { BlockConfig } from '@/blocks/types'
+import {
+  resolveTableColumnSelection,
+  TABLE_OUTPUT_COLUMNS_INPUT,
+  TABLE_OUTPUT_COLUMNS_SUBBLOCK,
+} from '@/blocks/utils'
 import { TABLE_ID_PARAM } from '@/tools/table/params'
 import { getTrigger } from '@/triggers'
 
@@ -90,28 +95,6 @@ function resolveOrder(params: TableBlockParams): SortSpec | undefined {
   return (parsed as SortSpec | undefined) || undefined
 }
 
-/**
- * Resolves the "Columns to Return" selection. Only an absent value (`undefined`,
- * `null`, the dependsOn-cleared `''`, or `[]`) means "all columns"; anything
- * else must be a list of column ids/names. A malformed value fails fast rather
- * than silently widening a deliberately narrowed query to every column — the
- * same stance `parseOptionalLimit` takes for a malformed limit. A JSON string is
- * accepted because an agent-authored block config serializes arrays that way.
- */
-function resolveOutputColumns(raw: unknown): string[] | undefined {
-  if (raw === undefined || raw === null || raw === '') return undefined
-  const parsed = typeof raw === 'string' ? parseJSON(raw, 'Columns to Return') : raw
-  if (!Array.isArray(parsed)) {
-    throw new Error('Invalid Columns to Return: expected a list of column ids or names')
-  }
-  if (parsed.length === 0) return undefined
-  const columns = parsed.map((column) => (typeof column === 'string' ? column.trim() : ''))
-  if (columns.some((column) => column.length === 0)) {
-    throw new Error('Invalid Columns to Return: every entry must be a non-empty column id or name')
-  }
-  return columns
-}
-
 interface ParsedParams {
   tableId?: string
   rowId?: string
@@ -200,9 +183,7 @@ const paramTransformers: Record<string, (params: TableBlockParams) => ParsedPara
 
   query_rows: (params) => ({
     tableId: params.tableId,
-    // The canvas selection wins; a model-supplied `columns` (Agent tool use) is the
-    // fallback so it is not clobbered by an empty canvas value when merged.
-    columns: resolveOutputColumns(params.outputColumns) ?? resolveOutputColumns(params.columns),
+    columns: resolveTableColumnSelection(params),
     filter: resolveFilter(params),
     order: resolveOrder(params),
     // Omitted limit returns the entire result (server fails fast over 5MB);
@@ -342,19 +323,7 @@ export const TableV2Block: BlockConfig = {
       required: true,
     },
 
-    {
-      id: 'outputColumns',
-      title: 'Columns to Return',
-      type: 'dropdown',
-      selectorKey: 'table.outputColumns',
-      multiSelect: true,
-      searchable: true,
-      preserveLabelCase: true,
-      placeholder: 'All columns',
-      description: 'Only include these columns in each returned row. Leave empty for all columns.',
-      dependsOn: { any: ['tableSelector', 'manualTableId'] },
-      condition: { field: 'operation', value: 'query_rows' },
-    },
+    TABLE_OUTPUT_COLUMNS_SUBBLOCK,
 
     {
       id: 'rowId',
@@ -590,10 +559,7 @@ Return ONLY the JSON object:`,
     data: { type: 'json', description: 'Row data for insert/update' },
     rows: { type: 'array', description: 'Array of row data for batch insert' },
     rowId: { type: 'string', description: 'Row identifier for ID-based operations' },
-    outputColumns: {
-      type: 'array',
-      description: 'Table columns to include in each queried row; omit for all columns',
-    },
+    outputColumns: TABLE_OUTPUT_COLUMNS_INPUT,
     filterInput: {
       type: 'json',
       description:

apps/sim/blocks/utils.ts

@@ -703,6 +703,60 @@ export function parseOptionalJsonInput<T = unknown>(value: unknown, label: strin
   return value as T
 }
 
+const COLUMN_SELECTION_LABEL = 'Columns to Return'
+
+function parseColumnSelection(raw: unknown): string[] | undefined {
+  const parsed = parseOptionalJsonInput<unknown>(raw, COLUMN_SELECTION_LABEL)
+  if (parsed === undefined || parsed === null) return undefined
+  if (!Array.isArray(parsed)) {
+    throw new Error(`Invalid ${COLUMN_SELECTION_LABEL}: expected a list of column ids or names`)
+  }
+  if (parsed.length === 0) return undefined
+  const columns = parsed.map((column) => (typeof column === 'string' ? column.trim() : ''))
+  if (columns.some((column) => column.length === 0)) {
+    throw new Error(
+      `Invalid ${COLUMN_SELECTION_LABEL}: every entry must be a non-empty column id or name`
+    )
+  }
+  return columns
+}
+
+/**
+ * Resolves a Table block's "Columns to Return" for a query. The canvas selection wins; a
+ * model-supplied `columns` (Agent tool use) is the fallback, so an empty canvas value does not
+ * clobber it. Only an absent value (`undefined`, `null`, `''`, or `[]`) means "all columns";
+ * anything else must be a list of column ids/names, and a malformed value fails fast rather than
+ * silently widening a deliberately narrowed query. A JSON string is accepted because an
+ * agent-authored block config serializes arrays that way.
+ */
+export function resolveTableColumnSelection(params: {
+  outputColumns?: unknown
+  columns?: unknown
+}): string[] | undefined {
+  return parseColumnSelection(params.outputColumns) ?? parseColumnSelection(params.columns)
+}
+
+/** The Table blocks' "Columns to Return" picker, shown for Query Rows. */
+export const TABLE_OUTPUT_COLUMNS_SUBBLOCK: SubBlockConfig = {
+  id: 'outputColumns',
+  title: COLUMN_SELECTION_LABEL,
+  type: 'dropdown',
+  selectorKey: 'table.outputColumns',
+  multiSelect: true,
+  searchable: true,
+  preserveLabelCase: true,
+  placeholder: 'All columns',
+  description: 'Only include these columns in each returned row. Leave empty for all columns.',
+  dependsOn: { any: ['tableSelector', 'manualTableId'] },
+  condition: { field: 'operation', value: 'query_rows' },
+}
+
+/** Block input entry for {@link TABLE_OUTPUT_COLUMNS_SUBBLOCK}. */
+export const TABLE_OUTPUT_COLUMNS_INPUT = {
+  type: 'array',
+  description: 'Table columns to include in each queried row; omit for all columns',
+} as const
+
 /**
  * Parses an optional numeric block input value.
  * Returns `undefined` for empty values and throws when the provided value is not a valid number.

apps/sim/lib/api/contracts/tables.ts

@@ -783,6 +783,23 @@ const jsonQueryValue = <S extends z.ZodType>(schema: S) =>
     }
   }, schema)
 
+/** Columns to keep in each returned row's data, by stable id or name; omitted or empty means all. */
+const rowColumnSelectionSchema = z
+  .array(
+    requiredFieldSchema('Column reference must not be empty').max(
+      MAX_ID_LENGTH,
+      'Column reference is too long'
+    )
+  )
+  .max(
+    TABLE_LIMITS.MAX_COLUMNS_PER_TABLE,
+    `Cannot select more than ${TABLE_LIMITS.MAX_COLUMNS_PER_TABLE} columns`
+  )
+  .optional()
+  .describe(
+    'Stable column identifiers or column names to include. Omit or pass an empty array for all columns; a reference that matches no column is ignored.'
+  )
+
 /** Unrefined base so v1 contracts can `.extend()` — consumers use {@link tableRowsQuerySchema}. */
 export const tableRowsQueryBaseSchema = z.object({
   workspaceId: workspaceIdSchema,
@@ -849,7 +866,10 @@ const unboundedTableRowsLimitSchema = z.preprocess(
 )
 
 export const tableRowsQuerySchema = tableRowsQueryBaseSchema
-  .extend({ limit: unboundedTableRowsLimitSchema })
+  .extend({
+    limit: unboundedTableRowsLimitSchema,
+    columns: jsonQueryValue(rowColumnSelectionSchema).optional(),
+  })
   .refine((data) => !(data.after && data.sort), {
     message: 'after cursor cannot be combined with sort — cursors paginate the default order',
   })
@@ -1074,21 +1094,7 @@ export const rowQueryBodySchema = z.object({
   workspaceId: z.string().min(1, 'Workspace ID is required'),
   predicate: predicateInputSchema.optional(),
   sort: sortSpecSchema.optional(),
-  columns: z
-    .array(
-      requiredFieldSchema('Column reference must not be empty').max(
-        MAX_ID_LENGTH,
-        'Column reference is too long'
-      )
-    )
-    .max(
-      TABLE_LIMITS.MAX_COLUMNS_PER_TABLE,
-      `Cannot select more than ${TABLE_LIMITS.MAX_COLUMNS_PER_TABLE} columns`
-    )
-    .optional()
-    .describe(
-      'Stable column identifiers or column names to include. Omit or pass an empty array for all columns; a reference that matches no column is ignored.'
-    ),
+  columns: rowColumnSelectionSchema,
   // Omitted limit returns the ENTIRE matching result, failing fast (400) when
   // it exceeds the response byte budget. An explicit limit caps the page row
   // count; the byte budget may still end a page early with nextCursor set.

apps/sim/lib/internal/table/execute-tool.test.ts

@@ -40,6 +40,7 @@ vi.mock('@/lib/internal/table/operations', () => ({
 }))
 
 import { executeTableTool } from '@/lib/internal/table/execute-tool'
+import { tableQueryRowsTool } from '@/tools/table/query_rows'
 
 const { mockCreateExecutorPrincipalFromExecutionContext } = executorPrincipalMockFns
 
@@ -144,6 +145,30 @@ describe('executeTableTool', () => {
     })
   })
 
+  it('carries a legacy query tool column selection through the query-string contract', async () => {
+    const input = tableQueryRowsTool.operation.input?.({
+      tableId: 'table-1',
+      columns: ['Email', 'col_status'],
+      limit: 10,
+      _context: { workspaceId: 'workspace-1' },
+    } as never)
+
+    const response = await executeTableTool({
+      toolId: 'table_query_rows',
+      input,
+      headers: new Headers(),
+      context: CONTEXT,
+      requestId: 'request-1',
+    })
+
+    expect(response.status).toBe(200)
+    expect(mocks.queryRows).toHaveBeenCalledWith(
+      'table-1',
+      expect.objectContaining({ columns: ['Email', 'col_status'] }),
+      expect.anything()
+    )
+  })
+
   it('projects a stale workflow binding as authentication failure', async () => {
     mockCreateExecutorPrincipalFromExecutionContext.mockRejectedValueOnce(
       new InvalidInternalDelegationBindingError()

apps/sim/lib/internal/table/operations.ts

@@ -233,6 +233,7 @@ export async function executeTableQueryRows(
       legacyKeying: 'names',
       limit: query.limit,
       offset: query.offset,
+      columns: query.columns,
       includeTotal: query.includeTotal,
       includeRunState: query.limit !== undefined && query.limit <= TABLE_LIMITS.MAX_QUERY_LIMIT,
       allowExpandedLimit: true,