triggerdotdev--trigger.dev/apps/webapp/app/services/realtime/streamBasinProvisioner.server.ts

/**
 * Per-org S2 basin provisioning. Gated by
 * `REALTIME_STREAMS_PER_ORG_BASINS_ENABLED`: when off, all orgs share
 * `REALTIME_STREAMS_S2_BASIN` and this module no-ops.
 *
 * Pure retention-string in / S2-call out. Plan vocabulary lives in the
 * cloud billing app, which calls into the admin sync route to drive
 * provisioning + reconfiguration.
 */
import type { PrismaClientOrTransaction } from "~/db.server";
import { prisma } from "~/db.server";
import { env } from "~/env.server";
import { logger } from "~/services/logger.server";
import { controlPlaneResolver } from "~/v3/runOpsMigration/controlPlaneResolver.server";
import { parseDuration } from "./duration.server";

export function isPerOrgBasinsEnabled(): boolean {
  return env.REALTIME_STREAMS_PER_ORG_BASINS_ENABLED === "true";
}

export function defaultRetention(): string {
  return env.REALTIME_STREAMS_BASIN_DEFAULT_RETENTION;
}

// Org id is a cuid — fixed-length and stable, so the basin name is
// collision-free without truncation. Slugs are user-editable and would
// drift.
export function basinNameForOrg(org: { id: string }): string {
  const prefix = env.REALTIME_STREAMS_BASIN_NAME_PREFIX;
  const envName = env.REALTIME_STREAMS_BASIN_NAME_ENV;
  return `${prefix}-${envName}-org-${org.id}`;
}

type ProvisionInput = {
  id: string;
  retention?: string;
  streamBasinName: string | null | undefined;
};

type ProvisionResult =
  | { kind: "skipped"; reason: "feature-disabled" | "already-provisioned"; basin: string | null }
  | { kind: "provisioned"; basin: string; retention: string };

// Idempotent. Treats S2 409 as success (race with another caller, or
// previous run that crashed after S2 ack but before the column write).
export async function provisionBasinForOrg(
  org: ProvisionInput,
  prismaClient: PrismaClientOrTransaction = prisma
): Promise<ProvisionResult> {
  if (!isPerOrgBasinsEnabled()) {
    return { kind: "skipped", reason: "feature-disabled", basin: null };
  }

  if (org.streamBasinName) {
    return { kind: "skipped", reason: "already-provisioned", basin: org.streamBasinName };
  }

  const accessToken = env.REALTIME_STREAMS_S2_ACCESS_TOKEN;
  if (!accessToken) {
    throw new Error(
      "REALTIME_STREAMS_S2_ACCESS_TOKEN must be set when REALTIME_STREAMS_PER_ORG_BASINS_ENABLED=true"
    );
  }

  const basin = basinNameForOrg(org);
  const retention = org.retention ?? defaultRetention();

  await s2CreateBasin(basin, {
    accessToken,
    retentionPolicy: retention,
    storageClass: env.REALTIME_STREAMS_BASIN_STORAGE_CLASS,
    deleteOnEmptyMinAge: env.REALTIME_STREAMS_BASIN_DELETE_ON_EMPTY_MIN_AGE,
  });

  await prismaClient.organization.update({
    where: { id: org.id },
    data: { streamBasinName: basin },
  });

  // streamBasinName is embedded in every env of the org; drop all its cached env rows.
  controlPlaneResolver.invalidateOrganization(org.id);

  logger.info("[streamBasinProvisioner] provisioned basin for org", {
    orgId: org.id,
    basin,
    retention,
  });

  return { kind: "provisioned", basin, retention };
}

export async function reconfigureBasinForOrg(orgId: string, retention: string): Promise<void> {
  if (!isPerOrgBasinsEnabled()) return;

  const accessToken = env.REALTIME_STREAMS_S2_ACCESS_TOKEN;
  if (!accessToken) {
    throw new Error(
      "REALTIME_STREAMS_S2_ACCESS_TOKEN must be set when REALTIME_STREAMS_PER_ORG_BASINS_ENABLED=true"
    );
  }

  const org = await prisma.organization.findFirst({
    where: { id: orgId },
    select: { id: true, streamBasinName: true },
  });
  if (!org?.streamBasinName) return;

  await s2ReconfigureBasin(org.streamBasinName, { accessToken, retentionPolicy: retention });

  logger.info("[streamBasinProvisioner] reconfigured basin retention", {
    orgId,
    basin: org.streamBasinName,
    retention,
  });
}

