fix(selectors): drain Snowflake result partitions

apps/sim/lib/selectors/server/providers/snowflake.test.ts

@@ -0,0 +1,217 @@
+/**
+ * @vitest-environment node
+ */
+import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
+
+const { mockFetch, mockResolveCredentialBundle } = vi.hoisted(() => ({
+  mockFetch: vi.fn(),
+  mockResolveCredentialBundle: vi.fn(),
+}))
+
+vi.mock('@/lib/selectors/server/providers/credential-bundle', () => ({
+  resolveSelectorCredentialBundle: mockResolveCredentialBundle,
+}))
+
+import { createSelectorProtectedValues } from '@/lib/selectors/server/protected-values'
+import { snowflakeSelectorAttachments } from '@/lib/selectors/server/providers/snowflake'
+import type { ExecuteServerSelectorArgs } from '@/lib/selectors/server/types'
+
+const STATEMENT_HANDLE = '019c06a4-0000-df4f-0000-00100006589e'
+
+function jsonResponse(body: unknown, status = 200): Response {
+  return new Response(JSON.stringify(body), {
+    status,
+    headers: { 'Content-Type': 'application/json' },
+  })
+}
+
+function tableArgs(signal?: AbortSignal): ExecuteServerSelectorArgs {
+  return {
+    selectorKey: 'snowflake.tables',
+    context: {
+      oauthCredential: 'credential-1',
+      database: 'ANALYTICS',
+      schema: 'PUBLIC',
+    },
+    request: { kind: 'list' },
+    scope: { kind: 'workspace', workspaceId: 'workspace-1' },
+    workspaceId: 'workspace-1',
+    principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' },
+    requesterUserId: 'user-1',
+    credential: { suppliedId: 'credential-1' },
+    references: new Map(),
+    signal,
+    protectedValues: createSelectorProtectedValues(),
+  }
+}
+
+describe('Snowflake server selector adapter', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    vi.stubGlobal('fetch', mockFetch)
+    mockResolveCredentialBundle.mockResolvedValue({
+      accessToken: 'server-only-token',
+      domain: 'acme.snowflakecomputing.com',
+    })
+  })
+
+  afterAll(() => vi.unstubAllGlobals())
+
+  it('returns every advertised result partition in order', async () => {
+    mockFetch
+      .mockResolvedValueOnce(
+        jsonResponse({
+          statementHandle: STATEMENT_HANDLE,
+          data: [['ALPHA', 'first']],
+          resultSetMetaData: {
+            numRows: 3,
+            partitionInfo: [{ rowCount: 1 }, { rowCount: 2 }],
+          },
+        })
+      )
+      .mockResolvedValueOnce(
+        jsonResponse({
+          data: [
+            ['BETA', null],
+            ['GAMMA', 'third'],
+          ],
+        })
+      )
+
+    await expect(
+      snowflakeSelectorAttachments['snowflake.tables'].execute(tableArgs())
+    ).resolves.toEqual({
+      kind: 'list',
+      items: [
+        {
+          id: 'ALPHA',
+          label: 'ALPHA — first',
+          meta: { name: 'ALPHA', detail: 'first' },
+        },
+        { id: 'BETA', label: 'BETA', meta: { name: 'BETA' } },
+        {
+          id: 'GAMMA',
+          label: 'GAMMA — third',
+          meta: { name: 'GAMMA', detail: 'third' },
+        },
+      ],
+    })
+
+    expect(mockFetch).toHaveBeenCalledTimes(2)
+    expect(String(mockFetch.mock.calls[1]?.[0])).toBe(
+      `https://acme.snowflakecomputing.com/api/v2/statements/${STATEMENT_HANDLE}?partition=1`
+    )
+    expect(mockFetch.mock.calls[1]?.[1]).toMatchObject({ method: 'GET', redirect: 'error' })
+  })
+
+  it('rejects the whole selector when a later partition fails', async () => {
+    mockFetch
+      .mockResolvedValueOnce(
+        jsonResponse({
+          statementHandle: STATEMENT_HANDLE,
+          data: [['ALPHA', null]],
+          resultSetMetaData: {
+            numRows: 2,
+            partitionInfo: [{ rowCount: 1 }, { rowCount: 1 }],
+          },
+        })
+      )
+      .mockResolvedValueOnce(jsonResponse({ message: 'private provider payload' }, 500))
+
+    await expect(
+      snowflakeSelectorAttachments['snowflake.tables'].execute(tableArgs())
+    ).rejects.toMatchObject({
+      name: 'SelectorOptionsUnavailableError',
+      message: 'Options unavailable',
+      status: 502,
+    })
+    expect(mockFetch).toHaveBeenCalledTimes(2)
+  })
+
+  it('preserves caller cancellation during a later partition', async () => {
+    const controller = new AbortController()
+    const abortError = new DOMException('The operation was aborted', 'AbortError')
+    let markLaterFetchStarted: (() => void) | undefined
+    const laterFetchStarted = new Promise<void>((resolve) => {
+      markLaterFetchStarted = resolve
+    })
+    mockFetch
+      .mockResolvedValueOnce(
+        jsonResponse({
+          statementHandle: STATEMENT_HANDLE,
+          data: [['ALPHA', null]],
+          resultSetMetaData: {
+            numRows: 3,
+            partitionInfo: [{ rowCount: 1 }, { rowCount: 1 }, { rowCount: 1 }],
+          },
+        })
+      )
+      .mockImplementationOnce((_input: RequestInfo | URL, init?: RequestInit) => {
+        markLaterFetchStarted?.()
+        return new Promise<Response>((_resolve, reject) => {
+          init?.signal?.addEventListener('abort', () => reject(init.signal?.reason), { once: true })
+        })
+      })
+
+    const execution = snowflakeSelectorAttachments['snowflake.tables'].execute(
+      tableArgs(controller.signal)
+    )
+    await laterFetchStarted
+    controller.abort(abortError)
+
+    await expect(execution).rejects.toBe(abortError)
+    expect(mockFetch).toHaveBeenCalledTimes(2)
+  })
+
+  it.each([
+    {
+      name: 'missing partition metadata',
+      body: {
+        statementHandle: STATEMENT_HANDLE,
+        data: [['ALPHA', null]],
+        resultSetMetaData: { numRows: 1 },
+      },
+    },
+    {
+      name: 'more than 16 partitions',
+      body: {
+        statementHandle: STATEMENT_HANDLE,
+        data: [['ALPHA', null]],
+        resultSetMetaData: {
+          numRows: 1,
+          partitionInfo: Array.from({ length: 17 }, () => ({ rowCount: 0 })),
+        },
+      },
+    },
+    {
+      name: 'more than 1,000 rows',
+      body: {
+        statementHandle: STATEMENT_HANDLE,
+        data: [['ALPHA', null]],
+        resultSetMetaData: { numRows: 1_001, partitionInfo: [{ rowCount: 1 }] },
+      },
+    },
+    {
+      name: 'an invalid handle for a partitioned result',
+      body: {
+        statementHandle: '../untrusted-handle',
+        data: [['ALPHA', null]],
+        resultSetMetaData: {
+          numRows: 2,
+          partitionInfo: [{ rowCount: 1 }, { rowCount: 1 }],
+        },
+      },
+    },
+  ])('rejects $name before requesting more data', async ({ body }) => {
+    mockFetch.mockResolvedValueOnce(jsonResponse(body))
+
+    await expect(
+      snowflakeSelectorAttachments['snowflake.tables'].execute(tableArgs())
+    ).rejects.toMatchObject({
+      name: 'SelectorOptionsUnavailableError',
+      message: 'Options unavailable',
+      status: 502,
+    })
+    expect(mockFetch).toHaveBeenCalledTimes(1)
+  })
+})

