fix(data-drains): address PR review (update schema, S3 date partition, self-hosted gate)

- update body schema: drop the discriminated-union-with-.optional() that silently required destinationType for any non-undefined body. The route already validates destination payloads against the typed configSchema/credentialsSchema for the existing drain, so the contract is now a flat partial — clients can send {enabled:false} without supplying destinationType
- S3 buildKey: partition by run startedAt instead of new Date() per chunk so a single run that crosses midnight still lands under one YYYY/MM/DD prefix
- self-hosted gate: wire DATA_DRAINS_ENABLED into authorizeDrainAccess and the cron dispatcher route so the docs claim ("reserved for server-side feature gating") is actually enforced — mutating endpoints 404 and the dispatcher no-ops when unset

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Author: waleedlatif1

apps/docs/content/docs/en/enterprise/data-drains.mdx

@@ -162,6 +162,6 @@ DATA_DRAINS_ENABLED=true
 NEXT_PUBLIC_DATA_DRAINS_ENABLED=true
 ```
 
-`NEXT_PUBLIC_DATA_DRAINS_ENABLED` shows the **Settings → Enterprise → Data Drains** page in the UI. `DATA_DRAINS_ENABLED` is reserved for server-side feature gating on self-hosted deployments. Both should be set to `true` together.
+`NEXT_PUBLIC_DATA_DRAINS_ENABLED` shows the **Settings → Enterprise → Data Drains** page in the UI. `DATA_DRAINS_ENABLED` gates the server-side mutating endpoints and the cron dispatcher — when unset on a self-hosted deployment, drain create/update/delete/run requests return `404` and the dispatcher is a no-op. Both should be set to `true` together.
 
 Data Drains otherwise rely on the standard Trigger.dev background job infrastructure used elsewhere in Sim — no additional setup is required. The cron dispatcher runs hourly and fans out due drains as background jobs.

apps/sim/app/api/cron/run-data-drains/route.ts

@@ -2,6 +2,7 @@ import { createLogger } from '@sim/logger'
 import { toError } from '@sim/utils/errors'
 import { type NextRequest, NextResponse } from 'next/server'
 import { verifyCronAuth } from '@/lib/auth/internal'
+import { isBillingEnabled, isDataDrainsEnabled } from '@/lib/core/config/feature-flags'
 import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
 import { dispatchDueDrains } from '@/lib/data-drains/dispatcher'
 
@@ -11,6 +12,13 @@ export const GET = withRouteHandler(async (request: NextRequest) => {
   const authError = verifyCronAuth(request, 'Data drain dispatcher')
   if (authError) return authError
 
+  // Self-hosted opt-in: skip dispatch entirely when the deployment hasn't
+  // enabled drains. Sim Cloud (billing enabled) gates per-org by enterprise
+  // plan inside the dispatcher's join.
+  if (!isBillingEnabled && !isDataDrainsEnabled) {
+    return NextResponse.json({ success: true, dispatched: 0, skipped: 'disabled' })
+  }
+
   try {
     const result = await dispatchDueDrains()
     logger.info('Data drain dispatcher run complete', result)

apps/sim/lib/api/contracts/data-drains.ts

@@ -73,23 +73,20 @@ export const createDataDrainBodySchema = z.intersection(
   dataDrainDestinationBodySchema
 )
 
-export const updateDataDrainBodySchema = z.intersection(
-  drainCommonBodyFieldsSchema.partial(),
-  z
-    .discriminatedUnion('destinationType', [
-      z.object({
-        destinationType: z.literal('s3'),
-        destinationConfig: s3ConfigBodySchema.optional(),
-        destinationCredentials: s3CredentialsBodySchema.optional(),
-      }),
-      z.object({
-        destinationType: z.literal('webhook'),
-        destinationConfig: webhookConfigBodySchema.optional(),
-        destinationCredentials: webhookCredentialsBodySchema.optional(),
-      }),
-    ])
-    .optional()
-)
+/**
+ * Update bodies are partial — every field is optional. We deliberately don't
+ * use a discriminated union here: clients sending `{ enabled: false }` should
+ * not be forced to also send `destinationType`. The route validates the
+ * destination payloads against the typed `configSchema` / `credentialsSchema`
+ * for the existing drain's destination type before persisting, so the
+ * structural shape is still enforced — just at the route layer rather than at
+ * the contract boundary.
+ */
+export const updateDataDrainBodySchema = drainCommonBodyFieldsSchema.partial().extend({
+  destinationType: dataDrainDestinationTypeSchema.optional(),
+  destinationConfig: z.record(z.string(), z.unknown()).optional(),
+  destinationCredentials: z.record(z.string(), z.unknown()).optional(),
+})
 
 const drainDestinationResponseSchema = z.discriminatedUnion('destinationType', [
   z.object({

apps/sim/lib/data-drains/access.ts

@@ -4,7 +4,14 @@ 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 } from '@/lib/core/config/feature-flags'
+import { isBillingEnabled, isDataDrainsEnabled } from '@/lib/core/config/feature-flags'
+
+/**
+ * On Sim Cloud (billing enabled), enterprise plan is the gate. On self-hosted
+ * deployments, owners opt in by setting `DATA_DRAINS_ENABLED=true` — and
+ * routes 404 until they do, so enterprise customers don't accidentally expose
+ * mutating endpoints just by deploying a newer image.
+ */
 
 export interface DrainAccessSession {
   user: {
@@ -53,6 +60,15 @@ export async function authorizeDrainAccess(
   }
 
   if (options.requireMutating) {
+    if (!isBillingEnabled && !isDataDrainsEnabled) {
+      return {
+        ok: false,
+        response: NextResponse.json(
+          { error: 'Data Drains are not enabled on this deployment' },
+          { status: 404 }
+        ),
+      }
+    }
     if (memberEntry.role !== 'owner' && memberEntry.role !== 'admin') {
       return {
         ok: false,

apps/sim/lib/data-drains/destinations/s3.test.ts

@@ -47,6 +47,7 @@ describe('s3Destination openSession', () => {
       source: 'workflow_logs' as const,
       sequence,
       rowCount: 1,
+      runStartedAt: new Date('2025-06-15T12:00:00Z'),
     })
     const signal = new AbortController().signal
 
@@ -90,7 +91,14 @@ describe('s3Destination openSession', () => {
     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 },
+      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(
@@ -111,7 +119,14 @@ describe('s3Destination openSession', () => {
       session.deliver({
         body: Buffer.from('x'),
         contentType: 'application/x-ndjson',
-        metadata: { drainId: 'd', runId: 'r', source: 'audit_logs', sequence: 0, rowCount: 1 },
+        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/)

apps/sim/lib/data-drains/destinations/s3.ts

@@ -68,12 +68,20 @@ function normalizePrefix(raw: string | undefined): string {
 
 function buildKey(
   config: S3DestinationConfig,
-  metadata: { drainId: string; runId: string; source: string; sequence: number }
+  metadata: {
+    drainId: string
+    runId: string
+    source: string
+    sequence: number
+    runStartedAt: Date
+  }
 ): string {
-  const now = new Date()
-  const yyyy = now.getUTCFullYear().toString().padStart(4, '0')
-  const mm = (now.getUTCMonth() + 1).toString().padStart(2, '0')
-  const dd = now.getUTCDate().toString().padStart(2, '0')
+  // Partition by the run's start time so all chunks from one run share a
+  // single date prefix even if delivery crosses a midnight boundary.
+  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')
   const prefix = normalizePrefix(config.prefix)
   return `${prefix}${metadata.source}/${metadata.drainId}/${yyyy}/${mm}/${dd}/${metadata.runId}-${seq}.ndjson`

apps/sim/lib/data-drains/service.ts

@@ -101,6 +101,7 @@ export async function runDrain(
           source: drain.source,
           sequence,
           rowCount: chunk.length,
+          runStartedAt: startedAt,
         },
         signal,
       })

apps/sim/lib/data-drains/types.ts

@@ -57,6 +57,12 @@ export interface DeliveryMetadata {
   /** 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
 }
 
 export interface DeliveryResult {