type EnsureResult =
  | { kind: "skipped"; reason: "feature-disabled" | "org-not-found" }
  | { kind: "provisioned"; basin: string; retention: string }
  | { kind: "reconfigured"; basin: string; retention: string };

// Idempotent: provisions if the org has no basin, PATCHes retention if
// it does. The single entrypoint the cloud billing app drives — both
// for the live plan-change path and the bulk backfill.
export async function ensureBasinForOrg(orgId: string, retention: string): Promise<EnsureResult> {
  if (!isPerOrgBasinsEnabled()) {
    return { kind: "skipped", reason: "feature-disabled" };
  }

  const org = await prisma.organization.findFirst({
    where: { id: orgId },
    select: { id: true, streamBasinName: true },
  });
  if (!org) return { kind: "skipped", reason: "org-not-found" };

  if (!org.streamBasinName) {
    const result = await provisionBasinForOrg({ id: org.id, streamBasinName: null, retention });
    if (result.kind === "provisioned") {
      return { kind: "provisioned", basin: result.basin, retention: result.retention };
    }
    return { kind: "skipped", reason: "feature-disabled" };
  }

  await reconfigureBasinForOrg(org.id, retention);
  return { kind: "reconfigured", basin: org.streamBasinName, retention };
}

// Inverse of ensureBasinForOrg: nulls the column so future runs/sessions
// land in the shared global basin. The S2 basin lingers; existing streams
// age out on their original retention.
export async function deprovisionBasinForOrg(
  orgId: string
): Promise<{ kind: "deprovisioned" } | { kind: "skipped"; reason: "no-basin" }> {
  const org = await prisma.organization.findFirst({
    where: { id: orgId },
    select: { id: true, streamBasinName: true },
  });
  if (!org?.streamBasinName) return { kind: "skipped", reason: "no-basin" };

  await prisma.organization.update({
    where: { id: org.id },
    data: { streamBasinName: null },
  });

  // streamBasinName is embedded in every env of the org; drop all its cached env rows.
  controlPlaneResolver.invalidateOrganization(org.id);

  logger.info("[streamBasinProvisioner] deprovisioned basin for org", {
    orgId,
    previousBasin: org.streamBasinName,
  });

  return { kind: "deprovisioned" };
}

// S2 REST: POST /v1/basins to create, PATCH /v1/basins/{name} to
// reconfigure. Wire shape takes integer seconds; we accept human strings
// like "7d" / "1y" as env-var ergonomics and parse them here.

type CreateBasinOptions = {
  accessToken: string;
  retentionPolicy: string;
  storageClass: "express" | "standard";
  deleteOnEmptyMinAge: string;
};

async function s2CreateBasin(name: string, opts: CreateBasinOptions): Promise<void> {
  const url = `https://aws.s2.dev/v1/basins`;
  const body = {
    basin: name,
    config: {
      create_stream_on_append: true,
      create_stream_on_read: true,
      default_stream_config: {
        storage_class: opts.storageClass,
        retention_policy: { age: parseDuration(opts.retentionPolicy) },
        delete_on_empty: { min_age_secs: parseDuration(opts.deleteOnEmptyMinAge) },
      },
    },
  };

  const res = await fetch(url, {
    signal: AbortSignal.timeout(10_000),
    method: "POST",
    headers: {
      Authorization: `Bearer ${opts.accessToken}`,
      "Content-Type": "application/json",
    },
    body: JSON.stringify(body),
  });

  // 409 = basin already exists; treat as success (idempotent).
  if (res.ok || res.status === 409) return;

  const text = await res.text().catch(() => "");
  throw new Error(`S2 createBasin failed: ${res.status} ${res.statusText} ${text}`);
}

type ReconfigureBasinOptions = {
  accessToken: string;
  retentionPolicy: string;
};

async function s2ReconfigureBasin(name: string, opts: ReconfigureBasinOptions): Promise<void> {
  const url = `https://aws.s2.dev/v1/basins/${encodeURIComponent(name)}`;
  const body = {
    default_stream_config: {
      retention_policy: { age: parseDuration(opts.retentionPolicy) },
    },
  };

  const res = await fetch(url, {
    signal: AbortSignal.timeout(10_000),
    method: "PATCH",
    headers: {
      Authorization: `Bearer ${opts.accessToken}`,
      "Content-Type": "application/json",
    },
    body: JSON.stringify(body),
  });

  if (res.ok) return;

  const text = await res.text().catch(() => "");
  throw new Error(`S2 reconfigureBasin failed: ${res.status} ${res.statusText} ${text}`);
}