chore: import upstream snapshot with attribution
CI / Migrate Dev DB (push) Has been skipped
CI / Detect Version (push) Has been cancelled
CI / Migrate DB (push) Has been cancelled
CI / Build Dev ECR (./docker/app.Dockerfile, ECR_APP) (push) Has been cancelled
CI / Build Dev ECR (./docker/db.Dockerfile, ECR_MIGRATIONS) (push) Has been cancelled
CI / Build Dev ECR (./docker/pii.Dockerfile, ECR_PII) (push) Has been cancelled
CI / Build Dev ECR (./docker/realtime.Dockerfile, ECR_REALTIME) (push) Has been cancelled
CI / Deploy Trigger.dev (Dev) (push) Has been cancelled
CI / Build AMD64 (./docker/app.Dockerfile, ECR_APP, ghcr.io/simstudioai/simstudio) (push) Has been cancelled
CI / Build AMD64 (./docker/db.Dockerfile, ECR_MIGRATIONS, ghcr.io/simstudioai/migrations) (push) Has been cancelled
CI / Build AMD64 (./docker/pii.Dockerfile, ECR_PII, ghcr.io/simstudioai/pii) (push) Has been cancelled
CI / Build AMD64 (./docker/realtime.Dockerfile, ECR_REALTIME, ghcr.io/simstudioai/realtime) (push) Has been cancelled
CI / Build ARM64 (GHCR Only) (./docker/app.Dockerfile, ghcr.io/simstudioai/simstudio) (push) Has been cancelled
CI / Build ARM64 (GHCR Only) (./docker/db.Dockerfile, ghcr.io/simstudioai/migrations) (push) Has been cancelled
CI / Build ARM64 (GHCR Only) (./docker/pii.Dockerfile, ghcr.io/simstudioai/pii) (push) Has been cancelled
CI / Build ARM64 (GHCR Only) (./docker/realtime.Dockerfile, ghcr.io/simstudioai/realtime) (push) Has been cancelled
CI / Create GHCR Manifests (ghcr.io/simstudioai/migrations) (push) Has been cancelled
CI / Create GHCR Manifests (ghcr.io/simstudioai/pii) (push) Has been cancelled
CI / Create GHCR Manifests (ghcr.io/simstudioai/realtime) (push) Has been cancelled
CI / Create GHCR Manifests (ghcr.io/simstudioai/simstudio) (push) Has been cancelled
CI / Check Docs Changes (push) Has been cancelled
CI / Process Docs (push) Has been cancelled
CI / Create GitHub Release (push) Has been cancelled
CI / Test and Build (push) Has been cancelled
Publish CLI Package / publish-npm (push) Has been cancelled
Publish Python SDK / publish-pypi (push) Has been cancelled
Publish TypeScript SDK / publish-npm (push) Has been cancelled

