diff --git a/apps/sim/lib/selectors/server/providers/snowflake.test.ts b/apps/sim/lib/selectors/server/providers/snowflake.test.ts new file mode 100644 index 00000000000..b3d9fb08df3 --- /dev/null +++ b/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((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((_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) + }) +}) diff --git a/apps/sim/lib/selectors/server/providers/snowflake.ts b/apps/sim/lib/selectors/server/providers/snowflake.ts index fc51f3015a6..3708f047c1e 100644 --- a/apps/sim/lib/selectors/server/providers/snowflake.ts +++ b/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 { + 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() } diff --git a/apps/sim/tools/snowflake/utils.test.ts b/apps/sim/tools/snowflake/utils.test.ts index 85a7dd92709..e0331305c47 100644 --- a/apps/sim/tools/snowflake/utils.test.ts +++ b/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', () => { diff --git a/apps/sim/tools/snowflake/utils.ts b/apps/sim/tools/snowflake/utils.ts index 0fc1c3417b7..d6af74c7710 100644 --- a/apps/sim/tools/snowflake/utils.ts +++ b/apps/sim/tools/snowflake/utils.ts @@ -97,6 +97,9 @@ interface SnowflakeResponseOptions { partitionCount?: number canceled?: boolean fallbackStatementHandle?: string + signal?: AbortSignal + /** Remaining decoded response bytes available to a multi-response consumer. */ + byteBudget?: { remainingBytes: number } } interface SnowflakeStatementBodyOptions { @@ -291,7 +294,7 @@ export async function readSnowflakeResult( response: Response, options: SnowflakeResponseOptions = {} ): Promise { - const data = await readSnowflakeJson(response) + const data = await readSnowflakeJson(response, options) const pending = response.status === 202 const cancelRequest = options.canceled === true assertSnowflakeSuccess(response, data, cancelRequest) @@ -367,11 +370,20 @@ export function transformSnowflakeResult

( }) } -async function readSnowflakeJson(response: Response): Promise { +async function readSnowflakeJson( + response: Response, + options: Pick +): Promise { + const remainingBytes = options.byteBudget?.remainingBytes ?? SNOWFLAKE_MAX_RESPONSE_BYTES + if (!Number.isSafeInteger(remainingBytes) || remainingBytes < 0) { + throw new Error('Snowflake response byte budget is invalid') + } const body = await readResponseTextWithLimit(response, { - maxBytes: SNOWFLAKE_MAX_RESPONSE_BYTES, + maxBytes: Math.min(SNOWFLAKE_MAX_RESPONSE_BYTES, remainingBytes), label: 'Snowflake response body', + signal: options.signal, }) + if (options.byteBudget) options.byteBudget.remainingBytes -= Buffer.byteLength(body) let data: unknown try { data = JSON.parse(body)