Skip to content

Commit 11b10fb

Browse files
BillLeoutsakosvl346Bill Leoutsakos
BillLeoutsakosvl346
and
Bill Leoutsakos
authored
fix(selectors): drain Snowflake result partitions (#7357)
Co-authored-by: Bill Leoutsakos <billleoutsakos@Bills-MacBook-Pro.local>
1 parent 349f40a commit 11b10fb

4 files changed

Lines changed: 340 additions & 20 deletions

File tree

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

Lines changed: 217 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,217 @@
1+
/**
2+
* @vitest-environment node
3+
*/
4+
import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
5+
6+
const { mockFetch, mockResolveCredentialBundle } = vi.hoisted(() => ({
7+
mockFetch: vi.fn(),
8+
mockResolveCredentialBundle: vi.fn(),
9+
}))
10+
11+
vi.mock('@/lib/selectors/server/providers/credential-bundle', () => ({
12+
resolveSelectorCredentialBundle: mockResolveCredentialBundle,
13+
}))
14+
15+
import { createSelectorProtectedValues } from '@/lib/selectors/server/protected-values'
16+
import { snowflakeSelectorAttachments } from '@/lib/selectors/server/providers/snowflake'
17+
import type { ExecuteServerSelectorArgs } from '@/lib/selectors/server/types'
18+
19+
const STATEMENT_HANDLE = '019c06a4-0000-df4f-0000-00100006589e'
20+
21+
function jsonResponse(body: unknown, status = 200): Response {
22+
return new Response(JSON.stringify(body), {
23+
status,
24+
headers: { 'Content-Type': 'application/json' },
25+
})
26+
}
27+
28+
function tableArgs(signal?: AbortSignal): ExecuteServerSelectorArgs {
29+
return {
30+
selectorKey: 'snowflake.tables',
31+
context: {
32+
oauthCredential: 'credential-1',
33+
database: 'ANALYTICS',
34+
schema: 'PUBLIC',
35+
},
36+
request: { kind: 'list' },
37+
scope: { kind: 'workspace', workspaceId: 'workspace-1' },
38+
workspaceId: 'workspace-1',
39+
principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' },
40+
requesterUserId: 'user-1',
41+
credential: { suppliedId: 'credential-1' },
42+
references: new Map(),
43+
signal,
44+
protectedValues: createSelectorProtectedValues(),
45+
}
46+
}
47+
48+
describe('Snowflake server selector adapter', () => {
49+
beforeEach(() => {
50+
vi.clearAllMocks()
51+
vi.stubGlobal('fetch', mockFetch)
52+
mockResolveCredentialBundle.mockResolvedValue({
53+
accessToken: 'server-only-token',
54+
domain: 'acme.snowflakecomputing.com',
55+
})
56+
})
57+
58+
afterAll(() => vi.unstubAllGlobals())
59+
60+
it('returns every advertised result partition in order', async () => {
61+
mockFetch
62+
.mockResolvedValueOnce(
63+
jsonResponse({
64+
statementHandle: STATEMENT_HANDLE,
65+
data: [['ALPHA', 'first']],
66+
resultSetMetaData: {
67+
numRows: 3,
68+
partitionInfo: [{ rowCount: 1 }, { rowCount: 2 }],
69+
},
70+
})
71+
)
72+
.mockResolvedValueOnce(
73+
jsonResponse({
74+
data: [
75+
['BETA', null],
76+
['GAMMA', 'third'],
77+
],
78+
})
79+
)
80+
81+
await expect(
82+
snowflakeSelectorAttachments['snowflake.tables'].execute(tableArgs())
83+
).resolves.toEqual({
84+
kind: 'list',
85+
items: [
86+
{
87+
id: 'ALPHA',
88+
label: 'ALPHA — first',
89+
meta: { name: 'ALPHA', detail: 'first' },
90+
},
91+
{ id: 'BETA', label: 'BETA', meta: { name: 'BETA' } },
92+
{
93+
id: 'GAMMA',
94+
label: 'GAMMA — third',
95+
meta: { name: 'GAMMA', detail: 'third' },
96+
},
97+
],
98+
})
99+
100+
expect(mockFetch).toHaveBeenCalledTimes(2)
101+
expect(String(mockFetch.mock.calls[1]?.[0])).toBe(
102+
`https://acme.snowflakecomputing.com/api/v2/statements/${STATEMENT_HANDLE}?partition=1`
103+
)
104+
expect(mockFetch.mock.calls[1]?.[1]).toMatchObject({ method: 'GET', redirect: 'error' })
105+
})
106+
107+
it('rejects the whole selector when a later partition fails', async () => {
108+
mockFetch
109+
.mockResolvedValueOnce(
110+
jsonResponse({
111+
statementHandle: STATEMENT_HANDLE,
112+
data: [['ALPHA', null]],
113+
resultSetMetaData: {
114+
numRows: 2,
115+
partitionInfo: [{ rowCount: 1 }, { rowCount: 1 }],
116+
},
117+
})
118+
)
119+
.mockResolvedValueOnce(jsonResponse({ message: 'private provider payload' }, 500))
120+
121+
await expect(
122+
snowflakeSelectorAttachments['snowflake.tables'].execute(tableArgs())
123+
).rejects.toMatchObject({
124+
name: 'SelectorOptionsUnavailableError',
125+
message: 'Options unavailable',
126+
status: 502,
127+
})
128+
expect(mockFetch).toHaveBeenCalledTimes(2)
129+
})
130+
131+
it('preserves caller cancellation during a later partition', async () => {
132+
const controller = new AbortController()
133+
const abortError = new DOMException('The operation was aborted', 'AbortError')
134+
let markLaterFetchStarted: (() => void) | undefined
135+
const laterFetchStarted = new Promise<void>((resolve) => {
136+
markLaterFetchStarted = resolve
137+
})
138+
mockFetch
139+
.mockResolvedValueOnce(
140+
jsonResponse({
141+
statementHandle: STATEMENT_HANDLE,
142+
data: [['ALPHA', null]],
143+
resultSetMetaData: {
144+
numRows: 3,
145+
partitionInfo: [{ rowCount: 1 }, { rowCount: 1 }, { rowCount: 1 }],
146+
},
147+
})
148+
)
149+
.mockImplementationOnce((_input: RequestInfo | URL, init?: RequestInit) => {
150+
markLaterFetchStarted?.()
151+
return new Promise<Response>((_resolve, reject) => {
152+
init?.signal?.addEventListener('abort', () => reject(init.signal?.reason), { once: true })
153+
})
154+
})
155+
156+
const execution = snowflakeSelectorAttachments['snowflake.tables'].execute(
157+
tableArgs(controller.signal)
158+
)
159+
await laterFetchStarted
160+
controller.abort(abortError)
161+
162+
await expect(execution).rejects.toBe(abortError)
163+
expect(mockFetch).toHaveBeenCalledTimes(2)
164+
})
165+
166+
it.each([
167+
{
168+
name: 'missing partition metadata',
169+
body: {
170+
statementHandle: STATEMENT_HANDLE,
171+
data: [['ALPHA', null]],
172+
resultSetMetaData: { numRows: 1 },
173+
},
174+
},
175+
{
176+
name: 'more than 16 partitions',
177+
body: {
178+
statementHandle: STATEMENT_HANDLE,
179+
data: [['ALPHA', null]],
180+
resultSetMetaData: {
181+
numRows: 1,
182+
partitionInfo: Array.from({ length: 17 }, () => ({ rowCount: 0 })),
183+
},
184+
},
185+
},
186+
{
187+
name: 'more than 1,000 rows',
188+
body: {
189+
statementHandle: STATEMENT_HANDLE,
190+
data: [['ALPHA', null]],
191+
resultSetMetaData: { numRows: 1_001, partitionInfo: [{ rowCount: 1 }] },
192+
},
193+
},
194+
{
195+
name: 'an invalid handle for a partitioned result',
196+
body: {
197+
statementHandle: '../untrusted-handle',
198+
data: [['ALPHA', null]],
199+
resultSetMetaData: {
200+
numRows: 2,
201+
partitionInfo: [{ rowCount: 1 }, { rowCount: 1 }],
202+
},
203+
},
204+
},
205+
])('rejects $name before requesting more data', async ({ body }) => {
206+
mockFetch.mockResolvedValueOnce(jsonResponse(body))
207+
208+
await expect(
209+
snowflakeSelectorAttachments['snowflake.tables'].execute(tableArgs())
210+
).rejects.toMatchObject({
211+
name: 'SelectorOptionsUnavailableError',
212+
message: 'Options unavailable',
213+
status: 502,
214+
})
215+
expect(mockFetch).toHaveBeenCalledTimes(1)
216+
})
217+
})

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

Lines changed: 92 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import { MAX_SELECTOR_OPTIONS } from '@/lib/selectors/limits'
12
import type { ServerSelectorKey } from '@/lib/selectors/manifest'
23
import {
34
SelectorConnectionUnavailableError,
@@ -41,6 +42,10 @@ const SNOWFLAKE_SELECTOR_SPECS = {
4142
const SELECTOR_ROW_LIMIT = 1_000
4243
const SELECTOR_TIMEOUT_SECONDS = 20
4344
const SELECTOR_FETCH_TIMEOUT_MS = (SELECTOR_TIMEOUT_SECONDS + 10) * 1_000
45+
const SELECTOR_MAX_PARTITIONS = 16
46+
const SELECTOR_MAX_AGGREGATE_RESPONSE_BYTES = 16 * 1024 * 1024
47+
const SNOWFLAKE_STATEMENT_HANDLE_PATTERN =
48+
/^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i
4449

4550
interface SnowflakeObject {
4651
name: string
@@ -52,6 +57,46 @@ interface SnowflakeDestination {
5257
baseUrl: string
5358
}
5459

60+
function requirePartitionCount(value: number | null): number {
61+
if (
62+
typeof value !== 'number' ||
63+
!Number.isSafeInteger(value) ||
64+
value < 1 ||
65+
value > SELECTOR_MAX_PARTITIONS
66+
) {
67+
throw new SelectorOptionsUnavailableError()
68+
}
69+
return value
70+
}
71+
72+
function requireTotalRows(value: number | null): number {
73+
if (
74+
typeof value !== 'number' ||
75+
!Number.isSafeInteger(value) ||
76+
value < 0 ||
77+
value > SELECTOR_ROW_LIMIT
78+
) {
79+
throw new SelectorOptionsUnavailableError()
80+
}
81+
return value
82+
}
83+
84+
function requireStatementHandle(value: string): string {
85+
if (!SNOWFLAKE_STATEMENT_HANDLE_PATTERN.test(value)) {
86+
throw new SelectorOptionsUnavailableError()
87+
}
88+
return value
89+
}
90+
91+
async function fetchSnowflakeResponse(url: string, init: RequestInit): Promise<Response> {
92+
const response = await fetch(url, { ...init, redirect: 'error' })
93+
if (!response.ok) {
94+
await response.body?.cancel().catch(() => undefined)
95+
throw selectorProviderStatusError(response.status)
96+
}
97+
return response
98+
}
99+
55100
function parseAvailableRoles(cellValue: string | null | undefined): SnowflakeObject[] {
56101
if (!cellValue) return []
57102
let parsed: unknown
@@ -117,42 +162,72 @@ async function executeSnowflake(
117162

118163
const timeoutSignal = AbortSignal.timeout(SELECTOR_FETCH_TIMEOUT_MS)
119164
const signal = args.signal ? AbortSignal.any([args.signal, timeoutSignal]) : timeoutSignal
120-
let response: Response
165+
const headers = buildSnowflakeAuthHeaders(destination.accessToken)
166+
const byteBudget = { remainingBytes: SELECTOR_MAX_AGGREGATE_RESPONSE_BYTES }
121167
try {
122-
response = await fetch(`${destination.baseUrl}/api/v2/statements`, {
168+
const response = await fetchSnowflakeResponse(`${destination.baseUrl}/api/v2/statements`, {
123169
method: 'POST',
124-
headers: buildSnowflakeAuthHeaders(destination.accessToken),
170+
headers,
125171
body: JSON.stringify({
126172
statement,
127173
timeout: SELECTOR_TIMEOUT_SECONDS,
128174
parameters: { rows_per_resultset: SELECTOR_ROW_LIMIT },
129175
}),
130176
signal,
131-
redirect: 'error',
132177
})
133-
} catch (error) {
134-
if (args.signal?.aborted) throw error
135-
throw new SelectorOptionsUnavailableError()
136-
}
178+
const output = await readSnowflakeResult(response, { signal, byteBudget })
179+
if (output.status !== 'SUCCEEDED' || !output.result) {
180+
throw new SelectorOptionsUnavailableError()
181+
}
137182

138-
if (!response.ok) {
139-
await response.body?.cancel().catch(() => undefined)
140-
throw selectorProviderStatusError(response.status)
141-
}
183+
const partitionCount = requirePartitionCount(output.result.partitionCount)
184+
const totalRows = requireTotalRows(output.result.totalRows)
185+
const rows = [...output.result.rows]
186+
if (rows.length > SELECTOR_ROW_LIMIT) throw new SelectorOptionsUnavailableError()
142187

143-
try {
144-
const result = await readSnowflakeResult(response)
145-
if (!result.result) throw new SelectorOptionsUnavailableError()
188+
if (partitionCount > 1) {
189+
const statementHandle = requireStatementHandle(output.statementHandle)
190+
for (let partition = 1; partition < partitionCount; partition += 1) {
191+
signal.throwIfAborted()
192+
const partitionResponse = await fetchSnowflakeResponse(
193+
`${destination.baseUrl}/api/v2/statements/${encodeURIComponent(statementHandle)}?partition=${partition}`,
194+
{ method: 'GET', headers, signal }
195+
)
196+
const partitionOutput = await readSnowflakeResult(partitionResponse, {
197+
currentPartition: partition,
198+
partitionCount,
199+
fallbackStatementHandle: statementHandle,
200+
signal,
201+
byteBudget,
202+
})
203+
if (
204+
partitionOutput.status !== 'SUCCEEDED' ||
205+
partitionOutput.statementHandle !== statementHandle ||
206+
!partitionOutput.result ||
207+
partitionOutput.result.partitionCount !== partitionCount
208+
) {
209+
throw new SelectorOptionsUnavailableError()
210+
}
211+
rows.push(...partitionOutput.result.rows)
212+
if (rows.length > SELECTOR_ROW_LIMIT) throw new SelectorOptionsUnavailableError()
213+
}
214+
}
215+
216+
signal.throwIfAborted()
217+
if (rows.length !== totalRows) throw new SelectorOptionsUnavailableError()
146218
const objects: SnowflakeObject[] =
147219
spec.kind === 'roles'
148-
? parseAvailableRoles(result.result.rows[0]?.[0])
149-
: result.result.rows.flatMap((row) => {
220+
? parseAvailableRoles(rows[0]?.[0])
221+
: rows.flatMap((row) => {
150222
const name = row[0]
151223
if (typeof name !== 'string' || !name) return []
152224
return [{ name, detail: typeof row[1] === 'string' ? row[1] : null }]
153225
})
226+
if (objects.length > MAX_SELECTOR_OPTIONS) throw new SelectorOptionsUnavailableError()
154227
return flatSelectorResult(args.request, objects.map(toOption), true)
155228
} catch (error) {
229+
if (args.signal?.aborted) throw error
230+
if (error instanceof SelectorConnectionUnavailableError) throw error
156231
if (error instanceof SelectorOptionsUnavailableError) throw error
157232
throw new SelectorOptionsUnavailableError()
158233
}

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -871,6 +871,22 @@ describe('Snowflake SQL API transport', () => {
871871
)
872872
expect(canceled).toBe(true)
873873
expect(emittedBytes).toBeLessThan(SNOWFLAKE_MAX_RESPONSE_BYTES * 2)
874+
875+
const budgetedBody = JSON.stringify({
876+
statementHandle: 'budgeted',
877+
data: [],
878+
resultSetMetaData: { numRows: 0, partitionInfo: [{ rowCount: 0 }] },
879+
})
880+
const budgetedBodyBytes = Buffer.byteLength(budgetedBody)
881+
const byteBudget = { remainingBytes: budgetedBodyBytes + 5 }
882+
await expect(
883+
readSnowflakeResult(new Response(budgetedBody), { byteBudget })
884+
).resolves.toMatchObject({ statementHandle: 'budgeted', status: 'SUCCEEDED' })
885+
expect(byteBudget.remainingBytes).toBe(5)
886+
887+
await expect(readSnowflakeResult(new Response(budgetedBody), { byteBudget })).rejects.toThrow(
888+
'Snowflake response body exceeds maximum size'
889+
)
874890
})
875891

876892
it('rejects session-context names that are not Snowflake identifiers', () => {

0 commit comments

Comments
 (0)