This commit is contained in:
wehub-resource-sync
2026-07-13 13:20:55 +08:00
commit d25d482dc2
13754 changed files with 4996608 additions and 0 deletions
+111
View File
@@ -0,0 +1,111 @@
import { db } from '@sim/db'
import { dataDrains, member } from '@sim/db/schema'
import { and, eq } from 'drizzle-orm'
import { NextResponse } from 'next/server'
import { getSession } from '@/lib/auth'
import { isOrganizationOnEnterprisePlan } from '@/lib/billing/core/subscription'
import { isBillingEnabled, isDataDrainsEnabled } from '@/lib/core/config/env-flags'
interface DrainAccessSession {
user: {
id: string
name?: string | null
email?: string | null
}
membership: {
role: string
}
}
export type DrainAccessResult =
| { ok: true; session: DrainAccessSession }
| { ok: false; response: NextResponse }
/**
* Auth + membership + role + enterprise-plan gate shared by every data-drain
* route. Owner/admin role is required for reads as well as writes since drain
* configs expose customer bucket names and webhook URLs. On Sim Cloud the
* gate is the Enterprise plan; on self-hosted it's `DATA_DRAINS_ENABLED`,
* which 404s when unset so a newer image doesn't silently expose drains.
*/
export async function authorizeDrainAccess(
organizationId: string,
options: { requireMutating: boolean }
): Promise<DrainAccessResult> {
const session = await getSession()
if (!session?.user?.id) {
return { ok: false, response: NextResponse.json({ error: 'Unauthorized' }, { status: 401 }) }
}
const [memberEntry] = await db
.select({ role: member.role })
.from(member)
.where(and(eq(member.organizationId, organizationId), eq(member.userId, session.user.id)))
.limit(1)
if (!memberEntry) {
return {
ok: false,
response: NextResponse.json(
{ error: 'Forbidden - Not a member of this organization' },
{ status: 403 }
),
}
}
if (!isBillingEnabled && !isDataDrainsEnabled) {
return {
ok: false,
response: NextResponse.json(
{ error: 'Data Drains are not enabled on this deployment' },
{ status: 404 }
),
}
}
if (isBillingEnabled) {
const hasEnterprise = await isOrganizationOnEnterprisePlan(organizationId)
if (!hasEnterprise) {
return {
ok: false,
response: NextResponse.json(
{ error: 'Data Drains are available on Enterprise plans only' },
{ status: 403 }
),
}
}
}
if (memberEntry.role !== 'owner' && memberEntry.role !== 'admin') {
return {
ok: false,
response: NextResponse.json(
{
error: options.requireMutating
? 'Forbidden - Only organization owners and admins can manage data drains'
: 'Forbidden - Only organization owners and admins can view data drains',
},
{ status: 403 }
),
}
}
return {
ok: true,
session: {
user: {
id: session.user.id,
name: session.user.name ?? null,
email: session.user.email ?? null,
},
membership: { role: memberEntry.role },
},
}
}
export async function loadDrain(organizationId: string, drainId: string) {
const [drain] = await db
.select()
.from(dataDrains)
.where(and(eq(dataDrains.id, drainId), eq(dataDrains.organizationId, organizationId)))
.limit(1)
return drain ?? null
}
@@ -0,0 +1,186 @@
/**
* @vitest-environment node
*/
import { beforeEach, describe, expect, it, vi } from 'vitest'
const { mockUpload, mockDeleteIfExists, BlobServiceClientCtor, StorageSharedKeyCredentialCtor } =
vi.hoisted(() => {
const mockUpload = vi.fn(async () => ({}))
const mockDeleteIfExists = vi.fn(async () => ({ succeeded: true }))
const blockBlobClient = { upload: mockUpload, deleteIfExists: mockDeleteIfExists }
const containerClient = { getBlockBlobClient: vi.fn(() => blockBlobClient) }
return {
mockUpload,
mockDeleteIfExists,
BlobServiceClientCtor: vi.fn().mockImplementation(
class {
getContainerClient = vi.fn(() => containerClient)
}
),
StorageSharedKeyCredentialCtor: vi.fn().mockImplementation(class {}),
}
})
vi.mock('@azure/storage-blob', () => ({
BlobServiceClient: BlobServiceClientCtor,
StorageSharedKeyCredential: StorageSharedKeyCredentialCtor,
}))
import { azureBlobDestination } from '@/lib/data-drains/destinations/azure_blob'
const config = { accountName: 'simstore', containerName: 'drains', prefix: 'sim/' }
// Realistic 88-char base64 (64-byte) Azure storage key shape.
const credentials = {
accountKey:
'YWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWE=',
}
beforeEach(() => {
vi.clearAllMocks()
})
describe('azureBlobDestination openSession', () => {
it('uploads via BlockBlobClient and returns an azure:// locator', async () => {
const session = azureBlobDestination.openSession({ config, credentials })
const body = Buffer.from('row\n', 'utf8')
const result = await session.deliver({
body,
contentType: 'application/x-ndjson',
metadata: {
drainId: 'd1',
runId: 'r1',
source: 'workflow_logs',
sequence: 0,
rowCount: 1,
runStartedAt: new Date('2025-06-15T12:00:00Z'),
},
signal: new AbortController().signal,
})
expect(result.locator).toMatch(
/^azure:\/\/simstore\/drains\/sim\/workflow_logs\/d1\/\d{4}\/\d{2}\/\d{2}\/r1-00000\.ndjson$/
)
expect(mockUpload).toHaveBeenCalledTimes(1)
const [calledBody, calledLength, opts] = mockUpload.mock.calls[0] as [
Buffer,
number,
{ metadata?: Record<string, string> },
]
expect(calledBody).toBe(body)
expect(calledLength).toBe(body.byteLength)
expect(opts.metadata?.simdrainid).toBe('d1')
expect(opts.metadata?.simsequence).toBe('0')
const fullOpts = opts as {
blobHTTPHeaders?: { blobContentType?: string }
abortSignal?: AbortSignal
}
expect(fullOpts.blobHTTPHeaders?.blobContentType).toBe('application/x-ndjson')
expect(fullOpts.abortSignal).toBeDefined()
expect(StorageSharedKeyCredentialCtor).toHaveBeenCalledWith('simstore', credentials.accountKey)
expect(BlobServiceClientCtor).toHaveBeenCalledWith(
'https://simstore.blob.core.windows.net',
expect.anything(),
expect.objectContaining({ retryOptions: expect.any(Object) })
)
await session.close()
})
it('routes to a sovereign-cloud endpoint suffix when configured', async () => {
const session = azureBlobDestination.openSession({
config: { ...config, endpointSuffix: 'blob.core.usgovcloudapi.net' },
credentials,
})
await session.deliver({
body: Buffer.from('row\n', 'utf8'),
contentType: 'application/x-ndjson',
metadata: {
drainId: 'd',
runId: 'r',
source: 'workflow_logs',
sequence: 0,
rowCount: 1,
runStartedAt: new Date('2025-06-15T12:00:00Z'),
},
signal: new AbortController().signal,
})
expect(BlobServiceClientCtor).toHaveBeenCalledWith(
'https://simstore.blob.core.usgovcloudapi.net',
expect.anything(),
expect.objectContaining({ retryOptions: expect.any(Object) })
)
await session.close()
})
it('surfaces Azure REST errors', async () => {
mockUpload.mockRejectedValueOnce(
Object.assign(new Error('Forbidden'), {
statusCode: 403,
code: 'AuthenticationFailed',
})
)
const session = azureBlobDestination.openSession({ config, credentials })
await expect(
session.deliver({
body: Buffer.from('x'),
contentType: 'application/x-ndjson',
metadata: {
drainId: 'd',
runId: 'r',
source: 'audit_logs',
sequence: 0,
rowCount: 1,
runStartedAt: new Date('2025-06-15T12:00:00Z'),
},
signal: new AbortController().signal,
})
).rejects.toThrow(/AuthenticationFailed 403/)
await session.close()
})
})
describe('azureBlobDestination test()', () => {
it('writes a probe blob then attempts cleanup', async () => {
await azureBlobDestination.test!({
config,
credentials,
signal: new AbortController().signal,
})
expect(mockUpload).toHaveBeenCalled()
expect(mockDeleteIfExists).toHaveBeenCalled()
})
})
describe('azureBlobDestination config schema', () => {
it('rejects invalid account names', () => {
const result = azureBlobDestination.configSchema.safeParse({
accountName: 'BAD-NAME',
containerName: 'drains',
})
expect(result.success).toBe(false)
})
it('rejects invalid container names', () => {
const result = azureBlobDestination.configSchema.safeParse({
accountName: 'simstore',
containerName: '--bad--',
})
expect(result.success).toBe(false)
})
})
describe('azureBlobDestination credentials schema', () => {
it('rejects non-base64 account keys', () => {
const padded = 'a'.repeat(70)
const result = azureBlobDestination.credentialsSchema.safeParse({
accountKey: `${padded}!@#$`,
})
expect(result.success).toBe(false)
})
it('rejects keys that are too short', () => {
const result = azureBlobDestination.credentialsSchema.safeParse({ accountKey: 'YQ==' })
expect(result.success).toBe(false)
})
})
@@ -0,0 +1,203 @@
import { createLogger } from '@sim/logger'
import { generateShortId } from '@sim/utils/id'
import { z } from 'zod'
import { buildObjectKey, normalizePrefix } from '@/lib/data-drains/destinations/utils'
import type { DrainDestination } from '@/lib/data-drains/types'
const logger = createLogger('DataDrainAzureBlobDestination')
/**
* Azure storage account names: 3-24 chars, lowercase letters and digits only.
* https://learn.microsoft.com/en-us/azure/storage/common/storage-account-overview#storage-account-name
*/
const ACCOUNT_NAME_RE = /^[a-z0-9]{3,24}$/
/**
* Azure container names: 3-63 chars, lowercase letters, digits, single hyphens
* (no leading/trailing/double hyphens).
* https://learn.microsoft.com/en-us/rest/api/storageservices/naming-and-referencing-containers--blobs--and-metadata
*/
const CONTAINER_NAME_RE = /^[a-z0-9]([a-z0-9]|-(?!-))+[a-z0-9]$/
/** Azure storage account keys are 64 raw bytes => exactly 88 base64 chars (0-2 trailing `=`). */
const ACCOUNT_KEY_RE = /^[A-Za-z0-9+/]+={0,2}$/
/** Public cloud default; sovereign clouds (Gov/China/legacy DE) are validated via allowlist. */
const DEFAULT_ENDPOINT_SUFFIX = 'blob.core.windows.net'
/**
* Allowlist of Azure Storage endpoint suffixes. URL host must end with one of these
* (after the account name + dot). Reject anything else to prevent SSRF via attacker-controlled
* endpoint suffix.
*/
const ALLOWED_ENDPOINT_SUFFIXES = [
'blob.core.windows.net',
'blob.core.usgovcloudapi.net',
'blob.core.chinacloudapi.cn',
'blob.core.cloudapi.de',
] as const
const azureBlobConfigSchema = z.object({
accountName: z
.string()
.min(1, 'accountName is required')
.refine((value) => ACCOUNT_NAME_RE.test(value), {
message: 'accountName must be 3-24 lowercase letters or digits',
}),
containerName: z
.string()
.min(3, 'containerName must be 3-63 characters')
.max(63)
.refine((value) => CONTAINER_NAME_RE.test(value), {
message: 'containerName must use lowercase letters, digits, or single hyphens',
}),
/** Optional prefix; trailing slash is added automatically when assembling blob names. */
prefix: z.string().max(512).optional(),
/** Storage endpoint suffix. Must be one of the known Azure Storage suffixes (public/Gov/China/DE). */
endpointSuffix: z
.string()
.refine((v) => (ALLOWED_ENDPOINT_SUFFIXES as readonly string[]).includes(v), {
message: `endpointSuffix must be one of: ${ALLOWED_ENDPOINT_SUFFIXES.join(', ')}`,
})
.optional(),
})
const azureBlobCredentialsSchema = z.object({
accountKey: z
.string()
.length(88, 'accountKey must be exactly 88 base64 characters (64-byte Azure storage key)')
.refine((v) => ACCOUNT_KEY_RE.test(v), {
message: 'accountKey must be a base64-encoded Azure storage account key',
}),
})
export type AzureBlobDestinationConfig = z.infer<typeof azureBlobConfigSchema>
export type AzureBlobDestinationCredentials = z.infer<typeof azureBlobCredentialsSchema>
interface BlobClients {
containerClient: import('@azure/storage-blob').ContainerClient
}
async function buildClients(
config: AzureBlobDestinationConfig,
credentials: AzureBlobDestinationCredentials
): Promise<BlobClients> {
const { BlobServiceClient, StorageSharedKeyCredential } = await import('@azure/storage-blob')
const sharedKeyCredential = new StorageSharedKeyCredential(
config.accountName,
credentials.accountKey
)
const suffix = config.endpointSuffix ?? DEFAULT_ENDPOINT_SUFFIX
/**
* Bound per-attempt timeout. SDK default `tryTimeoutInMs` is "infinite" / OS
* connection-idle limits, so a hung receiver could pin a delivery indefinitely.
* 30s per try + 5 tries (~ exponential 0.5s → 30s) caps the worst-case wall time.
*/
const blobServiceClient = new BlobServiceClient(
`https://${config.accountName}.${suffix}`,
sharedKeyCredential,
{
retryOptions: {
tryTimeoutInMs: 30_000,
maxTries: 5,
retryDelayInMs: 500,
maxRetryDelayInMs: 30_000,
},
}
)
return { containerClient: blobServiceClient.getContainerClient(config.containerName) }
}
interface AzureRestErrorLike {
statusCode?: number
code?: string
message?: string
}
function isAzureRestError(error: unknown): error is AzureRestErrorLike {
return typeof error === 'object' && error !== null && ('statusCode' in error || 'code' in error)
}
async function withAzureErrorContext<T>(action: string, fn: () => Promise<T>): Promise<T> {
try {
return await fn()
} catch (error) {
if (isAzureRestError(error)) {
const status = error.statusCode
const code = error.code
logger.warn('Azure Blob operation failed', { action, code, status })
throw new Error(
`Azure Blob ${action} failed (${code ?? 'Error'}${status ? ` ${status}` : ''}): ${error.message ?? ''}`,
{ cause: error }
)
}
throw error
}
}
export const azureBlobDestination: DrainDestination<
AzureBlobDestinationConfig,
AzureBlobDestinationCredentials
> = {
type: 'azure_blob',
displayName: 'Azure Blob Storage',
configSchema: azureBlobConfigSchema,
credentialsSchema: azureBlobCredentialsSchema,
async test({ config, credentials, signal }) {
const { containerClient } = await buildClients(config, credentials)
const probeName = `${normalizePrefix(config.prefix)}.sim-drain-write-probe/${generateShortId(12)}`
const blockBlobClient = containerClient.getBlockBlobClient(probeName)
await withAzureErrorContext('test-put', () =>
blockBlobClient.upload(Buffer.alloc(0), 0, {
blobHTTPHeaders: { blobContentType: 'application/octet-stream' },
abortSignal: signal,
})
)
try {
await blockBlobClient.deleteIfExists({ abortSignal: signal })
} catch (cleanupError) {
logger.debug('Azure Blob test write probe cleanup failed (non-fatal)', {
accountName: config.accountName,
containerName: config.containerName,
blobName: probeName,
error: cleanupError,
})
}
},
openSession({ config, credentials }) {
let clientsPromise: Promise<BlobClients> | null = null
return {
async deliver({ body, contentType, metadata, signal }) {
if (clientsPromise === null) clientsPromise = buildClients(config, credentials)
const { containerClient } = await clientsPromise
const blobName = buildObjectKey(config.prefix, metadata)
const blockBlobClient = containerClient.getBlockBlobClient(blobName)
await withAzureErrorContext('put-object', () =>
blockBlobClient.upload(body, body.byteLength, {
blobHTTPHeaders: { blobContentType: contentType },
metadata: {
simdrainid: metadata.drainId,
simrunid: metadata.runId,
simsource: metadata.source,
simsequence: metadata.sequence.toString(),
simrowcount: metadata.rowCount.toString(),
},
abortSignal: signal,
})
)
logger.debug('Azure Blob chunk delivered', {
accountName: config.accountName,
containerName: config.containerName,
blobName,
bytes: body.byteLength,
})
return {
locator: `azure://${config.accountName}/${config.containerName}/${blobName}`,
}
},
async close() {},
}
},
}
@@ -0,0 +1,276 @@
/**
* @vitest-environment node
*/
import { beforeEach, describe, expect, it, vi } from 'vitest'
const { mockGetAccessToken, JWTCtor, loggerInstance } = vi.hoisted(() => {
const mockGetAccessToken = vi.fn(async () => ({ token: 'bq-token' }))
const loggerInstance = {
info: vi.fn(),
warn: vi.fn(),
error: vi.fn(),
debug: vi.fn(),
trace: vi.fn(),
fatal: vi.fn(),
child: vi.fn(),
withMetadata: vi.fn(),
}
return {
mockGetAccessToken,
JWTCtor: vi.fn().mockImplementation(
class {
getAccessToken = mockGetAccessToken
}
),
loggerInstance,
}
})
vi.mock('google-auth-library', () => ({ JWT: JWTCtor }))
vi.mock('@sim/logger', () => ({
createLogger: () => loggerInstance,
logger: loggerInstance,
runWithRequestContext: <T>(_ctx: unknown, fn: () => T): T => fn(),
getRequestContext: () => undefined,
}))
vi.mock('@sim/utils/helpers', () => ({
sleep: vi.fn(async () => {}),
}))
const fetchMock = vi.fn(
async () =>
new Response(JSON.stringify({}), {
status: 200,
headers: { 'content-type': 'application/json' },
})
)
vi.stubGlobal('fetch', fetchMock)
import { bigqueryDestination } from '@/lib/data-drains/destinations/bigquery'
const config = { projectId: 'my-proj', datasetId: 'logs', tableId: 'workflow' }
const credentials = {
serviceAccountJson: JSON.stringify({
client_email: 'sa@p.iam.gserviceaccount.com',
private_key: '-----BEGIN PRIVATE KEY-----\nfake\n-----END PRIVATE KEY-----\n',
}),
}
const meta = {
drainId: 'd',
runId: 'r',
source: 'workflow_logs' as const,
sequence: 0,
rowCount: 2,
runStartedAt: new Date('2025-06-15T12:00:00Z'),
}
beforeEach(() => {
vi.clearAllMocks()
fetchMock.mockResolvedValue(
new Response(JSON.stringify({}), {
status: 200,
headers: { 'content-type': 'application/json' },
})
)
mockGetAccessToken.mockResolvedValue({ token: 'bq-token' })
})
describe('bigqueryDestination', () => {
it('posts rows with stable insertIds for dedup', async () => {
const session = bigqueryDestination.openSession({ config, credentials })
const body = Buffer.from(
`${JSON.stringify({ id: 'a' })}\n${JSON.stringify({ id: 'b' })}\n`,
'utf8'
)
const result = await session.deliver({
body,
contentType: 'application/x-ndjson',
metadata: meta,
signal: new AbortController().signal,
})
expect(fetchMock).toHaveBeenCalledTimes(1)
const [url, init] = fetchMock.mock.calls[0] as [string, RequestInit]
expect(url).toContain('/projects/my-proj/datasets/logs/tables/workflow/insertAll')
const payload = JSON.parse(init.body as string)
expect(payload.rows).toHaveLength(2)
expect(payload.rows[0].insertId).toBe('d-r-0-0')
expect(payload.rows[0].json).toEqual({ id: 'a' })
expect(payload.rows[1].insertId).toBe('d-r-0-1')
expect(result.locator).toBe('bigquery://my-proj/logs/workflow#r-0')
await session.close()
})
it('throws with row indices and warns on partial-failure insertErrors', async () => {
fetchMock.mockResolvedValueOnce(
new Response(
JSON.stringify({
insertErrors: [
{ index: 1, errors: [{ message: 'invalid', reason: 'invalid' }] },
{ index: 2, errors: [{ message: 'bad', reason: 'invalid' }] },
],
}),
{ status: 200, headers: { 'content-type': 'application/json' } }
)
)
const session = bigqueryDestination.openSession({ config, credentials })
const body = Buffer.from(
`${JSON.stringify({ x: 1 })}\n${JSON.stringify({ x: 2 })}\n${JSON.stringify({ x: 3 })}\n`
)
await expect(
session.deliver({
body,
contentType: 'application/x-ndjson',
metadata: meta,
signal: new AbortController().signal,
})
).rejects.toThrow(/partial failure.*1,2.*dedup-keyed by insertId/s)
expect(loggerInstance.warn).toHaveBeenCalledWith(
expect.stringContaining('partial failure'),
expect.objectContaining({
partialFailure: true,
succeededRows: 1,
failedRows: 2,
})
)
await session.close()
})
it('test() probes table existence with a GET', async () => {
await bigqueryDestination.test!({
config,
credentials,
signal: new AbortController().signal,
})
const [url, init] = fetchMock.mock.calls[0] as [string, RequestInit]
expect(url).toContain('?fields=id')
expect(init.method).toBeUndefined()
})
it('throws a clear error when an NDJSON line is malformed', async () => {
const session = bigqueryDestination.openSession({ config, credentials })
const body = Buffer.from(`${JSON.stringify({ id: 'a' })}\n{not json}\n`, 'utf8')
await expect(
session.deliver({
body,
contentType: 'application/x-ndjson',
metadata: meta,
signal: new AbortController().signal,
})
).rejects.toThrow(/NDJSON parse failed at line 2/)
expect(fetchMock).not.toHaveBeenCalled()
await session.close()
})
it('throws when an NDJSON row is not a JSON object', async () => {
const session = bigqueryDestination.openSession({ config, credentials })
const body = Buffer.from(`${JSON.stringify({ id: 'a' })}\n42\n`, 'utf8')
await expect(
session.deliver({
body,
contentType: 'application/x-ndjson',
metadata: meta,
signal: new AbortController().signal,
})
).rejects.toThrow(/NDJSON row at line 2 is not an object/)
await session.close()
})
it('parses NDJSON with CRLF line endings', async () => {
const session = bigqueryDestination.openSession({ config, credentials })
const body = Buffer.from(
`${JSON.stringify({ id: 'a' })}\r\n${JSON.stringify({ id: 'b' })}\r\n`,
'utf8'
)
await session.deliver({
body,
contentType: 'application/x-ndjson',
metadata: meta,
signal: new AbortController().signal,
})
const [, init] = fetchMock.mock.calls[0] as [string, RequestInit]
const payload = JSON.parse(init.body as string)
expect(payload.rows).toHaveLength(2)
expect(payload.rows[0].json).toEqual({ id: 'a' })
expect(payload.rows[1].json).toEqual({ id: 'b' })
await session.close()
})
it('insertId includes drainId prefix to avoid cross-drain collisions', async () => {
const session = bigqueryDestination.openSession({ config, credentials })
const body = Buffer.from(`${JSON.stringify({ id: 'a' })}\n`, 'utf8')
await session.deliver({
body,
contentType: 'application/x-ndjson',
metadata: { ...meta, drainId: 'drain-xyz', runId: 'run-1', sequence: 7 },
signal: new AbortController().signal,
})
const [, init] = fetchMock.mock.calls[0] as [string, RequestInit]
const payload = JSON.parse(init.body as string)
expect(payload.rows[0].insertId).toBe('drain-xyz-run-1-7-0')
await session.close()
})
it('retries 5xx responses with backoff, then succeeds', async () => {
fetchMock.mockResolvedValueOnce(new Response('boom', { status: 503 })).mockResolvedValueOnce(
new Response(JSON.stringify({}), {
status: 200,
headers: { 'content-type': 'application/json' },
})
)
const session = bigqueryDestination.openSession({ config, credentials })
await session.deliver({
body: Buffer.from(`${JSON.stringify({ id: 'a' })}\n`),
contentType: 'application/x-ndjson',
metadata: meta,
signal: new AbortController().signal,
})
expect(fetchMock).toHaveBeenCalledTimes(2)
expect(loggerInstance.warn).toHaveBeenCalledWith(
expect.stringContaining('transient error'),
expect.objectContaining({ status: 503, attempt: 1 })
)
await session.close()
})
it('accepts domain-scoped project IDs', () => {
const result = bigqueryDestination.configSchema.safeParse({
projectId: 'example.com:my-project',
datasetId: 'logs',
tableId: 'workflow',
})
expect(result.success).toBe(true)
const standard = bigqueryDestination.configSchema.safeParse({
projectId: 'my-proj',
datasetId: 'logs',
tableId: 'workflow',
})
expect(standard.success).toBe(true)
})
it('test() throws when serviceAccountJson is missing required fields', async () => {
await expect(
bigqueryDestination.test!({
config,
credentials: {
serviceAccountJson: JSON.stringify({
client_email: 'sa@p.iam.gserviceaccount.com',
}),
},
signal: new AbortController().signal,
})
).rejects.toThrow(/missing private_key/)
await expect(
bigqueryDestination.test!({
config,
credentials: {
serviceAccountJson: JSON.stringify({
private_key: '-----BEGIN PRIVATE KEY-----\nfake\n-----END PRIVATE KEY-----\n',
}),
},
signal: new AbortController().signal,
})
).rejects.toThrow(/missing client_email/)
})
})
@@ -0,0 +1,333 @@
import { createHash } from 'node:crypto'
import { createLogger } from '@sim/logger'
import { toError } from '@sim/utils/errors'
import { backoffWithJitter, parseRetryAfter } from '@sim/utils/retry'
import { JWT } from 'google-auth-library'
import { z } from 'zod'
import {
type ParsedServiceAccount,
parseNdjsonObjects,
parseServiceAccount,
refineServiceAccountJson,
sleepUntilAborted,
} from '@/lib/data-drains/destinations/utils'
import type { DeliveryMetadata, DrainDestination } from '@/lib/data-drains/types'
const logger = createLogger('DataDrainBigQueryDestination')
/**
* Uses the legacy `tabledata.insertAll` streaming endpoint. The Storage Write
* API offers exactly-once semantics and lower pricing but requires gRPC; we
* stay on insertAll for simplicity and direct HTTP support.
*/
/** `insertdata` for streaming inserts; `readonly` for the `tables.get` probe in `test()`. */
const SCOPES = [
'https://www.googleapis.com/auth/bigquery.insertdata',
'https://www.googleapis.com/auth/bigquery.readonly',
]
/** Standard project IDs are 6-30 chars; the optional `domain.tld:` prefix supports legacy domain-scoped projects. */
const PROJECT_ID_RE = /^([a-z][a-z0-9.-]{0,61}[a-z0-9]:)?[a-z][a-z0-9-]{4,28}[a-z0-9]$/
const DATASET_RE = /^[A-Za-z0-9_]{1,1024}$/
const TABLE_RE = /^[\p{L}\p{M}\p{N}\p{Pc}\p{Pd} ]{1,1024}$/u
const USER_AGENT = 'sim-data-drain/1.0'
/** Per-request streaming limits: 10 MB body, 50,000 rows, 1 MB per row. */
const MAX_REQUEST_BYTES = 10 * 1024 * 1024
const MAX_ROWS_PER_REQUEST = 50_000
const MAX_ROW_BYTES = 1024 * 1024
/** `insertId` is capped at 128 characters (encoded length). */
const MAX_INSERT_ID_LENGTH = 128
const PER_ATTEMPT_TIMEOUT_MS = 60_000
const bigqueryConfigSchema = z.object({
projectId: z
.string()
.min(6, 'projectId is required')
.refine((v) => PROJECT_ID_RE.test(v), {
message: 'projectId must match Google Cloud project ID rules',
}),
datasetId: z
.string()
.min(1, 'datasetId is required')
.refine((v) => DATASET_RE.test(v), {
message: 'datasetId may only contain ASCII letters, digits, and underscores (max 1024 chars)',
}),
tableId: z
.string()
.min(1, 'tableId is required')
.refine((v) => TABLE_RE.test(v), {
message:
'tableId may only contain Unicode letters/marks/numbers, connectors, dashes, and spaces (max 1024 chars)',
})
.refine((v) => Buffer.byteLength(v, 'utf8') <= 1024, {
message: 'tableId must be at most 1024 bytes when UTF-8 encoded',
}),
})
const bigqueryCredentialsSchema = z
.object({
serviceAccountJson: z.string().min(1, 'serviceAccountJson is required'),
})
.superRefine(refineServiceAccountJson)
export type BigQueryDestinationConfig = z.infer<typeof bigqueryConfigSchema>
export type BigQueryDestinationCredentials = z.infer<typeof bigqueryCredentialsSchema>
function buildJwt(account: ParsedServiceAccount): JWT {
return new JWT({ email: account.clientEmail, key: account.privateKey, scopes: SCOPES })
}
async function getAccessToken(jwt: JWT, forceRefresh = false): Promise<string> {
if (forceRefresh) {
/** Clearing `credentials` forces `getAccessToken` to mint a new token instead of returning the cached one. */
jwt.credentials = {}
}
const { token } = await jwt.getAccessToken()
if (!token) throw new Error('Failed to obtain BigQuery access token')
return token
}
interface InsertAllInput {
config: BigQueryDestinationConfig
rows: Record<string, unknown>[]
metadata: DeliveryMetadata
jwt: JWT
signal: AbortSignal
}
interface InsertAllError {
index: number
errors: Array<{ reason?: string; message?: string; location?: string }>
}
async function postInsertAll(
input: InsertAllInput,
url: string,
body: string,
forceRefresh = false
): Promise<Response> {
const token = await getAccessToken(input.jwt, forceRefresh)
const perAttempt = AbortSignal.any([input.signal, AbortSignal.timeout(PER_ATTEMPT_TIMEOUT_MS)])
try {
return await fetch(url, {
method: 'POST',
headers: {
Authorization: `Bearer ${token}`,
'Content-Type': 'application/json',
'User-Agent': USER_AGENT,
},
body,
signal: perAttempt,
})
} catch (error) {
logger.warn('BigQuery request failed', {
table: `${input.config.projectId}.${input.config.datasetId}.${input.config.tableId}`,
error: toError(error).message,
})
throw error
}
}
/**
* Builds a stable `insertId` for best-effort dedup (~60s window). Prefixed
* with `drainId` so (runId, sequence) collisions across drains do not
* accidentally dedupe each other's rows. With UUID drain/run IDs the raw
* form fits well under 128 chars; if anything pushes it over (e.g. a future
* non-UUID id), hash the prefix and keep the row-distinguishing `index`
* suffix intact so BigQuery does not silently dedupe distinct rows.
*/
function buildInsertId(metadata: DeliveryMetadata, index: number): string {
const raw = `${metadata.drainId}-${metadata.runId}-${metadata.sequence}-${index}`
if (raw.length <= MAX_INSERT_ID_LENGTH) return raw
const prefixHash = createHash('sha1')
.update(`${metadata.drainId}-${metadata.runId}-${metadata.sequence}`)
.digest('hex')
return `${prefixHash}-${index}`
}
const RETRYABLE_STATUSES = new Set([408, 429, 500, 502, 503, 504])
const MAX_RETRY_ATTEMPTS = 3
const BASE_RETRY_DELAY_MS = 250
/**
* Streams a chunk of rows to `tabledata.insertAll`.
*
* Partial-success caveat: BigQuery may return HTTP 200 with a non-empty
* `insertErrors` array. Rows not listed there are inserted and dedup-keyed by
* `insertId` for ~60s. We throw on any `insertErrors`; retries within the
* dedup window are safe, but retries after it may duplicate succeeded rows.
*/
async function insertAll(input: InsertAllInput): Promise<void> {
if (input.rows.length > MAX_ROWS_PER_REQUEST) {
throw new Error(
`BigQuery insertAll chunk has ${input.rows.length} rows, exceeds the ${MAX_ROWS_PER_REQUEST} per-request limit`
)
}
const url = `https://bigquery.googleapis.com/bigquery/v2/projects/${encodeURIComponent(input.config.projectId)}/datasets/${encodeURIComponent(input.config.datasetId)}/tables/${encodeURIComponent(input.config.tableId)}/insertAll`
/**
* `skipInvalidRows: false` and `ignoreUnknownValues: false` surface schema
* mismatches as `insertErrors` instead of silently dropping data — drains
* should fail loudly so operators notice the schema drift.
*/
const payload = {
skipInvalidRows: false,
ignoreUnknownValues: false,
rows: input.rows.map((row, index) => {
const rowBytes = Buffer.byteLength(JSON.stringify(row), 'utf8')
if (rowBytes > MAX_ROW_BYTES) {
throw new Error(
`BigQuery row at index ${index} is ${rowBytes} bytes, exceeds the ${MAX_ROW_BYTES}-byte per-row limit`
)
}
return {
insertId: buildInsertId(input.metadata, index),
json: row,
}
}),
}
const body = JSON.stringify(payload)
const byteLength = Buffer.byteLength(body, 'utf8')
if (byteLength > MAX_REQUEST_BYTES) {
throw new Error(
`BigQuery insertAll body is ${byteLength} bytes, exceeds the ${MAX_REQUEST_BYTES}-byte per-request limit`
)
}
let attempt = 0
let response: Response | undefined
let refreshedOnce = false
while (true) {
attempt++
try {
response = await postInsertAll(input, url, body)
/** A 401 retry doesn't count against the 5xx/429 budget — token refresh is a one-shot recovery. */
if (response.status === 401 && !refreshedOnce) {
refreshedOnce = true
logger.debug('BigQuery returned 401; refreshing access token and retrying once')
/** Drain the 401 body before discarding so undici can return the socket to the keep-alive pool. */
await response.text().catch(() => '')
response = await postInsertAll(input, url, body, true)
}
if (!RETRYABLE_STATUSES.has(response.status)) break
if (attempt >= MAX_RETRY_ATTEMPTS) break
const retryAfterHeaderMs = parseRetryAfter(response.headers.get('retry-after'))
const retryAfterMs = backoffWithJitter(attempt, retryAfterHeaderMs, {
baseMs: BASE_RETRY_DELAY_MS,
})
logger.warn('BigQuery insertAll transient error; retrying', {
status: response.status,
attempt,
retryAfterMs,
})
/** Drain the body so the keep-alive connection can be reused. */
await response.text().catch(() => '')
await sleepUntilAborted(retryAfterMs, input.signal)
if (input.signal.aborted) throw input.signal.reason ?? new Error('Aborted')
} catch (error) {
/**
* Connection-level failures (DNS, socket reset, timeout) never produce
* a Response — treat them like 5xx and retry with backoff. Re-throw
* aborts unwrapped so callers see the cancellation reason.
*/
if (input.signal.aborted) throw input.signal.reason ?? error
if (attempt >= MAX_RETRY_ATTEMPTS) throw error
const retryAfterMs = backoffWithJitter(attempt, null, { baseMs: BASE_RETRY_DELAY_MS })
logger.warn('BigQuery insertAll network error; retrying', {
attempt,
retryAfterMs,
error: toError(error).message,
})
await sleepUntilAborted(retryAfterMs, input.signal)
if (input.signal.aborted) throw input.signal.reason ?? new Error('Aborted')
}
}
if (!response) throw new Error('BigQuery insertAll failed: no response')
if (!response.ok) {
const text = await response.text().catch(() => '')
throw new Error(`BigQuery insertAll failed (HTTP ${response.status}): ${text}`)
}
const result = (await response.json().catch(() => ({}))) as {
insertErrors?: InsertAllError[]
}
if (result.insertErrors && result.insertErrors.length > 0) {
const failedIndices = result.insertErrors.map((e) => e.index)
const total = input.rows.length
const failed = result.insertErrors.length
const succeeded = total - failed
logger.warn('BigQuery insertAll returned partial failure', {
partialFailure: true,
table: `${input.config.projectId}.${input.config.datasetId}.${input.config.tableId}`,
succeededRows: succeeded,
failedRows: failed,
failedIndices: failedIndices.slice(0, 20),
})
const summary = result.insertErrors
.slice(0, 3)
.map(
(e) =>
`row ${e.index}: ${e.errors.map((er) => er.message ?? er.reason ?? 'unknown').join('; ')}`
)
.join(' | ')
throw new Error(
`BigQuery insertAll partial failure: ${failed} of ${total} rows failed (indices: ${failedIndices.slice(0, 20).join(',')}${failedIndices.length > 20 ? ',...' : ''}); ${succeeded} rows were inserted and are dedup-keyed by insertId for ~60s — retries within that window are safe, but retries after the window may duplicate the succeeded rows. First errors: ${summary}`
)
}
}
export const bigqueryDestination: DrainDestination<
BigQueryDestinationConfig,
BigQueryDestinationCredentials
> = {
type: 'bigquery',
displayName: 'Google BigQuery',
configSchema: bigqueryConfigSchema,
credentialsSchema: bigqueryCredentialsSchema,
/**
* Probes table existence, IAM access, and credential validity in a single
* `tables.get` call. `fields=id` minimises response size — we only care
* whether the call succeeds, not the payload.
*/
async test({ config, credentials, signal }) {
const account = parseServiceAccount(credentials.serviceAccountJson)
const jwt = buildJwt(account)
const token = await getAccessToken(jwt)
const url = `https://bigquery.googleapis.com/bigquery/v2/projects/${encodeURIComponent(config.projectId)}/datasets/${encodeURIComponent(config.datasetId)}/tables/${encodeURIComponent(config.tableId)}?fields=id`
const perAttempt = AbortSignal.any([signal, AbortSignal.timeout(PER_ATTEMPT_TIMEOUT_MS)])
const response = await fetch(url, {
headers: { Authorization: `Bearer ${token}` },
signal: perAttempt,
})
if (!response.ok) {
const text = await response.text().catch(() => '')
throw new Error(`BigQuery probe failed (HTTP ${response.status}): ${text}`)
}
/** Drain the success body so undici can return the socket to the keep-alive pool. */
await response.text().catch(() => '')
},
openSession({ config, credentials }) {
const account = parseServiceAccount(credentials.serviceAccountJson)
const jwt = buildJwt(account)
return {
async deliver({ body, metadata, signal }) {
const rows = parseNdjsonObjects(body, { requireObject: true }) as Record<string, unknown>[]
if (rows.length === 0) {
return {
locator: `bigquery://${config.projectId}/${config.datasetId}/${config.tableId}#${metadata.runId}-${metadata.sequence}`,
}
}
await insertAll({ config, rows, metadata, jwt, signal })
logger.debug('BigQuery chunk delivered', {
table: `${config.projectId}.${config.datasetId}.${config.tableId}`,
rows: rows.length,
})
return {
locator: `bigquery://${config.projectId}/${config.datasetId}/${config.tableId}#${metadata.runId}-${metadata.sequence}`,
}
},
async close() {},
}
},
}
@@ -0,0 +1,211 @@
/**
* @vitest-environment node
*/
import { gunzipSync } from 'node:zlib'
import { beforeEach, describe, expect, it, vi } from 'vitest'
const fetchMock = vi.fn(async () => new Response(null, { status: 202 }))
vi.stubGlobal('fetch', fetchMock)
import { datadogDestination } from '@/lib/data-drains/destinations/datadog'
const config = { site: 'us1' as const, service: 'sim', tags: 'env:prod' }
const credentials = { apiKey: 'dd-key' }
const meta = (sequence: number) => ({
drainId: 'd1',
runId: 'r1',
source: 'workflow_logs' as const,
sequence,
rowCount: 2,
runStartedAt: new Date('2025-06-15T12:00:00Z'),
})
beforeEach(() => {
vi.clearAllMocks()
fetchMock.mockResolvedValue(new Response(null, { status: 202 }))
})
describe('datadogDestination', () => {
it('parses NDJSON and POSTs a JSON array of log entries', async () => {
const session = datadogDestination.openSession({ config, credentials })
const body = Buffer.from(
`${JSON.stringify({ id: 'a', name: 'one' })}\n${JSON.stringify({ id: 'b', name: 'two' })}\n`,
'utf8'
)
const result = await session.deliver({
body,
contentType: 'application/x-ndjson',
metadata: meta(0),
signal: new AbortController().signal,
})
expect(fetchMock).toHaveBeenCalledTimes(1)
const [url, init] = fetchMock.mock.calls[0] as [string, RequestInit]
expect(url).toBe('https://http-intake.logs.datadoghq.com/api/v2/logs')
const headers = init.headers as Record<string, string>
expect(headers['DD-API-KEY']).toBe('dd-key')
expect(headers['Content-Type']).toBe('application/json')
expect(headers.Accept).toBe('application/json')
expect(headers['User-Agent']).toBe('sim-data-drain/1.0')
expect(headers['Content-Encoding']).toBeUndefined()
const payload = JSON.parse(init.body as string)
expect(payload).toHaveLength(2)
expect(payload[0].ddsource).toBe('sim')
expect(payload[0].service).toBe('sim')
expect(payload[0].ddtags).toContain('sim_drain_id:d1')
expect(payload[0].ddtags).toContain('env:prod')
expect(payload[0].id).toBe('a')
expect(payload[0].name).toBe('one')
expect(payload[0].attributes).toBeUndefined()
expect(result.locator).toMatch(/^datadog:\/\/us1#r1-0/)
await session.close()
})
it('retries 5xx responses then surfaces the final error', async () => {
vi.useFakeTimers()
try {
fetchMock.mockResolvedValue(new Response('boom', { status: 503 }))
const session = datadogDestination.openSession({ config, credentials })
const promise = session.deliver({
body: Buffer.from(`${JSON.stringify({ x: 1 })}\n`),
contentType: 'application/x-ndjson',
metadata: meta(0),
signal: new AbortController().signal,
})
// Attach a handler so Node doesn't flag the in-flight rejection while
// we advance fake timers; we still assert via the original promise below.
const settled = promise.catch((e) => e)
await vi.runAllTimersAsync()
await expect(settled).resolves.toMatchObject({ message: expect.stringMatching(/HTTP 503/) })
expect(fetchMock).toHaveBeenCalledTimes(4)
await session.close()
} finally {
vi.useRealTimers()
}
})
it('does not retry on non-retryable 4xx (e.g. invalid API key)', async () => {
fetchMock.mockResolvedValueOnce(new Response('forbidden', { status: 403 }))
const session = datadogDestination.openSession({ config, credentials })
await expect(
session.deliver({
body: Buffer.from(`${JSON.stringify({ x: 1 })}\n`),
contentType: 'application/x-ndjson',
metadata: meta(0),
signal: new AbortController().signal,
})
).rejects.toThrow(/HTTP 403/)
expect(fetchMock).toHaveBeenCalledTimes(1)
await session.close()
})
it('routes to the EU site host', async () => {
const session = datadogDestination.openSession({
config: { ...config, site: 'eu1' },
credentials,
})
await session.deliver({
body: Buffer.from(`${JSON.stringify({ x: 1 })}\n`),
contentType: 'application/x-ndjson',
metadata: meta(0),
signal: new AbortController().signal,
})
expect(fetchMock.mock.calls[0]?.[0]).toBe('https://http-intake.logs.datadoghq.eu/api/v2/logs')
await session.close()
})
it('routes to the AP2 site host', async () => {
const session = datadogDestination.openSession({
config: { ...config, site: 'ap2' },
credentials,
})
await session.deliver({
body: Buffer.from(`${JSON.stringify({ x: 1 })}\n`),
contentType: 'application/x-ndjson',
metadata: meta(0),
signal: new AbortController().signal,
})
expect(fetchMock.mock.calls[0]?.[0]).toBe(
'https://http-intake.logs.ap2.datadoghq.com/api/v2/logs'
)
await session.close()
})
it('throws with the entry index when a single entry exceeds 1 MB', async () => {
const session = datadogDestination.openSession({ config, credentials })
// Two entries; the second exceeds the 1 MB per-entry limit.
const huge = 'x'.repeat(1024 * 1024 + 10)
const body = Buffer.from(
`${JSON.stringify({ id: 'small' })}\n${JSON.stringify({ blob: huge })}\n`,
'utf8'
)
await expect(
session.deliver({
body,
contentType: 'application/x-ndjson',
metadata: meta(0),
signal: new AbortController().signal,
})
).rejects.toThrow(/entry at index 1 is .* exceeds the 1048576-byte per-entry limit/)
expect(fetchMock).not.toHaveBeenCalled()
await session.close()
})
it('gzips payloads larger than 1KB and sets Content-Encoding: gzip', async () => {
const session = datadogDestination.openSession({ config, credentials })
// Build > 1KB raw payload; padding string is JSON-safe.
const padding = 'a'.repeat(2048)
const body = Buffer.from(`${JSON.stringify({ id: 'a', big: padding })}\n`, 'utf8')
await session.deliver({
body,
contentType: 'application/x-ndjson',
metadata: meta(0),
signal: new AbortController().signal,
})
const init = fetchMock.mock.calls[0]?.[1] as RequestInit
const headers = init.headers as Record<string, string>
expect(headers['Content-Encoding']).toBe('gzip')
expect(init.body).toBeInstanceOf(Uint8Array)
expect(typeof init.body).not.toBe('string')
const decoded = JSON.parse(gunzipSync(init.body as Uint8Array).toString('utf8'))
expect(decoded).toHaveLength(1)
expect(decoded[0].id).toBe('a')
expect(decoded[0].big).toBe(padding)
await session.close()
})
it('locator includes the dd-request-id header when present', async () => {
fetchMock.mockResolvedValueOnce(
new Response(null, { status: 202, headers: { 'dd-request-id': 'req-abc-123' } })
)
const session = datadogDestination.openSession({ config, credentials })
const result = await session.deliver({
body: Buffer.from(`${JSON.stringify({ x: 1 })}\n`),
contentType: 'application/x-ndjson',
metadata: meta(7),
signal: new AbortController().signal,
})
expect(result.locator).toBe('datadog://us1#r1-7@req-abc-123')
await session.close()
})
})
describe('datadogDestination test()', () => {
it('sends a single probe entry', async () => {
await datadogDestination.test!({
config,
credentials,
signal: new AbortController().signal,
})
const init = fetchMock.mock.calls[0]?.[1] as RequestInit
const headers = init.headers as Record<string, string>
expect(headers.Accept).toBe('application/json')
expect(headers['User-Agent']).toBe('sim-data-drain/1.0')
const payload = JSON.parse(init.body as string)
expect(payload).toHaveLength(1)
expect(payload[0].message).toContain('connection test')
})
})
@@ -0,0 +1,270 @@
import { gzipSync } from 'node:zlib'
import { createLogger } from '@sim/logger'
import { toError } from '@sim/utils/errors'
import { backoffWithJitter, parseRetryAfter } from '@sim/utils/retry'
import { z } from 'zod'
import { parseNdjsonObjects, sleepUntilAborted } from '@/lib/data-drains/destinations/utils'
import type { DeliveryMetadata, DrainDestination } from '@/lib/data-drains/types'
const logger = createLogger('DataDrainDatadogDestination')
const DATADOG_SITES = ['us1', 'us3', 'us5', 'eu1', 'ap1', 'ap2', 'gov'] as const
type DatadogSite = (typeof DATADOG_SITES)[number]
const SITE_HOSTS: Record<DatadogSite, string> = {
us1: 'datadoghq.com',
us3: 'us3.datadoghq.com',
us5: 'us5.datadoghq.com',
eu1: 'datadoghq.eu',
ap1: 'ap1.datadoghq.com',
ap2: 'ap2.datadoghq.com',
gov: 'ddog-gov.com',
}
const MAX_ATTEMPTS = 4
const PER_ATTEMPT_TIMEOUT_MS = 30_000
const MAX_UNCOMPRESSED_BYTES = 5 * 1024 * 1024
const MAX_WIRE_BYTES = 6 * 1024 * 1024
const MAX_ENTRY_BYTES = 1024 * 1024
const MAX_ENTRIES_PER_REQUEST = 1000
const GZIP_THRESHOLD_BYTES = 1024
/**
* Datadog tag format: comma-separated `key:value` pairs. Each key must start
* with a letter and contain only [A-Za-z0-9_:./-]. Validating here so the
* `ddtags` header we emit can't be mangled by user-supplied free-form input.
*/
const DATADOG_TAG_PAIR_RE = /^[A-Za-z][A-Za-z0-9_./-]*:[^,\s][^,]*$/
const datadogConfigSchema = z.object({
site: z.enum(DATADOG_SITES),
service: z.string().min(1).max(100).optional(),
tags: z
.string()
.min(1)
.max(1024)
.refine(
(v) =>
v
.split(',')
.map((t) => t.trim())
.filter((t) => t.length > 0)
.every((t) => DATADOG_TAG_PAIR_RE.test(t)),
{ message: 'tags must be comma-separated key:value pairs' }
)
.optional(),
})
const datadogCredentialsSchema = z.object({
apiKey: z.string().min(1, 'apiKey is required'),
})
export type DatadogDestinationConfig = z.infer<typeof datadogConfigSchema>
export type DatadogDestinationCredentials = z.infer<typeof datadogCredentialsSchema>
interface DatadogLogEntry {
ddsource: string
service: string
ddtags: string
message: string
[attribute: string]: unknown
}
function buildEndpoint(site: DatadogSite): string {
return `https://http-intake.logs.${SITE_HOSTS[site]}/api/v2/logs`
}
function buildEntries(
rows: unknown[],
config: DatadogDestinationConfig,
metadata: DeliveryMetadata
): DatadogLogEntry[] {
const ddtags = [
`sim_drain_id:${metadata.drainId}`,
`sim_run_id:${metadata.runId}`,
`sim_source:${metadata.source}`,
...(config.tags ? [config.tags] : []),
].join(',')
const service = config.service ?? 'sim'
return rows.map((row) => {
const attrs = typeof row === 'object' && row !== null ? (row as Record<string, unknown>) : {}
let message: string
if (typeof row === 'string') {
message = row
} else if (typeof attrs.message === 'string') {
message = attrs.message
} else {
message = JSON.stringify(row)
}
/**
* Spread user attributes first, then force all four reserved fields the
* drain owns: `ddsource`, `service`, `ddtags`, and `message`. Per
* https://docs.datadoghq.com/logs/log_configuration/pipelines/#service-and-source,
* Datadog uses `service` + `ddsource` to pick the processing pipeline, so
* letting a row field clobber them would silently re-route a drain.
*/
return {
...attrs,
ddsource: 'sim',
service,
ddtags,
message,
}
})
}
function isRetryableStatus(status: number): boolean {
return status === 408 || status === 429 || status >= 500
}
interface PreparedBody {
body: Uint8Array | string
headers: Record<string, string>
wireBytes: number
rawBytes: number
}
interface PostInput {
url: string
prepared: PreparedBody
signal: AbortSignal
}
function buildRequestBody(payload: string, apiKey: string): PreparedBody {
const headers: Record<string, string> = {
'Content-Type': 'application/json',
'DD-API-KEY': apiKey,
Accept: 'application/json',
'User-Agent': 'sim-data-drain/1.0',
}
const rawBytes = Buffer.byteLength(payload, 'utf8')
if (rawBytes > GZIP_THRESHOLD_BYTES) {
const compressed = gzipSync(payload)
headers['Content-Encoding'] = 'gzip'
const view = new Uint8Array(compressed.buffer, compressed.byteOffset, compressed.byteLength)
return { body: view, headers, wireBytes: view.byteLength, rawBytes }
}
return { body: payload, headers, wireBytes: rawBytes, rawBytes }
}
async function postWithRetries(input: PostInput): Promise<Response> {
const { body, headers } = input.prepared
let lastError: unknown
for (let attempt = 1; attempt <= MAX_ATTEMPTS; attempt++) {
if (input.signal.aborted) throw input.signal.reason ?? new Error('Aborted')
const perAttempt = AbortSignal.any([input.signal, AbortSignal.timeout(PER_ATTEMPT_TIMEOUT_MS)])
let retryAfterMs: number | null = null
let response: Response | undefined
try {
response = await fetch(input.url, {
method: 'POST',
// double-cast-allowed: Uint8Array is a valid runtime BodyInit but the DOM lib types only enumerate Blob/FormData/string/etc.
body: body as unknown as BodyInit,
headers,
signal: perAttempt,
})
} catch (error) {
lastError = error
logger.debug('Datadog request failed', { attempt, error: toError(error).message })
}
if (response) {
if (response.ok) {
/** Drain the success body so undici can return the socket to the keep-alive pool. Headers remain readable after consumption. */
await response.text().catch(() => '')
return response
}
if (!isRetryableStatus(response.status)) {
const text = await response.text().catch(() => '')
throw new Error(`Datadog responded with HTTP ${response.status}: ${text}`)
}
lastError = new Error(`Datadog responded with HTTP ${response.status}`)
retryAfterMs = parseRetryAfter(response.headers.get('retry-after'))
/** Drain the retryable response body so undici can return the socket to the keep-alive pool. */
await response.text().catch(() => '')
}
if (attempt < MAX_ATTEMPTS) {
await sleepUntilAborted(backoffWithJitter(attempt, retryAfterMs), input.signal)
}
}
throw lastError instanceof Error ? lastError : new Error('Datadog delivery failed after retries')
}
export const datadogDestination: DrainDestination<
DatadogDestinationConfig,
DatadogDestinationCredentials
> = {
type: 'datadog',
displayName: 'Datadog',
configSchema: datadogConfigSchema,
credentialsSchema: datadogCredentialsSchema,
async test({ config, credentials, signal }) {
const probe = [
{
ddsource: 'sim',
service: config.service ?? 'sim',
ddtags: `sim_probe:1${config.tags ? `,${config.tags}` : ''}`,
message: 'sim-data-drain connection test',
},
]
await postWithRetries({
url: buildEndpoint(config.site),
prepared: buildRequestBody(JSON.stringify(probe), credentials.apiKey),
signal,
})
},
openSession({ config, credentials }) {
const url = buildEndpoint(config.site)
return {
async deliver({ body, metadata, signal }) {
const rows = parseNdjsonObjects(body)
const entries = buildEntries(rows, config, metadata)
if (entries.length > MAX_ENTRIES_PER_REQUEST) {
throw new Error(
`Datadog chunk has ${entries.length} entries, exceeds the ${MAX_ENTRIES_PER_REQUEST} per-request limit`
)
}
for (let i = 0; i < entries.length; i++) {
const entryBytes = Buffer.byteLength(JSON.stringify(entries[i]), 'utf8')
if (entryBytes > MAX_ENTRY_BYTES) {
throw new Error(
`Datadog entry at index ${i} is ${entryBytes} bytes, exceeds the ${MAX_ENTRY_BYTES}-byte per-entry limit`
)
}
}
const payload = JSON.stringify(entries)
const prepared = buildRequestBody(payload, credentials.apiKey)
if (prepared.rawBytes > MAX_UNCOMPRESSED_BYTES) {
throw new Error(
`Datadog payload is ${prepared.rawBytes} bytes uncompressed, exceeds the ${MAX_UNCOMPRESSED_BYTES}-byte per-request limit`
)
}
if (prepared.wireBytes > MAX_WIRE_BYTES) {
throw new Error(
`Datadog payload is ${prepared.wireBytes} bytes on the wire, exceeds the ${MAX_WIRE_BYTES}-byte defensive wire-size cap`
)
}
const response = await postWithRetries({
url,
prepared,
signal,
})
const requestId = response.headers.get('dd-request-id') ?? null
logger.debug('Datadog chunk delivered', {
site: config.site,
rows: entries.length,
rawBytes: prepared.rawBytes,
wireBytes: prepared.wireBytes,
})
return {
locator: requestId
? `datadog://${config.site}#${metadata.runId}-${metadata.sequence}@${requestId}`
: `datadog://${config.site}#${metadata.runId}-${metadata.sequence}`,
}
},
async close() {},
}
},
}
@@ -0,0 +1,192 @@
/**
* @vitest-environment node
*/
import { beforeEach, describe, expect, it, vi } from 'vitest'
const { mockGetAccessToken, JWTCtor } = vi.hoisted(() => {
const mockGetAccessToken = vi.fn(async () => ({ token: 'fake-access-token' }))
return {
mockGetAccessToken,
JWTCtor: vi.fn().mockImplementation(
class {
getAccessToken = mockGetAccessToken
}
),
}
})
vi.mock('google-auth-library', () => ({ JWT: JWTCtor }))
const fetchMock = vi.fn(async () => new Response(null, { status: 200 }))
vi.stubGlobal('fetch', fetchMock)
import { gcsDestination } from '@/lib/data-drains/destinations/gcs'
const config = { bucket: 'my-bucket', prefix: 'sim/' }
const credentials = {
serviceAccountJson: JSON.stringify({
client_email: 'sa@project.iam.gserviceaccount.com',
private_key: '-----BEGIN PRIVATE KEY-----\nfake\n-----END PRIVATE KEY-----\n',
}),
}
beforeEach(() => {
vi.clearAllMocks()
fetchMock.mockResolvedValue(new Response(null, { status: 200 }))
})
describe('gcsDestination openSession', () => {
it('uploads via the JSON API and returns a gs:// locator', async () => {
const session = gcsDestination.openSession({ config, credentials })
const body = Buffer.from('row\n', 'utf8')
const result = await session.deliver({
body,
contentType: 'application/x-ndjson',
metadata: {
drainId: 'd1',
runId: 'r1',
source: 'workflow_logs',
sequence: 0,
rowCount: 1,
runStartedAt: new Date('2025-06-15T12:00:00Z'),
},
signal: new AbortController().signal,
})
expect(result.locator).toMatch(
/^gs:\/\/my-bucket\/sim\/workflow_logs\/d1\/\d{4}\/\d{2}\/\d{2}\/r1-00000\.ndjson$/
)
expect(fetchMock).toHaveBeenCalledTimes(1)
const [url, init] = fetchMock.mock.calls[0] as [string, RequestInit]
expect(url).toContain('/upload/storage/v1/b/my-bucket/o')
expect(url).toContain('uploadType=media')
const headers = init.headers as Record<string, string>
expect(headers.Authorization).toBe('Bearer fake-access-token')
expect(headers['Content-Type']).toBe('application/x-ndjson')
expect(headers['x-goog-meta-sim-drain-id']).toBe('d1')
expect(headers['x-goog-meta-sim-sequence']).toBe('0')
await session.close()
})
it('surfaces non-2xx responses as errors', async () => {
fetchMock.mockResolvedValueOnce(
new Response('Permission denied', { status: 403, statusText: 'Forbidden' })
)
const session = gcsDestination.openSession({ config, credentials })
await expect(
session.deliver({
body: Buffer.from('x'),
contentType: 'application/x-ndjson',
metadata: {
drainId: 'd',
runId: 'r',
source: 'audit_logs',
sequence: 0,
rowCount: 1,
runStartedAt: new Date('2025-06-15T12:00:00Z'),
},
signal: new AbortController().signal,
})
).rejects.toThrow(/HTTP 403/)
await session.close()
})
})
describe('gcsDestination test()', () => {
it('writes a probe object then attempts cleanup', async () => {
await gcsDestination.test!({
config,
credentials,
signal: new AbortController().signal,
})
expect(fetchMock).toHaveBeenCalledTimes(2)
const [, deleteCall] = fetchMock.mock.calls
expect((deleteCall[1] as RequestInit).method).toBe('DELETE')
})
})
describe('gcsDestination credentials schema', () => {
it('rejects invalid JSON', () => {
const result = gcsDestination.credentialsSchema.safeParse({ serviceAccountJson: 'not-json' })
expect(result.success).toBe(false)
})
it('rejects JSON missing client_email', () => {
const result = gcsDestination.credentialsSchema.safeParse({
serviceAccountJson: JSON.stringify({ private_key: 'k' }),
})
expect(result.success).toBe(false)
})
})
describe('gcsDestination config schema', () => {
it('accepts a 3-character bucket name', () => {
const result = gcsDestination.configSchema.safeParse({ bucket: 'abc' })
expect(result.success).toBe(true)
})
it('rejects bucket names beginning with goog or containing google', () => {
expect(gcsDestination.configSchema.safeParse({ bucket: 'goog-prefixed' }).success).toBe(false)
expect(gcsDestination.configSchema.safeParse({ bucket: 'my-google-bucket' }).success).toBe(
false
)
expect(gcsDestination.configSchema.safeParse({ bucket: 'g00gle-bucket' }).success).toBe(false)
})
})
describe('gcsDestination upload headers', () => {
it('does not set a Content-Length header on uploads', async () => {
const session = gcsDestination.openSession({ config, credentials })
await session.deliver({
body: Buffer.from('row\n', 'utf8'),
contentType: 'application/x-ndjson',
metadata: {
drainId: 'd1',
runId: 'r1',
source: 'workflow_logs',
sequence: 0,
rowCount: 1,
runStartedAt: new Date('2025-06-15T12:00:00Z'),
},
signal: new AbortController().signal,
})
const [, init] = fetchMock.mock.calls[0] as [string, RequestInit]
const headers = init.headers as Record<string, string>
const headerKeys = Object.keys(headers).map((k) => k.toLowerCase())
expect(headerKeys).not.toContain('content-length')
expect(headers['User-Agent']).toBe('sim-data-drain/1.0')
await session.close()
})
})
describe('gcsDestination deleteObject behavior', () => {
it('treats 404 as success on delete during test() cleanup', async () => {
fetchMock.mockResolvedValueOnce(new Response(null, { status: 200 })) // upload probe
fetchMock.mockResolvedValueOnce(new Response(null, { status: 404 })) // delete probe
await expect(
gcsDestination.test!({
config,
credentials,
signal: new AbortController().signal,
})
).resolves.toBeUndefined()
expect(fetchMock).toHaveBeenCalledTimes(2)
})
it('retries DELETE on 503 then succeeds on 200', async () => {
fetchMock.mockResolvedValueOnce(new Response(null, { status: 200 })) // upload probe
fetchMock.mockResolvedValueOnce(new Response('busy', { status: 503 })) // first delete attempt
fetchMock.mockResolvedValueOnce(new Response(null, { status: 200 })) // retry succeeds
await gcsDestination.test!({
config,
credentials,
signal: new AbortController().signal,
})
expect(fetchMock).toHaveBeenCalledTimes(3)
const [, deleteCall1] = fetchMock.mock.calls[1] as [string, RequestInit]
const [, deleteCall2] = fetchMock.mock.calls[2] as [string, RequestInit]
expect(deleteCall1.method).toBe('DELETE')
expect(deleteCall2.method).toBe('DELETE')
})
})
@@ -0,0 +1,320 @@
import { createLogger } from '@sim/logger'
import { toError } from '@sim/utils/errors'
import { generateShortId } from '@sim/utils/id'
import { backoffWithJitter, parseRetryAfter } from '@sim/utils/retry'
import { JWT } from 'google-auth-library'
import { z } from 'zod'
import {
buildObjectKey,
normalizePrefix,
type ParsedServiceAccount,
parseServiceAccount,
refineServiceAccountJson,
sleepUntilAborted,
} from '@/lib/data-drains/destinations/utils'
import type { DrainDestination } from '@/lib/data-drains/types'
const logger = createLogger('DataDrainGCSDestination')
const SCOPE = 'https://www.googleapis.com/auth/devstorage.read_write'
const GCS_HOST = 'https://storage.googleapis.com'
const USER_AGENT = 'sim-data-drain/1.0'
const MAX_ATTEMPTS = 4
const PER_ATTEMPT_TIMEOUT_MS = 60_000
/** GCS caps total custom metadata at 8 KiB per object (sum of key + value bytes). */
const MAX_CUSTOM_METADATA_BYTES = 8 * 1024
/** GCS object names are at most 1024 bytes when UTF-8 encoded (flat-namespace buckets). */
const MAX_OBJECT_NAME_BYTES = 1024
const GCS_BUCKET_COMPONENT_RE = /^[a-z0-9]([a-z0-9_-]*[a-z0-9])?$/
const IPV4_LIKE_RE = /^(\d{1,3}\.){3}\d{1,3}$/
const GOOGLE_RESERVED_RE = /^(goog|google|g00gle)/i
const GOOGLE_CONTAINS_RE = /(google|g00gle)/i
function validateGcsBucketComponents(v: string): string | null {
if (v.length < 3 || v.length > 222) return 'bucket must be 3-222 characters'
const components = v.split('.')
for (const c of components) {
if (c.length < 1 || c.length > 63) {
return 'each dot-separated component must be 1-63 characters'
}
if (!GCS_BUCKET_COMPONENT_RE.test(c)) {
return 'each component must be lowercase, start/end alphanumeric, letters/digits/_/- only'
}
}
return null
}
const gcsConfigSchema = z.object({
bucket: z
.string()
.min(3, 'bucket must be 3-222 characters')
.max(222, 'bucket must be 3-222 characters')
.superRefine((v, ctx) => {
const err = validateGcsBucketComponents(v)
if (err) ctx.addIssue({ code: z.ZodIssueCode.custom, message: err })
})
.refine((v) => !IPV4_LIKE_RE.test(v), {
message: 'bucket must not look like an IP address',
})
.refine((v) => !v.includes('..'), { message: 'bucket must not contain consecutive dots' })
.refine((v) => !v.includes('-.') && !v.includes('.-'), {
message: 'bucket must not contain a dash adjacent to a dot',
})
.refine((v) => !GOOGLE_RESERVED_RE.test(v) && !GOOGLE_CONTAINS_RE.test(v), {
message: 'bucket name cannot begin with "goog" or contain "google" / close misspellings',
}),
prefix: z
.string()
.max(512)
.refine((v) => Buffer.byteLength(v, 'utf8') <= 512, {
message: 'prefix must be at most 512 bytes (UTF-8)',
})
.refine((v) => !v.startsWith('.well-known/acme-challenge/'), {
message: 'prefix must not start with ".well-known/acme-challenge/" (reserved by GCS)',
})
.optional(),
})
const gcsCredentialsSchema = z
.object({
serviceAccountJson: z.string().min(1, 'serviceAccountJson is required'),
})
.superRefine(refineServiceAccountJson)
export type GCSDestinationConfig = z.infer<typeof gcsConfigSchema>
export type GCSDestinationCredentials = z.infer<typeof gcsCredentialsSchema>
function buildJwt(account: ParsedServiceAccount): JWT {
return new JWT({ email: account.clientEmail, key: account.privateKey, scopes: [SCOPE] })
}
async function getAccessToken(jwt: JWT): Promise<string> {
const { token } = await jwt.getAccessToken()
if (!token) throw new Error('Failed to obtain GCS access token')
return token
}
interface UploadInput {
bucket: string
objectName: string
body: Buffer
contentType: string
metadata: Record<string, string>
signal: AbortSignal
jwt: JWT
}
function isRetryableStatus(status: number): boolean {
return (
status === 408 ||
status === 429 ||
status === 500 ||
status === 502 ||
status === 503 ||
status === 504
)
}
interface RetryRequestInput {
action: string
bucket: string
url: string
method: string
/**
* Built per attempt so the OAuth access token is refreshed if it expired
* between retries (google-auth-library caches and refreshes on demand).
*/
buildHeaders: () => Promise<Record<string, string>>
body?: BodyInit | Buffer
signal: AbortSignal
/** HTTP statuses to treat as success in addition to 2xx. */
successStatuses?: number[]
}
async function fetchWithRetry(input: RetryRequestInput): Promise<void> {
let lastError: unknown
for (let attempt = 1; attempt <= MAX_ATTEMPTS; attempt++) {
if (input.signal.aborted) throw input.signal.reason ?? new Error('Aborted')
const perAttempt = AbortSignal.any([input.signal, AbortSignal.timeout(PER_ATTEMPT_TIMEOUT_MS)])
let response: Response
try {
const headers = await input.buildHeaders()
response = await fetch(input.url, {
method: input.method,
body: input.body as BodyInit | undefined,
headers,
signal: perAttempt,
})
} catch (error) {
lastError = error
logger.debug('GCS request failed', {
action: input.action,
attempt,
bucket: input.bucket,
error: toError(error).message,
})
if (attempt < MAX_ATTEMPTS) {
await sleepUntilAborted(backoffWithJitter(attempt, null), input.signal)
continue
}
throw error
}
if (response.ok || input.successStatuses?.includes(response.status)) {
/** Drain the success body so undici can return the socket to the keep-alive pool. */
await response.text().catch(() => '')
return
}
if (!isRetryableStatus(response.status) || attempt === MAX_ATTEMPTS) {
const text = await response.text().catch(() => '')
logger.warn('GCS operation failed', {
action: input.action,
bucket: input.bucket,
status: response.status,
})
throw new Error(
`GCS ${input.action} failed (HTTP ${response.status}): ${text || response.statusText}`
)
}
lastError = new Error(`GCS ${input.action} responded with HTTP ${response.status}`)
const retryAfterMs = parseRetryAfter(response.headers.get('retry-after'))
/** Drain the retryable response body so undici can return the socket to the keep-alive pool. */
await response.text().catch(() => '')
await sleepUntilAborted(backoffWithJitter(attempt, retryAfterMs), input.signal)
}
throw lastError instanceof Error
? lastError
: new Error(`GCS ${input.action} failed after retries`)
}
/** GCS uses HTTP headers (x-goog-meta-*) to carry custom metadata; the spec forbids non-ASCII. */
const ASCII_ONLY_RE = /^[\x20-\x7e]*$/
async function uploadObject(action: string, input: UploadInput): Promise<void> {
const objectNameBytes = Buffer.byteLength(input.objectName, 'utf8')
if (objectNameBytes < 1 || objectNameBytes > MAX_OBJECT_NAME_BYTES) {
throw new Error(
`GCS object name is ${objectNameBytes} bytes, must be 1-${MAX_OBJECT_NAME_BYTES} bytes (UTF-8)`
)
}
let metadataBytes = 0
for (const [key, value] of Object.entries(input.metadata)) {
if (!ASCII_ONLY_RE.test(key) || !ASCII_ONLY_RE.test(value)) {
throw new Error(`GCS custom metadata key/value must be ASCII printable: ${key}`)
}
metadataBytes += Buffer.byteLength(key, 'utf8') + Buffer.byteLength(value, 'utf8')
}
if (metadataBytes > MAX_CUSTOM_METADATA_BYTES) {
throw new Error(
`GCS custom metadata is ${metadataBytes} bytes, exceeds the ${MAX_CUSTOM_METADATA_BYTES}-byte per-object limit`
)
}
const url = `${GCS_HOST}/upload/storage/v1/b/${encodeURIComponent(input.bucket)}/o?uploadType=media&name=${encodeURIComponent(input.objectName)}`
await fetchWithRetry({
action,
bucket: input.bucket,
url,
method: 'POST',
buildHeaders: async () => {
const token = await getAccessToken(input.jwt)
const headers: Record<string, string> = {
Authorization: `Bearer ${token}`,
'Content-Type': input.contentType,
'User-Agent': USER_AGENT,
}
for (const [key, value] of Object.entries(input.metadata)) {
headers[`x-goog-meta-${key}`] = value
}
return headers
},
body: input.body,
signal: input.signal,
})
}
async function deleteObject(input: {
bucket: string
objectName: string
jwt: JWT
signal: AbortSignal
}): Promise<void> {
const url = `${GCS_HOST}/storage/v1/b/${encodeURIComponent(input.bucket)}/o/${encodeURIComponent(input.objectName)}`
await fetchWithRetry({
action: 'delete-object',
bucket: input.bucket,
url,
method: 'DELETE',
buildHeaders: async () => {
const token = await getAccessToken(input.jwt)
return {
Authorization: `Bearer ${token}`,
'User-Agent': USER_AGENT,
}
},
signal: input.signal,
successStatuses: [404],
})
}
export const gcsDestination: DrainDestination<GCSDestinationConfig, GCSDestinationCredentials> = {
type: 'gcs',
displayName: 'Google Cloud Storage',
configSchema: gcsConfigSchema,
credentialsSchema: gcsCredentialsSchema,
async test({ config, credentials, signal }) {
const account = parseServiceAccount(credentials.serviceAccountJson)
const jwt = buildJwt(account)
const probeName = `${normalizePrefix(config.prefix)}.sim-drain-write-probe/${generateShortId(12)}`
await uploadObject('test-put', {
bucket: config.bucket,
objectName: probeName,
body: Buffer.alloc(0),
contentType: 'application/octet-stream',
metadata: {},
signal,
jwt,
})
try {
await deleteObject({ bucket: config.bucket, objectName: probeName, jwt, signal })
} catch (cleanupError) {
logger.debug('GCS test write probe cleanup failed (non-fatal)', {
bucket: config.bucket,
objectName: probeName,
error: cleanupError,
})
}
},
openSession({ config, credentials }) {
const account = parseServiceAccount(credentials.serviceAccountJson)
const jwt = buildJwt(account)
return {
async deliver({ body, contentType, metadata, signal }) {
const objectName = buildObjectKey(config.prefix, metadata)
await uploadObject('put-object', {
bucket: config.bucket,
objectName,
body,
contentType,
metadata: {
'sim-drain-id': metadata.drainId,
'sim-run-id': metadata.runId,
'sim-source': metadata.source,
'sim-sequence': metadata.sequence.toString(),
'sim-row-count': metadata.rowCount.toString(),
},
signal,
jwt,
})
logger.debug('GCS chunk delivered', {
bucket: config.bucket,
objectName,
bytes: body.byteLength,
})
return { locator: `gs://${config.bucket}/${objectName}` }
},
async close() {},
}
},
}
@@ -0,0 +1,22 @@
import { azureBlobDestination } from '@/lib/data-drains/destinations/azure_blob'
import { bigqueryDestination } from '@/lib/data-drains/destinations/bigquery'
import { datadogDestination } from '@/lib/data-drains/destinations/datadog'
import { gcsDestination } from '@/lib/data-drains/destinations/gcs'
import { s3Destination } from '@/lib/data-drains/destinations/s3'
import { snowflakeDestination } from '@/lib/data-drains/destinations/snowflake'
import { webhookDestination } from '@/lib/data-drains/destinations/webhook'
import type { DestinationType, DrainDestination } from '@/lib/data-drains/types'
export const DESTINATION_REGISTRY = {
s3: s3Destination,
gcs: gcsDestination,
azure_blob: azureBlobDestination,
datadog: datadogDestination,
bigquery: bigqueryDestination,
snowflake: snowflakeDestination,
webhook: webhookDestination,
} as const satisfies Record<DestinationType, DrainDestination>
export function getDestination(type: DestinationType): DrainDestination {
return DESTINATION_REGISTRY[type]
}
@@ -0,0 +1,178 @@
/**
* @vitest-environment node
*/
import { beforeEach, describe, expect, it, vi } from 'vitest'
const { mockSend, mockDestroy, S3ClientCtor, PutObjectCommandCtor, DeleteObjectCommandCtor } =
vi.hoisted(() => {
const mockSend = vi.fn(async () => ({}))
const mockDestroy = vi.fn()
return {
mockSend,
mockDestroy,
S3ClientCtor: vi.fn().mockImplementation(
class {
send = mockSend
destroy = mockDestroy
}
),
PutObjectCommandCtor: vi.fn().mockImplementation(
class {
__cmd = 'put'
args: unknown
constructor(args: unknown) {
this.args = args
}
}
),
DeleteObjectCommandCtor: vi.fn().mockImplementation(
class {
__cmd = 'delete'
args: unknown
constructor(args: unknown) {
this.args = args
}
}
),
}
})
vi.mock('@aws-sdk/client-s3', () => ({
S3Client: S3ClientCtor,
PutObjectCommand: PutObjectCommandCtor,
DeleteObjectCommand: DeleteObjectCommandCtor,
}))
import { s3Destination } from '@/lib/data-drains/destinations/s3'
const config = {
bucket: 'my-bucket',
region: 'us-east-1',
prefix: 'sim/',
}
const credentials = { accessKeyId: 'AKID', secretAccessKey: 'SECRET' }
beforeEach(() => {
vi.clearAllMocks()
})
describe('s3Destination openSession', () => {
it('reuses one S3Client across multiple deliveries and destroys on close', async () => {
const session = s3Destination.openSession({ config, credentials })
expect(S3ClientCtor).toHaveBeenCalledTimes(1)
const body = Buffer.from('row\n', 'utf8')
const meta = (sequence: number) => ({
drainId: 'd1',
runId: 'r1',
source: 'workflow_logs' as const,
sequence,
rowCount: 1,
runStartedAt: new Date('2025-06-15T12:00:00Z'),
})
const signal = new AbortController().signal
const res1 = await session.deliver({
body,
contentType: 'application/x-ndjson',
metadata: meta(0),
signal,
})
const res2 = await session.deliver({
body,
contentType: 'application/x-ndjson',
metadata: meta(1),
signal,
})
expect(S3ClientCtor).toHaveBeenCalledTimes(1)
expect(mockSend).toHaveBeenCalledTimes(2)
expect(res1.locator).toMatch(
/^s3:\/\/my-bucket\/sim\/workflow_logs\/d1\/\d{4}\/\d{2}\/\d{2}\/r1-00000\.ndjson$/
)
expect(res2.locator).toMatch(/r1-00001\.ndjson$/)
const putArgs = (PutObjectCommandCtor.mock.calls[0]?.[0] ?? {}) as Record<string, unknown>
expect(putArgs.Bucket).toBe('my-bucket')
expect(putArgs.Body).toBe(body)
expect(putArgs.ContentType).toBe('application/x-ndjson')
expect((putArgs.Metadata as Record<string, string>)['sim-drain-id']).toBe('d1')
expect((putArgs.Metadata as Record<string, string>)['sim-sequence']).toBe('0')
await session.close()
expect(mockDestroy).toHaveBeenCalledTimes(1)
})
it('omits the prefix segment when prefix is empty', async () => {
const session = s3Destination.openSession({
config: { bucket: 'b', region: 'us-east-1' },
credentials,
})
const result = await session.deliver({
body: Buffer.from('x'),
contentType: 'application/x-ndjson',
metadata: {
drainId: 'd',
runId: 'r',
source: 'audit_logs',
sequence: 0,
rowCount: 1,
runStartedAt: new Date('2025-06-15T12:00:00Z'),
},
signal: new AbortController().signal,
})
expect(result.locator).toMatch(
/^s3:\/\/b\/audit_logs\/d\/\d{4}\/\d{2}\/\d{2}\/r-00000\.ndjson$/
)
await session.close()
})
it('surfaces AWS error code in delivery errors', async () => {
mockSend.mockRejectedValueOnce(
Object.assign(new Error('Access Denied'), {
name: 'AccessDenied',
$metadata: { httpStatusCode: 403, requestId: 'req-1' },
})
)
const session = s3Destination.openSession({ config, credentials })
await expect(
session.deliver({
body: Buffer.from('x'),
contentType: 'application/x-ndjson',
metadata: {
drainId: 'd',
runId: 'r',
source: 'audit_logs',
sequence: 0,
rowCount: 1,
runStartedAt: new Date('2025-06-15T12:00:00Z'),
},
signal: new AbortController().signal,
})
).rejects.toThrow(/AccessDenied 403/)
await session.close()
})
})
describe('s3Destination test()', () => {
it('writes a probe object then attempts cleanup', async () => {
await s3Destination.test!({
config,
credentials,
signal: new AbortController().signal,
})
expect(PutObjectCommandCtor).toHaveBeenCalled()
expect(DeleteObjectCommandCtor).toHaveBeenCalled()
expect(mockDestroy).toHaveBeenCalled()
})
it('still returns success when cleanup delete fails', async () => {
mockSend
.mockResolvedValueOnce({}) // put probe
.mockRejectedValueOnce(new Error('no delete perms')) // cleanup
await expect(
s3Destination.test!({ config, credentials, signal: new AbortController().signal })
).resolves.toBeUndefined()
})
})
+259
View File
@@ -0,0 +1,259 @@
import {
DeleteObjectCommand,
PutObjectCommand,
S3Client,
type S3ServiceException,
} from '@aws-sdk/client-s3'
import { createLogger } from '@sim/logger'
import { generateShortId } from '@sim/utils/id'
import { z } from 'zod'
import { validateExternalUrl } from '@/lib/core/security/input-validation'
import { validateUrlWithDNS } from '@/lib/core/security/input-validation.server'
import { buildObjectKey, normalizePrefix } from '@/lib/data-drains/destinations/utils'
import type { DrainDestination } from '@/lib/data-drains/types'
const logger = createLogger('DataDrainS3Destination')
/** https://docs.aws.amazon.com/AmazonS3/latest/userguide/bucketnamingrules.html */
const S3_BUCKET_NAME_RE = /^[a-z0-9][a-z0-9.-]{1,61}[a-z0-9]$/
const S3_IPV4_LIKE_RE = /^(\d{1,3}\.){3}\d{1,3}$/
/** Matches standard and 4-segment ISO partition codes (e.g. `us-iso-east-1`). */
const AWS_REGION_RE = /^[a-z]{2,}(-[a-z]+)+-\d+$/
/** Cap is over key + value bytes only (no `x-amz-meta-` prefix). */
const MAX_S3_METADATA_BYTES = 2 * 1024
const MAX_S3_KEY_BYTES = 1024
const s3BucketSchema = z
.string()
.min(3, 'bucket must be 3-63 characters')
.max(63, 'bucket must be 3-63 characters')
.refine((v) => S3_BUCKET_NAME_RE.test(v), {
message:
'bucket must be lowercase, 3-63 chars, start/end alphanumeric, only letters/digits/./-',
})
.refine((v) => !v.includes('..'), { message: 'bucket must not contain consecutive dots' })
.refine((v) => !v.includes('-.') && !v.includes('.-'), {
message: 'bucket must not contain a dash adjacent to a dot',
})
.refine((v) => !S3_IPV4_LIKE_RE.test(v), { message: 'bucket must not look like an IP address' })
.refine((v) => !v.startsWith('xn--'), { message: 'bucket must not start with "xn--"' })
.refine((v) => !v.startsWith('sthree-'), { message: 'bucket must not start with "sthree-"' })
.refine((v) => !v.startsWith('amzn-s3-demo-'), {
message: 'bucket must not start with "amzn-s3-demo-" (reserved by AWS)',
})
.refine((v) => !v.endsWith('-s3alias') && !v.endsWith('--ol-s3') && !v.endsWith('.mrap'), {
message: 'bucket must not end with reserved suffix (-s3alias, --ol-s3, .mrap)',
})
.refine((v) => !v.endsWith('--x-s3'), {
message:
'bucket must not end with "--x-s3" (reserved for S3 Express One Zone directory buckets)',
})
.refine((v) => !v.endsWith('--table-s3'), {
message: 'bucket must not end with "--table-s3" (reserved for S3 Tables)',
})
const s3RegionSchema = z
.string()
.min(1, 'region is required')
.max(32, 'region is too long')
.refine((v) => AWS_REGION_RE.test(v), {
message: 'region must look like an AWS region code, e.g. us-east-1',
})
const s3ConfigSchema = z.object({
bucket: s3BucketSchema,
region: s3RegionSchema,
/**
* Optional prefix; trailing slash is added automatically when assembling keys.
* Bounded by UTF-8 byte length (not code units) so non-ASCII prefixes can't
* push assembled keys past S3's 1024-byte object key limit.
*/
prefix: z
.string()
.max(512)
.refine((v) => Buffer.byteLength(v, 'utf8') <= 512, {
message: 'prefix must be at most 512 bytes (UTF-8)',
})
.optional(),
/**
* Optional override for non-AWS S3-compatible providers (MinIO, R2, GCS interop, etc.).
* SSRF-validated: HTTPS-only, must not resolve syntactically to a private,
* loopback, or cloud-metadata address. The AWS SDK will issue requests to
* this host, so we reject internal targets at the schema boundary.
*/
endpoint: z
.string()
.url()
.refine((v) => v.startsWith('https://'), { message: 'endpoint must use https://' })
.refine((value) => validateExternalUrl(value, 'endpoint').isValid, {
message: 'endpoint must be HTTPS and not point at a private, loopback, or metadata address',
})
.optional(),
/**
* Force path-style addressing. Set `true` for MinIO / Ceph RGW; defaults
* to `false` for AWS S3 and Cloudflare R2.
*/
forcePathStyle: z.boolean().optional(),
})
const s3CredentialsSchema = z.object({
accessKeyId: z.string().min(1, 'accessKeyId is required'),
secretAccessKey: z.string().min(1, 'secretAccessKey is required'),
})
export type S3DestinationConfig = z.infer<typeof s3ConfigSchema>
export type S3DestinationCredentials = z.infer<typeof s3CredentialsSchema>
function buildClient(config: S3DestinationConfig, credentials: S3DestinationCredentials): S3Client {
return new S3Client({
region: config.region,
credentials: {
accessKeyId: credentials.accessKeyId,
secretAccessKey: credentials.secretAccessKey,
},
endpoint: config.endpoint,
forcePathStyle: config.forcePathStyle ?? false,
})
}
function isS3ServiceException(error: unknown): error is S3ServiceException {
return (
typeof error === 'object' &&
error !== null &&
'$metadata' in error &&
typeof (error as { name?: unknown }).name === 'string'
)
}
/** DNS-aware SSRF check: catches hostnames that resolve to internal IPs (the schema check only catches IP literals). */
async function assertEndpointIsPublic(endpoint: string | undefined): Promise<void> {
if (!endpoint) return
const result = await validateUrlWithDNS(endpoint, 'endpoint')
if (!result.isValid) {
throw new Error(result.error ?? 'S3 endpoint failed SSRF validation')
}
}
/**
* Surfaces actionable S3 SDK error codes (`AccessDenied`, `NoSuchBucket`,
* `InvalidAccessKeyId`, `SignatureDoesNotMatch`, ...) and preserves the
* original error as `cause` so callers can still branch on `code`/`$metadata`.
*/
async function withS3ErrorContext<T>(action: string, fn: () => Promise<T>): Promise<T> {
try {
return await fn()
} catch (error) {
if (isS3ServiceException(error)) {
const code = error.name
const status = error.$metadata?.httpStatusCode
const requestId = error.$metadata?.requestId
logger.warn('S3 operation failed', { action, code, status, requestId })
/** Preserve SDK error as `cause` so callers can still branch on `code` / `$metadata`. */
throw new Error(
`S3 ${action} failed (${code}${status ? ` ${status}` : ''}): ${error.message}`,
{ cause: error }
)
}
throw error
}
}
export const s3Destination: DrainDestination<S3DestinationConfig, S3DestinationCredentials> = {
type: 's3',
displayName: 'Amazon S3',
configSchema: s3ConfigSchema,
credentialsSchema: s3CredentialsSchema,
async test({ config, credentials, signal }) {
await assertEndpointIsPublic(config.endpoint)
const client = buildClient(config, credentials)
/** Real write probe so write-only IAM policies surface here, not at first run. */
const probeKey = `${normalizePrefix(config.prefix)}.sim-drain-write-probe/${generateShortId(12)}`
try {
await withS3ErrorContext('test-put', () =>
client.send(
new PutObjectCommand({
Bucket: config.bucket,
Key: probeKey,
Body: Buffer.alloc(0),
ContentType: 'application/octet-stream',
ServerSideEncryption: 'AES256',
}),
{ abortSignal: signal }
)
)
/** Best-effort cleanup: write was already proven, so a missing s3:DeleteObject must not fail the test. */
try {
await client.send(new DeleteObjectCommand({ Bucket: config.bucket, Key: probeKey }), {
abortSignal: signal,
})
} catch (cleanupError) {
logger.debug('S3 test write probe cleanup failed (non-fatal)', {
bucket: config.bucket,
key: probeKey,
error: cleanupError,
})
}
} finally {
client.destroy()
}
},
openSession({ config, credentials }) {
const client = buildClient(config, credentials)
/**
* Lazy + cached DNS-aware endpoint check. SDK manages its own connections
* so we can't pin the IP, but failing the first deliver still rejects
* hostnames that resolve to internal targets. Lazy init avoids an
* unhandled rejection when the source yields no chunks.
*/
let endpointCheck: Promise<void> | null = null
return {
async deliver({ body, contentType, metadata, signal }) {
if (endpointCheck === null) endpointCheck = assertEndpointIsPublic(config.endpoint)
await endpointCheck
const key = buildObjectKey(config.prefix, metadata)
const keyBytes = Buffer.byteLength(key, 'utf8')
if (keyBytes > MAX_S3_KEY_BYTES) {
throw new Error(
`S3 object key is ${keyBytes} bytes, exceeds the ${MAX_S3_KEY_BYTES}-byte limit`
)
}
const userMetadata: Record<string, string> = {
'sim-drain-id': metadata.drainId,
'sim-run-id': metadata.runId,
'sim-source': metadata.source,
'sim-sequence': metadata.sequence.toString(),
'sim-row-count': metadata.rowCount.toString(),
}
let metadataBytes = 0
for (const [k, v] of Object.entries(userMetadata)) {
metadataBytes += Buffer.byteLength(k, 'utf8') + Buffer.byteLength(v, 'utf8')
}
if (metadataBytes > MAX_S3_METADATA_BYTES) {
throw new Error(
`S3 user metadata is ${metadataBytes} bytes, exceeds the ${MAX_S3_METADATA_BYTES}-byte per-object limit`
)
}
await withS3ErrorContext('put-object', () =>
client.send(
new PutObjectCommand({
Bucket: config.bucket,
Key: key,
Body: body,
ContentType: contentType,
ServerSideEncryption: 'AES256',
Metadata: userMetadata,
}),
{ abortSignal: signal }
)
)
logger.debug('S3 chunk delivered', { bucket: config.bucket, key, bytes: body.byteLength })
return { locator: `s3://${config.bucket}/${key}` }
},
async close() {
client.destroy()
},
}
},
}
@@ -0,0 +1,210 @@
/**
* @vitest-environment node
*/
import { generateKeyPairSync } from 'node:crypto'
import { decodeJwt } from 'jose'
import { beforeEach, describe, expect, it, vi } from 'vitest'
const fetchMock = vi.fn(async () => new Response(JSON.stringify({}), { status: 200 }))
vi.stubGlobal('fetch', fetchMock)
import { snowflakeDestination } from '@/lib/data-drains/destinations/snowflake'
const { privateKey } = generateKeyPairSync('rsa', {
modulusLength: 2048,
privateKeyEncoding: { type: 'pkcs8', format: 'pem' },
publicKeyEncoding: { type: 'spki', format: 'pem' },
})
const config = {
account: 'orgname-acct',
user: 'sim_user',
warehouse: 'WH',
database: 'DB',
schema: 'PUBLIC',
table: 'DRAINS',
}
const credentials = { privateKey }
const meta = {
drainId: 'd',
runId: 'r',
source: 'workflow_logs' as const,
sequence: 0,
rowCount: 2,
runStartedAt: new Date('2025-06-15T12:00:00Z'),
}
beforeEach(() => {
vi.clearAllMocks()
fetchMock.mockResolvedValue(new Response(JSON.stringify({}), { status: 200 }))
})
describe('snowflakeDestination', () => {
it('posts a multi-row INSERT with TEXT bindings and a Bearer JWT', async () => {
const session = snowflakeDestination.openSession({ config, credentials })
const body = Buffer.from(
`${JSON.stringify({ id: 'a' })}\n${JSON.stringify({ id: 'b' })}\n`,
'utf8'
)
await session.deliver({
body,
contentType: 'application/x-ndjson',
metadata: meta,
signal: new AbortController().signal,
})
expect(fetchMock).toHaveBeenCalledTimes(1)
const [url, init] = fetchMock.mock.calls[0] as [string, RequestInit]
expect(url).toMatch(
/^https:\/\/orgname-acct\.snowflakecomputing\.com\/api\/v2\/statements\?requestId=[0-9a-f-]+$/
)
const headers = init.headers as Record<string, string>
expect(headers.Authorization).toMatch(/^Bearer ey/)
expect(headers['X-Snowflake-Authorization-Token-Type']).toBe('KEYPAIR_JWT')
const payload = JSON.parse(init.body as string)
expect(payload.statement).toContain('INSERT INTO "DB"."PUBLIC"."DRAINS"')
expect(payload.statement).toContain('VALUES (PARSE_JSON(?)), (PARSE_JSON(?))')
expect(payload.statement.match(/PARSE_JSON\(\?\)/g)).toHaveLength(2)
expect(payload.bindings['1']).toEqual({ type: 'TEXT', value: JSON.stringify({ id: 'a' }) })
expect(payload.bindings['2']).toEqual({ type: 'TEXT', value: JSON.stringify({ id: 'b' }) })
expect(payload.warehouse).toBe('WH')
await session.close()
})
it('uses the configured column when provided', async () => {
const session = snowflakeDestination.openSession({
config: { ...config, column: 'payload' },
credentials,
})
await session.deliver({
body: Buffer.from(`${JSON.stringify({ x: 1 })}\n`),
contentType: 'application/x-ndjson',
metadata: meta,
signal: new AbortController().signal,
})
const init = fetchMock.mock.calls[0]?.[1] as RequestInit
const payload = JSON.parse(init.body as string)
expect(payload.statement).toContain('("payload")')
await session.close()
})
it('strips region/cloud suffix from the JWT iss/sub', async () => {
const session = snowflakeDestination.openSession({
config: { ...config, account: 'orgname-acct.us-east-1.aws' },
credentials,
})
await session.deliver({
body: Buffer.from(`${JSON.stringify({ x: 1 })}\n`),
contentType: 'application/x-ndjson',
metadata: meta,
signal: new AbortController().signal,
})
const init = fetchMock.mock.calls[0]?.[1] as RequestInit
const headers = init.headers as Record<string, string>
const token = headers.Authorization.replace(/^Bearer /, '')
const claims = decodeJwt(token)
expect(claims.sub).toBe('ORGNAME-ACCT.SIM_USER')
expect(claims.iss).toMatch(/^ORGNAME-ACCT\.SIM_USER\.SHA256:/)
await session.close()
})
it('polls /statements/{handle} when Snowflake returns 202', async () => {
fetchMock.mockResolvedValueOnce(
new Response(JSON.stringify({ statementHandle: 'h-1' }), { status: 202 })
)
fetchMock.mockResolvedValueOnce(new Response(JSON.stringify({}), { status: 202 }))
fetchMock.mockResolvedValueOnce(new Response(JSON.stringify({}), { status: 200 }))
const session = snowflakeDestination.openSession({ config, credentials })
await session.deliver({
body: Buffer.from(`${JSON.stringify({ x: 1 })}\n`),
contentType: 'application/x-ndjson',
metadata: meta,
signal: new AbortController().signal,
})
expect(fetchMock).toHaveBeenCalledTimes(3)
const pollUrl = fetchMock.mock.calls[1]?.[0]
expect(pollUrl).toBe('https://orgname-acct.snowflakecomputing.com/api/v2/statements/h-1')
await session.close()
})
it('parses CRLF NDJSON bodies correctly', async () => {
const session = snowflakeDestination.openSession({ config, credentials })
const body = Buffer.from(
`${JSON.stringify({ id: 'a' })}\r\n${JSON.stringify({ id: 'b' })}\r\n`,
'utf8'
)
await session.deliver({
body,
contentType: 'application/x-ndjson',
metadata: meta,
signal: new AbortController().signal,
})
const init = fetchMock.mock.calls[0]?.[1] as RequestInit
const payload = JSON.parse(init.body as string)
expect(payload.bindings['1']).toEqual({ type: 'TEXT', value: JSON.stringify({ id: 'a' }) })
expect(payload.bindings['2']).toEqual({ type: 'TEXT', value: JSON.stringify({ id: 'b' }) })
await session.close()
})
it('retries the POST on 5xx and succeeds on the next attempt', async () => {
fetchMock.mockResolvedValueOnce(new Response('boom', { status: 503 }))
fetchMock.mockResolvedValueOnce(new Response(JSON.stringify({}), { status: 200 }))
const session = snowflakeDestination.openSession({ config, credentials })
await session.deliver({
body: Buffer.from(`${JSON.stringify({ x: 1 })}\n`),
contentType: 'application/x-ndjson',
metadata: meta,
signal: new AbortController().signal,
})
expect(fetchMock).toHaveBeenCalledTimes(2)
await session.close()
})
it('honors Retry-After (delta seconds) on 429 before retrying', async () => {
fetchMock.mockResolvedValueOnce(
new Response('slow down', { status: 429, headers: { 'Retry-After': '1' } })
)
fetchMock.mockResolvedValueOnce(new Response(JSON.stringify({}), { status: 200 }))
const session = snowflakeDestination.openSession({ config, credentials })
const start = Date.now()
await session.deliver({
body: Buffer.from(`${JSON.stringify({ x: 1 })}\n`),
contentType: 'application/x-ndjson',
metadata: meta,
signal: new AbortController().signal,
})
const elapsed = Date.now() - start
expect(fetchMock).toHaveBeenCalledTimes(2)
expect(elapsed).toBeGreaterThanOrEqual(900)
await session.close()
})
it('throws a clear error when a binding exceeds the 16 MiB VARIANT limit', async () => {
const session = snowflakeDestination.openSession({ config, credentials })
const huge = `"${'a'.repeat(16 * 1024 * 1024 + 1)}"`
const body = Buffer.from(`${huge}\n`, 'utf8')
await expect(
session.deliver({
body,
contentType: 'application/x-ndjson',
metadata: meta,
signal: new AbortController().signal,
})
).rejects.toThrow(/16 MB/)
expect(fetchMock).not.toHaveBeenCalled()
await session.close()
})
it('test() runs SELECT 1', async () => {
await snowflakeDestination.test!({
config,
credentials,
signal: new AbortController().signal,
})
const init = fetchMock.mock.calls[0]?.[1] as RequestInit
const payload = JSON.parse(init.body as string)
expect(payload.statement).toBe('SELECT 1')
})
})
@@ -0,0 +1,515 @@
import { createHash, createPublicKey } from 'node:crypto'
import { createLogger } from '@sim/logger'
import { toError } from '@sim/utils/errors'
import { generateId } from '@sim/utils/id'
import { backoffWithJitter, parseRetryAfter } from '@sim/utils/retry'
import { importPKCS8, SignJWT } from 'jose'
import { z } from 'zod'
import { sleepUntilAborted } from '@/lib/data-drains/destinations/utils'
import type { DrainDestination } from '@/lib/data-drains/types'
const logger = createLogger('DataDrainSnowflakeDestination')
/**
* Snowflake account identifier formats (https://docs.snowflake.com/en/user-guide/admin-account-identifier):
* - Org-account: `<orgname>-<acctname>` — alphanumerics/underscore, hyphen-separated, no dots.
* - Legacy account locator: `<locator>` or `<locator>.<region>[.<cloud>]` — dots allowed.
*/
const ACCOUNT_ORG_RE = /^[A-Za-z0-9][A-Za-z0-9_]*(?:-[A-Za-z0-9_]+)+$/
/**
* First segment allows hyphens so org-account identifiers (`<orgname>-<acctname>`)
* carrying a legacy region/cloud suffix (e.g. `myorg-acct.us-east-1.aws`) match.
* `normalizeAccountForJwt` strips the dotted suffix for JWT `iss`/`sub`.
*/
const ACCOUNT_LOCATOR_RE = /^[A-Za-z0-9][A-Za-z0-9_-]*(?:\.[A-Za-z0-9][A-Za-z0-9_-]*){0,2}$/
function isValidAccount(v: string): boolean {
return ACCOUNT_ORG_RE.test(v) || ACCOUNT_LOCATOR_RE.test(v)
}
const IDENTIFIER_RE = /^[A-Za-z_][A-Za-z0-9_$]{0,254}$/
/** JWT lifetime; Snowflake caps server-side enforcement at 60 minutes regardless of `exp`. */
const JWT_LIFETIME_SECONDS = 55 * 60
/** Safety margin (in seconds) subtracted from the JWT exp when caching. */
const JWT_CACHE_SAFETY_MARGIN_SECONDS = 300
const PER_ATTEMPT_TIMEOUT_MS = 60_000
const POLL_INITIAL_INTERVAL_MS = 500
const POLL_MAX_INTERVAL_MS = 5_000
const POLL_DEADLINE_MS = 10 * 60_000
/**
* Cap on consecutive failed poll attempts (network errors or retryable HTTP
* statuses). Independent of the 10-minute wall-clock deadline so that
* persistent failures surface in seconds, not minutes — matches the
* MAX_ATTEMPTS shape used by `executeStatement`. Reset to 0 on a successful
* 202 (still-executing) response.
*/
const POLL_MAX_CONSECUTIVE_RETRIES = 8
/** Maximum number of attempts (including the initial attempt) for retryable POST failures. */
const EXECUTE_MAX_ATTEMPTS = 3
const EXECUTE_RETRY_BASE_DELAY_MS = 500
const EXECUTE_RETRY_MAX_DELAY_MS = 5_000
/** Conservative pre-2025_03 BCR VARIANT max size (16 MiB) so the same value works on every account. */
const VARIANT_MAX_BYTES = 16 * 1024 * 1024
/** Server-side statement execution timeout (seconds) sent in the SQL API request body. */
const SQL_API_TIMEOUT_SECONDS = 600
/**
* Snowflake JWT `iss`/`sub` require the bare account identifier without any
* region/cloud suffix. For account-locator format `xy12345.us-east-1.aws`,
* only `XY12345` is valid; for org-account format `myorg-acct.us-east-1`,
* only `MYORG-ACCT` is valid. Strip everything after the first dot.
*/
function normalizeAccountForJwt(account: string): string {
const dot = account.indexOf('.')
return (dot === -1 ? account : account.slice(0, dot)).toUpperCase()
}
const snowflakeConfigSchema = z.object({
/**
* Snowflake account identifier. Accepted formats:
* - Org-account (preferred): `<orgname>-<acctname>` (no dots), e.g. `myorg-acct`
* - Account locator: `<locator>` (no dots), e.g. `xy12345`
* - Legacy regional locator: `<locator>.<region>.<cloud>`, e.g. `xy12345.us-east-1.aws`
*
* Do not include the `.snowflakecomputing.com` suffix. Modern org-account
* identifiers must not contain dots; only legacy locator URLs use dots.
*/
account: z.string().min(3, 'account is required').refine(isValidAccount, {
message: 'account must be the Snowflake account identifier (e.g. orgname-accountname)',
}),
user: z
.string()
.min(1, 'user is required')
.refine((v) => IDENTIFIER_RE.test(v), {
message: 'user must be a valid Snowflake identifier',
}),
warehouse: z
.string()
.min(1)
.refine((v) => IDENTIFIER_RE.test(v), {
message: 'warehouse must be a valid Snowflake identifier',
}),
database: z
.string()
.min(1)
.refine((v) => IDENTIFIER_RE.test(v), {
message: 'database must be a valid Snowflake identifier',
}),
schema: z
.string()
.min(1)
.refine((v) => IDENTIFIER_RE.test(v), {
message: 'schema must be a valid Snowflake identifier',
}),
table: z
.string()
.min(1)
.refine((v) => IDENTIFIER_RE.test(v), {
message: 'table must be a valid Snowflake identifier',
}),
/** Target VARIANT column. Defaults to `DATA` (uppercase, matching Snowflake's unquoted identifier folding). */
column: z
.string()
.min(1)
.refine((v) => IDENTIFIER_RE.test(v), {
message: 'column must be a valid Snowflake identifier',
})
.optional(),
/** Optional Snowflake role to assume for the insert. */
role: z
.string()
.min(1)
.refine((v) => IDENTIFIER_RE.test(v), {
message: 'role must be a valid Snowflake identifier',
})
.optional(),
})
const snowflakeCredentialsSchema = z.object({
/** PKCS8-encoded RSA private key (PEM). The matching public key must be registered on the user. */
privateKey: z.string().min(1, 'privateKey is required'),
})
export type SnowflakeDestinationConfig = z.infer<typeof snowflakeConfigSchema>
export type SnowflakeDestinationCredentials = z.infer<typeof snowflakeCredentialsSchema>
/**
* Computes the SHA256:<base64> fingerprint of the public key derived from the
* given private key. Snowflake encodes this in the JWT issuer claim so the
* server can match the signature against the registered public key.
* Reference: https://docs.snowflake.com/en/developer-guide/sql-api/authenticating
*/
function computePublicKeyFingerprint(privateKeyPem: string): string {
const publicKey = createPublicKey({ key: privateKeyPem, format: 'pem' })
const spkiDer = publicKey.export({ type: 'spki', format: 'der' })
return `SHA256:${createHash('sha256').update(spkiDer).digest('base64')}`
}
interface JwtCacheEntry {
token: string
expiresAt: number
}
async function buildJwt(
account: string,
user: string,
privateKeyPem: string
): Promise<JwtCacheEntry> {
const fingerprint = computePublicKeyFingerprint(privateKeyPem)
const accountForJwt = normalizeAccountForJwt(account)
const userUpper = user.toUpperCase()
const issuer = `${accountForJwt}.${userUpper}.${fingerprint}`
const subject = `${accountForJwt}.${userUpper}`
const now = Math.floor(Date.now() / 1000)
const exp = now + JWT_LIFETIME_SECONDS
let privateKey: Awaited<ReturnType<typeof importPKCS8>>
try {
privateKey = await importPKCS8(privateKeyPem, 'RS256')
} catch (error) {
throw new Error(
`privateKey must be an unencrypted PKCS#8 PEM (-----BEGIN PRIVATE KEY-----). ` +
`Convert PKCS#1 with: openssl pkcs8 -topk8 -nocrypt -in rsa.pem -out pkcs8.pem. ` +
`Decrypt encrypted PEMs first. Underlying error: ${toError(error).message}`
)
}
const token = await new SignJWT({})
.setProtectedHeader({ alg: 'RS256', typ: 'JWT' })
.setIssuer(issuer)
.setSubject(subject)
.setIssuedAt(now)
.setExpirationTime(exp)
.sign(privateKey)
return { token, expiresAt: exp - JWT_CACHE_SAFETY_MARGIN_SECONDS }
}
/**
* Quotes a Snowflake identifier so that whatever case the user typed is
* preserved exactly. Without quoting, Snowflake folds unquoted identifiers
* to uppercase, which silently breaks any table whose canonical name was
* created with quoted mixed-case. Embedded `"` is escaped as `""`.
*/
function quoteIdentifier(name: string): string {
return `"${name.replace(/"/g, '""')}"`
}
function buildStatement(config: SnowflakeDestinationConfig, rowCount: number): string {
const column = quoteIdentifier(config.column ?? 'DATA')
const target = `${quoteIdentifier(config.database)}.${quoteIdentifier(config.schema)}.${quoteIdentifier(config.table)}`
const placeholders = Array.from({ length: rowCount }, () => '(PARSE_JSON(?))').join(', ')
return `INSERT INTO ${target} (${column}) VALUES ${placeholders}`
}
function isRetryableStatus(status: number): boolean {
return status === 408 || status === 429 || (status >= 500 && status <= 599)
}
interface ExecuteInput {
config: SnowflakeDestinationConfig
getJwt: () => Promise<string>
statement: string
bindings: string[]
signal: AbortSignal
}
async function executeStatement(input: ExecuteInput): Promise<void> {
for (const value of input.bindings) {
const bytes = Buffer.byteLength(value, 'utf8')
if (bytes > VARIANT_MAX_BYTES) {
throw new Error(
`Snowflake VARIANT value exceeds 16 MB limit (got ${bytes} bytes); split the row before delivery`
)
}
}
const baseUrl = `https://${input.config.account}.snowflakecomputing.com/api/v2/statements`
const bindings: Record<string, { type: 'TEXT'; value: string }> = {}
input.bindings.forEach((value, index) => {
bindings[(index + 1).toString()] = { type: 'TEXT', value }
})
const body = {
statement: input.statement,
timeout: SQL_API_TIMEOUT_SECONDS,
warehouse: input.config.warehouse,
role: input.config.role,
bindings,
}
const serializedBody = JSON.stringify(body)
/** Stable per-request UUID enables idempotent retries via `retry=true` on subsequent attempts. */
const requestId = generateId()
let lastError: unknown
for (let attempt = 1; attempt <= EXECUTE_MAX_ATTEMPTS; attempt++) {
if (input.signal.aborted) throw input.signal.reason ?? new Error('Aborted')
/** Acquire JWT before starting the per-attempt timer so token signing doesn't eat the network budget (mirrors pollStatement). */
const jwt = await input.getJwt()
const perAttempt = AbortSignal.any([input.signal, AbortSignal.timeout(PER_ATTEMPT_TIMEOUT_MS)])
const params = new URLSearchParams({ requestId })
if (attempt > 1) params.set('retry', 'true')
const url = `${baseUrl}?${params.toString()}`
let response: Response
try {
response = await fetch(url, {
method: 'POST',
headers: {
Authorization: `Bearer ${jwt}`,
'Content-Type': 'application/json',
Accept: 'application/json',
'X-Snowflake-Authorization-Token-Type': 'KEYPAIR_JWT',
'User-Agent': 'sim-data-drain/1.0',
},
body: serializedBody,
signal: perAttempt,
})
} catch (error) {
lastError = error
logger.warn('Snowflake request failed', {
attempt,
error: toError(error).message,
})
if (input.signal.aborted || attempt === EXECUTE_MAX_ATTEMPTS) throw error
await sleepUntilAborted(
backoffWithJitter(attempt, null, {
baseMs: EXECUTE_RETRY_BASE_DELAY_MS,
maxMs: EXECUTE_RETRY_MAX_DELAY_MS,
}),
input.signal
)
continue
}
if (response.status === 202) {
const json = (await response.json().catch(() => ({}))) as { statementHandle?: string }
if (!json.statementHandle) {
throw new Error('Snowflake returned 202 without a statementHandle')
}
await pollStatement({
account: input.config.account,
getJwt: input.getJwt,
handle: json.statementHandle,
signal: input.signal,
})
return
}
if (response.ok) {
/**
* Synchronous completions return 200 — same statement-level error envelope as
* the polled 200 path, so check `sqlState` here too instead of silently passing
* failures. Consuming the body also lets undici reuse the socket.
*/
const completion = (await response.json().catch(() => ({}))) as {
code?: string
sqlState?: string
message?: string
}
if (completion.sqlState && completion.sqlState !== '00000') {
throw new Error(
`Snowflake statement failed (sqlState ${completion.sqlState}${completion.code ? `, code ${completion.code}` : ''}): ${completion.message ?? ''}`
)
}
return
}
const text = await response.text().catch(() => '')
const error = new Error(`Snowflake responded with HTTP ${response.status}: ${text}`)
if (!isRetryableStatus(response.status) || attempt === EXECUTE_MAX_ATTEMPTS) throw error
lastError = error
const retryAfterMs = parseRetryAfter(response.headers.get('Retry-After'))
const delay = backoffWithJitter(attempt, retryAfterMs, {
baseMs: EXECUTE_RETRY_BASE_DELAY_MS,
maxMs: EXECUTE_RETRY_MAX_DELAY_MS,
})
logger.warn('Snowflake request retrying after retryable status', {
attempt,
status: response.status,
delayMs: delay,
})
await sleepUntilAborted(delay, input.signal)
}
throw lastError ?? new Error('Snowflake request failed after retries')
}
interface PollInput {
account: string
/** Thunk so long polls (past 55min) refresh the JWT instead of dying with 401. */
getJwt: () => Promise<string>
handle: string
signal: AbortSignal
}
/** Snowflake returns 202 while still executing and 200 on completion (async statement-handle semantics). */
async function pollStatement(input: PollInput): Promise<void> {
const url = `https://${input.account}.snowflakecomputing.com/api/v2/statements/${encodeURIComponent(input.handle)}`
const deadline = Date.now() + POLL_DEADLINE_MS
let interval = POLL_INITIAL_INTERVAL_MS
let skipIntervalSleep = true
let retryAttempt = 0
while (Date.now() < deadline) {
if (input.signal.aborted) throw input.signal.reason ?? new Error('Aborted')
if (!skipIntervalSleep) {
await sleepUntilAborted(interval, input.signal)
}
skipIntervalSleep = false
const jwt = await input.getJwt()
const perAttempt = AbortSignal.any([input.signal, AbortSignal.timeout(PER_ATTEMPT_TIMEOUT_MS)])
let response: Response
try {
response = await fetch(url, {
headers: {
Authorization: `Bearer ${jwt}`,
Accept: 'application/json',
'X-Snowflake-Authorization-Token-Type': 'KEYPAIR_JWT',
},
signal: perAttempt,
})
} catch (error) {
if (input.signal.aborted) throw error
retryAttempt++
if (retryAttempt > POLL_MAX_CONSECUTIVE_RETRIES) throw error
const delay = backoffWithJitter(retryAttempt, null, {
baseMs: EXECUTE_RETRY_BASE_DELAY_MS,
maxMs: EXECUTE_RETRY_MAX_DELAY_MS,
})
logger.warn('Snowflake poll request failed, retrying', {
attempt: retryAttempt,
delayMs: delay,
error: toError(error).message,
})
await sleepUntilAborted(delay, input.signal)
skipIntervalSleep = true
continue
}
if (response.status === 202) {
/** Drain the body so undici can return the socket to the keep-alive pool between polls. */
await response.text().catch(() => '')
retryAttempt = 0
interval = Math.min(interval * 2, POLL_MAX_INTERVAL_MS)
continue
}
if (isRetryableStatus(response.status)) {
retryAttempt++
if (retryAttempt > POLL_MAX_CONSECUTIVE_RETRIES) {
/** Drain the body so undici can return the socket to the keep-alive pool. */
const text = await response.text().catch(() => '')
throw new Error(
`Snowflake poll failed after ${POLL_MAX_CONSECUTIVE_RETRIES} consecutive retries (HTTP ${response.status}): ${text}`
)
}
const retryAfterMs = parseRetryAfter(response.headers.get('Retry-After'))
const delay = backoffWithJitter(retryAttempt, retryAfterMs, {
baseMs: EXECUTE_RETRY_BASE_DELAY_MS,
maxMs: EXECUTE_RETRY_MAX_DELAY_MS,
})
logger.warn('Snowflake poll retrying after retryable status', {
attempt: retryAttempt,
status: response.status,
delayMs: delay,
})
/** Drain the body so undici can return the socket to the keep-alive pool between retries. */
await response.text().catch(() => '')
await sleepUntilAborted(delay, input.signal)
skipIntervalSleep = true
continue
}
if (!response.ok) {
const text = await response.text().catch(() => '')
throw new Error(`Snowflake poll failed (HTTP ${response.status}): ${text}`)
}
/**
* Snowflake SQL API can return 200 with a statement-level error envelope
* (`code` / `sqlState` / `message`). Successful completions return
* `code === "090001"` ("statement executed successfully") or omit `code`,
* while statement errors come back with codes like `"002032"` and
* a populated `sqlState`. Treat anything with a `sqlState` as a failure.
*/
const completion = (await response.json().catch(() => ({}))) as {
code?: string
sqlState?: string
message?: string
}
if (completion.sqlState && completion.sqlState !== '00000') {
throw new Error(
`Snowflake statement failed (sqlState ${completion.sqlState}${completion.code ? `, code ${completion.code}` : ''}): ${completion.message ?? ''}`
)
}
return
}
throw new Error('Snowflake statement did not complete within the polling deadline')
}
export const snowflakeDestination: DrainDestination<
SnowflakeDestinationConfig,
SnowflakeDestinationCredentials
> = {
type: 'snowflake',
displayName: 'Snowflake',
configSchema: snowflakeConfigSchema,
credentialsSchema: snowflakeCredentialsSchema,
async test({ config, credentials, signal }) {
let cached: JwtCacheEntry | null = null
async function getJwt(): Promise<string> {
const now = Math.floor(Date.now() / 1000)
if (cached && cached.expiresAt > now) return cached.token
cached = await buildJwt(config.account, config.user, credentials.privateKey)
return cached.token
}
await executeStatement({
config,
getJwt,
statement: 'SELECT 1',
bindings: [],
signal,
})
},
openSession({ config, credentials }) {
let cached: JwtCacheEntry | null = null
async function getJwt(): Promise<string> {
const now = Math.floor(Date.now() / 1000)
if (cached && cached.expiresAt > now) return cached.token
cached = await buildJwt(config.account, config.user, credentials.privateKey)
return cached.token
}
return {
async deliver({ body, metadata, signal }) {
/**
* Bind the original line bytes — not `JSON.stringify(JSON.parse(line))` —
* so JSON numbers outside the JS safe-integer range (e.g. Snowflake
* NUMBER columns past 2^53-1) survive into VARIANT intact. We still
* parse each line so a malformed payload fails fast at the runner.
*/
const text = body.toString('utf8')
const rows: string[] = []
const lines = text.split(/\r?\n/)
for (let i = 0; i < lines.length; i++) {
const line = lines[i]
if (line.length === 0) continue
try {
JSON.parse(line)
} catch (error) {
throw new Error(
`Snowflake NDJSON parse failed at line ${i + 1}: ${toError(error).message}`
)
}
rows.push(line)
}
if (rows.length === 0) {
return {
locator: `snowflake://${config.account}/${config.database}.${config.schema}.${config.table}#${metadata.runId}-${metadata.sequence}`,
}
}
await executeStatement({
config,
getJwt,
statement: buildStatement(config, rows.length),
bindings: rows,
signal,
})
logger.debug('Snowflake chunk delivered', {
account: config.account,
table: `${config.database}.${config.schema}.${config.table}`,
rows: rows.length,
})
return {
locator: `snowflake://${config.account}/${config.database}.${config.schema}.${config.table}#${metadata.runId}-${metadata.sequence}`,
}
},
async close() {},
}
},
}
@@ -0,0 +1,164 @@
import { toError } from '@sim/utils/errors'
import { z } from 'zod'
/**
* Sleep for `ms` milliseconds, resolving early if `signal` aborts. Used by
* destination retry/poll loops so cancelled drain runs do not hang waiting on
* a `setTimeout` that ignores the abort signal.
*/
export function sleepUntilAborted(ms: number, signal: AbortSignal): Promise<void> {
if (signal.aborted) return Promise.resolve()
return new Promise((resolve) => {
const onAbort = () => {
clearTimeout(timeoutId)
resolve()
}
const timeoutId = setTimeout(() => {
signal.removeEventListener('abort', onAbort)
resolve()
}, ms)
signal.addEventListener('abort', onAbort, { once: true })
})
}
export function normalizePrefix(raw: string | undefined): string {
if (!raw) return ''
const trimmed = raw.replace(/^\/+/, '').replace(/\/+$/, '')
return trimmed.length === 0 ? '' : `${trimmed}/`
}
export interface ObjectKeyMetadata {
drainId: string
runId: string
source: string
sequence: number
runStartedAt: Date
}
/**
* Builds a date-partitioned NDJSON object key for blob-store destinations.
* Layout: `<prefix><source>/<drainId>/<YYYY>/<MM>/<DD>/<runId>-<seq>.ndjson`.
* Partition uses the run's start time so all chunks from a run share one
* date prefix even if delivery crosses a UTC midnight boundary.
*/
export function buildObjectKey(prefix: string | undefined, metadata: ObjectKeyMetadata): string {
const partition = metadata.runStartedAt
const yyyy = partition.getUTCFullYear().toString().padStart(4, '0')
const mm = (partition.getUTCMonth() + 1).toString().padStart(2, '0')
const dd = partition.getUTCDate().toString().padStart(2, '0')
const seq = metadata.sequence.toString().padStart(5, '0')
return `${normalizePrefix(prefix)}${metadata.source}/${metadata.drainId}/${yyyy}/${mm}/${dd}/${metadata.runId}-${seq}.ndjson`
}
export interface ParsedServiceAccount {
clientEmail: string
privateKey: string
}
/**
* Parses a Google service-account JSON key, returning the only two fields
* that destinations need (client email + private key). Shared by GCS and
* BigQuery so a fix in one place applies to both.
*/
export function parseServiceAccount(json: string): ParsedServiceAccount {
let parsed: unknown
try {
parsed = JSON.parse(json)
} catch (error) {
throw new Error(`serviceAccountJson is not valid JSON: ${toError(error).message}`)
}
if (typeof parsed !== 'object' || parsed === null) {
throw new Error('serviceAccountJson must be a JSON object')
}
const obj = parsed as Record<string, unknown>
const clientEmail = obj.client_email
const privateKey = obj.private_key
if (typeof clientEmail !== 'string' || clientEmail.length === 0) {
throw new Error('serviceAccountJson is missing client_email')
}
if (typeof privateKey !== 'string' || privateKey.length === 0) {
throw new Error('serviceAccountJson is missing private_key')
}
return { clientEmail, privateKey }
}
/**
* Zod `superRefine` helper that validates a service-account JSON key string
* is parseable and carries `client_email` + `private_key`. Used by both
* `gcsCredentialsSchema` and `bigqueryCredentialsSchema`.
*/
export interface ParseNdjsonObjectsOptions {
/** When true, throw if a parsed value is not a plain object. */
requireObject?: boolean
}
/**
* Parses an NDJSON buffer into per-row JSON values. Error messages use
* 1-indexed line numbers so they line up with how editors and `Content-Range`
* headers reference NDJSON payloads.
*/
export function parseNdjsonObjects(
body: Buffer,
options: ParseNdjsonObjectsOptions = {}
): unknown[] {
const text = body.toString('utf8')
const rows: unknown[] = []
const lines = text.split(/\r?\n/)
for (let i = 0; i < lines.length; i++) {
const line = lines[i]
if (line.length === 0) continue
let parsed: unknown
try {
parsed = JSON.parse(line)
} catch (error) {
throw new Error(`NDJSON parse failed at line ${i + 1}: ${toError(error).message}`)
}
if (options.requireObject) {
if (typeof parsed !== 'object' || parsed === null || Array.isArray(parsed)) {
throw new Error(`NDJSON row at line ${i + 1} is not an object`)
}
}
rows.push(parsed)
}
return rows
}
export function refineServiceAccountJson(
value: { serviceAccountJson: string },
ctx: z.RefinementCtx
): void {
let parsed: unknown
try {
parsed = JSON.parse(value.serviceAccountJson)
} catch {
ctx.addIssue({
code: z.ZodIssueCode.custom,
path: ['serviceAccountJson'],
message: 'serviceAccountJson must be valid JSON',
})
return
}
if (typeof parsed !== 'object' || parsed === null) {
ctx.addIssue({
code: z.ZodIssueCode.custom,
path: ['serviceAccountJson'],
message: 'serviceAccountJson must be a JSON object',
})
return
}
const obj = parsed as Record<string, unknown>
if (typeof obj.client_email !== 'string' || obj.client_email.length === 0) {
ctx.addIssue({
code: z.ZodIssueCode.custom,
path: ['serviceAccountJson'],
message: 'serviceAccountJson is missing client_email',
})
}
if (typeof obj.private_key !== 'string' || obj.private_key.length === 0) {
ctx.addIssue({
code: z.ZodIssueCode.custom,
path: ['serviceAccountJson'],
message: 'serviceAccountJson is missing private_key',
})
}
}
@@ -0,0 +1,176 @@
/**
* @vitest-environment node
*/
import { createHmac } from 'node:crypto'
import { inputValidationMock, inputValidationMockFns } from '@sim/testing'
import { beforeEach, describe, expect, it, vi } from 'vitest'
vi.mock('@/lib/core/security/input-validation.server', () => inputValidationMock)
import { webhookDestination } from '@/lib/data-drains/destinations/webhook'
const config = { url: 'https://example.com/hook' }
const credentials = { signingSecret: 'super-secret-key' }
const metadata = {
drainId: 'd1',
runId: 'r1',
source: 'workflow_logs' as const,
sequence: 3,
rowCount: 5,
}
function mockPinnedFetchOnce(response: { ok: boolean; status: number; headers?: Headers }) {
inputValidationMockFns.mockSecureFetchWithPinnedIP.mockResolvedValueOnce({
ok: response.ok,
status: response.status,
statusText: '',
headers: response.headers ?? new Headers(),
text: async () => '',
json: async () => ({}),
arrayBuffer: async () => new ArrayBuffer(0),
})
}
beforeEach(() => {
vi.clearAllMocks()
inputValidationMockFns.mockValidateUrlWithDNS.mockResolvedValue({
isValid: true,
resolvedIP: '93.184.216.34',
originalHostname: 'example.com',
})
})
describe('webhookDestination openSession', () => {
it('signs the body with HMAC-SHA256 over `<ts>.<body>`', async () => {
mockPinnedFetchOnce({ ok: true, status: 200 })
const session = webhookDestination.openSession({ config, credentials })
const body = Buffer.from('{"id":1}\n', 'utf8')
await session.deliver({
body,
contentType: 'application/x-ndjson',
metadata,
signal: new AbortController().signal,
})
const call = inputValidationMockFns.mockSecureFetchWithPinnedIP.mock.calls[0]
const [calledUrl, pinnedIP, init] = call
expect(calledUrl).toBe('https://example.com/hook')
expect(pinnedIP).toBe('93.184.216.34')
const headers = init.headers as Record<string, string>
expect(headers['Content-Type']).toBe('application/x-ndjson')
expect(headers['X-Sim-Drain-Id']).toBe('d1')
expect(headers['X-Sim-Run-Id']).toBe('r1')
expect(headers['X-Sim-Sequence']).toBe('3')
expect(headers['Idempotency-Key']).toBe('r1-3')
const sig = headers['X-Sim-Signature']
const tsPart = sig.match(/t=(\d+)/)![1]
const v1Part = sig.match(/v1=([0-9a-f]+)/)![1]
const expected = createHmac('sha256', credentials.signingSecret)
.update(`${tsPart}.`)
.update(body)
.digest('hex')
expect(v1Part).toBe(expected)
await session.close()
})
it('retries on 5xx and succeeds', async () => {
mockPinnedFetchOnce({ ok: false, status: 503 })
mockPinnedFetchOnce({ ok: true, status: 200 })
vi.spyOn(global, 'setTimeout').mockImplementation(((fn: () => void) => {
fn()
return 0 as unknown as NodeJS.Timeout
}) as never)
const session = webhookDestination.openSession({ config, credentials })
const result = await session.deliver({
body: Buffer.from('x'),
contentType: 'application/x-ndjson',
metadata,
signal: new AbortController().signal,
})
expect(result.locator).toContain('https://example.com/hook')
expect(inputValidationMockFns.mockSecureFetchWithPinnedIP).toHaveBeenCalledTimes(2)
})
it('does not retry on 4xx (other than 408/429)', async () => {
mockPinnedFetchOnce({ ok: false, status: 401 })
const session = webhookDestination.openSession({ config, credentials })
await expect(
session.deliver({
body: Buffer.from('x'),
contentType: 'application/x-ndjson',
metadata,
signal: new AbortController().signal,
})
).rejects.toThrow(/HTTP 401/)
expect(inputValidationMockFns.mockSecureFetchWithPinnedIP).toHaveBeenCalledTimes(1)
})
it('rejects when DNS resolves to a blocked IP', async () => {
inputValidationMockFns.mockValidateUrlWithDNS.mockResolvedValueOnce({
isValid: false,
error: 'url resolves to a blocked IP address',
})
const session = webhookDestination.openSession({ config, credentials })
await expect(
session.deliver({
body: Buffer.from('x'),
contentType: 'application/x-ndjson',
metadata,
signal: new AbortController().signal,
})
).rejects.toThrow(/blocked IP/)
expect(inputValidationMockFns.mockSecureFetchWithPinnedIP).not.toHaveBeenCalled()
})
it('reuses the same pinned IP across deliveries (no DNS rebinding window)', async () => {
mockPinnedFetchOnce({ ok: true, status: 200 })
mockPinnedFetchOnce({ ok: true, status: 200 })
const session = webhookDestination.openSession({ config, credentials })
const signal = new AbortController().signal
await session.deliver({
body: Buffer.from('x'),
contentType: 'application/x-ndjson',
metadata,
signal,
})
await session.deliver({
body: Buffer.from('y'),
contentType: 'application/x-ndjson',
metadata: { ...metadata, sequence: 4 },
signal,
})
expect(inputValidationMockFns.mockValidateUrlWithDNS).toHaveBeenCalledTimes(1)
const calls = inputValidationMockFns.mockSecureFetchWithPinnedIP.mock.calls
expect(calls[0][1]).toBe('93.184.216.34')
expect(calls[1][1]).toBe('93.184.216.34')
})
it('rejects every header buildHeaders writes when reused as signatureHeader (drift guard)', async () => {
mockPinnedFetchOnce({ ok: true, status: 200 })
const session = webhookDestination.openSession({ config, credentials })
await session.deliver({
body: Buffer.from('x'),
contentType: 'application/x-ndjson',
metadata,
signal: new AbortController().signal,
})
const init = inputValidationMockFns.mockSecureFetchWithPinnedIP.mock.calls[0][2]
const writtenHeaders = Object.keys(init.headers as Record<string, string>)
for (const name of writtenHeaders) {
const result = webhookDestination.configSchema.safeParse({
url: 'https://example.com/hook',
signatureHeader: name,
})
expect(
result.success,
`expected signatureHeader="${name}" to be rejected (it is written by buildHeaders)`
).toBe(false)
}
})
})
@@ -0,0 +1,246 @@
import { createHmac } from 'node:crypto'
import { createLogger } from '@sim/logger'
import { toError } from '@sim/utils/errors'
import { backoffWithJitter, parseRetryAfter } from '@sim/utils/retry'
import { z } from 'zod'
import { validateExternalUrl } from '@/lib/core/security/input-validation'
import {
secureFetchWithPinnedIP,
validateUrlWithDNS,
} from '@/lib/core/security/input-validation.server'
import { sleepUntilAborted } from '@/lib/data-drains/destinations/utils'
import type { DeliveryMetadata, DrainDestination } from '@/lib/data-drains/types'
const logger = createLogger('DataDrainWebhookDestination')
/** Initial attempt + 3 retries — 500ms/1s/2s backoff sequence. */
const MAX_ATTEMPTS = 4
const PER_ATTEMPT_TIMEOUT_MS = 30_000
/** Cap responder reply so a misbehaving receiver can't OOM the runner. */
const MAX_RESPONSE_BYTES = 256 * 1024
const SIGNATURE_VERSION = 'v1'
const USER_AGENT = 'Sim-DataDrain/1.0'
/**
* Headers `buildHeaders` emits. Callers cannot override these via
* `signatureHeader`. Keep in sync with `buildHeaders` — the drift-guard test
* enforces this by parsing the schema against every key the function writes.
*/
const RESERVED_SIGNATURE_HEADER_NAMES = new Set([
'authorization',
'content-type',
'user-agent',
'idempotency-key',
'x-sim-timestamp',
'x-sim-signature-version',
'x-sim-drain-id',
'x-sim-run-id',
'x-sim-source',
'x-sim-sequence',
'x-sim-row-count',
'x-sim-probe',
'x-sim-signature',
])
/** CR/LF/NUL would let a bearer token smuggle additional response headers. */
const HEADER_INJECTION_PATTERN = /[\r\n\0]/
async function resolvePublicTarget(url: string): Promise<string> {
const result = await validateUrlWithDNS(url, 'url')
if (!result.isValid || !result.resolvedIP) {
throw new Error(result.error ?? 'Webhook URL failed SSRF validation')
}
return result.resolvedIP
}
const webhookConfigSchema = z.object({
url: z
.string()
.url('url must be a valid URL')
.max(2048, 'url must be at most 2048 characters')
.refine((value) => validateExternalUrl(value, 'url').isValid, {
message: 'url must be HTTPS and not point at a private, loopback, or metadata address',
}),
signatureHeader: z
.string()
.min(1)
.max(128)
.refine((value) => !RESERVED_SIGNATURE_HEADER_NAMES.has(value.toLowerCase()), {
message: 'signatureHeader cannot reuse a reserved Sim header name',
})
.refine((value) => !HEADER_INJECTION_PATTERN.test(value) && /^[A-Za-z0-9\-_]+$/.test(value), {
message: 'signatureHeader must contain only letters, digits, hyphens, and underscores',
})
.optional(),
})
const webhookCredentialsSchema = z.object({
signingSecret: z
.string()
.min(32, 'signingSecret must be at least 32 characters')
.max(512, 'signingSecret must be at most 512 characters'),
bearerToken: z
.string()
.min(1)
.max(4096, 'bearerToken must be at most 4096 characters')
.refine((value) => !HEADER_INJECTION_PATTERN.test(value), {
message: 'bearerToken cannot contain CR, LF, or NUL characters',
})
.optional(),
})
export type WebhookDestinationConfig = z.infer<typeof webhookConfigSchema>
export type WebhookDestinationCredentials = z.infer<typeof webhookCredentialsSchema>
/**
* Stripe-style signature: HMAC-SHA256 over `${unixSeconds}.${body}` rendered as
* `t=<unixSeconds>,v1=<hex>`. Verifiers reject stale timestamps (~5 min skew)
* to block replay; we re-sign per attempt so long backoffs don't fall outside
* that window.
*/
function sign(body: Buffer, secret: string, timestamp: number): string {
const hmac = createHmac('sha256', secret).update(`${timestamp}.`).update(body).digest('hex')
return `t=${timestamp},${SIGNATURE_VERSION}=${hmac}`
}
function isRetryableStatus(status: number): boolean {
return status === 408 || status === 429 || status >= 500
}
function buildHeaders(input: {
config: WebhookDestinationConfig
credentials: WebhookDestinationCredentials
body: Buffer
contentType: string
metadata?: DeliveryMetadata
isProbe?: boolean
}): Record<string, string> {
const timestamp = Math.floor(Date.now() / 1000)
const headers: Record<string, string> = {
'Content-Type': input.contentType,
'User-Agent': USER_AGENT,
'X-Sim-Timestamp': timestamp.toString(),
'X-Sim-Signature-Version': SIGNATURE_VERSION,
[input.config.signatureHeader ?? 'X-Sim-Signature']: sign(
input.body,
input.credentials.signingSecret,
timestamp
),
}
if (input.metadata) {
headers['X-Sim-Drain-Id'] = input.metadata.drainId
headers['X-Sim-Run-Id'] = input.metadata.runId
headers['X-Sim-Source'] = input.metadata.source
headers['X-Sim-Sequence'] = input.metadata.sequence.toString()
headers['X-Sim-Row-Count'] = input.metadata.rowCount.toString()
// Stable across retries of the same chunk so receivers can dedupe.
headers['Idempotency-Key'] = `${input.metadata.runId}-${input.metadata.sequence}`
}
if (input.isProbe) {
headers['X-Sim-Probe'] = '1'
}
if (input.credentials.bearerToken) {
headers.Authorization = `Bearer ${input.credentials.bearerToken}`
}
return headers
}
export const webhookDestination: DrainDestination<
WebhookDestinationConfig,
WebhookDestinationCredentials
> = {
type: 'webhook',
displayName: 'HTTPS Webhook',
configSchema: webhookConfigSchema,
credentialsSchema: webhookCredentialsSchema,
async test({ config, credentials, signal }) {
const resolvedIP = await resolvePublicTarget(config.url)
const probe = Buffer.from('{"sim":"connection-test"}\n', 'utf8')
const headers = buildHeaders({
config,
credentials,
body: probe,
contentType: 'application/x-ndjson',
isProbe: true,
})
const response = await secureFetchWithPinnedIP(config.url, resolvedIP, {
method: 'POST',
body: new Uint8Array(probe),
headers,
signal,
timeout: PER_ATTEMPT_TIMEOUT_MS,
maxResponseBytes: MAX_RESPONSE_BYTES,
})
if (!response.ok) {
throw new Error(`Webhook probe failed: HTTP ${response.status}`)
}
},
openSession({ config, credentials }) {
let resolvedIP: string | null = null
return {
async deliver({ body, contentType, metadata, signal }) {
// Resolve once per session and pin across retries to defeat DNS rebinding (TOCTOU).
if (resolvedIP === null) {
resolvedIP = await resolvePublicTarget(config.url)
}
let lastError: unknown
for (let attempt = 1; attempt <= MAX_ATTEMPTS; attempt++) {
if (signal.aborted) throw signal.reason ?? new Error('Aborted')
const headers = buildHeaders({ config, credentials, body, contentType, metadata })
let retryAfterMs: number | null = null
let response: Awaited<ReturnType<typeof secureFetchWithPinnedIP>> | undefined
try {
response = await secureFetchWithPinnedIP(config.url, resolvedIP, {
method: 'POST',
body: new Uint8Array(body),
headers,
signal,
timeout: PER_ATTEMPT_TIMEOUT_MS,
maxResponseBytes: MAX_RESPONSE_BYTES,
})
} catch (error) {
lastError = error
logger.debug('Webhook delivery attempt failed', {
url: config.url,
attempt,
error: toError(error).message,
})
}
if (response) {
if (response.ok) {
const requestId =
response.headers.get('x-request-id') ??
response.headers.get('x-amzn-trace-id') ??
null
logger.debug('Webhook chunk delivered', {
url: config.url,
attempt,
status: response.status,
bytes: body.byteLength,
})
return {
locator: requestId
? `${config.url}#${metadata.runId}-${metadata.sequence}@${requestId}`
: `${config.url}#${metadata.runId}-${metadata.sequence}`,
}
}
if (!isRetryableStatus(response.status)) {
throw new Error(`Webhook responded with HTTP ${response.status}`)
}
lastError = new Error(`Webhook responded with HTTP ${response.status}`)
retryAfterMs = parseRetryAfter(response.headers.get('retry-after'))
}
if (attempt < MAX_ATTEMPTS) {
await sleepUntilAborted(backoffWithJitter(attempt, retryAfterMs), signal)
}
}
throw lastError instanceof Error
? lastError
: new Error('Webhook delivery failed after retries')
},
async close() {},
}
},
}
+117
View File
@@ -0,0 +1,117 @@
/**
* @vitest-environment node
*/
import { dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing'
import { beforeEach, describe, expect, it, vi } from 'vitest'
vi.mock('@sim/db', () => dbChainMock)
const { mockIsEnterprise, mockEnqueue, mockGetJobQueue } = vi.hoisted(() => {
const mockEnqueue = vi.fn(async () => 'job-id')
return {
mockIsEnterprise: vi.fn(),
mockEnqueue,
mockGetJobQueue: vi.fn(async () => ({ enqueue: mockEnqueue })),
}
})
vi.mock('@/lib/billing/core/subscription', () => ({
isOrganizationOnEnterprisePlan: mockIsEnterprise,
}))
vi.mock('@/lib/core/async-jobs', () => ({ getJobQueue: mockGetJobQueue }))
vi.mock('@/lib/core/config/env-flags', () => ({ isBillingEnabled: true }))
import { dispatchDueDrains, reapOrphanedRuns } from '@/lib/data-drains/dispatcher'
function mockCandidates(rows: Array<{ id: string; organizationId: string }>) {
// db.select().from().where() — override `from` so awaiting `.where(pred)`
// resolves with the candidate rows.
dbChainMockFns.from.mockReturnValueOnce({
where: vi.fn().mockResolvedValueOnce(rows),
} as never)
}
beforeEach(() => {
vi.clearAllMocks()
resetDbChainMock()
})
describe('reapOrphanedRuns', () => {
it('returns the count of rows updated to failed', async () => {
dbChainMockFns.returning.mockResolvedValueOnce([{ id: 'run-1' }, { id: 'run-2' }])
const result = await reapOrphanedRuns(new Date('2026-01-01T12:00:00.000Z'))
expect(result).toEqual({ reaped: 2 })
expect(dbChainMockFns.set).toHaveBeenCalledWith(
expect.objectContaining({ status: 'failed', error: expect.stringContaining('Orphaned') })
)
})
it('returns 0 when nothing is stuck', async () => {
dbChainMockFns.returning.mockResolvedValueOnce([])
expect(await reapOrphanedRuns()).toEqual({ reaped: 0 })
})
})
describe('dispatchDueDrains', () => {
it('returns early when no candidates are due', async () => {
dbChainMockFns.returning.mockResolvedValueOnce([]) // reaper
mockCandidates([])
const result = await dispatchDueDrains()
expect(result).toEqual({ candidates: 0, dispatched: 0, skipped: 0, reaped: 0 })
expect(mockGetJobQueue).not.toHaveBeenCalled()
})
it('skips drains for orgs not on enterprise plan', async () => {
dbChainMockFns.returning.mockResolvedValueOnce([]) // reaper
mockCandidates([{ id: 'd1', organizationId: 'org-a' }])
mockIsEnterprise.mockResolvedValueOnce(false)
const result = await dispatchDueDrains()
expect(result).toMatchObject({ candidates: 1, dispatched: 0, skipped: 1 })
expect(mockEnqueue).not.toHaveBeenCalled()
})
it('claims and enqueues a job per due drain', async () => {
dbChainMockFns.returning
.mockResolvedValueOnce([]) // reaper
.mockResolvedValueOnce([{ id: 'd1' }]) // claim succeeds
mockCandidates([{ id: 'd1', organizationId: 'org-a' }])
mockIsEnterprise.mockResolvedValueOnce(true)
const result = await dispatchDueDrains()
expect(result).toMatchObject({ candidates: 1, dispatched: 1, skipped: 0 })
expect(mockEnqueue).toHaveBeenCalledWith(
'run-data-drain',
{ drainId: 'd1', trigger: 'cron' },
{ concurrencyKey: 'data-drain:d1' }
)
})
it('does not enqueue when claim loses the race', async () => {
dbChainMockFns.returning
.mockResolvedValueOnce([]) // reaper
.mockResolvedValueOnce([]) // claim returns nothing — lost the race
mockCandidates([{ id: 'd1', organizationId: 'org-a' }])
mockIsEnterprise.mockResolvedValueOnce(true)
const result = await dispatchDueDrains()
expect(result.dispatched).toBe(0)
expect(mockEnqueue).not.toHaveBeenCalled()
})
it('caches enterprise check across drains in the same org', async () => {
dbChainMockFns.returning
.mockResolvedValueOnce([]) // reaper
.mockResolvedValueOnce([{ id: 'd1' }])
.mockResolvedValueOnce([{ id: 'd2' }])
mockCandidates([
{ id: 'd1', organizationId: 'org-a' },
{ id: 'd2', organizationId: 'org-a' },
])
mockIsEnterprise.mockResolvedValue(true)
await dispatchDueDrains()
expect(mockIsEnterprise).toHaveBeenCalledTimes(1)
})
})
+186
View File
@@ -0,0 +1,186 @@
import { db } from '@sim/db'
import { dataDrainRuns, dataDrains } from '@sim/db/schema'
import { createLogger } from '@sim/logger'
import { toError } from '@sim/utils/errors'
import { and, eq, isNull, lt, or } from 'drizzle-orm'
import { isOrganizationOnEnterprisePlan } from '@/lib/billing/core/subscription'
import { getJobQueue } from '@/lib/core/async-jobs'
import { isBillingEnabled } from '@/lib/core/config/env-flags'
const logger = createLogger('DataDrainsDispatcher')
const HOUR_MS = 60 * 60 * 1000
const DAY_MS = 24 * HOUR_MS
/**
* Cron fires hourly. Without a buffer, a drain that finishes a few minutes
* after the tick (lastRunAt = 10:05) won't satisfy `lastRunAt < now - cadence`
* at the next tick (10:05 < 10:00 is false), so an "hourly" drain effectively
* runs every two hours. Subtracting a small buffer from the cadence absorbs
* normal run duration plus cron jitter without allowing back-to-back runs
* within the same tick.
*/
const CADENCE_BUFFER_MS = 5 * 60 * 1000
/**
* Maximum wall-clock duration any single drain run is allowed before its
* `data_drain_runs` row is considered orphaned. Runs that exceed this are
* almost certainly the result of a Trigger.dev worker crash mid-run — there
* is no live process still updating them.
*/
const ORPHAN_THRESHOLD_MS = 60 * 60 * 1000
/**
* Marks `running` rows older than the orphan threshold as `failed`. Without
* this, a worker crash leaves run history permanently misleading and (worse)
* the drain row's `lastRunAt` reflects a successful claim that never finished
* — but the drain `cursor` never advanced, so re-running is safe.
*/
export async function reapOrphanedRuns(now: Date = new Date()): Promise<{ reaped: number }> {
const cutoff = new Date(now.getTime() - ORPHAN_THRESHOLD_MS)
const reaped = await db
.update(dataDrainRuns)
.set({
status: 'failed',
finishedAt: now,
error: `Orphaned run reaped after exceeding ${ORPHAN_THRESHOLD_MS / 60_000}m without completion`,
})
.where(and(eq(dataDrainRuns.status, 'running'), lt(dataDrainRuns.startedAt, cutoff)))
.returning({ id: dataDrainRuns.id })
if (reaped.length > 0) {
logger.warn('Reaped orphaned data drain runs', { count: reaped.length })
}
return { reaped: reaped.length }
}
/**
* Selects every enabled drain whose schedule is due (or has never run) and
* fans out one `run-data-drain` job per drain. Each drain is atomically
* claimed via a conditional UPDATE before being enqueued — two concurrent
* dispatcher invocations cannot both win the same row, and a manual run that
* lands between the SELECT and the UPDATE will lose the race cleanly. Drains
* belonging to orgs that have lapsed off the enterprise plan are skipped.
*/
export async function dispatchDueDrains(now: Date = new Date()): Promise<{
candidates: number
dispatched: number
skipped: number
reaped: number
}> {
const { reaped } = await reapOrphanedRuns(now)
const hourlyCutoff = new Date(now.getTime() - HOUR_MS + CADENCE_BUFFER_MS)
const dailyCutoff = new Date(now.getTime() - DAY_MS + CADENCE_BUFFER_MS)
const duePredicate = and(
eq(dataDrains.enabled, true),
or(
isNull(dataDrains.lastRunAt),
and(eq(dataDrains.scheduleCadence, 'hourly'), lt(dataDrains.lastRunAt, hourlyCutoff)),
and(eq(dataDrains.scheduleCadence, 'daily'), lt(dataDrains.lastRunAt, dailyCutoff))
)
)
const candidates = await db
.select({
id: dataDrains.id,
organizationId: dataDrains.organizationId,
lastRunAt: dataDrains.lastRunAt,
})
.from(dataDrains)
.where(duePredicate)
if (candidates.length === 0) {
return { candidates: 0, dispatched: 0, skipped: 0, reaped }
}
// Self-hosted deployments have no subscription infra; `DATA_DRAINS_ENABLED`
// is the global on/off there. Cache per-org so a multi-drain org pays one
// billing lookup.
const enterpriseCache = new Map<string, boolean>()
const isEnterprise = async (orgId: string): Promise<boolean> => {
if (!isBillingEnabled) return true
const cached = enterpriseCache.get(orgId)
if (cached !== undefined) return cached
const result = await isOrganizationOnEnterprisePlan(orgId)
enterpriseCache.set(orgId, result)
return result
}
const queue = await getJobQueue()
let dispatched = 0
let skipped = 0
for (const candidate of candidates) {
let enterprise: boolean
try {
enterprise = await isEnterprise(candidate.organizationId)
} catch (error) {
// A billing-API failure for one org must not abort the whole batch —
// skip this drain and let the next cron tick retry it.
logger.warn('Enterprise check failed; skipping drain', {
drainId: candidate.id,
organizationId: candidate.organizationId,
error,
})
skipped++
continue
}
if (!enterprise) {
skipped++
continue
}
// Conditional claim — re-asserts the due predicate to lose to any other
// dispatcher or manual-run path that's already moved this drain forward.
const claimed = await db
.update(dataDrains)
.set({ lastRunAt: now, updatedAt: now })
.where(and(eq(dataDrains.id, candidate.id), duePredicate))
.returning({ id: dataDrains.id })
if (claimed.length === 0) continue
try {
// concurrencyKey serializes runs of the same drain on the job queue, so
// a manual run-now racing a cron claim can never execute in parallel.
await queue.enqueue(
'run-data-drain',
{ drainId: candidate.id, trigger: 'cron' },
{ concurrencyKey: `data-drain:${candidate.id}` }
)
dispatched++
} catch (error) {
// Roll back the claim so a transient queue outage doesn't delay this
// drain by a full cadence. Scoped to our own claim timestamp so it
// can't trample a concurrent advance. The rollback itself is guarded
// so a DB error here doesn't abort the rest of the batch.
try {
await db
.update(dataDrains)
.set({ lastRunAt: candidate.lastRunAt, updatedAt: now })
.where(and(eq(dataDrains.id, candidate.id), eq(dataDrains.lastRunAt, now)))
} catch (rollbackError) {
logger.error('Failed to roll back data-drain claim after enqueue failure', {
drainId: candidate.id,
enqueueError: toError(error).message,
rollbackError: toError(rollbackError).message,
})
continue
}
logger.error('Failed to enqueue data-drain job; rolled back claim', {
drainId: candidate.id,
error,
})
}
}
logger.info('Data drain dispatch complete', {
candidates: candidates.length,
dispatched,
skipped,
reaped,
})
return { candidates: candidates.length, dispatched, skipped, reaped }
}
+21
View File
@@ -0,0 +1,21 @@
import { decryptSecret, encryptSecret } from '@/lib/core/security/encryption'
/**
* Encrypts an arbitrary JSON-serializable credentials object into a single
* `iv:ciphertext:authTag` string suitable for storage in
* `data_drains.destination_credentials`. Wraps the shared AES-256-GCM helper.
*/
export async function encryptCredentials<T>(plaintext: T): Promise<string> {
const { encrypted } = await encryptSecret(JSON.stringify(plaintext))
return encrypted
}
/**
* Decrypts the inverse of `encryptCredentials`. The caller is expected to run
* the destination's `credentialsSchema` on the result to defend against
* encryption-format drift.
*/
export async function decryptCredentials<T>(ciphertext: string): Promise<T> {
const { decrypted } = await decryptSecret(ciphertext)
return JSON.parse(decrypted) as T
}
+54
View File
@@ -0,0 +1,54 @@
import type { dataDrainRuns, dataDrains } from '@sim/db/schema'
import { type DataDrain, type DataDrainRun, dataDrainSchema } from '@/lib/api/contracts/data-drains'
import { getDestination } from '@/lib/data-drains/destinations/registry'
type DataDrainRow = typeof dataDrains.$inferSelect
type DataDrainRunRow = typeof dataDrainRuns.$inferSelect
/**
* Projects a DB row into the public `DataDrain` wire shape. Strips the
* encrypted credentials column and normalizes timestamps to ISO strings so
* clients receive a stable, schema-validated payload.
*
* The stored `destinationConfig` is JSONB and is re-validated against the
* destination's typed config schema before serialization so unexpected shapes
* surface as errors instead of leaking through the response.
*/
export function serializeDrain(row: DataDrainRow): DataDrain {
const destinationConfig = getDestination(row.destinationType).configSchema.parse(
row.destinationConfig
)
return dataDrainSchema.parse({
id: row.id,
organizationId: row.organizationId,
name: row.name,
source: row.source,
scheduleCadence: row.scheduleCadence,
enabled: row.enabled,
cursor: row.cursor,
lastRunAt: row.lastRunAt ? row.lastRunAt.toISOString() : null,
lastSuccessAt: row.lastSuccessAt ? row.lastSuccessAt.toISOString() : null,
createdBy: row.createdBy,
createdAt: row.createdAt.toISOString(),
updatedAt: row.updatedAt.toISOString(),
destinationType: row.destinationType,
destinationConfig,
})
}
export function serializeDrainRun(row: DataDrainRunRow): DataDrainRun {
return {
id: row.id,
drainId: row.drainId,
status: row.status,
trigger: row.trigger,
startedAt: row.startedAt.toISOString(),
finishedAt: row.finishedAt ? row.finishedAt.toISOString() : null,
rowsExported: row.rowsExported,
bytesWritten: row.bytesWritten,
cursorBefore: row.cursorBefore,
cursorAfter: row.cursorAfter,
error: row.error,
locators: row.locators ?? [],
}
}
+159
View File
@@ -0,0 +1,159 @@
/**
* @vitest-environment node
*/
import { dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing'
import { beforeEach, describe, expect, it, vi } from 'vitest'
vi.mock('@sim/db', () => dbChainMock)
const { mockGetSource, mockGetDestination, mockDecryptCredentials } = vi.hoisted(() => ({
mockGetSource: vi.fn(),
mockGetDestination: vi.fn(),
mockDecryptCredentials: vi.fn(),
}))
vi.mock('@/lib/data-drains/sources/registry', () => ({ getSource: mockGetSource }))
vi.mock('@/lib/data-drains/destinations/registry', () => ({ getDestination: mockGetDestination }))
vi.mock('@/lib/data-drains/encryption', () => ({ decryptCredentials: mockDecryptCredentials }))
import { runDrain } from '@/lib/data-drains/service'
type Row = { id: string; ts: string }
function makeSource(pages: Row[][]) {
return {
type: 'workflow_logs' as const,
displayName: 'Test',
pages: vi.fn(async function* () {
for (const page of pages) yield page
}),
serialize: vi.fn((row: Row) => row),
cursorAfter: vi.fn((row: Row) => JSON.stringify({ ts: row.ts, id: row.id })),
}
}
function makeDestination(
opts: { deliver?: ReturnType<typeof vi.fn>; close?: ReturnType<typeof vi.fn> } = {}
) {
const deliver =
opts.deliver ??
vi.fn(async ({ metadata }: { metadata: { sequence: number } }) => ({
locator: `loc-${metadata.sequence}`,
}))
const close = opts.close ?? vi.fn(async () => {})
return {
type: 's3' as const,
displayName: 'Test',
configSchema: { parse: (v: unknown) => v },
credentialsSchema: { parse: (v: unknown) => v },
openSession: vi.fn(() => ({ deliver, close })),
_deliver: deliver,
_close: close,
}
}
const baseDrain = {
id: 'drain-1',
organizationId: 'org-1',
enabled: true,
source: 'workflow_logs',
destinationType: 's3',
destinationConfig: {},
destinationCredentials: 'enc:blob',
cursor: null,
}
beforeEach(() => {
vi.clearAllMocks()
resetDbChainMock()
mockDecryptCredentials.mockResolvedValue({})
})
describe('runDrain', () => {
it('returns skipped when drain is disabled', async () => {
dbChainMockFns.limit.mockResolvedValueOnce([{ ...baseDrain, enabled: false }])
const result = await runDrain('drain-1', 'manual')
expect(result.status).toBe('skipped')
expect(result.rowsExported).toBe(0)
expect(mockGetSource).not.toHaveBeenCalled()
})
it('throws when drain does not exist', async () => {
dbChainMockFns.limit.mockResolvedValueOnce([])
await expect(runDrain('drain-1', 'manual')).rejects.toThrow(/not found/)
})
it('delivers each page and advances cursor on success', async () => {
dbChainMockFns.limit.mockResolvedValueOnce([baseDrain])
const source = makeSource([
[
{ id: 'r1', ts: '2026-01-01T00:00:00.000Z' },
{ id: 'r2', ts: '2026-01-01T00:00:01.000Z' },
],
[{ id: 'r3', ts: '2026-01-01T00:00:02.000Z' }],
])
const destination = makeDestination()
mockGetSource.mockReturnValue(source)
mockGetDestination.mockReturnValue(destination)
const result = await runDrain('drain-1', 'cron')
expect(result.status).toBe('success')
expect(result.rowsExported).toBe(3)
expect(destination._deliver).toHaveBeenCalledTimes(2)
expect(destination._close).toHaveBeenCalledTimes(1)
expect(result.cursorAfter).toBe(JSON.stringify({ ts: '2026-01-01T00:00:02.000Z', id: 'r3' }))
expect(result.locators).toEqual(['loc-0', 'loc-1'])
// Drain row updated with new cursor; transaction was used.
expect(dbChainMockFns.transaction).toHaveBeenCalled()
const drainUpdate = dbChainMockFns.set.mock.calls.find(
(call) => (call[0] as { cursor?: unknown }).cursor !== undefined
)
expect(drainUpdate?.[0]).toMatchObject({ cursor: result.cursorAfter })
})
it('does not advance drain cursor when delivery fails', async () => {
dbChainMockFns.limit.mockResolvedValueOnce([{ ...baseDrain, cursor: 'prior' }])
const source = makeSource([[{ id: 'r1', ts: '2026-01-01T00:00:00.000Z' }]])
const destination = makeDestination({
deliver: vi.fn(async () => {
throw new Error('boom')
}),
})
mockGetSource.mockReturnValue(source)
mockGetDestination.mockReturnValue(destination)
await expect(runDrain('drain-1', 'cron')).rejects.toThrow('boom')
// Run row updated with status=failed and cursorAfter equal to prior cursor.
const failedUpdate = dbChainMockFns.set.mock.calls.find(
(call) => (call[0] as { status?: unknown }).status === 'failed'
)
expect(failedUpdate?.[0]).toMatchObject({ status: 'failed', cursorAfter: 'prior' })
// No drain-row update with a new cursor field.
const cursorAdvanced = dbChainMockFns.set.mock.calls.some(
(call) => 'cursor' in (call[0] as object)
)
expect(cursorAdvanced).toBe(false)
expect(destination._close).toHaveBeenCalledTimes(1)
})
it('closes session even if close throws', async () => {
dbChainMockFns.limit.mockResolvedValueOnce([baseDrain])
const source = makeSource([])
const destination = makeDestination({
close: vi.fn(async () => {
throw new Error('close-failed')
}),
})
mockGetSource.mockReturnValue(source)
mockGetDestination.mockReturnValue(destination)
const result = await runDrain('drain-1', 'manual')
expect(result.status).toBe('success')
expect(destination._close).toHaveBeenCalled()
})
})
+227
View File
@@ -0,0 +1,227 @@
import { db } from '@sim/db'
import { dataDrainRuns, dataDrains } from '@sim/db/schema'
import { createLogger } from '@sim/logger'
import { toError } from '@sim/utils/errors'
import { generateId } from '@sim/utils/id'
import { eq } from 'drizzle-orm'
import { getDestination } from '@/lib/data-drains/destinations/registry'
import { decryptCredentials } from '@/lib/data-drains/encryption'
import { getSource } from '@/lib/data-drains/sources/registry'
import type { Cursor, RunTrigger } from '@/lib/data-drains/types'
const logger = createLogger('DataDrainsService')
const CHUNK_SIZE = 1000
export interface RunDrainResult {
drainId: string
runId: string
status: 'success' | 'failed' | 'skipped'
rowsExported: number
bytesWritten: number
cursorBefore: Cursor
cursorAfter: Cursor
locators: string[]
error?: string
}
/**
* Orchestrates one drain export. Source-/destination-agnostic — talks only to
* the registry interfaces. The drain's cursor is advanced only when the entire
* run completes successfully so consumers see at-least-once delivery and can
* dedupe on the per-row `id` field.
*/
export async function runDrain(
drainId: string,
trigger: RunTrigger,
options: { signal?: AbortSignal } = {}
): Promise<RunDrainResult> {
const signal = options.signal ?? new AbortController().signal
const [drain] = await db.select().from(dataDrains).where(eq(dataDrains.id, drainId)).limit(1)
if (!drain) {
throw new Error(`Data drain not found: ${drainId}`)
}
if (!drain.enabled) {
return {
drainId,
runId: '',
status: 'skipped',
rowsExported: 0,
bytesWritten: 0,
cursorBefore: drain.cursor,
cursorAfter: drain.cursor,
locators: [],
}
}
const source = getSource(drain.source)
const destination = getDestination(drain.destinationType)
const runId = generateId()
const startedAt = new Date()
await db.insert(dataDrainRuns).values({
id: runId,
drainId,
status: 'running',
trigger,
startedAt,
cursorBefore: drain.cursor,
})
const cursorBefore = drain.cursor
let cursor: Cursor = drain.cursor
let rowsExported = 0
let bytesWritten = 0
let sequence = 0
const locators: string[] = []
/**
* Schema-parse and decrypt happen *after* the run row is created so failures
* in either (e.g. encryption-key rotation, schema drift across versions)
* surface as a `failed` run row in the UI rather than vanishing into the
* background-job logs while `lastRunAt` quietly advances.
*/
let session: ReturnType<typeof destination.openSession> | null = null
try {
const config = destination.configSchema.parse(drain.destinationConfig)
const credentials = destination.credentialsSchema.parse(
await decryptCredentials(drain.destinationCredentials)
)
session = destination.openSession({ config, credentials })
for await (const chunk of source.pages({
organizationId: drain.organizationId,
cursor,
chunkSize: CHUNK_SIZE,
signal,
})) {
const ndjson = `${chunk.map((row) => JSON.stringify(source.serialize(row))).join('\n')}\n`
const body = Buffer.from(ndjson, 'utf8')
const result = await session.deliver({
body,
contentType: 'application/x-ndjson',
metadata: {
drainId,
runId,
source: drain.source,
sequence,
rowCount: chunk.length,
runStartedAt: startedAt,
},
signal,
})
locators.push(result.locator)
rowsExported += chunk.length
bytesWritten += body.byteLength
cursor = source.cursorAfter(chunk[chunk.length - 1])
sequence++
}
if (signal.aborted) {
throw new Error('Data drain run cancelled')
}
const finishedAt = new Date()
await db.transaction(async (tx) => {
await tx
.update(dataDrains)
.set({
cursor,
lastRunAt: finishedAt,
lastSuccessAt: finishedAt,
updatedAt: finishedAt,
})
.where(eq(dataDrains.id, drainId))
await tx
.update(dataDrainRuns)
.set({
status: 'success',
finishedAt,
rowsExported,
bytesWritten,
cursorAfter: cursor,
locators,
error: null,
})
.where(eq(dataDrainRuns.id, runId))
})
logger.info('Data drain run succeeded', {
drainId,
runId,
source: drain.source,
destinationType: drain.destinationType,
rowsExported,
bytesWritten,
chunks: sequence,
})
return {
drainId,
runId,
status: 'success',
rowsExported,
bytesWritten,
cursorBefore,
cursorAfter: cursor,
locators,
}
} catch (error) {
const finishedAt = new Date()
const message = toError(error).message
try {
await db.transaction(async (tx) => {
await tx
.update(dataDrains)
.set({ lastRunAt: finishedAt, updatedAt: finishedAt })
.where(eq(dataDrains.id, drainId))
await tx
.update(dataDrainRuns)
.set({
status: 'failed',
finishedAt,
rowsExported,
bytesWritten,
cursorAfter: cursorBefore,
locators,
error: message.slice(0, 4000),
})
.where(eq(dataDrainRuns.id, runId))
})
} catch (statusError) {
// Best-effort status write — the reaper repairs stuck rows. Log so DB
// outages don't hide behind the original delivery error.
logger.error('Failed to record data drain failure status', {
drainId,
runId,
deliveryError: message,
statusError: toError(statusError).message,
})
}
logger.error('Data drain run failed', {
drainId,
runId,
source: drain.source,
destinationType: drain.destinationType,
error: message,
})
throw error
} finally {
if (session) {
try {
await session.close()
} catch (closeError) {
logger.warn('Data drain session close failed', {
drainId,
runId,
error: toError(closeError).message,
})
}
}
}
}
@@ -0,0 +1,79 @@
import { dbReplica } from '@sim/db'
import { auditLog } from '@sim/db/schema'
import { and, inArray, isNull, or, sql } from 'drizzle-orm'
import {
decodeTimeCursor,
encodeTimeCursor,
timeCursorOrderBy,
timeCursorPredicate,
timeCursorStabilityBound,
} from '@/lib/data-drains/sources/cursor'
import { getOrganizationWorkspaceIds } from '@/lib/data-drains/sources/helpers'
import type { Cursor, DrainSource, SourcePageInput } from '@/lib/data-drains/types'
type AuditLogRow = typeof auditLog.$inferSelect
/**
* Drains audit events scoped to the organization: rows from any of the org's
* workspaces, plus org-level rows (`workspace_id IS NULL`) where
* `metadata->>'organizationId'` matches. Audit-log writers consistently set
* `metadata.organizationId` for org-scoped actions even though the table has
* no dedicated FK column.
*/
async function* pages(input: SourcePageInput): AsyncIterable<AuditLogRow[]> {
const workspaceIds = await getOrganizationWorkspaceIds(input.organizationId)
const orgScopedClause = and(
isNull(auditLog.workspaceId),
sql`${auditLog.metadata}->>'organizationId' = ${input.organizationId}`
)
const scopeClause =
workspaceIds.length === 0
? orgScopedClause
: or(inArray(auditLog.workspaceId, workspaceIds), orgScopedClause)
let cursor = decodeTimeCursor(input.cursor)
while (!input.signal.aborted) {
const cursorClause = timeCursorPredicate(auditLog.createdAt, auditLog.id, cursor)
const rows = await dbReplica
.select()
.from(auditLog)
.where(and(scopeClause, timeCursorStabilityBound(auditLog.createdAt), cursorClause))
.orderBy(...timeCursorOrderBy(auditLog.createdAt, auditLog.id))
.limit(input.chunkSize)
if (rows.length === 0) return
yield rows
const last = rows[rows.length - 1]
cursor = { ts: last.createdAt.toISOString(), id: last.id }
if (rows.length < input.chunkSize) return
}
}
export const auditLogsSource: DrainSource<AuditLogRow> = {
type: 'audit_logs',
displayName: 'Audit logs',
pages,
serialize(row) {
return {
id: row.id,
workspaceId: row.workspaceId,
actorId: row.actorId,
actorName: row.actorName,
actorEmail: row.actorEmail,
action: row.action,
resourceType: row.resourceType,
resourceId: row.resourceId,
resourceName: row.resourceName,
description: row.description,
metadata: row.metadata,
ipAddress: row.ipAddress,
userAgent: row.userAgent,
createdAt: row.createdAt.toISOString(),
}
},
cursorAfter(row): Cursor {
return encodeTimeCursor({ ts: row.createdAt.toISOString(), id: row.id })
},
}
@@ -0,0 +1,131 @@
import { dbReplica } from '@sim/db'
import { copilotChats, copilotMessages } from '@sim/db/schema'
import { and, asc, inArray, isNull, sql } from 'drizzle-orm'
import {
decodeTimeCursor,
encodeTimeCursor,
timeCursorOrderBy,
timeCursorPredicate,
timeCursorStabilityBound,
} from '@/lib/data-drains/sources/cursor'
import { getOrganizationWorkspaceIds } from '@/lib/data-drains/sources/helpers'
import type { Cursor, DrainSource, SourcePageInput } from '@/lib/data-drains/types'
/**
* The transcript no longer lives on `copilot_chats.messages` — it is assembled
* per page from the normalized `copilot_messages` table, so `messages` is the
* ordered list of message `content` objects rather than the DB column.
*/
type CopilotChatRow = Omit<typeof copilotChats.$inferSelect, 'messages'> & {
messages: unknown[]
}
/** Chat metadata columns, excluding the legacy `messages` JSONB. */
const chatColumns = {
id: copilotChats.id,
userId: copilotChats.userId,
workflowId: copilotChats.workflowId,
workspaceId: copilotChats.workspaceId,
type: copilotChats.type,
title: copilotChats.title,
model: copilotChats.model,
conversationId: copilotChats.conversationId,
previewYaml: copilotChats.previewYaml,
planArtifact: copilotChats.planArtifact,
config: copilotChats.config,
resources: copilotChats.resources,
lastSeenAt: copilotChats.lastSeenAt,
pinned: copilotChats.pinned,
createdAt: copilotChats.createdAt,
updatedAt: copilotChats.updatedAt,
} as const
/**
* Cursor is `createdAt` (immutable) but rows themselves are mutable —
* `messages`, `title`, `lastSeenAt`, etc. are updated in-place over the chat's
* lifetime. This means a chat exported once will not be re-exported when its
* messages change. Consumers who need the latest state should periodically
* full-refresh from a separate snapshot job; drains are append-mostly by
* design and `data-drains` is not a CDC pipeline.
*/
async function* pages(input: SourcePageInput): AsyncIterable<CopilotChatRow[]> {
const workspaceIds = await getOrganizationWorkspaceIds(input.organizationId)
if (workspaceIds.length === 0) return
let cursor = decodeTimeCursor(input.cursor)
while (!input.signal.aborted) {
const cursorClause = timeCursorPredicate(copilotChats.createdAt, copilotChats.id, cursor)
const metaRows = await dbReplica
.select(chatColumns)
.from(copilotChats)
.where(
and(
inArray(copilotChats.workspaceId, workspaceIds),
timeCursorStabilityBound(copilotChats.createdAt),
cursorClause
)
)
.orderBy(...timeCursorOrderBy(copilotChats.createdAt, copilotChats.id))
.limit(input.chunkSize)
if (metaRows.length === 0) return
const chatIds = metaRows.map((r) => r.id)
const messageRows = await dbReplica
.select({ chatId: copilotMessages.chatId, content: copilotMessages.content })
.from(copilotMessages)
.where(and(inArray(copilotMessages.chatId, chatIds), isNull(copilotMessages.deletedAt)))
.orderBy(
asc(copilotMessages.chatId),
sql`${copilotMessages.seq} asc nulls last`,
asc(copilotMessages.createdAt),
asc(copilotMessages.id)
)
const messagesByChat = new Map<string, unknown[]>()
for (const m of messageRows) {
const existing = messagesByChat.get(m.chatId)
if (existing) existing.push(m.content)
else messagesByChat.set(m.chatId, [m.content])
}
const rows: CopilotChatRow[] = metaRows.map((r) => ({
...r,
messages: messagesByChat.get(r.id) ?? [],
}))
yield rows
const last = metaRows[metaRows.length - 1]
cursor = { ts: last.createdAt.toISOString(), id: last.id }
if (metaRows.length < input.chunkSize) return
}
}
export const copilotChatsSource: DrainSource<CopilotChatRow> = {
type: 'copilot_chats',
displayName: 'Chats',
pages,
serialize(row) {
return {
id: row.id,
userId: row.userId,
workflowId: row.workflowId,
workspaceId: row.workspaceId,
type: row.type,
title: row.title,
messages: row.messages,
model: row.model,
conversationId: row.conversationId,
previewYaml: row.previewYaml,
planArtifact: row.planArtifact,
config: row.config,
resources: row.resources,
lastSeenAt: row.lastSeenAt ? row.lastSeenAt.toISOString() : null,
createdAt: row.createdAt.toISOString(),
updatedAt: row.updatedAt.toISOString(),
}
},
cursorAfter(row): Cursor {
return encodeTimeCursor({ ts: row.createdAt.toISOString(), id: row.id })
},
}
@@ -0,0 +1,79 @@
import { dbReplica } from '@sim/db'
import { copilotRuns } from '@sim/db/schema'
import { and, inArray, isNotNull } from 'drizzle-orm'
import {
decodeTimeCursor,
encodeTimeCursor,
timeCursorOrderBy,
timeCursorPredicate,
timeCursorStabilityBound,
} from '@/lib/data-drains/sources/cursor'
import { getOrganizationWorkspaceIds } from '@/lib/data-drains/sources/helpers'
import type { Cursor, DrainSource, SourcePageInput } from '@/lib/data-drains/types'
type CopilotRunRow = typeof copilotRuns.$inferSelect
/**
* Cursors on terminal `completedAt` so in-flight runs (mutable `status`,
* `error`, `completedAt`) are not exported until they reach a terminal state.
*/
async function* pages(input: SourcePageInput): AsyncIterable<CopilotRunRow[]> {
const workspaceIds = await getOrganizationWorkspaceIds(input.organizationId)
if (workspaceIds.length === 0) return
let cursor = decodeTimeCursor(input.cursor)
while (!input.signal.aborted) {
const cursorClause = timeCursorPredicate(copilotRuns.completedAt, copilotRuns.id, cursor)
const rows = await dbReplica
.select()
.from(copilotRuns)
.where(
and(
inArray(copilotRuns.workspaceId, workspaceIds),
isNotNull(copilotRuns.completedAt),
timeCursorStabilityBound(copilotRuns.completedAt),
cursorClause
)
)
.orderBy(...timeCursorOrderBy(copilotRuns.completedAt, copilotRuns.id))
.limit(input.chunkSize)
if (rows.length === 0) return
yield rows
const last = rows[rows.length - 1]
cursor = { ts: last.completedAt!.toISOString(), id: last.id }
if (rows.length < input.chunkSize) return
}
}
export const copilotRunsSource: DrainSource<CopilotRunRow> = {
type: 'copilot_runs',
displayName: 'Chat runs',
pages,
serialize(row) {
return {
id: row.id,
executionId: row.executionId,
parentRunId: row.parentRunId,
chatId: row.chatId,
userId: row.userId,
workflowId: row.workflowId,
workspaceId: row.workspaceId,
streamId: row.streamId,
agent: row.agent,
model: row.model,
provider: row.provider,
status: row.status,
requestContext: row.requestContext,
startedAt: row.startedAt.toISOString(),
completedAt: row.completedAt ? row.completedAt.toISOString() : null,
createdAt: row.createdAt.toISOString(),
updatedAt: row.updatedAt.toISOString(),
error: row.error,
}
},
cursorAfter(row): Cursor {
return encodeTimeCursor({ ts: row.completedAt!.toISOString(), id: row.id })
},
}
@@ -0,0 +1,26 @@
/**
* @vitest-environment node
*/
import { describe, expect, it } from 'vitest'
import { decodeTimeCursor, encodeTimeCursor } from '@/lib/data-drains/sources/cursor'
describe('time cursor encoding', () => {
it('round-trips a valid cursor', () => {
const value = { ts: '2026-01-01T00:00:00.000Z', id: 'row-1' }
expect(decodeTimeCursor(encodeTimeCursor(value))).toEqual(value)
})
it('returns null for null input', () => {
expect(decodeTimeCursor(null)).toBeNull()
})
it('returns null for malformed JSON', () => {
expect(decodeTimeCursor('not-json')).toBeNull()
})
it('returns null when shape is wrong', () => {
expect(decodeTimeCursor(JSON.stringify({ ts: 1, id: 'x' }))).toBeNull()
expect(decodeTimeCursor(JSON.stringify({ ts: '2026', id: 5 }))).toBeNull()
expect(decodeTimeCursor(JSON.stringify({}))).toBeNull()
})
})
@@ -0,0 +1,67 @@
import { type SQL, sql } from 'drizzle-orm'
import type { PgColumn } from 'drizzle-orm/pg-core'
import type { Cursor } from '@/lib/data-drains/types'
/**
* Composite cursor for time-ordered tables. Pairs a timestamp with the row's id
* so chunks split across rows that share a timestamp pick up cleanly without
* skipping or duplicating.
*/
export interface TimeCursor {
ts: string
id: string
}
export function encodeTimeCursor(value: TimeCursor): Cursor {
return JSON.stringify(value)
}
export function decodeTimeCursor(cursor: Cursor): TimeCursor | null {
if (!cursor) return null
try {
const parsed = JSON.parse(cursor) as TimeCursor
if (typeof parsed?.ts !== 'string' || typeof parsed?.id !== 'string') return null
return parsed
} catch {
return null
}
}
/**
* Builds a strict-greater-than predicate over a `(timestampCol, idCol)` pair.
*
* Postgres `timestamp` columns store microsecond precision but JS `Date`
* round-trips at millisecond precision, so the cursor only ever captures
* millisecond-truncated timestamps. We compare in millisecond buckets via
* `date_trunc('milliseconds', col)` so the predicate's notion of order matches
* `timeCursorOrderBy` exactly. If ORDER BY used raw microseconds while the
* predicate used millisecond buckets, a row sorted later by µs but with a
* lexicographically earlier id than the cursor row would be skipped forever.
*/
export function timeCursorPredicate(
timestampCol: PgColumn,
idCol: PgColumn,
cursor: TimeCursor | null
): SQL | undefined {
if (!cursor) return undefined
return sql`(date_trunc('milliseconds', ${timestampCol}), ${idCol}) > (${new Date(cursor.ts)}, ${cursor.id})`
}
/**
* ORDER BY fragments paired with `timeCursorPredicate`. Both must agree on
* millisecond bucketing so cursor advancement never skips rows.
*/
export function timeCursorOrderBy(timestampCol: PgColumn, idCol: PgColumn): [SQL, SQL] {
return [sql`date_trunc('milliseconds', ${timestampCol}) asc`, sql`${idCol} asc`]
}
/**
* Excludes rows newer than a short stability window. Timestamp cursors assume
* rows become visible in timestamp order, but out-of-order commits and replica
* lag can surface an earlier-stamped row after the cursor has advanced past it
* — permanently skipping it. Leaving the freshest rows for the next run bounds
* both.
*/
export function timeCursorStabilityBound(timestampCol: PgColumn): SQL {
return sql`${timestampCol} <= now() - interval '5 minutes'`
}
@@ -0,0 +1,15 @@
import { db } from '@sim/db'
import { workspace } from '@sim/db/schema'
import { eq } from 'drizzle-orm'
/**
* Returns the IDs of all workspaces belonging to the organization. Used by
* sources whose underlying tables are workspace-scoped rather than org-scoped.
*/
export async function getOrganizationWorkspaceIds(organizationId: string): Promise<string[]> {
const rows = await db
.select({ id: workspace.id })
.from(workspace)
.where(eq(workspace.organizationId, organizationId))
return rows.map((row) => row.id)
}
@@ -0,0 +1,74 @@
import { dbReplica } from '@sim/db'
import { jobExecutionLogs } from '@sim/db/schema'
import { and, inArray, isNotNull } from 'drizzle-orm'
import {
decodeTimeCursor,
encodeTimeCursor,
timeCursorOrderBy,
timeCursorPredicate,
timeCursorStabilityBound,
} from '@/lib/data-drains/sources/cursor'
import { getOrganizationWorkspaceIds } from '@/lib/data-drains/sources/helpers'
import type { Cursor, DrainSource, SourcePageInput } from '@/lib/data-drains/types'
type JobLogRow = typeof jobExecutionLogs.$inferSelect
/**
* Cursors on terminal `endedAt` so in-flight rows (mutable `status`, `endedAt`,
* `totalDurationMs`, `executionData`) are not exported until finalized.
*/
async function* pages(input: SourcePageInput): AsyncIterable<JobLogRow[]> {
const workspaceIds = await getOrganizationWorkspaceIds(input.organizationId)
if (workspaceIds.length === 0) return
let cursor = decodeTimeCursor(input.cursor)
while (!input.signal.aborted) {
const cursorClause = timeCursorPredicate(jobExecutionLogs.endedAt, jobExecutionLogs.id, cursor)
const rows = await dbReplica
.select()
.from(jobExecutionLogs)
.where(
and(
inArray(jobExecutionLogs.workspaceId, workspaceIds),
isNotNull(jobExecutionLogs.endedAt),
timeCursorStabilityBound(jobExecutionLogs.endedAt),
cursorClause
)
)
.orderBy(...timeCursorOrderBy(jobExecutionLogs.endedAt, jobExecutionLogs.id))
.limit(input.chunkSize)
if (rows.length === 0) return
yield rows
const last = rows[rows.length - 1]
cursor = { ts: last.endedAt!.toISOString(), id: last.id }
if (rows.length < input.chunkSize) return
}
}
export const jobLogsSource: DrainSource<JobLogRow> = {
type: 'job_logs',
displayName: 'Job execution logs',
pages,
serialize(row) {
return {
id: row.id,
executionId: row.executionId,
scheduleId: row.scheduleId,
workspaceId: row.workspaceId,
level: row.level,
status: row.status,
trigger: row.trigger,
startedAt: row.startedAt.toISOString(),
endedAt: row.endedAt ? row.endedAt.toISOString() : null,
totalDurationMs: row.totalDurationMs,
executionData: row.executionData,
cost: row.cost,
createdAt: row.createdAt.toISOString(),
}
},
cursorAfter(row): Cursor {
return encodeTimeCursor({ ts: row.endedAt!.toISOString(), id: row.id })
},
}
@@ -0,0 +1,18 @@
import { auditLogsSource } from '@/lib/data-drains/sources/audit-logs'
import { copilotChatsSource } from '@/lib/data-drains/sources/copilot-chats'
import { copilotRunsSource } from '@/lib/data-drains/sources/copilot-runs'
import { jobLogsSource } from '@/lib/data-drains/sources/job-logs'
import { workflowLogsSource } from '@/lib/data-drains/sources/workflow-logs'
import type { DrainSource, SourceType } from '@/lib/data-drains/types'
export const SOURCE_REGISTRY = {
workflow_logs: workflowLogsSource,
job_logs: jobLogsSource,
audit_logs: auditLogsSource,
copilot_chats: copilotChatsSource,
copilot_runs: copilotRunsSource,
} as const satisfies Record<SourceType, DrainSource>
export function getSource(type: SourceType): DrainSource {
return SOURCE_REGISTRY[type]
}
@@ -0,0 +1,103 @@
import { dbReplica } from '@sim/db'
import { workflowExecutionLogs } from '@sim/db/schema'
import { and, inArray, isNotNull } from 'drizzle-orm'
import { MATERIALIZE_CONCURRENCY, mapWithConcurrency } from '@/lib/core/utils/concurrency'
import {
decodeTimeCursor,
encodeTimeCursor,
timeCursorOrderBy,
timeCursorPredicate,
timeCursorStabilityBound,
} from '@/lib/data-drains/sources/cursor'
import { getOrganizationWorkspaceIds } from '@/lib/data-drains/sources/helpers'
import type { Cursor, DrainSource, SourcePageInput } from '@/lib/data-drains/types'
import { materializeExecutionData } from '@/lib/logs/execution/trace-store'
type WorkflowLogRow = typeof workflowExecutionLogs.$inferSelect
/**
* Cursors on `endedAt` (terminal timestamp) rather than `startedAt`. A running
* row's mutable fields (`endedAt`, `status`, `totalDurationMs`, `executionData`)
* would otherwise be exported mid-flight and never re-emitted with their final
* values. Filtering on `endedAt IS NOT NULL` guarantees rows are immutable
* once visible to the drain.
*/
async function* pages(input: SourcePageInput): AsyncIterable<WorkflowLogRow[]> {
const workspaceIds = await getOrganizationWorkspaceIds(input.organizationId)
if (workspaceIds.length === 0) return
let cursor = decodeTimeCursor(input.cursor)
while (!input.signal.aborted) {
const cursorClause = timeCursorPredicate(
workflowExecutionLogs.endedAt,
workflowExecutionLogs.id,
cursor
)
const rows = await dbReplica
.select()
.from(workflowExecutionLogs)
.where(
and(
inArray(workflowExecutionLogs.workspaceId, workspaceIds),
isNotNull(workflowExecutionLogs.endedAt),
timeCursorStabilityBound(workflowExecutionLogs.endedAt),
cursorClause
)
)
.orderBy(...timeCursorOrderBy(workflowExecutionLogs.endedAt, workflowExecutionLogs.id))
.limit(input.chunkSize)
if (rows.length === 0) return
// Heavy execution data may live in object storage; resolve pointers (bounded
// concurrency) so the drain exports full execution data, not the slim row.
// Use the order-preserving returned array (the util's documented contract)
// and write back, rather than mutating rows inside the mapper.
const materialized = await mapWithConcurrency(rows, MATERIALIZE_CONCURRENCY, (row) =>
materializeExecutionData(row.executionData as Record<string, unknown> | null, {
workspaceId: row.workspaceId,
workflowId: row.workflowId,
executionId: row.executionId,
})
)
for (let i = 0; i < rows.length; i++) {
rows[i].executionData = materialized[i] as WorkflowLogRow['executionData']
}
yield rows
const last = rows[rows.length - 1]
cursor = { ts: last.endedAt!.toISOString(), id: last.id }
if (rows.length < input.chunkSize) return
}
}
export const workflowLogsSource: DrainSource<WorkflowLogRow> = {
type: 'workflow_logs',
displayName: 'Workflow execution logs',
pages,
serialize(row) {
return {
id: row.id,
executionId: row.executionId,
workflowId: row.workflowId,
workspaceId: row.workspaceId,
stateSnapshotId: row.stateSnapshotId,
deploymentVersionId: row.deploymentVersionId,
level: row.level,
status: row.status,
trigger: row.trigger,
startedAt: row.startedAt.toISOString(),
endedAt: row.endedAt ? row.endedAt.toISOString() : null,
totalDurationMs: row.totalDurationMs,
executionData: row.executionData,
// cost_total projection of the usage_log ledger (not the deprecated jsonb).
cost: row.costTotal != null ? { total: Number(row.costTotal) } : null,
files: row.files,
createdAt: row.createdAt.toISOString(),
}
},
cursorAfter(row): Cursor {
return encodeTimeCursor({ ts: row.endedAt!.toISOString(), id: row.id })
},
}
+106
View File
@@ -0,0 +1,106 @@
import type { z } from 'zod'
export const SOURCE_TYPES = [
'workflow_logs',
'job_logs',
'audit_logs',
'copilot_chats',
'copilot_runs',
] as const
export type SourceType = (typeof SOURCE_TYPES)[number]
export const DESTINATION_TYPES = [
's3',
'gcs',
'azure_blob',
'datadog',
'bigquery',
'snowflake',
'webhook',
] as const
export type DestinationType = (typeof DESTINATION_TYPES)[number]
export const CADENCE_TYPES = ['hourly', 'daily'] as const
export type CadenceType = (typeof CADENCE_TYPES)[number]
export const RUN_TRIGGERS = ['cron', 'manual'] as const
export type RunTrigger = (typeof RUN_TRIGGERS)[number]
/**
* Opaque, source-defined cursor. Stored as text in `data_drains.cursor` and
* round-tripped untouched. Sources may encode timestamps, ULIDs, or composite
* keys — the runner never inspects it.
*/
export type Cursor = string | null
export interface SourcePageInput {
organizationId: string
cursor: Cursor
chunkSize: number
signal: AbortSignal
}
export interface DrainSource<TRow = unknown> {
readonly type: SourceType
readonly displayName: string
/**
* Pages rows strictly newer than `cursor` in cursor-ascending order.
* An empty iterator means no new rows.
*/
pages(input: SourcePageInput): AsyncIterable<TRow[]>
/** Stable JSON-safe shape sent to destinations. Public NDJSON contract. */
serialize(row: TRow): Record<string, unknown>
/** Returns the cursor that, when passed back, excludes `row` and everything before it. */
cursorAfter(row: TRow): Cursor
}
export interface DeliveryMetadata {
drainId: string
runId: string
source: SourceType
/** 0-based chunk index within the run. */
sequence: number
rowCount: number
/**
* Wall-clock start of the run. Destinations that partition by date (e.g. S3
* `YYYY/MM/DD` keys) should derive the partition from this so a single run
* lands under one prefix even when delivery crosses a midnight boundary.
*/
runStartedAt: Date
}
interface DeliveryResult {
/** Stable identifier for the written object: e.g. `s3://bucket/key` or `https://host/path`. */
locator: string
}
interface DrainDeliverySession {
deliver(input: {
body: Buffer
contentType: 'application/x-ndjson'
metadata: DeliveryMetadata
signal: AbortSignal
}): Promise<DeliveryResult>
close(): Promise<void>
}
export interface DrainDestination<TConfig = unknown, TCredentials = unknown> {
readonly type: DestinationType
readonly displayName: string
/** Validates non-secret config (bucket, region, prefix, url, ...) at the API boundary. */
readonly configSchema: z.ZodType<TConfig>
/** Validates secret payload separately so it can live in an encrypted column. */
readonly credentialsSchema: z.ZodType<TCredentials>
/** Optional reachability probe used by the "Test connection" UI button. */
test?(input: { config: TConfig; credentials: TCredentials; signal: AbortSignal }): Promise<void>
/**
* Opens a delivery session for one drain run. Lets destinations amortize
* expensive resources (S3Client, keep-alive connections) across all chunks
* in a run instead of rebuilding per chunk. Caller must `close()` when done.
*/
openSession(input: { config: TConfig; credentials: TCredentials }): DrainDeliverySession
}