apps/sim/lib/selectors/server/providers/snowflake.ts

@@ -1,3 +1,4 @@
+import { MAX_SELECTOR_OPTIONS } from '@/lib/selectors/limits'
 import type { ServerSelectorKey } from '@/lib/selectors/manifest'
 import {
   SelectorConnectionUnavailableError,
@@ -41,6 +42,10 @@ const SNOWFLAKE_SELECTOR_SPECS = {
 const SELECTOR_ROW_LIMIT = 1_000
 const SELECTOR_TIMEOUT_SECONDS = 20
 const SELECTOR_FETCH_TIMEOUT_MS = (SELECTOR_TIMEOUT_SECONDS + 10) * 1_000
+const SELECTOR_MAX_PARTITIONS = 16
+const SELECTOR_MAX_AGGREGATE_RESPONSE_BYTES = 16 * 1024 * 1024
+const SNOWFLAKE_STATEMENT_HANDLE_PATTERN =
+  /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i
 
 interface SnowflakeObject {
   name: string
@@ -52,6 +57,46 @@ interface SnowflakeDestination {
   baseUrl: string
 }
 
+function requirePartitionCount(value: number | null): number {
+  if (
+    typeof value !== 'number' ||
+    !Number.isSafeInteger(value) ||
+    value < 1 ||
+    value > SELECTOR_MAX_PARTITIONS
+  ) {
+    throw new SelectorOptionsUnavailableError()
+  }
+  return value
+}
+
+function requireTotalRows(value: number | null): number {
+  if (
+    typeof value !== 'number' ||
+    !Number.isSafeInteger(value) ||
+    value < 0 ||
+    value > SELECTOR_ROW_LIMIT
+  ) {
+    throw new SelectorOptionsUnavailableError()
+  }
+  return value
+}
+
+function requireStatementHandle(value: string): string {
+  if (!SNOWFLAKE_STATEMENT_HANDLE_PATTERN.test(value)) {
+    throw new SelectorOptionsUnavailableError()
+  }
+  return value
+}
+
+async function fetchSnowflakeResponse(url: string, init: RequestInit): Promise<Response> {
+  const response = await fetch(url, { ...init, redirect: 'error' })
+  if (!response.ok) {
+    await response.body?.cancel().catch(() => undefined)
+    throw selectorProviderStatusError(response.status)
+  }
+  return response
+}
+
 function parseAvailableRoles(cellValue: string | null | undefined): SnowflakeObject[] {
   if (!cellValue) return []
   let parsed: unknown
@@ -117,42 +162,72 @@ async function executeSnowflake(
 
   const timeoutSignal = AbortSignal.timeout(SELECTOR_FETCH_TIMEOUT_MS)
   const signal = args.signal ? AbortSignal.any([args.signal, timeoutSignal]) : timeoutSignal
-  let response: Response
+  const headers = buildSnowflakeAuthHeaders(destination.accessToken)
+  const byteBudget = { remainingBytes: SELECTOR_MAX_AGGREGATE_RESPONSE_BYTES }
   try {
-    response = await fetch(`${destination.baseUrl}/api/v2/statements`, {
+    const response = await fetchSnowflakeResponse(`${destination.baseUrl}/api/v2/statements`, {
       method: 'POST',
-      headers: buildSnowflakeAuthHeaders(destination.accessToken),
+      headers,
       body: JSON.stringify({
         statement,
         timeout: SELECTOR_TIMEOUT_SECONDS,
         parameters: { rows_per_resultset: SELECTOR_ROW_LIMIT },
       }),
       signal,
-      redirect: 'error',
     })
-  } catch (error) {
-    if (args.signal?.aborted) throw error
-    throw new SelectorOptionsUnavailableError()
-  }
+    const output = await readSnowflakeResult(response, { signal, byteBudget })
+    if (output.status !== 'SUCCEEDED' || !output.result) {
+      throw new SelectorOptionsUnavailableError()
+    }
 
-  if (!response.ok) {
-    await response.body?.cancel().catch(() => undefined)
-    throw selectorProviderStatusError(response.status)
-  }
+    const partitionCount = requirePartitionCount(output.result.partitionCount)
+    const totalRows = requireTotalRows(output.result.totalRows)
+    const rows = [...output.result.rows]
+    if (rows.length > SELECTOR_ROW_LIMIT) throw new SelectorOptionsUnavailableError()
 
-  try {
-    const result = await readSnowflakeResult(response)
-    if (!result.result) throw new SelectorOptionsUnavailableError()
+    if (partitionCount > 1) {
+      const statementHandle = requireStatementHandle(output.statementHandle)
+      for (let partition = 1; partition < partitionCount; partition += 1) {
+        signal.throwIfAborted()
+        const partitionResponse = await fetchSnowflakeResponse(
+          `${destination.baseUrl}/api/v2/statements/${encodeURIComponent(statementHandle)}?partition=${partition}`,
+          { method: 'GET', headers, signal }
+        )
+        const partitionOutput = await readSnowflakeResult(partitionResponse, {
+          currentPartition: partition,
+          partitionCount,
+          fallbackStatementHandle: statementHandle,
+          signal,
+          byteBudget,
+        })
+        if (
+          partitionOutput.status !== 'SUCCEEDED' ||
+          partitionOutput.statementHandle !== statementHandle ||
+          !partitionOutput.result ||
+          partitionOutput.result.partitionCount !== partitionCount
+        ) {
+          throw new SelectorOptionsUnavailableError()
+        }
+        rows.push(...partitionOutput.result.rows)
+        if (rows.length > SELECTOR_ROW_LIMIT) throw new SelectorOptionsUnavailableError()
+      }
+    }
+
+    signal.throwIfAborted()
+    if (rows.length !== totalRows) throw new SelectorOptionsUnavailableError()
     const objects: SnowflakeObject[] =
       spec.kind === 'roles'
-        ? parseAvailableRoles(result.result.rows[0]?.[0])
-        : result.result.rows.flatMap((row) => {
+        ? parseAvailableRoles(rows[0]?.[0])
+        : rows.flatMap((row) => {
             const name = row[0]
             if (typeof name !== 'string' || !name) return []
             return [{ name, detail: typeof row[1] === 'string' ? row[1] : null }]
           })
+    if (objects.length > MAX_SELECTOR_OPTIONS) throw new SelectorOptionsUnavailableError()
     return flatSelectorResult(args.request, objects.map(toOption), true)
   } catch (error) {
+    if (args.signal?.aborted) throw error
+    if (error instanceof SelectorConnectionUnavailableError) throw error
     if (error instanceof SelectorOptionsUnavailableError) throw error
     throw new SelectorOptionsUnavailableError()
   }

apps/sim/tools/snowflake/utils.test.ts

@@ -871,6 +871,22 @@ describe('Snowflake SQL API transport', () => {
     )
     expect(canceled).toBe(true)
     expect(emittedBytes).toBeLessThan(SNOWFLAKE_MAX_RESPONSE_BYTES * 2)
+
+    const budgetedBody = JSON.stringify({
+      statementHandle: 'budgeted',
+      data: [],
+      resultSetMetaData: { numRows: 0, partitionInfo: [{ rowCount: 0 }] },
+    })
+    const budgetedBodyBytes = Buffer.byteLength(budgetedBody)
+    const byteBudget = { remainingBytes: budgetedBodyBytes + 5 }
+    await expect(
+      readSnowflakeResult(new Response(budgetedBody), { byteBudget })
+    ).resolves.toMatchObject({ statementHandle: 'budgeted', status: 'SUCCEEDED' })
+    expect(byteBudget.remainingBytes).toBe(5)
+
+    await expect(readSnowflakeResult(new Response(budgetedBody), { byteBudget })).rejects.toThrow(
+      'Snowflake response body exceeds maximum size'
+    )
   })
 
   it('rejects session-context names that are not Snowflake identifiers', () => {