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
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:
@@ -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()
|
||||
})
|
||||
})
|
||||
@@ -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() {},
|
||||
}
|
||||
},
|
||||
}
|
||||
@@ -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)
|
||||
})
|
||||
})
|
||||
@@ -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 }
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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 ?? [],
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
})
|
||||
})
|
||||
@@ -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 })
|
||||
},
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user