diff --git a/.cursor/mcp.json b/.cursor/mcp.json index da39e4ffafe..c4b06a67630 100644 --- a/.cursor/mcp.json +++ b/.cursor/mcp.json @@ -1,3 +1,7 @@ { - "mcpServers": {} + "mcpServers": { + "linear": { + "url": "https://mcp.linear.app/mcp" + } + } } diff --git a/.server-changes/organization-scoped-clickhouse.md b/.server-changes/organization-scoped-clickhouse.md new file mode 100644 index 00000000000..874b9dc6026 --- /dev/null +++ b/.server-changes/organization-scoped-clickhouse.md @@ -0,0 +1,6 @@ +--- +area: webapp +type: feature +--- + +Organization-scoped ClickHouse routing enables customers with HIPAA and other data security requirements to use dedicated database instances diff --git a/CLAUDE.md b/CLAUDE.md index 0a54cced672..c43bddf323c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -66,6 +66,17 @@ containerTest("should use both", async ({ prisma, redisOptions }) => { }); ``` +## Code Style + +### Imports + +**Prefer static imports over dynamic imports.** Only use dynamic `import()` when: +- Circular dependencies cannot be resolved otherwise +- Code splitting is genuinely needed for performance +- The module must be loaded conditionally at runtime + +Dynamic imports add unnecessary overhead in hot paths and make code harder to analyze. If you find yourself using `await import()`, ask if a regular `import` statement would work instead. + ## Changesets and Server Changes When modifying any public package (`packages/*` or `integrations/*`), add a changeset: diff --git a/apps/webapp/app/env.server.ts b/apps/webapp/app/env.server.ts index 82b73b34853..de02c39a477 100644 --- a/apps/webapp/app/env.server.ts +++ b/apps/webapp/app/env.server.ts @@ -1301,6 +1301,9 @@ const EnvironmentSchema = z EVENTS_CLICKHOUSE_MAX_TRACE_DETAILED_SUMMARY_VIEW_COUNT: z.coerce.number().int().default(5_000), EVENTS_CLICKHOUSE_MAX_LIVE_RELOADING_SETTING: z.coerce.number().int().default(2000), + // Organization data stores registry + ORGANIZATION_DATA_STORES_RELOAD_INTERVAL_MS: z.coerce.number().int().default(5 * 60 * 1000), // 5 minutes + // LLM cost tracking LLM_COST_TRACKING_ENABLED: BoolEnv.default(true), LLM_PRICING_RELOAD_INTERVAL_MS: z.coerce.number().int().default(5 * 60 * 1000), // 5 minutes diff --git a/apps/webapp/app/presenters/v3/ApiRunListPresenter.server.ts b/apps/webapp/app/presenters/v3/ApiRunListPresenter.server.ts index 254ec18d1c0..357d66501b8 100644 --- a/apps/webapp/app/presenters/v3/ApiRunListPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/ApiRunListPresenter.server.ts @@ -9,7 +9,7 @@ import { type Project, type RuntimeEnvironment, type TaskRunStatus } from "@trig import assertNever from "assert-never"; import { z } from "zod"; import { API_VERSIONS, RunStatusUnspecifiedApiVersion } from "~/api/versions"; -import { clickhouseClient } from "~/services/clickhouseInstance.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { logger } from "~/services/logger.server"; import { CoercedDate } from "~/utils/zod"; import { ServiceValidationError } from "~/v3/services/baseService.server"; @@ -259,7 +259,8 @@ export class ApiRunListPresenter extends BasePresenter { options.machines = searchParams["filter[machine]"]; } - const presenter = new NextRunListPresenter(this._replica, clickhouseClient); + const clickhouse = await clickhouseFactory.getClickhouseForOrganization(organizationId, "standard"); + const presenter = new NextRunListPresenter(this._replica, clickhouse); logger.debug("Calling RunListPresenter", { options }); diff --git a/apps/webapp/app/presenters/v3/CreateBulkActionPresenter.server.ts b/apps/webapp/app/presenters/v3/CreateBulkActionPresenter.server.ts index acf511f0f5e..eeb5b3d871e 100644 --- a/apps/webapp/app/presenters/v3/CreateBulkActionPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/CreateBulkActionPresenter.server.ts @@ -1,6 +1,6 @@ import { type PrismaClient } from "@trigger.dev/database"; import { CreateBulkActionSearchParams } from "~/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.bulkaction"; -import { clickhouseClient } from "~/services/clickhouseInstance.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { RunsRepository } from "~/services/runsRepository/runsRepository.server"; import { getRunFiltersFromRequest } from "../RunFilters.server"; import { BasePresenter } from "./basePresenter.server"; @@ -24,8 +24,9 @@ export class CreateBulkActionPresenter extends BasePresenter { Object.fromEntries(new URL(request.url).searchParams) ); + const clickhouse = await clickhouseFactory.getClickhouseForOrganization(organizationId, "standard"); const runsRepository = new RunsRepository({ - clickhouse: clickhouseClient, + clickhouse, prisma: this._replica as PrismaClient, }); diff --git a/apps/webapp/app/presenters/v3/RunPresenter.server.ts b/apps/webapp/app/presenters/v3/RunPresenter.server.ts index 5e8dab2d0b6..cd437b07b91 100644 --- a/apps/webapp/app/presenters/v3/RunPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/RunPresenter.server.ts @@ -3,11 +3,11 @@ import { createTreeFromFlatItems, flattenTree } from "~/components/primitives/Tr import { prisma, type PrismaClient } from "~/db.server"; import { createTimelineSpanEventsFromSpanEvents } from "~/utils/timelineSpanEvents"; import { getUsername } from "~/utils/username"; -import { resolveEventRepositoryForStore } from "~/v3/eventRepository/index.server"; import { SpanSummary } from "~/v3/eventRepository/eventRepository.types"; import { getTaskEventStoreTableForRun } from "~/v3/taskEventStore.server"; import { isFinalRunStatus } from "~/v3/taskStatus"; import { env } from "~/env.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; type Result = Awaited>; export type Run = Result["run"]; @@ -145,10 +145,13 @@ export class RunPresenter { }; } - const eventRepository = resolveEventRepositoryForStore(run.taskEventStore); + const { repository } = await clickhouseFactory.getEventRepositoryForOrganization( + run.taskEventStore, + run.runtimeEnvironment.organizationId + ); // get the events - let traceSummary = await eventRepository.getTraceSummary( + let traceSummary = await repository.getTraceSummary( getTaskEventStoreTableForRun(run), run.runtimeEnvironment.id, run.traceId, @@ -272,7 +275,7 @@ export class RunPresenter { overridesBySpanId: traceSummary.overridesBySpanId, linkedRunIdBySpanId, }, - maximumLiveReloadingSetting: eventRepository.maximumLiveReloadingSetting, + maximumLiveReloadingSetting: repository.maximumLiveReloadingSetting, }; } } diff --git a/apps/webapp/app/presenters/v3/RunTagListPresenter.server.ts b/apps/webapp/app/presenters/v3/RunTagListPresenter.server.ts index e9de368eceb..44d6f2a0747 100644 --- a/apps/webapp/app/presenters/v3/RunTagListPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/RunTagListPresenter.server.ts @@ -1,6 +1,6 @@ import { RunsRepository } from "~/services/runsRepository/runsRepository.server"; import { BasePresenter } from "./basePresenter.server"; -import { clickhouseClient } from "~/services/clickhouseInstance.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { type PrismaClient } from "@trigger.dev/database"; import { timeFilters } from "~/components/runs/v3/SharedFilters"; @@ -37,8 +37,9 @@ export class RunTagListPresenter extends BasePresenter { }: TagListOptions) { const hasFilters = Boolean(name?.trim()); + const clickhouse = await clickhouseFactory.getClickhouseForOrganization(organizationId, "standard"); const runsRepository = new RunsRepository({ - clickhouse: clickhouseClient, + clickhouse, prisma: this._replica as PrismaClient, }); diff --git a/apps/webapp/app/presenters/v3/SpanPresenter.server.ts b/apps/webapp/app/presenters/v3/SpanPresenter.server.ts index 0ea9b37ab7f..353c9b25ebd 100644 --- a/apps/webapp/app/presenters/v3/SpanPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/SpanPresenter.server.ts @@ -21,7 +21,6 @@ import { isFailedRunStatus, isFinalRunStatus } from "~/v3/taskStatus"; import { BasePresenter } from "./basePresenter.server"; import { WaitpointPresenter } from "./WaitpointPresenter.server"; import { engine } from "~/v3/runEngine.server"; -import { resolveEventRepositoryForStore } from "~/v3/eventRepository/index.server"; import { IEventRepository, SpanDetail } from "~/v3/eventRepository/eventRepository.types"; import { safeJsonParse } from "~/utils/json"; import { @@ -30,6 +29,7 @@ import { extractAIToolCallData, extractAIEmbedData, } from "~/components/runs/v3/ai"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; export type PromptSpanData = { slug: string; @@ -42,9 +42,7 @@ export type PromptSpanData = { config?: string; }; -function extractPromptSpanData( - properties: Record -): PromptSpanData | undefined { +function extractPromptSpanData(properties: Record): PromptSpanData | undefined { // Properties come as an unflattened nested object from ClickHouse, // e.g. { prompt: { slug: "...", version: 3, ... } } const prompt = properties.prompt; @@ -132,14 +130,17 @@ export class SpanPresenter extends BasePresenter { const { traceId } = parentRun; - const eventRepository = resolveEventRepositoryForStore(parentRun.taskEventStore); + const { repository } = await clickhouseFactory.getEventRepositoryForOrganization( + parentRun.taskEventStore, + project.organizationId + ); const eventStore = getTaskEventStoreTableForRun(parentRun); const run = await this.getRun({ eventStore, traceId, - eventRepository, + eventRepository: repository, spanId, linkedRunId, createdAt: parentRun.createdAt, @@ -161,7 +162,7 @@ export class SpanPresenter extends BasePresenter { projectId: parentRun.projectId, createdAt: parentRun.createdAt, completedAt: parentRun.completedAt, - eventRepository, + eventRepository: repository, }); if (!span) { @@ -592,10 +593,7 @@ export class SpanPresenter extends BasePresenter { triggeredRuns, aiData: span.properties && typeof span.properties === "object" - ? extractAISpanData( - span.properties as Record, - span.duration / 1_000_000 - ) + ? extractAISpanData(span.properties as Record, span.duration / 1_000_000) : undefined, }; @@ -739,10 +737,7 @@ export class SpanPresenter extends BasePresenter { "ai.streamObject", ]; - if ( - typeof span.message === "string" && - AI_SUMMARY_MESSAGES.includes(span.message) - ) { + if (typeof span.message === "string" && AI_SUMMARY_MESSAGES.includes(span.message)) { const aiSummaryData = extractAISummarySpanData( span.properties as Record, span.duration / 1_000_000 diff --git a/apps/webapp/app/presenters/v3/TaskListPresenter.server.ts b/apps/webapp/app/presenters/v3/TaskListPresenter.server.ts index f1635f23375..8a41d0ef630 100644 --- a/apps/webapp/app/presenters/v3/TaskListPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/TaskListPresenter.server.ts @@ -4,7 +4,7 @@ import { type TaskTriggerSource, } from "@trigger.dev/database"; import { $replica } from "~/db.server"; -import { clickhouseClient } from "~/services/clickhouseInstance.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { type AverageDurations, ClickHouseEnvironmentMetricsRepository, @@ -25,10 +25,7 @@ export type TaskListItem = { export type TaskActivity = DailyTaskActivity[string]; export class TaskListPresenter { - constructor( - private readonly environmentMetricsRepository: EnvironmentMetricsRepository, - private readonly _replica: PrismaClientOrTransaction - ) {} + constructor(private readonly _replica: PrismaClientOrTransaction) {} public async call({ organizationId, @@ -76,9 +73,15 @@ export class TaskListPresenter { const slugs = tasks.map((t) => t.slug); + // Create org-specific environment metrics repository + const clickhouse = await clickhouseFactory.getClickhouseForOrganization(organizationId, "standard"); + const environmentMetricsRepository = new ClickHouseEnvironmentMetricsRepository({ + clickhouse, + }); + // IMPORTANT: Don't await these, we want to return the promises // so we can defer the loading of the data - const activity = this.environmentMetricsRepository.getDailyTaskActivity({ + const activity = environmentMetricsRepository.getDailyTaskActivity({ organizationId, projectId, environmentId, @@ -86,7 +89,7 @@ export class TaskListPresenter { tasks: slugs, }); - const runningStats = this.environmentMetricsRepository.getCurrentRunningStats({ + const runningStats = environmentMetricsRepository.getCurrentRunningStats({ organizationId, projectId, environmentId, @@ -94,7 +97,7 @@ export class TaskListPresenter { tasks: slugs, }); - const durations = this.environmentMetricsRepository.getAverageDurations({ + const durations = environmentMetricsRepository.getAverageDurations({ organizationId, projectId, environmentId, @@ -109,9 +112,5 @@ export class TaskListPresenter { export const taskListPresenter = singleton("taskListPresenter", setupTaskListPresenter); function setupTaskListPresenter() { - const environmentMetricsRepository = new ClickHouseEnvironmentMetricsRepository({ - clickhouse: clickhouseClient, - }); - - return new TaskListPresenter(environmentMetricsRepository, $replica); + return new TaskListPresenter($replica); } diff --git a/apps/webapp/app/presenters/v3/UsagePresenter.server.ts b/apps/webapp/app/presenters/v3/UsagePresenter.server.ts index 2fac95617a6..d312088b6d9 100644 --- a/apps/webapp/app/presenters/v3/UsagePresenter.server.ts +++ b/apps/webapp/app/presenters/v3/UsagePresenter.server.ts @@ -4,7 +4,7 @@ import { getUsage, getUsageSeries } from "~/services/platform.v3.server"; import { createTimeSeriesData } from "~/utils/graphs"; import { BasePresenter } from "./basePresenter.server"; import { DataPoint, linear } from "regression"; -import { clickhouseClient } from "~/services/clickhouseInstance.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; type Options = { organizationId: string; @@ -124,7 +124,8 @@ async function getTaskUsageByOrganization( endOfMonth: Date, replica: PrismaClientOrTransaction ) { - const [queryError, tasks] = await clickhouseClient.taskRuns.getTaskUsageByOrganization({ + const clickhouse = await clickhouseFactory.getClickhouseForOrganization(organizationId, "standard"); + const [queryError, tasks] = await clickhouse.taskRuns.getTaskUsageByOrganization({ startTime: startOfMonth.getTime(), endTime: endOfMonth.getTime(), organizationId, diff --git a/apps/webapp/app/presenters/v3/ViewSchedulePresenter.server.ts b/apps/webapp/app/presenters/v3/ViewSchedulePresenter.server.ts index f0e955fd04d..5341568e6f7 100644 --- a/apps/webapp/app/presenters/v3/ViewSchedulePresenter.server.ts +++ b/apps/webapp/app/presenters/v3/ViewSchedulePresenter.server.ts @@ -1,7 +1,7 @@ import { ScheduleObject } from "@trigger.dev/core/v3"; import { PrismaClient, prisma } from "~/db.server"; import { displayableEnvironment } from "~/models/runtimeEnvironment.server"; -import { clickhouseClient } from "~/services/clickhouseInstance.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { nextScheduledTimestamps } from "~/v3/utils/calculateNextSchedule.server"; import { NextRunListPresenter } from "./NextRunListPresenter.server"; import { scheduleWhereClause } from "~/models/schedules.server"; @@ -75,7 +75,8 @@ export class ViewSchedulePresenter { ? nextScheduledTimestamps(schedule.generatorExpression, schedule.timezone, new Date(), 5) : []; - const runPresenter = new NextRunListPresenter(this.#prismaClient, clickhouseClient); + const clickhouse = await clickhouseFactory.getClickhouseForOrganization(schedule.project.organizationId, "standard"); + const runPresenter = new NextRunListPresenter(this.#prismaClient, clickhouse); const { runs } = await runPresenter.call(schedule.project.organizationId, environmentId, { projectId: schedule.project.id, scheduleId: schedule.id, diff --git a/apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts b/apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts index 9abcdf32215..dc9bf3d1ef0 100644 --- a/apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts @@ -1,5 +1,5 @@ import { isWaitpointOutputTimeout, prettyPrintPacket } from "@trigger.dev/core/v3"; -import { clickhouseClient } from "~/services/clickhouseInstance.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { generateHttpCallbackUrl } from "~/services/httpCallback.server"; import { logger } from "~/services/logger.server"; import { BasePresenter } from "./basePresenter.server"; @@ -79,7 +79,8 @@ export class WaitpointPresenter extends BasePresenter { const connectedRuns: NextRunListItem[] = []; if (connectedRunIds.length > 0) { - const runPresenter = new NextRunListPresenter(this._prisma, clickhouseClient); + const clickhouse = await clickhouseFactory.getClickhouseForOrganization(waitpoint.environment.organizationId, "standard"); + const runPresenter = new NextRunListPresenter(this._prisma, clickhouse); const { runs } = await runPresenter.call( waitpoint.environment.organizationId, environmentId, diff --git a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.dashboards.$dashboardKey/route.tsx b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.dashboards.$dashboardKey/route.tsx index 57b6b71db6f..6c5eb8df842 100644 --- a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.dashboards.$dashboardKey/route.tsx +++ b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.dashboards.$dashboardKey/route.tsx @@ -32,7 +32,7 @@ import { MetricDashboardPresenter, } from "~/presenters/v3/MetricDashboardPresenter.server"; import { PromptPresenter } from "~/presenters/v3/PromptPresenter.server"; -import { clickhouseClient } from "~/services/clickhouseInstance.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { requireUser } from "~/services/session.server"; import { cn } from "~/utils/cn"; import { EnvironmentParamSchema } from "~/utils/pathBuilder"; @@ -75,10 +75,12 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => { const filters = dashboard.filters ?? ["tasks", "queues"]; + const clickhouse = await clickhouseFactory.getClickhouseForOrganization(project.organizationId, "standard"); + // Load distinct models from ClickHouse if the dashboard has a models filter let possibleModels: { model: string; system: string }[] = []; if (filters.includes("models")) { - const queryFn = clickhouseClient.reader.query({ + const queryFn = clickhouse.reader.query({ name: "getDistinctModels", query: `SELECT response_model, any(gen_ai_system) AS gen_ai_system FROM trigger_dev.llm_metrics_v1 WHERE organization_id = {organizationId: String} AND project_id = {projectId: String} AND environment_id = {environmentId: String} AND response_model != '' GROUP BY response_model ORDER BY response_model`, params: z.object({ @@ -98,7 +100,7 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => { } } - const promptPresenter = new PromptPresenter(clickhouseClient); + const promptPresenter = new PromptPresenter(clickhouse); const [possiblePrompts, possibleOperations, possibleProviders] = await Promise.all([ filters.includes("prompts") ? promptPresenter.getDistinctPromptSlugs(project.organizationId, project.id, environment.id) diff --git a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.errors.$fingerprint/route.tsx b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.errors.$fingerprint/route.tsx index f42c73b5ea3..491de3766cd 100644 --- a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.errors.$fingerprint/route.tsx +++ b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.errors.$fingerprint/route.tsx @@ -26,7 +26,11 @@ import { } from "~/presenters/v3/ErrorGroupPresenter.server"; import { type NextRunList } from "~/presenters/v3/NextRunListPresenter.server"; import { $replica } from "~/db.server"; -import { logsClickhouseClient, clickhouseClient } from "~/services/clickhouseInstance.server"; +import { + clickhouseFactory, + getDefaultClickhouseClient, + getDefaultLogsClickhouseClient, +} from "~/services/clickhouse/clickhouseFactory.server"; import { NavBar, PageTitle } from "~/components/primitives/PageHeader"; import { PageBody } from "~/components/layout/AppLayout"; import { @@ -163,6 +167,11 @@ export const action = async ({ request, params }: ActionFunctionArgs) => { let occurrenceCountAtIgnoreTime: number | undefined; if (submission.value.totalOccurrences) { + const clickhouseClient = await clickhouseFactory.getClickhouseForOrganization( + environment.organizationId, + "query" + ); + const qb = clickhouseClient.errors.listQueryBuilder(); qb.where("organization_id = {organizationId: String}", { organizationId: project.organizationId, @@ -236,7 +245,12 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => { const directionRaw = url.searchParams.get("direction") ?? undefined; const direction = directionRaw ? DirectionSchema.parse(directionRaw) : undefined; - const presenter = new ErrorGroupPresenter($replica, logsClickhouseClient, clickhouseClient); + const clickhouseClient = await clickhouseFactory.getClickhouseForOrganization( + environment.organizationId, + "query" + ); + + const presenter = new ErrorGroupPresenter($replica, clickhouseClient, clickhouseClient); const detailPromise = presenter .call(project.organizationId, environment.id, { diff --git a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.errors._index/route.tsx b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.errors._index/route.tsx index e92b5b34644..d7ed8c35798 100644 --- a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.errors._index/route.tsx +++ b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.errors._index/route.tsx @@ -70,7 +70,7 @@ import { type ErrorOccurrences, type ErrorsList as ErrorsListData, } from "~/presenters/v3/ErrorsListPresenter.server"; -import { logsClickhouseClient } from "~/services/clickhouseInstance.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { getCurrentPlan } from "~/services/platform.v3.server"; import { requireUser } from "~/services/session.server"; import { formatNumberCompact } from "~/utils/numberFormatter"; @@ -123,7 +123,11 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => { const plan = await getCurrentPlan(project.organizationId); const retentionLimitDays = plan?.v3Subscription?.plan?.limits.logRetentionDays.number ?? 30; - const presenter = new ErrorsListPresenter($replica, logsClickhouseClient); + const queryClickhouse = await clickhouseFactory.getClickhouseForOrganization( + project.organizationId, + "query" + ); + const presenter = new ErrorsListPresenter($replica, queryClickhouse); const listPromise = presenter .call(project.organizationId, environment.id, { diff --git a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.logs/route.tsx b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.logs/route.tsx index 80a5c6ef232..75155681c30 100644 --- a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.logs/route.tsx +++ b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.logs/route.tsx @@ -16,7 +16,7 @@ import { findEnvironmentBySlug } from "~/models/runtimeEnvironment.server"; import { LogsListPresenter, LogEntry } from "~/presenters/v3/LogsListPresenter.server"; import type { LogLevel } from "~/utils/logUtils"; import { $replica, prisma } from "~/db.server"; -import { logsClickhouseClient } from "~/services/clickhouseInstance.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { NavBar, PageTitle } from "~/components/primitives/PageHeader"; import { PageBody, PageContainer } from "~/components/layout/AppLayout"; import { Suspense, useCallback, useEffect, useMemo, useRef, useState, useTransition } from "react"; @@ -134,7 +134,8 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => { const plan = await getCurrentPlan(project.organizationId); const retentionLimitDays = plan?.v3Subscription?.plan?.limits.logRetentionDays.number ?? 30; - const presenter = new LogsListPresenter($replica, logsClickhouseClient); + const logsClickhouse = await clickhouseFactory.getClickhouseForOrganization(project.organizationId, "logs"); + const presenter = new LogsListPresenter($replica, logsClickhouse); const listPromise = presenter .call(project.organizationId, environment.id, { diff --git a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.models.$modelId/route.tsx b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.models.$modelId/route.tsx index 2d2a0ac850f..05a6cc7bdcd 100644 --- a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.models.$modelId/route.tsx +++ b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.models.$modelId/route.tsx @@ -28,7 +28,7 @@ import type { QueryWidgetConfig } from "~/components/metrics/QueryWidget"; import { findProjectBySlug } from "~/models/project.server"; import { findEnvironmentBySlug } from "~/models/runtimeEnvironment.server"; import { ModelRegistryPresenter } from "~/presenters/v3/ModelRegistryPresenter.server"; -import { clickhouseClient } from "~/services/clickhouseInstance.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { requireUserId } from "~/services/session.server"; import { useOrganization } from "~/hooks/useOrganizations"; import { useProject } from "~/hooks/useProject"; @@ -68,7 +68,8 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => { throw new Response("Environment not found", { status: 404 }); } - const presenter = new ModelRegistryPresenter(clickhouseClient); + const clickhouse = await clickhouseFactory.getClickhouseForOrganization(project.organizationId, "standard"); + const presenter = new ModelRegistryPresenter(clickhouse); const model = await presenter.getModelDetail(modelId); if (!model) { diff --git a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.models._index/route.tsx b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.models._index/route.tsx index 394530b6335..ffcd6ae2172 100644 --- a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.models._index/route.tsx +++ b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.models._index/route.tsx @@ -68,7 +68,7 @@ import { type PopularModel, ModelRegistryPresenter, } from "~/presenters/v3/ModelRegistryPresenter.server"; -import { clickhouseClient } from "~/services/clickhouseInstance.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { requireUserId } from "~/services/session.server"; import { useEnvironment } from "~/hooks/useEnvironment"; import { useOrganization } from "~/hooks/useOrganizations"; @@ -106,7 +106,8 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => { throw new Response("Environment not found", { status: 404 }); } - const presenter = new ModelRegistryPresenter(clickhouseClient); + const clickhouse = await clickhouseFactory.getClickhouseForOrganization(project.organizationId, "standard"); + const presenter = new ModelRegistryPresenter(clickhouse); const catalog = await presenter.getModelCatalog(); const now = new Date(); diff --git a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.models.compare/route.tsx b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.models.compare/route.tsx index 661fb294268..1306eb91943 100644 --- a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.models.compare/route.tsx +++ b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.models.compare/route.tsx @@ -20,7 +20,7 @@ import { type ModelComparisonItem, ModelRegistryPresenter, } from "~/presenters/v3/ModelRegistryPresenter.server"; -import { clickhouseClient } from "~/services/clickhouseInstance.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { requireUserId } from "~/services/session.server"; import { useOrganization } from "~/hooks/useOrganizations"; import { useProject } from "~/hooks/useProject"; @@ -55,7 +55,8 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => { return typedjson({ comparison: [] as ModelComparisonItem[], models: responseModels }); } - const presenter = new ModelRegistryPresenter(clickhouseClient); + const clickhouse = await clickhouseFactory.getClickhouseForOrganization(project.organizationId, "standard"); + const presenter = new ModelRegistryPresenter(clickhouse); const now = new Date(); const sevenDaysAgo = new Date(now.getTime() - 7 * 24 * 60 * 60 * 1000); diff --git a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.prompts.$promptSlug/route.tsx b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.prompts.$promptSlug/route.tsx index 5a953c0199b..de43ea46b23 100644 --- a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.prompts.$promptSlug/route.tsx +++ b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.prompts.$promptSlug/route.tsx @@ -70,7 +70,7 @@ import { findProjectBySlug } from "~/models/project.server"; import { findEnvironmentBySlug } from "~/models/runtimeEnvironment.server"; import { type GenerationRow, PromptPresenter } from "~/presenters/v3/PromptPresenter.server"; import { SpanView } from "~/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.spans.$spanParam/route"; -import { clickhouseClient } from "~/services/clickhouseInstance.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { getResizableSnapshot } from "~/services/resizablePanel.server"; import { requireUserId } from "~/services/session.server"; import { PromptService } from "~/v3/services/promptService.server"; @@ -242,7 +242,8 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => { const startTime = fromTime ? new Date(fromTime) : new Date(Date.now() - periodMs); const endTime = toTime ? new Date(toTime) : new Date(); - const presenter = new PromptPresenter(clickhouseClient); + const clickhouse = await clickhouseFactory.getClickhouseForOrganization(project.organizationId, "standard"); + const presenter = new PromptPresenter(clickhouse); let generations: Awaited>["generations"] = []; let generationsPagination: { next?: string } = {}; try { @@ -273,7 +274,7 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => { // Load distinct filter values and resizable snapshots in parallel const distinctQuery = (col: string, name: string) => - clickhouseClient.reader.query({ + clickhouse.reader.query({ name, query: `SELECT DISTINCT ${col} AS val FROM trigger_dev.llm_metrics_v1 WHERE environment_id = {environmentId: String} AND prompt_slug = {promptSlug: String} AND ${col} != '' ORDER BY val`, params: z.object({ environmentId: z.string(), promptSlug: z.string() }), diff --git a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.prompts._index/route.tsx b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.prompts._index/route.tsx index 02c7cc444b7..b44c5954ffe 100644 --- a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.prompts._index/route.tsx +++ b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.prompts._index/route.tsx @@ -22,7 +22,7 @@ import { useProject } from "~/hooks/useProject"; import { findProjectBySlug } from "~/models/project.server"; import { findEnvironmentBySlug } from "~/models/runtimeEnvironment.server"; import { PromptPresenter } from "~/presenters/v3/PromptPresenter.server"; -import { clickhouseClient } from "~/services/clickhouseInstance.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { requireUserId } from "~/services/session.server"; import { docsPath, EnvironmentParamSchema, v3PromptsPath } from "~/utils/pathBuilder"; import { LinkButton } from "~/components/primitives/Buttons"; @@ -46,7 +46,8 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => { throw new Response("Environment not found", { status: 404 }); } - const presenter = new PromptPresenter(clickhouseClient); + const clickhouse = await clickhouseFactory.getClickhouseForOrganization(project.organizationId, "standard"); + const presenter = new PromptPresenter(clickhouse); const prompts = await presenter.listPrompts(project.id, environment.id); const sparklines = await presenter.getUsageSparklines( diff --git a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam/route.tsx b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam/route.tsx index dc1f3fa2703..ee1f291d154 100644 --- a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam/route.tsx +++ b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam/route.tsx @@ -89,7 +89,7 @@ import { findProjectBySlug } from "~/models/project.server"; import { findEnvironmentBySlug } from "~/models/runtimeEnvironment.server"; import { NextRunListPresenter } from "~/presenters/v3/NextRunListPresenter.server"; import { RunEnvironmentMismatchError, RunPresenter } from "~/presenters/v3/RunPresenter.server"; -import { clickhouseClient } from "~/services/clickhouseInstance.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { getImpersonationId } from "~/services/impersonation.server"; import { logger } from "~/services/logger.server"; import { getResizableSnapshot } from "~/services/resizablePanel.server"; @@ -179,7 +179,8 @@ async function getRunsListFromTableState({ return null; } - const runsListPresenter = new NextRunListPresenter($replica, clickhouseClient); + const clickhouse = await clickhouseFactory.getClickhouseForOrganization(project.organizationId, "standard"); + const runsListPresenter = new NextRunListPresenter($replica, clickhouse); const currentPageResult = await runsListPresenter.call(project.organizationId, environment.id, { userId, projectId: project.id, diff --git a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs._index/route.tsx b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs._index/route.tsx index 9f8cf278bef..96025ea14a3 100644 --- a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs._index/route.tsx +++ b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs._index/route.tsx @@ -42,7 +42,7 @@ import { findProjectBySlug } from "~/models/project.server"; import { findEnvironmentBySlug } from "~/models/runtimeEnvironment.server"; import { getRunFiltersFromRequest } from "~/presenters/RunFilters.server"; import { NextRunListPresenter } from "~/presenters/v3/NextRunListPresenter.server"; -import { clickhouseClient } from "~/services/clickhouseInstance.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { setRootOnlyFilterPreference, uiPreferencesStorage, @@ -85,7 +85,8 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => { const filters = await getRunFiltersFromRequest(request); - const presenter = new NextRunListPresenter($replica, clickhouseClient); + const clickhouse = await clickhouseFactory.getClickhouseForOrganization(project.organizationId, "standard"); + const presenter = new NextRunListPresenter($replica, clickhouse); const list = presenter.call(project.organizationId, environment.id, { userId, projectId: project.id, diff --git a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.test.tasks.$taskParam/route.tsx b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.test.tasks.$taskParam/route.tsx index ee69419e1b7..69ce432f4e0 100644 --- a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.test.tasks.$taskParam/route.tsx +++ b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.test.tasks.$taskParam/route.tsx @@ -74,7 +74,7 @@ import { Dialog, DialogContent, DialogHeader, DialogTrigger } from "~/components import { DialogClose, DialogDescription } from "@radix-ui/react-dialog"; import { FormButtons } from "~/components/primitives/FormButtons"; import { $replica } from "~/db.server"; -import { clickhouseClient } from "~/services/clickhouseInstance.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { RegionsPresenter, type Region } from "~/presenters/v3/RegionsPresenter.server"; import { TestSidebarTabs } from "./TestSidebarTabs"; import { AIPayloadTabContent } from "./AIPayloadTabContent"; @@ -102,7 +102,8 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => { }); } - const presenter = new TestTaskPresenter($replica, clickhouseClient); + const clickhouse = await clickhouseFactory.getClickhouseForOrganization(project.organizationId, "standard"); + const presenter = new TestTaskPresenter($replica, clickhouse); try { const [result, regionsResult] = await Promise.all([ presenter.call({ diff --git a/apps/webapp/app/routes/admin.api.v1.runs-replication.create.ts b/apps/webapp/app/routes/admin.api.v1.runs-replication.create.ts index 483c2d219a1..b8167d462f2 100644 --- a/apps/webapp/app/routes/admin.api.v1.runs-replication.create.ts +++ b/apps/webapp/app/routes/admin.api.v1.runs-replication.create.ts @@ -2,8 +2,8 @@ import { ActionFunctionArgs, json } from "@remix-run/server-runtime"; import { prisma } from "~/db.server"; import { authenticateApiRequestWithPersonalAccessToken } from "~/services/personalAccessToken.server"; import { z } from "zod"; -import { ClickHouse } from "@internal/clickhouse"; import { env } from "~/env.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { RunsReplicationService } from "~/services/runsReplicationService.server"; import { getRunsReplicationGlobal, @@ -62,6 +62,8 @@ export async function action({ request }: ActionFunctionArgs) { const params = CreateRunReplicationServiceParams.parse(await request.json()); + await clickhouseFactory.isReady(); + const service = createRunReplicationService(params); setRunsReplicationGlobal(service); @@ -77,24 +79,23 @@ export async function action({ request }: ActionFunctionArgs) { } function createRunReplicationService(params: CreateRunReplicationServiceParams) { - const clickhouse = new ClickHouse({ - url: env.RUN_REPLICATION_CLICKHOUSE_URL, - name: params.name, - keepAlive: { - enabled: params.keepAliveEnabled, - idleSocketTtl: params.keepAliveIdleSocketTtl, - }, - logLevel: "debug", - compression: { - request: true, - }, - maxOpenConnections: params.maxOpenConnections, - }); + const { + name, + maxFlushConcurrency, + flushIntervalMs, + flushBatchSize, + leaderLockTimeoutMs, + leaderLockExtendIntervalMs, + leaderLockAcquireAdditionalTimeMs, + leaderLockRetryIntervalMs, + ackIntervalSeconds, + waitForAsyncInsert, + } = params; const service = new RunsReplicationService({ - clickhouse: clickhouse, + clickhouseFactory, pgConnectionUrl: env.DATABASE_URL, - serviceName: params.name, + serviceName: name, slotName: env.RUN_REPLICATION_SLOT_NAME, publicationName: env.RUN_REPLICATION_PUBLICATION_NAME, redisOptions: { @@ -106,16 +107,16 @@ function createRunReplicationService(params: CreateRunReplicationServiceParams) enableAutoPipelining: true, ...(env.RUN_REPLICATION_REDIS_TLS_DISABLED === "true" ? {} : { tls: {} }), }, - maxFlushConcurrency: params.maxFlushConcurrency, - flushIntervalMs: params.flushIntervalMs, - flushBatchSize: params.flushBatchSize, - leaderLockTimeoutMs: params.leaderLockTimeoutMs, - leaderLockExtendIntervalMs: params.leaderLockExtendIntervalMs, - leaderLockAcquireAdditionalTimeMs: params.leaderLockAcquireAdditionalTimeMs, - leaderLockRetryIntervalMs: params.leaderLockRetryIntervalMs, - ackIntervalSeconds: params.ackIntervalSeconds, + maxFlushConcurrency, + flushIntervalMs, + flushBatchSize, + leaderLockTimeoutMs, + leaderLockExtendIntervalMs, + leaderLockAcquireAdditionalTimeMs, + leaderLockRetryIntervalMs, + ackIntervalSeconds, logLevel: "debug", - waitForAsyncInsert: params.waitForAsyncInsert, + waitForAsyncInsert, }); return service; diff --git a/apps/webapp/app/routes/admin.api.v1.runs-replication.start.ts b/apps/webapp/app/routes/admin.api.v1.runs-replication.start.ts index a700c4d4f11..9f043f8ccd7 100644 --- a/apps/webapp/app/routes/admin.api.v1.runs-replication.start.ts +++ b/apps/webapp/app/routes/admin.api.v1.runs-replication.start.ts @@ -1,6 +1,7 @@ import { ActionFunctionArgs, json } from "@remix-run/server-runtime"; import { prisma } from "~/db.server"; import { authenticateApiRequestWithPersonalAccessToken } from "~/services/personalAccessToken.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { getRunsReplicationGlobal } from "~/services/runsReplicationGlobal.server"; import { runsReplicationInstance } from "~/services/runsReplicationInstance.server"; @@ -29,6 +30,8 @@ export async function action({ request }: ActionFunctionArgs) { try { const globalService = getRunsReplicationGlobal(); + await clickhouseFactory.isReady(); + if (globalService) { await globalService.start(); } else { diff --git a/apps/webapp/app/routes/admin.data-stores.tsx b/apps/webapp/app/routes/admin.data-stores.tsx new file mode 100644 index 00000000000..033b143151f --- /dev/null +++ b/apps/webapp/app/routes/admin.data-stores.tsx @@ -0,0 +1,400 @@ +import { useState } from "react"; +import { useFetcher } from "@remix-run/react"; +import type { ActionFunctionArgs, LoaderFunctionArgs } from "@remix-run/server-runtime"; +import { redirect } from "@remix-run/server-runtime"; +import { typedjson, useTypedLoaderData } from "remix-typedjson"; +import { z } from "zod"; +import { Button } from "~/components/primitives/Buttons"; +import { + Dialog, + DialogContent, + DialogFooter, + DialogHeader, + DialogTitle, +} from "~/components/primitives/Dialog"; +import { Input } from "~/components/primitives/Input"; +import { Paragraph } from "~/components/primitives/Paragraph"; +import { Popover, PopoverContent, PopoverTrigger } from "~/components/primitives/Popover"; +import { + Table, + TableBlankRow, + TableBody, + TableCell, + TableHeader, + TableHeaderCell, + TableRow, +} from "~/components/primitives/Table"; +import { prisma } from "~/db.server"; +import { requireUser } from "~/services/session.server"; +import { ClickhouseConnectionSchema } from "~/services/clickhouse/clickhouseSecretSchemas.server"; +import { organizationDataStoresRegistry } from "~/services/dataStores/organizationDataStoresRegistryInstance.server"; +import { tryCatch } from "@trigger.dev/core"; + +// --------------------------------------------------------------------------- +// Loader +// --------------------------------------------------------------------------- + +export const loader = async ({ request }: LoaderFunctionArgs) => { + const user = await requireUser(request); + if (!user.admin) throw redirect("/"); + + const dataStores = await prisma.organizationDataStore.findMany({ + orderBy: { createdAt: "desc" }, + }); + + return typedjson({ dataStores }); +}; + +// --------------------------------------------------------------------------- +// Action +// --------------------------------------------------------------------------- + +const AddSchema = z.object({ + _action: z.literal("add"), + key: z.string().min(1), + organizationIds: z.string().min(1), + connectionUrl: z.string().url(), +}); + +const UpdateSchema = z.object({ + _action: z.literal("update"), + key: z.string().min(1), + organizationIds: z.string().min(1), + connectionUrl: z.string().url().optional(), +}); + +const DeleteSchema = z.object({ + _action: z.literal("delete"), + key: z.string().min(1), +}); + +const FormSchema = z.discriminatedUnion("_action", [AddSchema, UpdateSchema, DeleteSchema]); + +export async function action({ request }: ActionFunctionArgs) { + const user = await requireUser(request); + if (!user.admin) throw redirect("/"); + + const formData = await request.formData(); + + const result = FormSchema.safeParse(Object.fromEntries(formData)); + + if (!result.success) { + return typedjson( + { error: result.error.issues.map((i) => i.message).join(", ") }, + { status: 400 } + ); + } + + switch (result.data._action) { + case "add": { + const { key, organizationIds: rawOrgIds, connectionUrl } = result.data; + const organizationIds = rawOrgIds + .split(",") + .map((s) => s.trim()) + .filter(Boolean); + + const config = ClickhouseConnectionSchema.parse({ url: connectionUrl }); + + const [error, _] = await tryCatch( + organizationDataStoresRegistry.addDataStore({ + key, + kind: "CLICKHOUSE", + organizationIds, + config, + }) + ); + + if (error) { + return typedjson({ error: error.message }, { status: 400 }); + } + + return typedjson({ success: true }); + } + case "update": { + const { key, organizationIds: rawOrgIds, connectionUrl } = result.data; + const organizationIds = rawOrgIds + .split(",") + .map((s) => s.trim()) + .filter(Boolean); + + const config = connectionUrl + ? ClickhouseConnectionSchema.parse({ url: connectionUrl }) + : undefined; + + const [error, _] = await tryCatch( + organizationDataStoresRegistry.updateDataStore({ + key, + kind: "CLICKHOUSE", + organizationIds, + config, + }) + ); + + if (error) { + return typedjson({ error: error.message }, { status: 400 }); + } + + return typedjson({ success: true }); + } + case "delete": { + const { key } = result.data; + + const [error, _] = await tryCatch( + organizationDataStoresRegistry.deleteDataStore({ + key, + kind: "CLICKHOUSE", + }) + ); + + if (error) { + return typedjson({ error: error.message }, { status: 400 }); + } + + return typedjson({ success: true }); + } + default: { + return typedjson({ error: "Unknown action" }, { status: 400 }); + } + } +} + +// --------------------------------------------------------------------------- +// Component +// --------------------------------------------------------------------------- + +export default function AdminDataStoresRoute() { + const { dataStores } = useTypedLoaderData(); + const [addOpen, setAddOpen] = useState(false); + + return ( +
+
+
+ + {dataStores.length} data store{dataStores.length !== 1 ? "s" : ""} + + +
+ + + + + Key + Kind + Organizations + Created + Updated + + Actions + + + + + {dataStores.length === 0 ? ( + + No data stores configured + + ) : ( + dataStores.map((ds) => ( + + + {ds.key} + + + + {ds.kind} + + + + + {ds.organizationIds.length} org{ds.organizationIds.length !== 1 ? "s" : ""} + + {ds.organizationIds.length > 0 && ( + + ({ds.organizationIds.slice(0, 2).join(", ")} + {ds.organizationIds.length > 2 + ? ` +${ds.organizationIds.length - 2} more` + : ""} + ) + + )} + + + + {new Date(ds.createdAt).toLocaleString()} + + + + + {new Date(ds.updatedAt).toLocaleString()} + + + + + + + )) + )} + +
+
+ + +
+ ); +} + +// --------------------------------------------------------------------------- +// Delete button with popover confirmation +// --------------------------------------------------------------------------- + +function DeleteButton({ name }: { name: string }) { + const [open, setOpen] = useState(false); + const fetcher = useFetcher<{ success?: boolean; error?: string }>(); + const isDeleting = fetcher.state !== "idle"; + + return ( + + + + + + + Delete {name}? + + + This will remove the data store and its secret. Organizations using it will fall back to + the default ClickHouse instance. + +
+ + setOpen(false)}> + + + + +
+ + + ); +} + +// --------------------------------------------------------------------------- +// Add data store dialog +// --------------------------------------------------------------------------- + +function AddDataStoreDialog({ + open, + onOpenChange, +}: { + open: boolean; + onOpenChange: (open: boolean) => void; +}) { + const fetcher = useFetcher<{ success?: boolean; error?: string }>(); + const isSubmitting = fetcher.state !== "idle"; + + // Close dialog on success + if (fetcher.data?.success && open) { + onOpenChange(false); + } + + return ( + + + + Add data store + + + + + +
+ + +

+ Unique identifier for this data store. Used as the secret key prefix. +

+
+ +
+ + +
+ +
+ + +

Comma-separated organization IDs.

+
+ +
+ + +

+ Stored encrypted in SecretStore. Never logged or displayed again. +

+
+ + {fetcher.data?.error &&

{fetcher.data.error}

} + + + + + + + + + ); +} diff --git a/apps/webapp/app/routes/admin.tsx b/apps/webapp/app/routes/admin.tsx index 4cd2deca533..3b0858f899a 100644 --- a/apps/webapp/app/routes/admin.tsx +++ b/apps/webapp/app/routes/admin.tsx @@ -44,6 +44,10 @@ export default function Page() { label: "Notifications", to: "/admin/notifications", }, + { + label: "Data Stores", + to: "/admin/data-stores", + }, ]} layoutId={"admin"} /> diff --git a/apps/webapp/app/routes/api.v1.prompts.$slug.ts b/apps/webapp/app/routes/api.v1.prompts.$slug.ts index 32ea1525c14..919133d3cbd 100644 --- a/apps/webapp/app/routes/api.v1.prompts.$slug.ts +++ b/apps/webapp/app/routes/api.v1.prompts.$slug.ts @@ -2,7 +2,7 @@ import { json } from "@remix-run/server-runtime"; import { z } from "zod"; import { prisma } from "~/db.server"; import { PromptPresenter } from "~/presenters/v3/PromptPresenter.server"; -import { clickhouseClient } from "~/services/clickhouseInstance.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { createActionApiRoute, createLoaderApiRoute, @@ -33,6 +33,13 @@ export const loader = createLoaderApiRoute( slug: params.slug, }, }, + include: { + project: { + select: { + organizationId: true, + }, + }, + }, }); }, authorization: { @@ -46,7 +53,8 @@ export const loader = createLoaderApiRoute( return json({ error: "Prompt not found" }, { status: 404 }); } - const presenter = new PromptPresenter(clickhouseClient); + const clickhouse = await clickhouseFactory.getClickhouseForOrganization(prompt.project.organizationId, "standard"); + const presenter = new PromptPresenter(clickhouse); const version = await presenter.resolveVersion(prompt.id, { version: searchParams.version, label: searchParams.label, @@ -117,7 +125,8 @@ const { action } = createActionApiRoute( return json({ error: "Prompt not found" }, { status: 404 }); } - const presenter = new PromptPresenter(clickhouseClient); + const clickhouse = await clickhouseFactory.getClickhouseForOrganization(authentication.environment.organizationId, "standard"); + const presenter = new PromptPresenter(clickhouse); const version = await presenter.resolveVersion(prompt.id, { version: body.version, label: body.label, diff --git a/apps/webapp/app/routes/api.v1.prompts.$slug.versions.ts b/apps/webapp/app/routes/api.v1.prompts.$slug.versions.ts index c40b3e62dbf..8f953d1a48c 100644 --- a/apps/webapp/app/routes/api.v1.prompts.$slug.versions.ts +++ b/apps/webapp/app/routes/api.v1.prompts.$slug.versions.ts @@ -2,7 +2,7 @@ import { json } from "@remix-run/server-runtime"; import { z } from "zod"; import { prisma } from "~/db.server"; import { PromptPresenter } from "~/presenters/v3/PromptPresenter.server"; -import { clickhouseClient } from "~/services/clickhouseInstance.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { createLoaderApiRoute } from "~/services/routeBuilders/apiBuilder.server"; const ParamsSchema = z.object({ @@ -23,6 +23,13 @@ export const loader = createLoaderApiRoute( slug: params.slug, }, }, + include: { + project: { + select: { + organizationId: true, + }, + }, + }, }); }, authorization: { @@ -36,7 +43,8 @@ export const loader = createLoaderApiRoute( return json({ error: "Prompt not found" }, { status: 404 }); } - const presenter = new PromptPresenter(clickhouseClient); + const clickhouse = await clickhouseFactory.getClickhouseForOrganization(prompt.project.organizationId, "standard"); + const presenter = new PromptPresenter(clickhouse); const versions = await presenter.listVersions(prompt.id); return json({ diff --git a/apps/webapp/app/routes/api.v1.prompts._index.ts b/apps/webapp/app/routes/api.v1.prompts._index.ts index ccbc0ec38d0..a6ad065c245 100644 --- a/apps/webapp/app/routes/api.v1.prompts._index.ts +++ b/apps/webapp/app/routes/api.v1.prompts._index.ts @@ -1,6 +1,6 @@ import { json } from "@remix-run/server-runtime"; import { PromptPresenter } from "~/presenters/v3/PromptPresenter.server"; -import { clickhouseClient } from "~/services/clickhouseInstance.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { createLoaderApiRoute } from "~/services/routeBuilders/apiBuilder.server"; export const loader = createLoaderApiRoute( @@ -15,7 +15,8 @@ export const loader = createLoaderApiRoute( }, }, async ({ authentication }) => { - const presenter = new PromptPresenter(clickhouseClient); + const clickhouse = await clickhouseFactory.getClickhouseForOrganization(authentication.environment.organizationId, "standard"); + const presenter = new PromptPresenter(clickhouse); const prompts = await presenter.listPrompts( authentication.environment.projectId, authentication.environment.id diff --git a/apps/webapp/app/routes/otel.v1.logs.ts b/apps/webapp/app/routes/otel.v1.logs.ts index a05ddd24cf2..1dc7c07c16c 100644 --- a/apps/webapp/app/routes/otel.v1.logs.ts +++ b/apps/webapp/app/routes/otel.v1.logs.ts @@ -4,12 +4,13 @@ import { otlpExporter } from "~/v3/otlpExporter.server"; export async function action({ request }: ActionFunctionArgs) { try { + const exporter = await otlpExporter; const contentType = request.headers.get("content-type")?.toLowerCase() ?? ""; if (contentType.startsWith("application/json")) { const body = await request.json(); - const exportResponse = await otlpExporter.exportLogs(body as ExportLogsServiceRequest); + const exportResponse = await exporter.exportLogs(body as ExportLogsServiceRequest); return json(exportResponse, { status: 200 }); } else if (contentType.startsWith("application/x-protobuf")) { @@ -17,7 +18,7 @@ export async function action({ request }: ActionFunctionArgs) { const exportRequest = ExportLogsServiceRequest.decode(new Uint8Array(buffer)); - const exportResponse = await otlpExporter.exportLogs(exportRequest); + const exportResponse = await exporter.exportLogs(exportRequest); return new Response(ExportLogsServiceResponse.encode(exportResponse).finish(), { status: 200, diff --git a/apps/webapp/app/routes/otel.v1.metrics.ts b/apps/webapp/app/routes/otel.v1.metrics.ts index 5529f9310ec..9a09cb18233 100644 --- a/apps/webapp/app/routes/otel.v1.metrics.ts +++ b/apps/webapp/app/routes/otel.v1.metrics.ts @@ -7,12 +7,13 @@ import { otlpExporter } from "~/v3/otlpExporter.server"; export async function action({ request }: ActionFunctionArgs) { try { + const exporter = await otlpExporter; const contentType = request.headers.get("content-type")?.toLowerCase() ?? ""; if (contentType.startsWith("application/json")) { const body = await request.json(); - const exportResponse = await otlpExporter.exportMetrics( + const exportResponse = await exporter.exportMetrics( body as ExportMetricsServiceRequest ); @@ -22,7 +23,7 @@ export async function action({ request }: ActionFunctionArgs) { const exportRequest = ExportMetricsServiceRequest.decode(new Uint8Array(buffer)); - const exportResponse = await otlpExporter.exportMetrics(exportRequest); + const exportResponse = await exporter.exportMetrics(exportRequest); return new Response(ExportMetricsServiceResponse.encode(exportResponse).finish(), { status: 200, diff --git a/apps/webapp/app/routes/otel.v1.traces.ts b/apps/webapp/app/routes/otel.v1.traces.ts index 609b72c0465..8e974c7b1dd 100644 --- a/apps/webapp/app/routes/otel.v1.traces.ts +++ b/apps/webapp/app/routes/otel.v1.traces.ts @@ -4,12 +4,13 @@ import { otlpExporter } from "~/v3/otlpExporter.server"; export async function action({ request }: ActionFunctionArgs) { try { + const exporter = await otlpExporter; const contentType = request.headers.get("content-type")?.toLowerCase() ?? ""; if (contentType.startsWith("application/json")) { const body = await request.json(); - const exportResponse = await otlpExporter.exportTraces(body as ExportTraceServiceRequest); + const exportResponse = await exporter.exportTraces(body as ExportTraceServiceRequest); return json(exportResponse, { status: 200 }); } else if (contentType.startsWith("application/x-protobuf")) { @@ -17,7 +18,7 @@ export async function action({ request }: ActionFunctionArgs) { const exportRequest = ExportTraceServiceRequest.decode(new Uint8Array(buffer)); - const exportResponse = await otlpExporter.exportTraces(exportRequest); + const exportResponse = await exporter.exportTraces(exportRequest); return new Response(ExportTraceServiceResponse.encode(exportResponse).finish(), { status: 200, diff --git a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.logs.$logId.tsx b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.logs.$logId.tsx index f862ced6b05..ae552e96eb9 100644 --- a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.logs.$logId.tsx +++ b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.logs.$logId.tsx @@ -1,7 +1,7 @@ import { type LoaderFunctionArgs } from "@remix-run/server-runtime"; import { typedjson } from "remix-typedjson"; import { z } from "zod"; -import { logsClickhouseClient } from "~/services/clickhouseInstance.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { requireUserId } from "~/services/session.server"; import { LogDetailPresenter } from "~/presenters/v3/LogDetailPresenter.server"; import { findProjectBySlug } from "~/models/project.server"; @@ -43,7 +43,8 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => { const [traceId, spanId, , startTime] = parts; - const presenter = new LogDetailPresenter($replica, logsClickhouseClient); + const logsClickhouse = await clickhouseFactory.getClickhouseForOrganization(project.organizationId, "logs"); + const presenter = new LogDetailPresenter($replica, logsClickhouse); let result; try { diff --git a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.logs.ts b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.logs.ts index 66ddebe4e2a..7bb0db0d54b 100644 --- a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.logs.ts +++ b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.logs.ts @@ -6,7 +6,7 @@ import { findProjectBySlug } from "~/models/project.server"; import { findEnvironmentBySlug } from "~/models/runtimeEnvironment.server"; import { LogsListPresenter, type LogLevel, LogsListOptionsSchema } from "~/presenters/v3/LogsListPresenter.server"; import { $replica } from "~/db.server"; -import { logsClickhouseClient } from "~/services/clickhouseInstance.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { getCurrentPlan } from "~/services/platform.v3.server"; // Valid log levels for filtering @@ -69,7 +69,8 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => { retentionLimitDays, }) as any; // Validated by LogsListOptionsSchema at runtime - const presenter = new LogsListPresenter($replica, logsClickhouseClient); + const logsClickhouse = await clickhouseFactory.getClickhouseForOrganization(project.organizationId, "logs"); + const presenter = new LogsListPresenter($replica, logsClickhouse); const result = await presenter.call(project.organizationId, environment.id, options); return json({ diff --git a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.prompts.$promptSlug.generations.ts b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.prompts.$promptSlug.generations.ts index 77a55ec3f0b..5188d8ccdfe 100644 --- a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.prompts.$promptSlug.generations.ts +++ b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.prompts.$promptSlug.generations.ts @@ -6,7 +6,7 @@ import { EnvironmentParamSchema } from "~/utils/pathBuilder"; import { parsePeriodToMs } from "~/utils/periods"; import { findProjectBySlug } from "~/models/project.server"; import { findEnvironmentBySlug } from "~/models/runtimeEnvironment.server"; -import { clickhouseClient } from "~/services/clickhouseInstance.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { PromptPresenter, type GenerationRow, @@ -59,7 +59,8 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => { const operations = url.searchParams.getAll("operations").filter(Boolean); const providers = url.searchParams.getAll("providers").filter(Boolean); - const presenter = new PromptPresenter(clickhouseClient); + const clickhouse = await clickhouseFactory.getClickhouseForOrganization(project.organizationId, "standard"); + const presenter = new PromptPresenter(clickhouse); const result = await presenter.listGenerations({ environmentId: environment.id, promptSlug, diff --git a/apps/webapp/app/runEngine/concerns/traceEvents.server.ts b/apps/webapp/app/runEngine/concerns/traceEvents.server.ts index cb2eaa30a58..8c9029e2d60 100644 --- a/apps/webapp/app/runEngine/concerns/traceEvents.server.ts +++ b/apps/webapp/app/runEngine/concerns/traceEvents.server.ts @@ -10,6 +10,7 @@ export class DefaultTraceEventsConcern implements TraceEventConcern { parentStore: string | undefined ): Promise<{ repository: IEventRepository; store: string }> { return await getEventRepository( + request.environment.organization.id, request.environment.organization.featureFlags as Record, parentStore ); @@ -162,18 +163,15 @@ export class DefaultTraceEventsConcern implements TraceEventConcern { }, async (event, traceContext, traceparent) => { // Log a message about the debounced trigger - await repository.recordEvent( - `Debounced: using existing run with key "${debounceKey}"`, - { - taskSlug: request.taskId, - environment: request.environment, - attributes: { - runId: existingRun.friendlyId, - }, - context: request.options?.traceContext, - parentId: event.spanId, - } - ); + await repository.recordEvent(`Debounced: using existing run with key "${debounceKey}"`, { + taskSlug: request.taskId, + environment: request.environment, + attributes: { + runId: existingRun.friendlyId, + }, + context: request.options?.traceContext, + parentId: event.spanId, + }); return await callback( { diff --git a/apps/webapp/app/runEngine/services/triggerFailedTask.server.ts b/apps/webapp/app/runEngine/services/triggerFailedTask.server.ts index cdcfa63ff0b..4e58862a7a2 100644 --- a/apps/webapp/app/runEngine/services/triggerFailedTask.server.ts +++ b/apps/webapp/app/runEngine/services/triggerFailedTask.server.ts @@ -68,6 +68,7 @@ export class TriggerFailedTaskService { try { const { repository, store } = await getEventRepository( + request.environment.organization.id, request.environment.organization.featureFlags as Record, undefined ); @@ -75,11 +76,11 @@ export class TriggerFailedTaskService { // Resolve parent run for rootTaskRunId and depth (same as triggerTask.server.ts) const parentRun = request.parentRunId ? await this.prisma.taskRun.findFirst({ - where: { - id: RunId.fromFriendlyId(request.parentRunId), - runtimeEnvironmentId: request.environment.id, - }, - }) + where: { + id: RunId.fromFriendlyId(request.parentRunId), + runtimeEnvironmentId: request.environment.id, + }, + }) : undefined; const depth = parentRun ? parentRun.depth + 1 : 0; @@ -110,18 +111,18 @@ export class TriggerFailedTaskService { // resolveQueueProperties requires the worker to be passed when lockToVersion is present. const lockedToBackgroundWorker = bodyOptions?.lockToVersion ? await this.prisma.backgroundWorker.findFirst({ - where: { - projectId: request.environment.projectId, - runtimeEnvironmentId: request.environment.id, - version: bodyOptions.lockToVersion, - }, - select: { - id: true, - version: true, - sdkVersion: true, - cliVersion: true, - }, - }) + where: { + projectId: request.environment.projectId, + runtimeEnvironmentId: request.environment.id, + version: bodyOptions.lockToVersion, + }, + select: { + id: true, + version: true, + sdkVersion: true, + cliVersion: true, + }, + }) : undefined; const resolved = await queueConcern.resolveQueueProperties( @@ -267,9 +268,7 @@ export class TriggerFailedTaskService { }, taskIdentifier: opts.taskId, payload: - typeof opts.payload === "string" - ? opts.payload - : JSON.stringify(opts.payload ?? ""), + typeof opts.payload === "string" ? opts.payload : JSON.stringify(opts.payload ?? ""), payloadType: opts.payloadType ?? "application/json", error: { type: "INTERNAL_ERROR" as const, diff --git a/apps/webapp/app/services/admin/missingLlmModels.server.ts b/apps/webapp/app/services/admin/missingLlmModels.server.ts index 7ce6bc2ab7e..07e6160ee03 100644 --- a/apps/webapp/app/services/admin/missingLlmModels.server.ts +++ b/apps/webapp/app/services/admin/missingLlmModels.server.ts @@ -1,4 +1,4 @@ -import { adminClickhouseClient } from "~/services/clickhouseInstance.server"; +import { getAdminClickhouse } from "~/services/clickhouse/clickhouseFactory.server"; import { llmPricingRegistry } from "~/v3/llmPricingRegistry.server"; export type MissingLlmModel = { @@ -13,8 +13,10 @@ export async function getMissingLlmModels(opts: { const lookbackHours = opts.lookbackHours ?? 24; const since = new Date(Date.now() - lookbackHours * 60 * 60 * 1000); + const adminClickhouse = getAdminClickhouse(); + // queryBuilderFast returns a factory function — call it to get the builder - const createBuilder = adminClickhouseClient.reader.queryBuilderFast<{ + const createBuilder = adminClickhouse.reader.queryBuilderFast<{ model: string; system: string; cnt: string; @@ -93,7 +95,9 @@ export async function getMissingModelSamples(opts: { const limit = opts.limit ?? 10; const since = new Date(Date.now() - lookbackHours * 60 * 60 * 1000); - const createBuilder = adminClickhouseClient.reader.queryBuilderFast({ + const adminClickhouse = getAdminClickhouse(); + + const createBuilder = adminClickhouse.reader.queryBuilderFast({ name: "missingModelSamples", table: "trigger_dev.task_events_v2", columns: [ diff --git a/apps/webapp/app/services/clickhouse/clickhouseFactory.server.ts b/apps/webapp/app/services/clickhouse/clickhouseFactory.server.ts new file mode 100644 index 00000000000..92df7de2b54 --- /dev/null +++ b/apps/webapp/app/services/clickhouse/clickhouseFactory.server.ts @@ -0,0 +1,389 @@ +import { ClickHouse } from "@internal/clickhouse"; +import { createHash } from "crypto"; +import { ClickhouseEventRepository } from "~/v3/eventRepository/clickhouseEventRepository.server"; +import { env } from "~/env.server"; +import { singleton } from "~/utils/singleton"; +import { organizationDataStoresRegistry } from "~/services/dataStores/organizationDataStoresRegistryInstance.server"; +import type { OrganizationDataStoresRegistry } from "~/services/dataStores/organizationDataStoresRegistry.server"; +import { type IEventRepository } from "~/v3/eventRepository/eventRepository.types"; + +// --------------------------------------------------------------------------- +// Default clients (singleton per process) +// --------------------------------------------------------------------------- + +const defaultClickhouseClient = singleton("clickhouseClient", initializeClickhouseClient); + +function initializeClickhouseClient() { + const url = new URL(env.CLICKHOUSE_URL); + url.searchParams.delete("secure"); + + console.log(`🗃️ Clickhouse service enabled to host ${url.host}`); + + return new ClickHouse({ + url: url.toString(), + name: "clickhouse-instance", + keepAlive: { + enabled: env.CLICKHOUSE_KEEP_ALIVE_ENABLED === "1", + idleSocketTtl: env.CLICKHOUSE_KEEP_ALIVE_IDLE_SOCKET_TTL_MS, + }, + logLevel: env.CLICKHOUSE_LOG_LEVEL, + compression: { request: true }, + maxOpenConnections: env.CLICKHOUSE_MAX_OPEN_CONNECTIONS, + }); +} + +const defaultLogsClickhouseClient = singleton( + "logsClickhouseClient", + initializeLogsClickhouseClient +); + +function initializeLogsClickhouseClient() { + if (!env.LOGS_CLICKHOUSE_URL) { + throw new Error("LOGS_CLICKHOUSE_URL is not set"); + } + + const url = new URL(env.LOGS_CLICKHOUSE_URL); + url.searchParams.delete("secure"); + + return new ClickHouse({ + url: url.toString(), + name: "logs-clickhouse", + keepAlive: { + enabled: env.CLICKHOUSE_KEEP_ALIVE_ENABLED === "1", + idleSocketTtl: env.CLICKHOUSE_KEEP_ALIVE_IDLE_SOCKET_TTL_MS, + }, + logLevel: env.CLICKHOUSE_LOG_LEVEL, + compression: { request: true }, + maxOpenConnections: env.CLICKHOUSE_MAX_OPEN_CONNECTIONS, + clickhouseSettings: { + max_memory_usage: env.CLICKHOUSE_LOGS_LIST_MAX_MEMORY_USAGE.toString(), + max_bytes_before_external_sort: + env.CLICKHOUSE_LOGS_LIST_MAX_BYTES_BEFORE_EXTERNAL_SORT.toString(), + max_threads: env.CLICKHOUSE_LOGS_LIST_MAX_THREADS, + ...(env.CLICKHOUSE_LOGS_LIST_MAX_ROWS_TO_READ && { + max_rows_to_read: env.CLICKHOUSE_LOGS_LIST_MAX_ROWS_TO_READ.toString(), + }), + ...(env.CLICKHOUSE_LOGS_LIST_MAX_EXECUTION_TIME && { + max_execution_time: env.CLICKHOUSE_LOGS_LIST_MAX_EXECUTION_TIME, + }), + }, + }); +} + +const defaultAdminClickhouseClient = singleton( + "adminClickhouseClient", + initializeAdminClickhouseClient +); + +function initializeAdminClickhouseClient() { + if (!env.ADMIN_CLICKHOUSE_URL) { + throw new Error("ADMIN_CLICKHOUSE_URL is not set"); + } + + const url = new URL(env.ADMIN_CLICKHOUSE_URL); + url.searchParams.delete("secure"); + + return new ClickHouse({ + url: url.toString(), + name: "admin-clickhouse", + keepAlive: { + enabled: env.CLICKHOUSE_KEEP_ALIVE_ENABLED === "1", + idleSocketTtl: env.CLICKHOUSE_KEEP_ALIVE_IDLE_SOCKET_TTL_MS, + }, + logLevel: env.CLICKHOUSE_LOG_LEVEL, + compression: { request: true }, + maxOpenConnections: env.CLICKHOUSE_MAX_OPEN_CONNECTIONS, + }); +} + +const defaultQueryClickhouseClient = singleton( + "queryClickhouseClient", + initializeQueryClickhouseClient +); + +function initializeQueryClickhouseClient() { + if (!env.QUERY_CLICKHOUSE_URL) { + throw new Error("QUERY_CLICKHOUSE_URL is not set"); + } + + const url = new URL(env.QUERY_CLICKHOUSE_URL); + url.searchParams.delete("secure"); + + return new ClickHouse({ + url: url.toString(), + name: "query-clickhouse", + keepAlive: { + enabled: env.CLICKHOUSE_KEEP_ALIVE_ENABLED === "1", + idleSocketTtl: env.CLICKHOUSE_KEEP_ALIVE_IDLE_SOCKET_TTL_MS, + }, + logLevel: env.CLICKHOUSE_LOG_LEVEL, + compression: { request: true }, + maxOpenConnections: env.CLICKHOUSE_MAX_OPEN_CONNECTIONS, + }); +} + +/** TaskRun replication to ClickHouse (`RUN_REPLICATION_CLICKHOUSE_URL`); not exported. */ +const defaultRunsReplicationClickhouseClient = singleton( + "runsReplicationClickhouseClient", + initializeRunsReplicationClickhouseClient +); + +function initializeRunsReplicationClickhouseClient(): ClickHouse { + if (!env.RUN_REPLICATION_CLICKHOUSE_URL) { + // Runs replication worker gates on this URL; factory may still resolve "replication" for tests. + return defaultClickhouseClient; + } + + const url = new URL(env.RUN_REPLICATION_CLICKHOUSE_URL); + url.searchParams.delete("secure"); + + return new ClickHouse({ + url: url.toString(), + name: "runs-replication", + keepAlive: { + enabled: env.RUN_REPLICATION_KEEP_ALIVE_ENABLED === "1", + idleSocketTtl: env.RUN_REPLICATION_KEEP_ALIVE_IDLE_SOCKET_TTL_MS, + }, + logLevel: env.RUN_REPLICATION_CLICKHOUSE_LOG_LEVEL, + compression: { request: true }, + maxOpenConnections: env.RUN_REPLICATION_MAX_OPEN_CONNECTIONS, + }); +} + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +function hashHostname(url: string): string { + const parsed = new URL(url); + return createHash("sha256").update(parsed.hostname).digest("hex"); +} + +export type ClientType = "standard" | "events" | "replication" | "logs" | "query" | "admin"; + +function buildOrgClickhouseClient(url: string, clientType: ClientType): ClickHouse { + const parsed = new URL(url); + parsed.searchParams.delete("secure"); + + return new ClickHouse({ + url: parsed.toString(), + name: `org-clickhouse-${clientType}`, + keepAlive: { + enabled: env.CLICKHOUSE_KEEP_ALIVE_ENABLED === "1", + idleSocketTtl: env.CLICKHOUSE_KEEP_ALIVE_IDLE_SOCKET_TTL_MS, + }, + logLevel: env.CLICKHOUSE_LOG_LEVEL, + compression: { request: true }, + maxOpenConnections: env.CLICKHOUSE_MAX_OPEN_CONNECTIONS, + }); +} + +// --------------------------------------------------------------------------- +// Factory class (injectable for testing) +// --------------------------------------------------------------------------- + +export class ClickhouseFactory { + /** ClickHouse clients keyed by hostname hash + clientType. */ + private readonly _clientCache = new Map(); + /** Event repositories keyed by hostname hash (stateful, must be reused). */ + private readonly _eventRepositoryCache = new Map(); + + constructor(private readonly _registry: OrganizationDataStoresRegistry) {} + + async isReady(): Promise { + if (!this._registry.isLoaded) { + await this._registry.isReady; + } + return true; + } + + async getClickhouseForOrganization( + organizationId: string, + clientType: ClientType + ): Promise { + if (!this._registry.isLoaded) { + await this._registry.isReady; + } + + return this.getClickhouseForOrganizationSync(organizationId, clientType); + } + + getClickhouseForOrganizationSync(organizationId: string, clientType: ClientType): ClickHouse { + const dataStore = this._registry.get(organizationId, "CLICKHOUSE"); + + if (!dataStore) { + switch (clientType) { + case "standard": + case "events": + return defaultClickhouseClient; + case "replication": + return defaultRunsReplicationClickhouseClient; + case "logs": + return defaultLogsClickhouseClient; + case "query": + return defaultQueryClickhouseClient; + case "admin": + return defaultAdminClickhouseClient; + } + } + + const hostnameHash = hashHostname(dataStore.url); + const cacheKey = `${hostnameHash}:${clientType}`; + let client = this._clientCache.get(cacheKey); + + if (!client) { + client = buildOrgClickhouseClient(dataStore.url, clientType); + this._clientCache.set(cacheKey, client); + } + + return client; + } + + async getEventRepositoryForOrganization( + store: string, + organizationId: string + ): Promise<{ key: string; repository: IEventRepository }> { + if (!this._registry.isLoaded) { + await this._registry.isReady; + } + + return this.getEventRepositoryForOrganizationSync(store, organizationId); + } + + getEventRepositoryForOrganizationSync( + store: string, + organizationId: string + ): { key: string; repository: IEventRepository } { + const dataStore = this._registry.get(organizationId, "CLICKHOUSE"); + + if (!dataStore) { + const defaultKey = `default:events:${store}`; + let defaultRepo = this._eventRepositoryCache.get(defaultKey); + if (!defaultRepo) { + const eventsClickhouse = getEventsClickhouseClient(); + defaultRepo = buildEventRepository(store, eventsClickhouse); + this._eventRepositoryCache.set(defaultKey, defaultRepo); + } + return { key: defaultKey, repository: defaultRepo }; + } + + const hostnameHash = hashHostname(dataStore.url); + const cacheKey = `${hostnameHash}:events:${store}`; + let repository = this._eventRepositoryCache.get(cacheKey); + + if (!repository) { + const client = this.getClickhouseForOrganizationSync(organizationId, "events"); + repository = buildEventRepository(store, client); + this._eventRepositoryCache.set(cacheKey, repository); + } + + return { key: cacheKey, repository: repository }; + } +} + +// --------------------------------------------------------------------------- +// Singleton factory instance +// --------------------------------------------------------------------------- + +export const clickhouseFactory = singleton( + "clickhouseFactory", + () => new ClickhouseFactory(organizationDataStoresRegistry) +); + +/** + * Get admin ClickHouse client for cross-organization queries. + * Only use for admin tools and analytics that need to query across all orgs. + */ +export function getAdminClickhouse(): ClickHouse { + return defaultAdminClickhouseClient; +} + +export function getDefaultClickhouseClient(): ClickHouse { + return defaultClickhouseClient; +} + +export function getDefaultLogsClickhouseClient(): ClickHouse { + return defaultLogsClickhouseClient; +} + +// --------------------------------------------------------------------------- +// Private helpers +// --------------------------------------------------------------------------- + +function getEventsClickhouseClient(): ClickHouse { + if (!env.EVENTS_CLICKHOUSE_URL) { + throw new Error("EVENTS_CLICKHOUSE_URL is not set"); + } + + const url = new URL(env.EVENTS_CLICKHOUSE_URL); + url.searchParams.delete("secure"); + + return new ClickHouse({ + url: url.toString(), + name: "task-events", + keepAlive: { + enabled: env.EVENTS_CLICKHOUSE_KEEP_ALIVE_ENABLED === "1", + idleSocketTtl: env.EVENTS_CLICKHOUSE_KEEP_ALIVE_IDLE_SOCKET_TTL_MS, + }, + logLevel: env.EVENTS_CLICKHOUSE_LOG_LEVEL, + compression: { + request: env.EVENTS_CLICKHOUSE_COMPRESSION_REQUEST === "1", + }, + maxOpenConnections: env.EVENTS_CLICKHOUSE_MAX_OPEN_CONNECTIONS, + }); +} + +function buildEventRepository(store: string, clickhouse: ClickHouse): ClickhouseEventRepository { + switch (store) { + case "clickhouse": { + return new ClickhouseEventRepository({ + clickhouse, + batchSize: env.EVENTS_CLICKHOUSE_BATCH_SIZE, + flushInterval: env.EVENTS_CLICKHOUSE_FLUSH_INTERVAL_MS, + maximumTraceSummaryViewCount: env.EVENTS_CLICKHOUSE_MAX_TRACE_SUMMARY_VIEW_COUNT, + maximumTraceDetailedSummaryViewCount: + env.EVENTS_CLICKHOUSE_MAX_TRACE_DETAILED_SUMMARY_VIEW_COUNT, + maximumLiveReloadingSetting: env.EVENTS_CLICKHOUSE_MAX_LIVE_RELOADING_SETTING, + insertStrategy: env.EVENTS_CLICKHOUSE_INSERT_STRATEGY, + waitForAsyncInsert: env.EVENTS_CLICKHOUSE_WAIT_FOR_ASYNC_INSERT === "1", + asyncInsertMaxDataSize: env.EVENTS_CLICKHOUSE_ASYNC_INSERT_MAX_DATA_SIZE, + asyncInsertBusyTimeoutMs: env.EVENTS_CLICKHOUSE_ASYNC_INSERT_BUSY_TIMEOUT_MS, + startTimeMaxAgeMs: env.EVENTS_CLICKHOUSE_START_TIME_MAX_AGE_MS, + llmMetricsBatchSize: env.LLM_METRICS_BATCH_SIZE, + llmMetricsFlushInterval: env.LLM_METRICS_FLUSH_INTERVAL_MS, + llmMetricsMaxBatchSize: env.LLM_METRICS_MAX_BATCH_SIZE, + llmMetricsMaxConcurrency: env.LLM_METRICS_MAX_CONCURRENCY, + otlpMetricsBatchSize: env.METRICS_CLICKHOUSE_BATCH_SIZE, + otlpMetricsFlushInterval: env.METRICS_CLICKHOUSE_FLUSH_INTERVAL_MS, + otlpMetricsMaxConcurrency: env.METRICS_CLICKHOUSE_MAX_CONCURRENCY, + version: "v1", + }); + } + case "clickhouse_v2": { + return new ClickhouseEventRepository({ + clickhouse: clickhouse, + batchSize: env.EVENTS_CLICKHOUSE_BATCH_SIZE, + flushInterval: env.EVENTS_CLICKHOUSE_FLUSH_INTERVAL_MS, + maximumTraceSummaryViewCount: env.EVENTS_CLICKHOUSE_MAX_TRACE_SUMMARY_VIEW_COUNT, + maximumTraceDetailedSummaryViewCount: + env.EVENTS_CLICKHOUSE_MAX_TRACE_DETAILED_SUMMARY_VIEW_COUNT, + maximumLiveReloadingSetting: env.EVENTS_CLICKHOUSE_MAX_LIVE_RELOADING_SETTING, + insertStrategy: env.EVENTS_CLICKHOUSE_INSERT_STRATEGY, + waitForAsyncInsert: env.EVENTS_CLICKHOUSE_WAIT_FOR_ASYNC_INSERT === "1", + asyncInsertMaxDataSize: env.EVENTS_CLICKHOUSE_ASYNC_INSERT_MAX_DATA_SIZE, + asyncInsertBusyTimeoutMs: env.EVENTS_CLICKHOUSE_ASYNC_INSERT_BUSY_TIMEOUT_MS, + llmMetricsBatchSize: env.LLM_METRICS_BATCH_SIZE, + llmMetricsFlushInterval: env.LLM_METRICS_FLUSH_INTERVAL_MS, + llmMetricsMaxBatchSize: env.LLM_METRICS_MAX_BATCH_SIZE, + llmMetricsMaxConcurrency: env.LLM_METRICS_MAX_CONCURRENCY, + otlpMetricsBatchSize: env.METRICS_CLICKHOUSE_BATCH_SIZE, + otlpMetricsFlushInterval: env.METRICS_CLICKHOUSE_FLUSH_INTERVAL_MS, + otlpMetricsMaxConcurrency: env.METRICS_CLICKHOUSE_MAX_CONCURRENCY, + version: "v2", + }); + } + default: { + throw new Error(`Unknown ClickHouse event repository store: ${store}`); + } + } +} diff --git a/apps/webapp/app/services/clickhouse/clickhouseSecretSchemas.server.ts b/apps/webapp/app/services/clickhouse/clickhouseSecretSchemas.server.ts new file mode 100644 index 00000000000..016eb717c18 --- /dev/null +++ b/apps/webapp/app/services/clickhouse/clickhouseSecretSchemas.server.ts @@ -0,0 +1,11 @@ +import { z } from "zod"; + +export const ClickhouseConnectionSchema = z.object({ + url: z.string().url(), +}); + +export type ClickhouseConnection = z.infer; + +export function getClickhouseSecretKey(orgId: string, clientType: string): string { + return `org:${orgId}:clickhouse:${clientType}`; +} diff --git a/apps/webapp/app/services/clickhouseInstance.server.ts b/apps/webapp/app/services/clickhouseInstance.server.ts deleted file mode 100644 index 9c4941671f3..00000000000 --- a/apps/webapp/app/services/clickhouseInstance.server.ts +++ /dev/null @@ -1,130 +0,0 @@ -import { ClickHouse } from "@internal/clickhouse"; -import { env } from "~/env.server"; -import { singleton } from "~/utils/singleton"; - -export const clickhouseClient = singleton("clickhouseClient", initializeClickhouseClient); - -function initializeClickhouseClient() { - const url = new URL(env.CLICKHOUSE_URL); - - // Remove secure param - url.searchParams.delete("secure"); - - console.log(`🗃️ Clickhouse service enabled to host ${url.host}`); - - const clickhouse = new ClickHouse({ - url: url.toString(), - name: "clickhouse-instance", - keepAlive: { - enabled: env.CLICKHOUSE_KEEP_ALIVE_ENABLED === "1", - idleSocketTtl: env.CLICKHOUSE_KEEP_ALIVE_IDLE_SOCKET_TTL_MS, - }, - logLevel: env.CLICKHOUSE_LOG_LEVEL, - compression: { - request: true, - }, - maxOpenConnections: env.CLICKHOUSE_MAX_OPEN_CONNECTIONS, - }); - - return clickhouse; -} - -export const logsClickhouseClient = singleton( - "logsClickhouseClient", - initializeLogsClickhouseClient -); - -function initializeLogsClickhouseClient() { - if (!env.LOGS_CLICKHOUSE_URL) { - throw new Error("LOGS_CLICKHOUSE_URL is not set"); - } - - const url = new URL(env.LOGS_CLICKHOUSE_URL); - - // Remove secure param - url.searchParams.delete("secure"); - - return new ClickHouse({ - url: url.toString(), - name: "logs-clickhouse", - keepAlive: { - enabled: env.CLICKHOUSE_KEEP_ALIVE_ENABLED === "1", - idleSocketTtl: env.CLICKHOUSE_KEEP_ALIVE_IDLE_SOCKET_TTL_MS, - }, - logLevel: env.CLICKHOUSE_LOG_LEVEL, - compression: { - request: true, - }, - maxOpenConnections: env.CLICKHOUSE_MAX_OPEN_CONNECTIONS, - clickhouseSettings: { - max_memory_usage: env.CLICKHOUSE_LOGS_LIST_MAX_MEMORY_USAGE.toString(), - max_bytes_before_external_sort: - env.CLICKHOUSE_LOGS_LIST_MAX_BYTES_BEFORE_EXTERNAL_SORT.toString(), - max_threads: env.CLICKHOUSE_LOGS_LIST_MAX_THREADS, - ...(env.CLICKHOUSE_LOGS_LIST_MAX_ROWS_TO_READ && { - max_rows_to_read: env.CLICKHOUSE_LOGS_LIST_MAX_ROWS_TO_READ.toString(), - }), - ...(env.CLICKHOUSE_LOGS_LIST_MAX_EXECUTION_TIME && { - max_execution_time: env.CLICKHOUSE_LOGS_LIST_MAX_EXECUTION_TIME, - }), - }, - }); -} - -export const adminClickhouseClient = singleton( - "adminClickhouseClient", - initializeAdminClickhouseClient -); - -function initializeAdminClickhouseClient() { - if (!env.ADMIN_CLICKHOUSE_URL) { - throw new Error("ADMIN_CLICKHOUSE_URL is not set"); - } - - const url = new URL(env.ADMIN_CLICKHOUSE_URL); - url.searchParams.delete("secure"); - - return new ClickHouse({ - url: url.toString(), - name: "admin-clickhouse", - keepAlive: { - enabled: env.CLICKHOUSE_KEEP_ALIVE_ENABLED === "1", - idleSocketTtl: env.CLICKHOUSE_KEEP_ALIVE_IDLE_SOCKET_TTL_MS, - }, - logLevel: env.CLICKHOUSE_LOG_LEVEL, - compression: { - request: true, - }, - maxOpenConnections: env.CLICKHOUSE_MAX_OPEN_CONNECTIONS, - }); -} - -export const queryClickhouseClient = singleton( - "queryClickhouseClient", - initializeQueryClickhouseClient -); - -function initializeQueryClickhouseClient() { - if (!env.QUERY_CLICKHOUSE_URL) { - throw new Error("QUERY_CLICKHOUSE_URL is not set"); - } - - const url = new URL(env.QUERY_CLICKHOUSE_URL); - - // Remove secure param - url.searchParams.delete("secure"); - - return new ClickHouse({ - url: url.toString(), - name: "query-clickhouse", - keepAlive: { - enabled: env.CLICKHOUSE_KEEP_ALIVE_ENABLED === "1", - idleSocketTtl: env.CLICKHOUSE_KEEP_ALIVE_IDLE_SOCKET_TTL_MS, - }, - logLevel: env.CLICKHOUSE_LOG_LEVEL, - compression: { - request: true, - }, - maxOpenConnections: env.CLICKHOUSE_MAX_OPEN_CONNECTIONS, - }); -} diff --git a/apps/webapp/app/services/dataStores/organizationDataStoreConfigSchemas.server.ts b/apps/webapp/app/services/dataStores/organizationDataStoreConfigSchemas.server.ts new file mode 100644 index 00000000000..a4a2491a2d6 --- /dev/null +++ b/apps/webapp/app/services/dataStores/organizationDataStoreConfigSchemas.server.ts @@ -0,0 +1,39 @@ +import { z } from "zod"; + +// --------------------------------------------------------------------------- +// ClickHouse config (kind = CLICKHOUSE) +// --------------------------------------------------------------------------- + +/** V1: single secret-store key that supplies the ClickHouse connection URL. */ +export const ClickhouseDataStoreConfigV1 = z.object({ + version: z.literal(1), + data: z.object({ + /** Key into the SecretStore that resolves to a ClickhouseConnection ({url}). */ + secretKey: z.string(), + }), +}); + +export type ClickhouseDataStoreConfigV1 = z.infer; + +/** Discriminated union over version — extend by adding new literals here. */ +export const ClickhouseDataStoreConfig = z.discriminatedUnion("version", [ + ClickhouseDataStoreConfigV1, +]); + +export type ClickhouseDataStoreConfig = z.infer; + +// --------------------------------------------------------------------------- +// Top-level per-kind union +// --------------------------------------------------------------------------- + +/** + * Secrets are resolved to URLs at registry load time so the factory never + * needs to touch the secret store on the hot path. + */ +export type ParsedClickhouseDataStore = { + kind: "CLICKHOUSE"; + url: string; +}; + +/** Union of all parsed data store types. Extend as new DataStoreKind values are added. */ +export type ParsedDataStore = ParsedClickhouseDataStore; diff --git a/apps/webapp/app/services/dataStores/organizationDataStoresRegistry.server.ts b/apps/webapp/app/services/dataStores/organizationDataStoresRegistry.server.ts new file mode 100644 index 00000000000..379084a0d74 --- /dev/null +++ b/apps/webapp/app/services/dataStores/organizationDataStoresRegistry.server.ts @@ -0,0 +1,167 @@ +import type { DataStoreKind, PrismaClient, PrismaReplicaClient } from "@trigger.dev/database"; +import { + ClickhouseDataStoreConfig, + type ParsedDataStore, +} from "./organizationDataStoreConfigSchemas.server"; +import { getSecretStore } from "../secrets/secretStore.server"; +import { ClickhouseConnectionSchema } from "../clickhouse/clickhouseSecretSchemas.server"; + +export class OrganizationDataStoresRegistry { + private _prisma: PrismaClient | PrismaReplicaClient; + /** Keyed by `${organizationId}:${kind}` */ + private _lookup: Map = new Map(); + private _loaded = false; + private _readyResolve!: () => void; + + /** Resolves once the initial `loadFromDatabase()` completes successfully. */ + readonly isReady: Promise; + + constructor(prisma: PrismaClient | PrismaReplicaClient) { + this._prisma = prisma; + this.isReady = new Promise((resolve) => { + this._readyResolve = resolve; + }); + } + + get isLoaded(): boolean { + return this._loaded; + } + + async loadFromDatabase(): Promise { + const rows = await this._prisma.organizationDataStore.findMany(); + const secretStore = getSecretStore("DATABASE", { prismaClient: this._prisma }); + + const lookup = new Map(); + + for (const row of rows) { + let parsed: ParsedDataStore | null = null; + + switch (row.kind) { + case "CLICKHOUSE": { + const result = ClickhouseDataStoreConfig.safeParse(row.config); + if (!result.success) { + console.warn( + `[OrganizationDataStoresRegistry] Invalid config for OrganizationDataStore "${row.key}" (kind=CLICKHOUSE): ${result.error.message}` + ); + continue; + } + + const connection = await secretStore.getSecret( + ClickhouseConnectionSchema, + result.data.data.secretKey + ); + + if (!connection) { + console.warn( + `[OrganizationDataStoresRegistry] Secret "${result.data.data.secretKey}" not found for OrganizationDataStore "${row.key}" — skipping` + ); + continue; + } + + parsed = { kind: "CLICKHOUSE", url: connection.url }; + break; + } + default: { + console.warn( + `[OrganizationDataStoresRegistry] Unknown kind "${row.kind}" for OrganizationDataStore "${row.key}" — skipping` + ); + continue; + } + } + + for (const orgId of row.organizationIds) { + const key = `${orgId}:${row.kind}`; + lookup.set(key, parsed); + } + } + + this._lookup = lookup; + + if (!this._loaded) { + this._loaded = true; + this._readyResolve(); + } + } + + async reload(): Promise { + await this.loadFromDatabase(); + } + + #secretKey(key: string, kind: DataStoreKind) { + return `data-store:${key}:${kind.toLocaleLowerCase()}`; + } + + async addDataStore({ + key, + kind, + organizationIds, + config, + }: { + key: string; + kind: DataStoreKind; + organizationIds: string[]; + config: any; + }) { + const secretKey = this.#secretKey(key, kind); + + const secretStore = getSecretStore("DATABASE", { prismaClient: this._prisma }); + await secretStore.setSecret(secretKey, config); + + return this._prisma.organizationDataStore.create({ + data: { + key, + organizationIds, + kind: "CLICKHOUSE", + config: { version: 1, data: { secretKey } }, + }, + }); + } + + async updateDataStore({ + key, + kind, + organizationIds, + config, + }: { + key: string; + kind: DataStoreKind; + organizationIds: string[]; + config?: any; + }) { + const secretKey = this.#secretKey(key, kind); + + if (config) { + const secretStore = getSecretStore("DATABASE", { prismaClient: this._prisma }); + await secretStore.setSecret(secretKey, config); + } + + return this._prisma.organizationDataStore.update({ + where: { + key, + }, + data: { + organizationIds, + kind: "CLICKHOUSE", + }, + }); + } + + async deleteDataStore({ key, kind }: { key: string; kind: DataStoreKind }) { + const secretKey = this.#secretKey(key, kind); + const secretStore = getSecretStore("DATABASE", { prismaClient: this._prisma }); + await secretStore.deleteSecret(secretKey).catch(() => { + // Secret may not exist — proceed with deletion + }); + + await this._prisma.organizationDataStore.delete({ where: { key } }); + } + + /** + * Returns the parsed data store config for the given organization and kind, + * or `null` if no override is configured (caller should use the default). + */ + get(organizationId: string, kind: DataStoreKind): ParsedDataStore | null { + if (!this._loaded) return null; + return this._lookup.get(`${organizationId}:${kind}`) ?? null; + } +} diff --git a/apps/webapp/app/services/dataStores/organizationDataStoresRegistryInstance.server.ts b/apps/webapp/app/services/dataStores/organizationDataStoresRegistryInstance.server.ts new file mode 100644 index 00000000000..f7d82da5682 --- /dev/null +++ b/apps/webapp/app/services/dataStores/organizationDataStoresRegistryInstance.server.ts @@ -0,0 +1,24 @@ +import { $replica } from "~/db.server"; +import { env } from "~/env.server"; +import { signalsEmitter } from "~/services/signals.server"; +import { singleton } from "~/utils/singleton"; +import { OrganizationDataStoresRegistry } from "./organizationDataStoresRegistry.server"; + +export const organizationDataStoresRegistry = singleton("organizationDataStoresRegistry", () => { + const registry = new OrganizationDataStoresRegistry($replica); + + registry.loadFromDatabase().catch((err) => { + console.error("[OrganizationDataStoresRegistry] Failed to initialize", err); + }); + + const interval = setInterval(() => { + registry.reload().catch((err) => { + console.error("[OrganizationDataStoresRegistry] Failed to reload", err); + }); + }, env.ORGANIZATION_DATA_STORES_RELOAD_INTERVAL_MS); + + signalsEmitter.on("SIGTERM", () => clearInterval(interval)); + signalsEmitter.on("SIGINT", () => clearInterval(interval)); + + return registry; +}); diff --git a/apps/webapp/app/services/queryService.server.ts b/apps/webapp/app/services/queryService.server.ts index 1f3bdbba18a..214232ebdf5 100644 --- a/apps/webapp/app/services/queryService.server.ts +++ b/apps/webapp/app/services/queryService.server.ts @@ -11,7 +11,7 @@ import type { TableSchema, WhereClauseCondition } from "@internal/tsql"; import { z } from "zod"; import { prisma } from "~/db.server"; import { env } from "~/env.server"; -import { queryClickhouseClient } from "./clickhouseInstance.server"; +import { clickhouseFactory } from "./clickhouse/clickhouseFactory.server"; import { queryConcurrencyLimiter, DEFAULT_ORG_CONCURRENCY_LIMIT, @@ -275,7 +275,8 @@ export async function executeQuery( environment: Object.fromEntries(environments.map((e) => [e.id, e.slug])), }; - const result = await executeTSQL(queryClickhouseClient.reader, { + const queryClickhouse = await clickhouseFactory.getClickhouseForOrganization(organizationId, "query"); + const result = await executeTSQL(queryClickhouse.reader, { ...baseOptions, schema: z.record(z.any()), tableSchema: querySchemas, diff --git a/apps/webapp/app/services/runsReplicationInstance.server.ts b/apps/webapp/app/services/runsReplicationInstance.server.ts index 0a8ab5e1bde..1a507aafd74 100644 --- a/apps/webapp/app/services/runsReplicationInstance.server.ts +++ b/apps/webapp/app/services/runsReplicationInstance.server.ts @@ -1,6 +1,6 @@ -import { ClickHouse } from "@internal/clickhouse"; import invariant from "tiny-invariant"; import { env } from "~/env.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { singleton } from "~/utils/singleton"; import { meter, provider } from "~/v3/tracer.server"; import { RunsReplicationService } from "./runsReplicationService.server"; @@ -22,22 +22,8 @@ function initializeRunsReplicationInstance() { console.log("🗃️ Runs replication service enabled"); - const clickhouse = new ClickHouse({ - url: env.RUN_REPLICATION_CLICKHOUSE_URL, - name: "runs-replication", - keepAlive: { - enabled: env.RUN_REPLICATION_KEEP_ALIVE_ENABLED === "1", - idleSocketTtl: env.RUN_REPLICATION_KEEP_ALIVE_IDLE_SOCKET_TTL_MS, - }, - logLevel: env.RUN_REPLICATION_CLICKHOUSE_LOG_LEVEL, - compression: { - request: true, - }, - maxOpenConnections: env.RUN_REPLICATION_MAX_OPEN_CONNECTIONS, - }); - const service = new RunsReplicationService({ - clickhouse: clickhouse, + clickhouseFactory, pgConnectionUrl: DATABASE_URL, serviceName: "runs-replication", slotName: env.RUN_REPLICATION_SLOT_NAME, @@ -72,8 +58,9 @@ function initializeRunsReplicationInstance() { }); if (env.RUN_REPLICATION_ENABLED === "1") { - service - .start() + clickhouseFactory + .isReady() + .then(() => service.start()) .then(() => { console.log("🗃️ Runs replication service started"); }) diff --git a/apps/webapp/app/services/runsReplicationService.server.ts b/apps/webapp/app/services/runsReplicationService.server.ts index 7930c05481f..bd128f4f315 100644 --- a/apps/webapp/app/services/runsReplicationService.server.ts +++ b/apps/webapp/app/services/runsReplicationService.server.ts @@ -1,5 +1,11 @@ -import type { ClickHouse, TaskRunInsertArray, PayloadInsertArray } from "@internal/clickhouse"; -import { getTaskRunField, getPayloadField } from "@internal/clickhouse"; +import type { ClickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; +import { + type ClickHouse, + type PayloadInsertArray, + type TaskRunInsertArray, + getPayloadField, + getTaskRunField, +} from "@internal/clickhouse"; import { type RedisOptions } from "@internal/redis"; import { LogicalReplicationClient, @@ -21,7 +27,10 @@ import { import { Logger, type LogLevel } from "@trigger.dev/core/logger"; import { tryCatch } from "@trigger.dev/core/utils"; import { parsePacketAsJson } from "@trigger.dev/core/v3/utils/ioSerialization"; -import { unsafeExtractIdempotencyKeyScope, unsafeExtractIdempotencyKeyUser } from "@trigger.dev/core/v3/serverOnly"; +import { + unsafeExtractIdempotencyKeyScope, + unsafeExtractIdempotencyKeyUser, +} from "@trigger.dev/core/v3/serverOnly"; import { RunAnnotations } from "@trigger.dev/core/v3"; import { type TaskRun } from "@trigger.dev/database"; import { nanoid } from "nanoid"; @@ -46,7 +55,7 @@ interface Transaction { } export type RunsReplicationServiceOptions = { - clickhouse: ClickHouse; + clickhouseFactory: ClickhouseFactory; pgConnectionUrl: string; serviceName: string; slotName: string; @@ -560,16 +569,50 @@ export class RunsReplicationService { const flushStartTime = performance.now(); await startSpan(this._tracer, "flushBatch", async (span) => { - const preparedInserts = await startSpan(this._tracer, "prepare_inserts", async (span) => { + const preparedInserts = await startSpan(this._tracer, "prepare_inserts", async () => { return await Promise.all(batch.map(this.#prepareRunInserts.bind(this))); }); - const taskRunInserts = preparedInserts - .map(({ taskRunInsert }) => taskRunInsert) - .filter((x): x is TaskRunInsertArray => Boolean(x)) - // batch inserts in clickhouse are more performant if the items - // are pre-sorted by the primary key - .sort((a, b) => { + const routeCache = new Map(); + const groups = new Map< + ClickHouse, + { taskRunInserts: TaskRunInsertArray[]; payloadInserts: PayloadInsertArray[] } + >(); + + for (let i = 0; i < batch.length; i++) { + const batchedRun = batch[i]!; + const prep = preparedInserts[i]!; + const { run } = batchedRun; + + if (!run.organizationId || !run.environmentType) { + continue; + } + + let client = routeCache.get(run.organizationId); + if (!client) { + client = this.options.clickhouseFactory.getClickhouseForOrganizationSync( + run.organizationId, + "replication" + ); + routeCache.set(run.organizationId, client); + } + + let group = groups.get(client); + if (!group) { + group = { taskRunInserts: [], payloadInserts: [] }; + groups.set(client, group); + } + + if (prep.taskRunInsert) { + group.taskRunInserts.push(prep.taskRunInsert); + } + if (prep.payloadInsert) { + group.payloadInserts.push(prep.payloadInsert); + } + } + + const sortTaskRunInserts = (rows: TaskRunInsertArray[]) => + rows.sort((a, b) => { const aOrgId = getTaskRunField(a, "organization_id"); const bOrgId = getTaskRunField(b, "organization_id"); if (aOrgId !== bOrgId) { @@ -596,41 +639,61 @@ export class RunsReplicationService { return aRunId < bRunId ? -1 : 1; }); - const payloadInserts = preparedInserts - .map(({ payloadInsert }) => payloadInsert) - .filter((x): x is PayloadInsertArray => Boolean(x)) - // batch inserts in clickhouse are more performant if the items - // are pre-sorted by the primary key - .sort((a, b) => { + const sortPayloadInserts = (rows: PayloadInsertArray[]) => + rows.sort((a, b) => { const aRunId = getPayloadField(a, "run_id"); const bRunId = getPayloadField(b, "run_id"); if (aRunId === bRunId) return 0; return aRunId < bRunId ? -1 : 1; }); - span.setAttribute("task_run_inserts", taskRunInserts.length); - span.setAttribute("payload_inserts", payloadInserts.length); + const combinedTaskRunInserts: TaskRunInsertArray[] = []; + const combinedPayloadInserts: PayloadInsertArray[] = []; + let taskRunError: Error | null = null; + let payloadError: Error | null = null; + + for (const [clickhouse, group] of groups) { + sortTaskRunInserts(group.taskRunInserts); + sortPayloadInserts(group.payloadInserts); + combinedTaskRunInserts.push(...group.taskRunInserts); + combinedPayloadInserts.push(...group.payloadInserts); + + const [trErr] = await this.#insertWithRetry( + (attempt) => this.#insertTaskRunInserts(clickhouse, group.taskRunInserts, attempt), + "task run inserts", + flushId + ); + if (trErr && !taskRunError) { + taskRunError = trErr; + } + + const [plErr] = await this.#insertWithRetry( + (attempt) => this.#insertPayloadInserts(clickhouse, group.payloadInserts, attempt), + "payload inserts", + flushId + ); + if (plErr && !payloadError) { + payloadError = plErr; + } + + if (!trErr) { + this._taskRunsInsertedCounter.add(group.taskRunInserts.length); + } + if (!plErr) { + this._payloadsInsertedCounter.add(group.payloadInserts.length); + } + } + + span.setAttribute("task_run_inserts", combinedTaskRunInserts.length); + span.setAttribute("payload_inserts", combinedPayloadInserts.length); this.logger.debug("Flushing inserts", { flushId, - taskRunInserts: taskRunInserts.length, - payloadInserts: payloadInserts.length, + taskRunInserts: combinedTaskRunInserts.length, + payloadInserts: combinedPayloadInserts.length, + clickhouseGroups: groups.size, }); - // Insert task runs and payloads with retry logic for connection errors - const [taskRunError, taskRunResult] = await this.#insertWithRetry( - (attempt) => this.#insertTaskRunInserts(taskRunInserts, attempt), - "task run inserts", - flushId - ); - - const [payloadError, payloadResult] = await this.#insertWithRetry( - (attempt) => this.#insertPayloadInserts(payloadInserts, attempt), - "payload inserts", - flushId - ); - - // Log any errors that occurred if (taskRunError) { this.logger.error("Error inserting task run inserts", { error: taskRunError, @@ -649,27 +712,22 @@ export class RunsReplicationService { this.logger.debug("Flushed inserts", { flushId, - taskRunInserts: taskRunInserts.length, - payloadInserts: payloadInserts.length, + taskRunInserts: combinedTaskRunInserts.length, + payloadInserts: combinedPayloadInserts.length, }); - this.events.emit("batchFlushed", { flushId, taskRunInserts, payloadInserts }); + this.events.emit("batchFlushed", { + flushId, + taskRunInserts: combinedTaskRunInserts, + payloadInserts: combinedPayloadInserts, + }); - // Record metrics const flushDurationMs = performance.now() - flushStartTime; const hasErrors = taskRunError !== null || payloadError !== null; this._batchSizeHistogram.record(batch.length); this._flushDurationHistogram.record(flushDurationMs); this._batchesFlushedCounter.add(1, { success: !hasErrors }); - - if (!taskRunError) { - this._taskRunsInsertedCounter.add(taskRunInserts.length); - } - - if (!payloadError) { - this._payloadsInsertedCounter.add(payloadInserts.length); - } }); } @@ -770,10 +828,17 @@ export class RunsReplicationService { }; } - async #insertTaskRunInserts(taskRunInserts: TaskRunInsertArray[], attempt: number) { + async #insertTaskRunInserts( + clickhouse: ClickHouse, + taskRunInserts: TaskRunInsertArray[], + attempt: number + ) { + if (taskRunInserts.length === 0) { + return; + } return await startSpan(this._tracer, "insertTaskRunsInserts", async (span) => { const [insertError, insertResult] = - await this.options.clickhouse.taskRuns.insertCompactArrays(taskRunInserts, { + await clickhouse.taskRuns.insertCompactArrays(taskRunInserts, { params: { clickhouse_settings: this.#getClickhouseInsertSettings(), }, @@ -793,10 +858,17 @@ export class RunsReplicationService { }); } - async #insertPayloadInserts(payloadInserts: PayloadInsertArray[], attempt: number) { + async #insertPayloadInserts( + clickhouse: ClickHouse, + payloadInserts: PayloadInsertArray[], + attempt: number + ) { + if (payloadInserts.length === 0) { + return; + } return await startSpan(this._tracer, "insertPayloadInserts", async (span) => { const [insertError, insertResult] = - await this.options.clickhouse.taskRuns.insertPayloadsCompactArrays(payloadInserts, { + await clickhouse.taskRuns.insertPayloadsCompactArrays(payloadInserts, { params: { clickhouse_settings: this.#getClickhouseInsertSettings(), }, @@ -860,12 +932,13 @@ export class RunsReplicationService { const errorData = { data: run.error }; // Calculate error fingerprint for failed runs - const errorFingerprint = ( + const errorFingerprint = !this._disableErrorFingerprinting && - ['SYSTEM_FAILURE', 'CRASHED', 'INTERRUPTED', 'COMPLETED_WITH_ERRORS', 'TIMED_OUT'].includes(run.status) - ) - ? calculateErrorFingerprint(run.error) - : ''; + ["SYSTEM_FAILURE", "CRASHED", "INTERRUPTED", "COMPLETED_WITH_ERRORS", "TIMED_OUT"].includes( + run.status + ) + ? calculateErrorFingerprint(run.error) + : ""; const annotations = this.#parseAnnotations(run.annotations); @@ -978,7 +1051,6 @@ export class RunsReplicationService { return { data: parsedData }; } - } export type ConcurrentFlushSchedulerConfig = { diff --git a/apps/webapp/app/v3/eventRepository/clickhouseEventRepository.server.ts b/apps/webapp/app/v3/eventRepository/clickhouseEventRepository.server.ts index 27b96ed7d37..5ed11b85540 100644 --- a/apps/webapp/app/v3/eventRepository/clickhouseEventRepository.server.ts +++ b/apps/webapp/app/v3/eventRepository/clickhouseEventRepository.server.ts @@ -1,6 +1,7 @@ import type { ClickHouse, LlmMetricsV1Input, + MetricsV1Input, TaskEventDetailedSummaryV1Result, TaskEventDetailsV1Result, TaskEventSummaryV1Result, @@ -91,6 +92,10 @@ export type ClickhouseEventRepositoryConfig = { llmMetricsFlushInterval?: number; llmMetricsMaxBatchSize?: number; llmMetricsMaxConcurrency?: number; + /** OTLP / task metrics_v1 flush scheduler config */ + otlpMetricsBatchSize?: number; + otlpMetricsFlushInterval?: number; + otlpMetricsMaxConcurrency?: number; }; /** @@ -102,6 +107,7 @@ export class ClickhouseEventRepository implements IEventRepository { private _config: ClickhouseEventRepositoryConfig; private readonly _flushScheduler: DynamicFlushScheduler; private readonly _llmMetricsFlushScheduler: DynamicFlushScheduler; + private readonly _otlpMetricsFlushScheduler: DynamicFlushScheduler; private _tracer: Tracer; private _version: "v1" | "v2"; @@ -137,6 +143,15 @@ export class ClickhouseEventRepository implements IEventRepository { memoryPressureThreshold: config.llmMetricsMaxBatchSize ?? 10000, loadSheddingEnabled: false, }); + + this._otlpMetricsFlushScheduler = new DynamicFlushScheduler({ + batchSize: config.otlpMetricsBatchSize ?? 10000, + flushInterval: config.otlpMetricsFlushInterval ?? 1000, + callback: this.#flushOtelMetricsBatch.bind(this), + minConcurrency: 1, + maxConcurrency: config.otlpMetricsMaxConcurrency ?? 3, + loadSheddingEnabled: false, + }); } get version() { @@ -236,7 +251,6 @@ export class ClickhouseEventRepository implements IEventRepository { } async #flushLlmMetricsBatch(flushId: string, rows: LlmMetricsV1Input[]) { - const [insertError] = await this._clickhouse.llmMetrics.insert(rows, { params: { clickhouse_settings: this.#getClickhouseInsertSettings(), @@ -252,6 +266,27 @@ export class ClickhouseEventRepository implements IEventRepository { }); } + async #flushOtelMetricsBatch(flushId: string, rows: MetricsV1Input[]) { + await startSpan(this._tracer, "flushOtelMetricsBatch", async (span) => { + span.setAttribute("flush_id", flushId); + span.setAttribute("row_count", rows.length); + + const [insertError] = await this._clickhouse.metrics.insert(rows, { + params: { + clickhouse_settings: this.#getClickhouseInsertSettings(), + }, + }); + + if (insertError) { + throw insertError; + } + + logger.info("ClickhouseEventRepository.flushOtelMetricsBatch Inserted OTLP metrics batch", { + rows: rows.length, + }); + }); + } + #createLlmMetricsInput(event: CreateEventInput): LlmMetricsV1Input { const llmMetrics = event._llmMetrics!; @@ -310,7 +345,7 @@ export class ClickhouseEventRepository implements IEventRepository { await tracePubSub.publish(events.map((e) => e.trace_id)); } - async insertMany(events: CreateEventInput[]): Promise { + insertMany(events: CreateEventInput[]): void { this.addToBatch(events.flatMap((event) => this.createEventToTaskEventV1Input(event))); // Dual-write LLM metrics records for spans with cost enrichment @@ -327,6 +362,11 @@ export class ClickhouseEventRepository implements IEventRepository { this.insertMany(events); } + insertManyMetrics(rows: MetricsV1Input[]): void { + if (rows.length === 0) return; + this._otlpMetricsFlushScheduler.addToBatch(rows); + } + private createEventToTaskEventV1Input(event: CreateEventInput): TaskEventV1Input[] { return [ { diff --git a/apps/webapp/app/v3/eventRepository/clickhouseEventRepositoryInstance.server.ts b/apps/webapp/app/v3/eventRepository/clickhouseEventRepositoryInstance.server.ts deleted file mode 100644 index d4e28c5841a..00000000000 --- a/apps/webapp/app/v3/eventRepository/clickhouseEventRepositoryInstance.server.ts +++ /dev/null @@ -1,112 +0,0 @@ -import { ClickHouse } from "@internal/clickhouse"; -import { env } from "~/env.server"; -import { singleton } from "~/utils/singleton"; -import { ClickhouseEventRepository } from "./clickhouseEventRepository.server"; - -export const clickhouseEventRepository = singleton( - "clickhouseEventRepository", - initializeClickhouseRepository -); - -export const clickhouseEventRepositoryV2 = singleton( - "clickhouseEventRepositoryV2", - initializeClickhouseRepositoryV2 -); - -function getClickhouseClient() { - if (!env.EVENTS_CLICKHOUSE_URL) { - throw new Error("EVENTS_CLICKHOUSE_URL is not set"); - } - - const url = new URL(env.EVENTS_CLICKHOUSE_URL); - url.searchParams.delete("secure"); - - return new ClickHouse({ - url: url.toString(), - name: "task-events", - keepAlive: { - enabled: env.EVENTS_CLICKHOUSE_KEEP_ALIVE_ENABLED === "1", - idleSocketTtl: env.EVENTS_CLICKHOUSE_KEEP_ALIVE_IDLE_SOCKET_TTL_MS, - }, - logLevel: env.EVENTS_CLICKHOUSE_LOG_LEVEL, - compression: { - request: env.EVENTS_CLICKHOUSE_COMPRESSION_REQUEST === "1", - }, - maxOpenConnections: env.EVENTS_CLICKHOUSE_MAX_OPEN_CONNECTIONS, - }); -} - -function initializeClickhouseRepository() { - if (!env.EVENTS_CLICKHOUSE_URL) { - throw new Error("EVENTS_CLICKHOUSE_URL is not set"); - } - - const url = new URL(env.EVENTS_CLICKHOUSE_URL); - url.searchParams.delete("secure"); - - const safeUrl = new URL(url.toString()); - safeUrl.password = "redacted"; - - console.log("🗃️ Initializing Clickhouse event repository (v1)", { url: safeUrl.toString() }); - - const clickhouse = getClickhouseClient(); - - const repository = new ClickhouseEventRepository({ - clickhouse: clickhouse, - batchSize: env.EVENTS_CLICKHOUSE_BATCH_SIZE, - flushInterval: env.EVENTS_CLICKHOUSE_FLUSH_INTERVAL_MS, - maximumTraceSummaryViewCount: env.EVENTS_CLICKHOUSE_MAX_TRACE_SUMMARY_VIEW_COUNT, - maximumTraceDetailedSummaryViewCount: - env.EVENTS_CLICKHOUSE_MAX_TRACE_DETAILED_SUMMARY_VIEW_COUNT, - maximumLiveReloadingSetting: env.EVENTS_CLICKHOUSE_MAX_LIVE_RELOADING_SETTING, - insertStrategy: env.EVENTS_CLICKHOUSE_INSERT_STRATEGY, - waitForAsyncInsert: env.EVENTS_CLICKHOUSE_WAIT_FOR_ASYNC_INSERT === "1", - asyncInsertMaxDataSize: env.EVENTS_CLICKHOUSE_ASYNC_INSERT_MAX_DATA_SIZE, - asyncInsertBusyTimeoutMs: env.EVENTS_CLICKHOUSE_ASYNC_INSERT_BUSY_TIMEOUT_MS, - startTimeMaxAgeMs: env.EVENTS_CLICKHOUSE_START_TIME_MAX_AGE_MS, - llmMetricsBatchSize: env.LLM_METRICS_BATCH_SIZE, - llmMetricsFlushInterval: env.LLM_METRICS_FLUSH_INTERVAL_MS, - llmMetricsMaxBatchSize: env.LLM_METRICS_MAX_BATCH_SIZE, - llmMetricsMaxConcurrency: env.LLM_METRICS_MAX_CONCURRENCY, - version: "v1", - }); - - return repository; -} - -function initializeClickhouseRepositoryV2() { - if (!env.EVENTS_CLICKHOUSE_URL) { - throw new Error("EVENTS_CLICKHOUSE_URL is not set"); - } - - const url = new URL(env.EVENTS_CLICKHOUSE_URL); - url.searchParams.delete("secure"); - - const safeUrl = new URL(url.toString()); - safeUrl.password = "redacted"; - - console.log("🗃️ Initializing Clickhouse event repository (v2)", { url: safeUrl.toString() }); - - const clickhouse = getClickhouseClient(); - - const repository = new ClickhouseEventRepository({ - clickhouse: clickhouse, - batchSize: env.EVENTS_CLICKHOUSE_BATCH_SIZE, - flushInterval: env.EVENTS_CLICKHOUSE_FLUSH_INTERVAL_MS, - maximumTraceSummaryViewCount: env.EVENTS_CLICKHOUSE_MAX_TRACE_SUMMARY_VIEW_COUNT, - maximumTraceDetailedSummaryViewCount: - env.EVENTS_CLICKHOUSE_MAX_TRACE_DETAILED_SUMMARY_VIEW_COUNT, - maximumLiveReloadingSetting: env.EVENTS_CLICKHOUSE_MAX_LIVE_RELOADING_SETTING, - insertStrategy: env.EVENTS_CLICKHOUSE_INSERT_STRATEGY, - waitForAsyncInsert: env.EVENTS_CLICKHOUSE_WAIT_FOR_ASYNC_INSERT === "1", - asyncInsertMaxDataSize: env.EVENTS_CLICKHOUSE_ASYNC_INSERT_MAX_DATA_SIZE, - asyncInsertBusyTimeoutMs: env.EVENTS_CLICKHOUSE_ASYNC_INSERT_BUSY_TIMEOUT_MS, - llmMetricsBatchSize: env.LLM_METRICS_BATCH_SIZE, - llmMetricsFlushInterval: env.LLM_METRICS_FLUSH_INTERVAL_MS, - llmMetricsMaxBatchSize: env.LLM_METRICS_MAX_BATCH_SIZE, - llmMetricsMaxConcurrency: env.LLM_METRICS_MAX_CONCURRENCY, - version: "v2", - }); - - return repository; -} diff --git a/apps/webapp/app/v3/eventRepository/eventRepository.server.ts b/apps/webapp/app/v3/eventRepository/eventRepository.server.ts index de2a19e395e..efc879b2f53 100644 --- a/apps/webapp/app/v3/eventRepository/eventRepository.server.ts +++ b/apps/webapp/app/v3/eventRepository/eventRepository.server.ts @@ -18,6 +18,7 @@ import { unflattenAttributes, } from "@trigger.dev/core/v3"; import { serializeTraceparent } from "@trigger.dev/core/v3/isomorphic"; +import type { MetricsV1Input } from "@internal/clickhouse"; import { Prisma, TaskEvent, TaskEventKind } from "@trigger.dev/database"; import { nanoid } from "nanoid"; import { Gauge } from "prom-client"; @@ -151,7 +152,7 @@ export class EventRepository implements IEventRepository { await this.#flushBatch(nanoid(), [this.#createableEventToPrismaEvent(event)]); } - async insertMany(events: CreateEventInput[]) { + insertMany(events: CreateEventInput[]) { this._flushScheduler.addToBatch(events.map(this.#createableEventToPrismaEvent)); } @@ -159,6 +160,8 @@ export class EventRepository implements IEventRepository { await this.#flushBatchWithReturn(nanoid(), events.map(this.#createableEventToPrismaEvent)); } + insertManyMetrics(_rows: MetricsV1Input[]): void {} + async completeSuccessfulRunEvent({ run, endTime }: { run: CompleteableTaskRun; endTime?: Date }) { const startTime = convertDateToNanoseconds(run.createdAt); diff --git a/apps/webapp/app/v3/eventRepository/eventRepository.types.ts b/apps/webapp/app/v3/eventRepository/eventRepository.types.ts index 0b45e536490..5c617a2482a 100644 --- a/apps/webapp/app/v3/eventRepository/eventRepository.types.ts +++ b/apps/webapp/app/v3/eventRepository/eventRepository.types.ts @@ -14,6 +14,7 @@ import type { TaskEventStatus, TaskRun, } from "@trigger.dev/database"; +import type { MetricsV1Input } from "@internal/clickhouse"; import type { DetailedTraceEvent, TaskEventStoreTable } from "../taskEventStore.server"; export type { ExceptionEventProperties }; @@ -345,8 +346,9 @@ export type TraceDetailedSummary = { export interface IEventRepository { maximumLiveReloadingSetting: number; // Event insertion methods - insertMany(events: CreateEventInput[]): Promise; + insertMany(events: CreateEventInput[]): void; insertManyImmediate(events: CreateEventInput[]): Promise; + insertManyMetrics(rows: MetricsV1Input[]): void; // Run event completion methods completeSuccessfulRunEvent(params: { run: CompleteableTaskRun; endTime?: Date }): Promise; diff --git a/apps/webapp/app/v3/eventRepository/index.server.ts b/apps/webapp/app/v3/eventRepository/index.server.ts index 70ea6440321..5c9026572c6 100644 --- a/apps/webapp/app/v3/eventRepository/index.server.ts +++ b/apps/webapp/app/v3/eventRepository/index.server.ts @@ -1,29 +1,12 @@ import { env } from "~/env.server"; import { eventRepository } from "./eventRepository.server"; -import { - clickhouseEventRepository, - clickhouseEventRepositoryV2, -} from "./clickhouseEventRepositoryInstance.server"; -import { IEventRepository, TraceEventOptions } from "./eventRepository.types"; +import { type IEventRepository, type TraceEventOptions } from "./eventRepository.types"; import { prisma } from "~/db.server"; import { logger } from "~/services/logger.server"; import { FEATURE_FLAG } from "../featureFlags"; import { flag } from "../featureFlags.server"; import { getTaskEventStore } from "../taskEventStore.server"; - -export function resolveEventRepositoryForStore(store: string | undefined): IEventRepository { - const taskEventStore = store ?? env.EVENT_REPOSITORY_DEFAULT_STORE; - - if (taskEventStore === "clickhouse_v2") { - return clickhouseEventRepositoryV2; - } - - if (taskEventStore === "clickhouse") { - return clickhouseEventRepository; - } - - return eventRepository; -} +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; export const EVENT_STORE_TYPES = { POSTGRES: "postgres", @@ -58,63 +41,68 @@ export async function getConfiguredEventRepository( (organization.featureFlags as Record | null) ?? undefined ); + const { repository: resolvedRepository } = await clickhouseFactory.getEventRepositoryForOrganization( + taskEventStore, + organizationId + ); + if (taskEventStore === EVENT_STORE_TYPES.CLICKHOUSE_V2) { - return { repository: clickhouseEventRepositoryV2, store: EVENT_STORE_TYPES.CLICKHOUSE_V2 }; + return { repository: resolvedRepository, store: EVENT_STORE_TYPES.CLICKHOUSE_V2 }; } if (taskEventStore === EVENT_STORE_TYPES.CLICKHOUSE) { - return { repository: clickhouseEventRepository, store: EVENT_STORE_TYPES.CLICKHOUSE }; + return { repository: resolvedRepository, store: EVENT_STORE_TYPES.CLICKHOUSE }; } return { repository: eventRepository, store: EVENT_STORE_TYPES.POSTGRES }; } export async function getEventRepository( + organizationId: string, featureFlags: Record | undefined, parentStore: string | undefined ): Promise<{ repository: IEventRepository; store: string }> { - if (typeof parentStore === "string") { - if (parentStore === "clickhouse_v2") { - return { repository: clickhouseEventRepositoryV2, store: "clickhouse_v2" }; - } - if (parentStore === "clickhouse") { - return { repository: clickhouseEventRepository, store: "clickhouse" }; - } else { - return { repository: eventRepository, store: getTaskEventStore() }; - } - } - - const taskEventRepository = await resolveTaskEventRepositoryFlag(featureFlags); + const taskEventStore = parentStore ?? (await resolveTaskEventRepositoryFlag(featureFlags)); + const { repository: resolvedRepository } = await clickhouseFactory.getEventRepositoryForOrganization( + taskEventStore, + organizationId + ); - if (taskEventRepository === "clickhouse_v2") { - return { repository: clickhouseEventRepositoryV2, store: "clickhouse_v2" }; + if (taskEventStore === "clickhouse_v2") { + return { repository: resolvedRepository, store: "clickhouse_v2" }; } - if (taskEventRepository === "clickhouse") { - return { repository: clickhouseEventRepository, store: "clickhouse" }; + if (taskEventStore === "clickhouse") { + return { repository: resolvedRepository, store: "clickhouse" }; } return { repository: eventRepository, store: getTaskEventStore() }; } export async function getV3EventRepository( + organizationId: string, parentStore: string | undefined ): Promise<{ repository: IEventRepository; store: string }> { if (typeof parentStore === "string") { - if (parentStore === "clickhouse_v2") { - return { repository: clickhouseEventRepositoryV2, store: "clickhouse_v2" }; - } - if (parentStore === "clickhouse") { - return { repository: clickhouseEventRepository, store: "clickhouse" }; - } else { - return { repository: eventRepository, store: getTaskEventStore() }; - } + const { repository: resolvedRepository } = await clickhouseFactory.getEventRepositoryForOrganization( + parentStore, + organizationId + ); + return { repository: resolvedRepository, store: parentStore }; } if (env.EVENT_REPOSITORY_DEFAULT_STORE === "clickhouse_v2") { - return { repository: clickhouseEventRepositoryV2, store: "clickhouse_v2" }; + const { repository: resolvedRepository } = await clickhouseFactory.getEventRepositoryForOrganization( + "clickhouse_v2", + organizationId + ); + return { repository: resolvedRepository, store: "clickhouse_v2" }; } else if (env.EVENT_REPOSITORY_DEFAULT_STORE === "clickhouse") { - return { repository: clickhouseEventRepository, store: "clickhouse" }; + const { repository: resolvedRepository } = await clickhouseFactory.getEventRepositoryForOrganization( + "clickhouse", + organizationId + ); + return { repository: resolvedRepository, store: "clickhouse" }; } else { return { repository: eventRepository, store: getTaskEventStore() }; } @@ -203,7 +191,10 @@ async function recordRunEvent( }; } - const $eventRepository = resolveEventRepositoryForStore(foundRun.taskEventStore); + const { repository: $eventRepository } = await clickhouseFactory.getEventRepositoryForOrganization( + foundRun.taskEventStore, + foundRun.runtimeEnvironment.organizationId + ); const { attributes, startTime, ...optionsRest } = options; diff --git a/apps/webapp/app/v3/otlpExporter.server.ts b/apps/webapp/app/v3/otlpExporter.server.ts index 7505693e3ab..37c5c2840bc 100644 --- a/apps/webapp/app/v3/otlpExporter.server.ts +++ b/apps/webapp/app/v3/otlpExporter.server.ts @@ -20,15 +20,12 @@ import { } from "@trigger.dev/otlp-importer"; import type { MetricsV1Input } from "@internal/clickhouse"; import { logger } from "~/services/logger.server"; -import { clickhouseClient } from "~/services/clickhouseInstance.server"; -import { DynamicFlushScheduler } from "./dynamicFlushScheduler.server"; -import { ClickhouseEventRepository } from "./eventRepository/clickhouseEventRepository.server"; import { - clickhouseEventRepository, - clickhouseEventRepositoryV2, -} from "./eventRepository/clickhouseEventRepositoryInstance.server"; + clickhouseFactory, + type ClickhouseFactory, +} from "~/services/clickhouse/clickhouseFactory.server"; + import { generateSpanId } from "./eventRepository/common.server"; -import { EventRepository, eventRepository } from "./eventRepository/eventRepository.server"; import type { CreatableEventKind, CreatableEventStatus, @@ -39,21 +36,25 @@ import { startSpan } from "./tracing.server"; import { enrichCreatableEvents } from "./utils/enrichCreatableEvents.server"; import { waitForLlmPricingReady } from "./llmPricingRegistry.server"; import { env } from "~/env.server"; -import { detectBadJsonStrings } from "~/utils/detectBadJsonStrings"; import { singleton } from "~/utils/singleton"; +type OTLPExporterConfig = { + clickhouseFactory: ClickhouseFactory; + verbose: boolean; + spanAttributeValueLengthLimit: number; +}; + class OTLPExporter { private _tracer: Tracer; + private readonly _clickhouseFactory: ClickhouseFactory; + private readonly _verbose: boolean; + private readonly _spanAttributeValueLengthLimit: number; - constructor( - private readonly _eventRepository: EventRepository, - private readonly _clickhouseEventRepository: ClickhouseEventRepository, - private readonly _clickhouseEventRepositoryV2: ClickhouseEventRepository, - private readonly _metricsFlushScheduler: DynamicFlushScheduler, - private readonly _verbose: boolean, - private readonly _spanAttributeValueLengthLimit: number - ) { + constructor(config: OTLPExporterConfig) { this._tracer = trace.getTracer("otlp-exporter"); + this._clickhouseFactory = config.clickhouseFactory; + this._verbose = config.verbose; + this._spanAttributeValueLengthLimit = config.spanAttributeValueLengthLimit; } async exportTraces(request: ExportTraceServiceRequest): Promise { @@ -74,23 +75,16 @@ class OTLPExporter { }); } - async exportMetrics( - request: ExportMetricsServiceRequest - ): Promise { + async exportMetrics(request: ExportMetricsServiceRequest): Promise { return await startSpan(this._tracer, "exportMetrics", async (span) => { - const rows = this.#filterResourceMetrics(request.resourceMetrics).flatMap( - (resourceMetrics) => { - return convertMetricsToClickhouseRows( - resourceMetrics, - this._spanAttributeValueLengthLimit - ); - } + const rows = this.#filterResourceMetrics(request.resourceMetrics).flatMap((resourceMetrics) => + convertMetricsToClickhouseRows(resourceMetrics, this._spanAttributeValueLengthLimit) ); span.setAttribute("metric_row_count", rows.length); if (rows.length > 0) { - this._metricsFlushScheduler.addToBatch(rows); + await this.#exportMetricRows(rows); } return ExportMetricsServiceResponse.create(); @@ -118,40 +112,73 @@ class OTLPExporter { async #exportEvents( eventsWithStores: { events: Array; taskEventStore: string }[] ) { - const eventsGroupedByStore = eventsWithStores.reduce((acc, { events, taskEventStore }) => { - acc[taskEventStore] = acc[taskEventStore] || []; - acc[taskEventStore].push(...events); - return acc; - }, {} as Record>); + await waitForLlmPricingReady(); + + // Group by unique event repositories + const routeCache = new Map(); + const groups = new Map(); + for (const { events, taskEventStore } of eventsWithStores) { + for (const event of events) { + const routeKey = `${event.organizationId}\0${taskEventStore}`; + let resolved = routeCache.get(routeKey); + if (!resolved) { + resolved = this._clickhouseFactory.getEventRepositoryForOrganizationSync( + taskEventStore, + event.organizationId + ); + routeCache.set(routeKey, resolved); + } - let eventCount = 0; + let group = groups.get(resolved.key); + if (!group) { + group = { repository: resolved.repository, events: [] }; + groups.set(resolved.key, group); + } + group.events.push(event); + } + } - for (const [store, events] of Object.entries(eventsGroupedByStore)) { - const eventRepository = this.#getEventRepositoryForStore(store); + let eventCount = 0; - await waitForLlmPricingReady(); + for (const [repoKey, { repository, events }] of groups) { const enrichedEvents = enrichCreatableEvents(events); - this.#logEventsVerbose(enrichedEvents, `exportEvents ${store}`); + this.#logEventsVerbose(enrichedEvents, `exportEvents ${repoKey}`); eventCount += enrichedEvents.length; - await eventRepository.insertMany(enrichedEvents); + repository.insertMany(enrichedEvents); } return eventCount; } - #getEventRepositoryForStore(store: string): IEventRepository { - if (store === "clickhouse") { - return this._clickhouseEventRepository; - } + async #exportMetricRows(rows: MetricsV1Input[]): Promise { + const routeCache = new Map(); + const groups = new Map(); + + for (const row of rows) { + const routeKey = row.organization_id; + let resolved = routeCache.get(routeKey); + if (!resolved) { + resolved = this._clickhouseFactory.getEventRepositoryForOrganizationSync( + "clickhouse_v2", + row.organization_id + ); + routeCache.set(routeKey, resolved); + } - if (store === "clickhouse_v2") { - return this._clickhouseEventRepositoryV2; + let group = groups.get(resolved.key); + if (!group) { + group = { repository: resolved.repository, rows: [] }; + groups.set(resolved.key, group); + } + group.rows.push(row); } - return this._eventRepository; + for (const [, { repository, rows: groupedRows }] of groups) { + repository.insertManyMetrics(groupedRows); + } } #logEventsVerbose(events: CreateEventInput[], prefix: string) { @@ -393,7 +420,10 @@ function convertSpansToCreateableEvents( SemanticInternalAttributes.METADATA ); - const runTags = extractArrayAttribute(span.attributes ?? [], SemanticInternalAttributes.RUN_TAGS); + const runTags = extractArrayAttribute( + span.attributes ?? [], + SemanticInternalAttributes.RUN_TAGS + ); const properties = truncateAttributes( @@ -464,7 +494,10 @@ function floorToTenSecondBucket(timeUnixNano: bigint | number): string { const flooredMs = Math.floor(epochMs / 10_000) * 10_000; const date = new Date(flooredMs); // Format as ClickHouse DateTime: YYYY-MM-DD HH:MM:SS - return date.toISOString().replace("T", " ").replace(/\.\d{3}Z$/, ""); + return date + .toISOString() + .replace("T", " ") + .replace(/\.\d{3}Z$/, ""); } function convertMetricsToClickhouseRows( @@ -584,8 +617,7 @@ function resolveDataPointContext( attributes: Record; } { const runId = - resourceCtx.runId ?? - extractStringAttribute(dpAttributes, SemanticInternalAttributes.RUN_ID); + resourceCtx.runId ?? extractStringAttribute(dpAttributes, SemanticInternalAttributes.RUN_ID); const taskSlug = resourceCtx.taskSlug ?? extractStringAttribute(dpAttributes, SemanticInternalAttributes.TASK_SLUG); @@ -1172,26 +1204,13 @@ function hasUnpairedSurrogateAtEnd(str: string): boolean { export const otlpExporter = singleton("otlpExporter", initializeOTLPExporter); -function initializeOTLPExporter() { - const metricsFlushScheduler = new DynamicFlushScheduler({ - batchSize: env.METRICS_CLICKHOUSE_BATCH_SIZE, - flushInterval: env.METRICS_CLICKHOUSE_FLUSH_INTERVAL_MS, - callback: async (_flushId, batch) => { - await clickhouseClient.metrics.insert(batch); - }, - minConcurrency: 1, - maxConcurrency: env.METRICS_CLICKHOUSE_MAX_CONCURRENCY, - loadSheddingEnabled: false, - }); - - return new OTLPExporter( - eventRepository, - clickhouseEventRepository, - clickhouseEventRepositoryV2, - metricsFlushScheduler, - process.env.OTLP_EXPORTER_VERBOSE === "1", - process.env.SERVER_OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT +async function initializeOTLPExporter() { + await clickhouseFactory.isReady(); + return new OTLPExporter({ + clickhouseFactory, + verbose: process.env.OTLP_EXPORTER_VERBOSE === "1", + spanAttributeValueLengthLimit: process.env.SERVER_OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT ? parseInt(process.env.SERVER_OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT, 10) - : 8192 - ); -} \ No newline at end of file + : 8192, + }); +} diff --git a/apps/webapp/app/v3/services/alerts/errorAlertEvaluator.server.ts b/apps/webapp/app/v3/services/alerts/errorAlertEvaluator.server.ts index e935a8a6911..40fe2d5f10d 100644 --- a/apps/webapp/app/v3/services/alerts/errorAlertEvaluator.server.ts +++ b/apps/webapp/app/v3/services/alerts/errorAlertEvaluator.server.ts @@ -7,7 +7,7 @@ import { } from "@trigger.dev/database"; import { $replica, prisma } from "~/db.server"; import { ErrorAlertConfig } from "~/models/projectAlert.server"; -import { clickhouseClient } from "~/services/clickhouseInstance.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { logger } from "~/services/logger.server"; import { alertsWorker } from "~/v3/alertsWorker.server"; @@ -45,8 +45,7 @@ const DEFAULT_INTERVAL_MS = 300_000; export class ErrorAlertEvaluator { constructor( protected readonly _prisma: PrismaClientOrTransaction = prisma, - protected readonly _replica: PrismaClientOrTransaction = $replica, - protected readonly _clickhouse: ClickHouse = clickhouseClient + protected readonly _replica: PrismaClientOrTransaction = $replica ) {} async evaluate(projectId: string, scheduledAt: number): Promise { @@ -245,10 +244,7 @@ export class ErrorAlertEvaluator { } } - if ( - state.ignoredUntilTotalOccurrences != null && - state.ignoredAtOccurrenceCount != null - ) { + if (state.ignoredUntilTotalOccurrences != null && state.ignoredAtOccurrenceCount != null) { const occurrencesSinceIgnored = context.totalOccurrenceCount - Number(state.ignoredAtOccurrenceCount); if (occurrencesSinceIgnored >= state.ignoredUntilTotalOccurrences) { @@ -335,7 +331,11 @@ export class ErrorAlertEvaluator { envIds: string[], scheduledAt: number ): Promise { - const qb = this._clickhouse.errors.activeErrorsSinceQueryBuilder(); + const queryClickhouse = await clickhouseFactory.getClickhouseForOrganization( + organizationId, + "query" + ); + const qb = queryClickhouse.errors.activeErrorsSinceQueryBuilder(); qb.where("organization_id = {organizationId: String}", { organizationId }); qb.where("project_id = {projectId: String}", { projectId }); qb.where("environment_id IN {envIds: Array(String)}", { envIds }); @@ -389,7 +389,11 @@ export class ErrorAlertEvaluator { occurrences_since: number; }> > { - const qb = this._clickhouse.errors.occurrenceCountsSinceQueryBuilder(); + const queryClickhouse = await clickhouseFactory.getClickhouseForOrganization( + organizationId, + "query" + ); + const qb = queryClickhouse.errors.occurrenceCountsSinceQueryBuilder(); qb.where("organization_id = {organizationId: String}", { organizationId }); qb.where("project_id = {projectId: String}", { projectId }); qb.where("environment_id IN {envIds: Array(String)}", { envIds }); diff --git a/apps/webapp/app/v3/services/bulk/BulkActionV2.server.ts b/apps/webapp/app/v3/services/bulk/BulkActionV2.server.ts index 156b68bff59..cf1f80165dc 100644 --- a/apps/webapp/app/v3/services/bulk/BulkActionV2.server.ts +++ b/apps/webapp/app/v3/services/bulk/BulkActionV2.server.ts @@ -7,7 +7,7 @@ import { } from "@trigger.dev/database"; import { getRunFiltersFromRequest } from "~/presenters/RunFilters.server"; import { type CreateBulkActionPayload } from "~/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.bulkaction"; -import { clickhouseClient } from "~/services/clickhouseInstance.server"; +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; import { parseRunListInputOptions, type RunListInputFilters, @@ -38,8 +38,9 @@ export class BulkActionService extends BaseService { const filters = await getFilters(payload, request); // Count the runs that will be affected by the bulk action + const clickhouse = await clickhouseFactory.getClickhouseForOrganization(organizationId, "standard"); const runsRepository = new RunsRepository({ - clickhouse: clickhouseClient, + clickhouse, prisma: this._replica as PrismaClient, }); const count = await runsRepository.countRuns({ @@ -147,8 +148,9 @@ export class BulkActionService extends BaseService { ...rawParams, }); + const clickhouse = await clickhouseFactory.getClickhouseForOrganization(group.project.organizationId, "standard"); const runsRepository = new RunsRepository({ - clickhouse: clickhouseClient, + clickhouse, prisma: this._replica as PrismaClient, }); diff --git a/apps/webapp/test/clickhouseFactory.test.ts b/apps/webapp/test/clickhouseFactory.test.ts new file mode 100644 index 00000000000..501462ab677 --- /dev/null +++ b/apps/webapp/test/clickhouseFactory.test.ts @@ -0,0 +1,130 @@ +import { describe, expect } from "vitest"; + +vi.mock("~/db.server", () => ({ prisma: {}, $replica: {} })); + +import { postgresTest } from "@internal/testcontainers"; +import { OrganizationDataStoresRegistry } from "~/services/dataStores/organizationDataStoresRegistry.server"; +import { ClickhouseConnectionSchema } from "~/services/clickhouse/clickhouseSecretSchemas.server"; +import { ClickhouseFactory } from "~/services/clickhouse/clickhouseFactory.server"; + +vi.setConfig({ testTimeout: 60_000 }); + +const TEST_URL = "https://default:password@ch-org.example.com:8443"; +const TEST_URL_2 = "https://default:password@ch-other.example.com:8443"; + +describe("ClickHouse Factory", () => { + postgresTest( + "returns default client when org has no data store", + async ({ prisma }) => { + const registry = new OrganizationDataStoresRegistry(prisma); + await registry.loadFromDatabase(); + + const factory = new ClickhouseFactory(registry); + const client = await factory.getClickhouseForOrganization("org-no-store", "standard"); + expect(client).toBeDefined(); + } + ); + + postgresTest( + "returns org-specific client when a data store is configured", + async ({ prisma }) => { + const registry = new OrganizationDataStoresRegistry(prisma); + + await registry.addDataStore({ + key: "factory-store", + kind: "CLICKHOUSE", + organizationIds: ["org-custom"], + config: ClickhouseConnectionSchema.parse({ url: TEST_URL }), + }); + + await registry.loadFromDatabase(); + + const factory = new ClickhouseFactory(registry); + const client = await factory.getClickhouseForOrganization("org-custom", "standard"); + expect(client).toBeDefined(); + + // Default client is a different instance from the org-specific one + const defaultClient = await factory.getClickhouseForOrganization("org-no-store", "standard"); + expect(client).not.toBe(defaultClient); + } + ); + + postgresTest( + "two orgs sharing the same data store get the same cached client", + async ({ prisma }) => { + const registry = new OrganizationDataStoresRegistry(prisma); + + await registry.addDataStore({ + key: "shared-factory-store", + kind: "CLICKHOUSE", + organizationIds: ["org-shared-1", "org-shared-2"], + config: ClickhouseConnectionSchema.parse({ url: TEST_URL }), + }); + + await registry.loadFromDatabase(); + + const factory = new ClickhouseFactory(registry); + const client1 = await factory.getClickhouseForOrganization("org-shared-1", "standard"); + const client2 = await factory.getClickhouseForOrganization("org-shared-2", "standard"); + + // Same hostname → same cached client instance + expect(client1).toBe(client2); + } + ); + + postgresTest( + "two data stores with different URLs produce different clients", + async ({ prisma }) => { + const registry = new OrganizationDataStoresRegistry(prisma); + + await registry.addDataStore({ + key: "store-a", + kind: "CLICKHOUSE", + organizationIds: ["org-a"], + config: ClickhouseConnectionSchema.parse({ url: TEST_URL }), + }); + + await registry.addDataStore({ + key: "store-b", + kind: "CLICKHOUSE", + organizationIds: ["org-b"], + config: ClickhouseConnectionSchema.parse({ url: TEST_URL_2 }), + }); + + await registry.loadFromDatabase(); + + const factory = new ClickhouseFactory(registry); + const clientA = await factory.getClickhouseForOrganization("org-a", "standard"); + const clientB = await factory.getClickhouseForOrganization("org-b", "standard"); + + expect(clientA).not.toBe(clientB); + } + ); + + postgresTest( + "after reload with a deleted store, org falls back to default", + async ({ prisma }) => { + const registry = new OrganizationDataStoresRegistry(prisma); + + await registry.addDataStore({ + key: "removable-store", + kind: "CLICKHOUSE", + organizationIds: ["org-removable"], + config: ClickhouseConnectionSchema.parse({ url: TEST_URL }), + }); + + await registry.loadFromDatabase(); + + const factory = new ClickhouseFactory(registry); + const before = await factory.getClickhouseForOrganization("org-removable", "standard"); + const defaultClient = await factory.getClickhouseForOrganization("org-no-store", "standard"); + expect(before).not.toBe(defaultClient); + + await registry.deleteDataStore({ key: "removable-store", kind: "CLICKHOUSE" }); + await registry.reload(); + + const after = await factory.getClickhouseForOrganization("org-removable", "standard"); + expect(after).toBe(defaultClient); + } + ); +}); diff --git a/apps/webapp/test/organizationDataStoresRegistry.test.ts b/apps/webapp/test/organizationDataStoresRegistry.test.ts new file mode 100644 index 00000000000..44fe851a404 --- /dev/null +++ b/apps/webapp/test/organizationDataStoresRegistry.test.ts @@ -0,0 +1,197 @@ +import { describe, expect } from "vitest"; + +vi.mock("~/db.server", () => ({ prisma: {}, $replica: {} })); + +import { postgresTest } from "@internal/testcontainers"; +import { OrganizationDataStoresRegistry } from "~/services/dataStores/organizationDataStoresRegistry.server"; +import { ClickhouseConnectionSchema } from "~/services/clickhouse/clickhouseSecretSchemas.server"; + +vi.setConfig({ testTimeout: 60_000 }); + +const TEST_URL = "https://default:password@clickhouse.example.com:8443"; +const TEST_URL_2 = "https://default:password@clickhouse2.example.com:8443"; + +describe("OrganizationDataStoresRegistry", () => { + postgresTest("isLoaded is false before loadFromDatabase", async ({ prisma }) => { + const registry = new OrganizationDataStoresRegistry(prisma); + expect(registry.isLoaded).toBe(false); + expect(registry.get("any-org", "CLICKHOUSE")).toBeNull(); + }); + + postgresTest("isLoaded is true after loadFromDatabase", async ({ prisma }) => { + const registry = new OrganizationDataStoresRegistry(prisma); + await registry.loadFromDatabase(); + expect(registry.isLoaded).toBe(true); + }); + + postgresTest("isReady resolves after loadFromDatabase", async ({ prisma }) => { + const registry = new OrganizationDataStoresRegistry(prisma); + let resolved = false; + registry.isReady.then(() => { + resolved = true; + }); + await registry.loadFromDatabase(); + await registry.isReady; + expect(resolved).toBe(true); + }); + + postgresTest("get returns null when no data stores exist", async ({ prisma }) => { + const registry = new OrganizationDataStoresRegistry(prisma); + await registry.loadFromDatabase(); + expect(registry.get("org-1", "CLICKHOUSE")).toBeNull(); + }); + + postgresTest("addDataStore creates a row and stores the secret", async ({ prisma }) => { + const registry = new OrganizationDataStoresRegistry(prisma); + + await registry.addDataStore({ + key: "test-store", + kind: "CLICKHOUSE", + organizationIds: ["org-1", "org-2"], + config: ClickhouseConnectionSchema.parse({ url: TEST_URL }), + }); + + const row = await prisma.organizationDataStore.findFirst({ where: { key: "test-store" } }); + expect(row).not.toBeNull(); + expect(row?.organizationIds).toEqual(["org-1", "org-2"]); + expect(row?.kind).toBe("CLICKHOUSE"); + + const secret = await prisma.secretStore.findFirst({ + where: { key: "data-store:test-store:clickhouse" }, + }); + expect(secret).not.toBeNull(); + }); + + postgresTest("loadFromDatabase resolves secrets and makes orgs available via get", async ({ prisma }) => { + const registry = new OrganizationDataStoresRegistry(prisma); + + await registry.addDataStore({ + key: "hipaa-store", + kind: "CLICKHOUSE", + organizationIds: ["org-hipaa"], + config: ClickhouseConnectionSchema.parse({ url: TEST_URL }), + }); + + await registry.loadFromDatabase(); + + const result = registry.get("org-hipaa", "CLICKHOUSE"); + expect(result).not.toBeNull(); + expect(result?.kind).toBe("CLICKHOUSE"); + expect(result?.url).toBe(TEST_URL); + }); + + postgresTest("get returns null for orgs not in any data store", async ({ prisma }) => { + const registry = new OrganizationDataStoresRegistry(prisma); + + await registry.addDataStore({ + key: "partial-store", + kind: "CLICKHOUSE", + organizationIds: ["org-a"], + config: ClickhouseConnectionSchema.parse({ url: TEST_URL }), + }); + + await registry.loadFromDatabase(); + + expect(registry.get("org-a", "CLICKHOUSE")).not.toBeNull(); + expect(registry.get("org-b", "CLICKHOUSE")).toBeNull(); + }); + + postgresTest("multiple orgs can share the same data store", async ({ prisma }) => { + const registry = new OrganizationDataStoresRegistry(prisma); + + await registry.addDataStore({ + key: "shared-store", + kind: "CLICKHOUSE", + organizationIds: ["org-x", "org-y", "org-z"], + config: ClickhouseConnectionSchema.parse({ url: TEST_URL }), + }); + + await registry.loadFromDatabase(); + + const x = registry.get("org-x", "CLICKHOUSE"); + const y = registry.get("org-y", "CLICKHOUSE"); + const z = registry.get("org-z", "CLICKHOUSE"); + + expect(x?.url).toBe(TEST_URL); + expect(y?.url).toBe(TEST_URL); + expect(z?.url).toBe(TEST_URL); + }); + + postgresTest("updateDataStore updates organizationIds and rotates the secret", async ({ prisma }) => { + const registry = new OrganizationDataStoresRegistry(prisma); + + await registry.addDataStore({ + key: "update-store", + kind: "CLICKHOUSE", + organizationIds: ["org-old"], + config: ClickhouseConnectionSchema.parse({ url: TEST_URL }), + }); + + await registry.updateDataStore({ + key: "update-store", + kind: "CLICKHOUSE", + organizationIds: ["org-new-1", "org-new-2"], + config: ClickhouseConnectionSchema.parse({ url: TEST_URL_2 }), + }); + + const row = await prisma.organizationDataStore.findFirst({ where: { key: "update-store" } }); + expect(row?.organizationIds).toEqual(["org-new-1", "org-new-2"]); + + await registry.loadFromDatabase(); + expect(registry.get("org-new-1", "CLICKHOUSE")?.url).toBe(TEST_URL_2); + expect(registry.get("org-old", "CLICKHOUSE")).toBeNull(); + }); + + postgresTest("reload picks up changes made after initial load", async ({ prisma }) => { + const registry = new OrganizationDataStoresRegistry(prisma); + await registry.loadFromDatabase(); + expect(registry.get("org-reload", "CLICKHOUSE")).toBeNull(); + + await registry.addDataStore({ + key: "reload-store", + kind: "CLICKHOUSE", + organizationIds: ["org-reload"], + config: ClickhouseConnectionSchema.parse({ url: TEST_URL }), + }); + + expect(registry.get("org-reload", "CLICKHOUSE")).toBeNull(); + + await registry.reload(); + expect(registry.get("org-reload", "CLICKHOUSE")?.url).toBe(TEST_URL); + }); + + postgresTest("deleteDataStore removes the row and its secret", async ({ prisma }) => { + const registry = new OrganizationDataStoresRegistry(prisma); + + await registry.addDataStore({ + key: "delete-store", + kind: "CLICKHOUSE", + organizationIds: ["org-delete"], + config: ClickhouseConnectionSchema.parse({ url: TEST_URL }), + }); + + await registry.deleteDataStore({ key: "delete-store", kind: "CLICKHOUSE" }); + + expect(await prisma.organizationDataStore.findFirst({ where: { key: "delete-store" } })).toBeNull(); + expect(await prisma.secretStore.findFirst({ where: { key: "data-store:delete-store:clickhouse" } })).toBeNull(); + }); + + postgresTest("after delete and reload, org no longer has a data store", async ({ prisma }) => { + const registry = new OrganizationDataStoresRegistry(prisma); + + await registry.addDataStore({ + key: "ephemeral-store", + kind: "CLICKHOUSE", + organizationIds: ["org-ephemeral"], + config: ClickhouseConnectionSchema.parse({ url: TEST_URL }), + }); + + await registry.loadFromDatabase(); + expect(registry.get("org-ephemeral", "CLICKHOUSE")?.url).toBe(TEST_URL); + + await registry.deleteDataStore({ key: "ephemeral-store", kind: "CLICKHOUSE" }); + await registry.reload(); + + expect(registry.get("org-ephemeral", "CLICKHOUSE")).toBeNull(); + }); +}); diff --git a/apps/webapp/test/runsBackfiller.test.ts b/apps/webapp/test/runsBackfiller.test.ts index 7051fb976f5..87bc3822d98 100644 --- a/apps/webapp/test/runsBackfiller.test.ts +++ b/apps/webapp/test/runsBackfiller.test.ts @@ -12,6 +12,7 @@ import { z } from "zod"; import { RunsBackfillerService } from "~/services/runsBackfiller.server"; import { RunsReplicationService } from "~/services/runsReplicationService.server"; import { createInMemoryTracing } from "./utils/tracing"; +import { TestReplicationClickhouseFactory } from "./utils/testReplicationClickhouseFactory"; vi.setConfig({ testTimeout: 60_000 }); @@ -30,7 +31,7 @@ describe("RunsBackfillerService", () => { const { tracer, exporter } = createInMemoryTracing(); const runsReplicationService = new RunsReplicationService({ - clickhouse, + clickhouseFactory: new TestReplicationClickhouseFactory(clickhouse), pgConnectionUrl: postgresContainer.getConnectionUri(), serviceName: "runs-replication", slotName: "task_runs_to_clickhouse_v1", diff --git a/apps/webapp/test/runsReplicationBenchmark.test.ts b/apps/webapp/test/runsReplicationBenchmark.test.ts index 5c04ef55cc8..e17d6b41212 100644 --- a/apps/webapp/test/runsReplicationBenchmark.test.ts +++ b/apps/webapp/test/runsReplicationBenchmark.test.ts @@ -7,6 +7,7 @@ import path from "node:path"; import { z } from "zod"; import { RunsReplicationService } from "~/services/runsReplicationService.server"; import { createInMemoryTracing, createInMemoryMetrics } from "./utils/tracing"; +import { TestReplicationClickhouseFactory } from "./utils/testReplicationClickhouseFactory"; // Extend test timeout for benchmarks vi.setConfig({ testTimeout: 300_000 }); // 5 minutes @@ -320,7 +321,7 @@ async function runBenchmark( // Create and start replication service const runsReplicationService = new RunsReplicationService({ - clickhouse, + clickhouseFactory: new TestReplicationClickhouseFactory(clickhouse), pgConnectionUrl: postgresContainer.getConnectionUri(), serviceName: `benchmark-${name}`, slotName: `benchmark_${name.replace(/-/g, "_")}`, diff --git a/apps/webapp/test/runsReplicationService.part1.test.ts b/apps/webapp/test/runsReplicationService.part1.test.ts index 715d4583dc2..d2a3c1b7627 100644 --- a/apps/webapp/test/runsReplicationService.part1.test.ts +++ b/apps/webapp/test/runsReplicationService.part1.test.ts @@ -5,6 +5,7 @@ import { z } from "zod"; import { TaskRunStatus } from "~/database-types"; import { RunsReplicationService } from "~/services/runsReplicationService.server"; import { createInMemoryTracing, createInMemoryMetrics } from "./utils/tracing"; +import { TestReplicationClickhouseFactory } from "./utils/testReplicationClickhouseFactory"; import superjson from "superjson"; vi.setConfig({ testTimeout: 60_000 }); @@ -27,7 +28,7 @@ describe("RunsReplicationService (part 1/2)", () => { const { tracer, exporter } = createInMemoryTracing(); const runsReplicationService = new RunsReplicationService({ - clickhouse, + clickhouseFactory: new TestReplicationClickhouseFactory(clickhouse), pgConnectionUrl: postgresContainer.getConnectionUri(), serviceName: "runs-replication", slotName: "task_runs_to_clickhouse_v1", @@ -151,7 +152,7 @@ describe("RunsReplicationService (part 1/2)", () => { const { tracer, exporter } = createInMemoryTracing(); const runsReplicationService = new RunsReplicationService({ - clickhouse, + clickhouseFactory: new TestReplicationClickhouseFactory(clickhouse), pgConnectionUrl: postgresContainer.getConnectionUri(), serviceName: "runs-replication", slotName: "task_runs_to_clickhouse_v1", @@ -289,7 +290,7 @@ describe("RunsReplicationService (part 1/2)", () => { const { tracer, exporter } = createInMemoryTracing(); const runsReplicationService = new RunsReplicationService({ - clickhouse, + clickhouseFactory: new TestReplicationClickhouseFactory(clickhouse), pgConnectionUrl: postgresContainer.getConnectionUri(), serviceName: "runs-replication", slotName: "task_runs_to_clickhouse_v1", @@ -359,7 +360,7 @@ describe("RunsReplicationService (part 1/2)", () => { }); const runsReplicationService = new RunsReplicationService({ - clickhouse, + clickhouseFactory: new TestReplicationClickhouseFactory(clickhouse), pgConnectionUrl: postgresContainer.getConnectionUri(), serviceName: "runs-replication-batching", slotName: "task_runs_to_clickhouse_v1", @@ -463,7 +464,7 @@ describe("RunsReplicationService (part 1/2)", () => { }); const runsReplicationService = new RunsReplicationService({ - clickhouse, + clickhouseFactory: new TestReplicationClickhouseFactory(clickhouse), pgConnectionUrl: postgresContainer.getConnectionUri(), serviceName: "runs-replication-payload", slotName: "task_runs_to_clickhouse_v1", @@ -564,7 +565,7 @@ describe("RunsReplicationService (part 1/2)", () => { }); const runsReplicationService = new RunsReplicationService({ - clickhouse, + clickhouseFactory: new TestReplicationClickhouseFactory(clickhouse), pgConnectionUrl: postgresContainer.getConnectionUri(), serviceName: "runs-replication-payload", slotName: "task_runs_to_clickhouse_v1", @@ -670,7 +671,7 @@ describe("RunsReplicationService (part 1/2)", () => { }); const runsReplicationService = new RunsReplicationService({ - clickhouse, + clickhouseFactory: new TestReplicationClickhouseFactory(clickhouse), pgConnectionUrl: postgresContainer.getConnectionUri(), serviceName: "runs-replication-update", slotName: "task_runs_to_clickhouse_v1", @@ -777,7 +778,7 @@ describe("RunsReplicationService (part 1/2)", () => { }); const runsReplicationService = new RunsReplicationService({ - clickhouse, + clickhouseFactory: new TestReplicationClickhouseFactory(clickhouse), pgConnectionUrl: postgresContainer.getConnectionUri(), serviceName: "runs-replication-delete", slotName: "task_runs_to_clickhouse_v1", @@ -878,7 +879,7 @@ describe("RunsReplicationService (part 1/2)", () => { // Service A const runsReplicationServiceA = new RunsReplicationService({ - clickhouse, + clickhouseFactory: new TestReplicationClickhouseFactory(clickhouse), pgConnectionUrl: postgresContainer.getConnectionUri(), serviceName: "runs-replication-shutdown-handover", slotName: "task_runs_to_clickhouse_v1", @@ -982,7 +983,7 @@ describe("RunsReplicationService (part 1/2)", () => { // Service B const runsReplicationServiceB = new RunsReplicationService({ - clickhouse, + clickhouseFactory: new TestReplicationClickhouseFactory(clickhouse), pgConnectionUrl: postgresContainer.getConnectionUri(), serviceName: "runs-replication-shutdown-handover", slotName: "task_runs_to_clickhouse_v1", @@ -1029,7 +1030,7 @@ describe("RunsReplicationService (part 1/2)", () => { // Service A const runsReplicationServiceA = new RunsReplicationService({ - clickhouse, + clickhouseFactory: new TestReplicationClickhouseFactory(clickhouse), pgConnectionUrl: postgresContainer.getConnectionUri(), serviceName: "runs-replication-shutdown-after-processed", slotName: "task_runs_to_clickhouse_v1", @@ -1131,7 +1132,7 @@ describe("RunsReplicationService (part 1/2)", () => { // Service B const runsReplicationServiceB = new RunsReplicationService({ - clickhouse, + clickhouseFactory: new TestReplicationClickhouseFactory(clickhouse), pgConnectionUrl: postgresContainer.getConnectionUri(), serviceName: "runs-replication-shutdown-after-processed", slotName: "task_runs_to_clickhouse_v1", @@ -1174,7 +1175,7 @@ describe("RunsReplicationService (part 1/2)", () => { const metricsHelper = createInMemoryMetrics(); const runsReplicationService = new RunsReplicationService({ - clickhouse, + clickhouseFactory: new TestReplicationClickhouseFactory(clickhouse), pgConnectionUrl: postgresContainer.getConnectionUri(), serviceName: "runs-replication-metrics", slotName: "task_runs_to_clickhouse_v1", diff --git a/apps/webapp/test/runsReplicationService.part2.test.ts b/apps/webapp/test/runsReplicationService.part2.test.ts index a93f99f246b..32bfb5ba73d 100644 --- a/apps/webapp/test/runsReplicationService.part2.test.ts +++ b/apps/webapp/test/runsReplicationService.part2.test.ts @@ -6,6 +6,7 @@ import { setTimeout } from "node:timers/promises"; import { z } from "zod"; import { RunsReplicationService } from "~/services/runsReplicationService.server"; import { detectBadJsonStrings } from "~/utils/detectBadJsonStrings"; +import { TestReplicationClickhouseFactory } from "./utils/testReplicationClickhouseFactory"; vi.setConfig({ testTimeout: 60_000 }); @@ -23,7 +24,7 @@ describe("RunsReplicationService (part 2/2)", () => { // Service A const runsReplicationServiceA = new RunsReplicationService({ - clickhouse, + clickhouseFactory: new TestReplicationClickhouseFactory(clickhouse), pgConnectionUrl: postgresContainer.getConnectionUri(), serviceName: "runs-replication-shutdown-handover", slotName: "task_runs_to_clickhouse_v1", @@ -43,7 +44,7 @@ describe("RunsReplicationService (part 2/2)", () => { // Service A const runsReplicationServiceB = new RunsReplicationService({ - clickhouse, + clickhouseFactory: new TestReplicationClickhouseFactory(clickhouse), pgConnectionUrl: postgresContainer.getConnectionUri(), serviceName: "runs-replication-shutdown-handover", slotName: "task_runs_to_clickhouse_v1", @@ -152,7 +153,7 @@ describe("RunsReplicationService (part 2/2)", () => { }); const runsReplicationService = new RunsReplicationService({ - clickhouse, + clickhouseFactory: new TestReplicationClickhouseFactory(clickhouse), pgConnectionUrl: postgresContainer.getConnectionUri(), serviceName: "runs-replication-stress-bulk-insert", slotName: "task_runs_to_clickhouse_v1", @@ -267,7 +268,7 @@ describe("RunsReplicationService (part 2/2)", () => { }); const runsReplicationService = new RunsReplicationService({ - clickhouse, + clickhouseFactory: new TestReplicationClickhouseFactory(clickhouse), pgConnectionUrl: postgresContainer.getConnectionUri(), serviceName: "runs-replication-stress-bulk-insert", slotName: "task_runs_to_clickhouse_v1", @@ -388,7 +389,7 @@ describe("RunsReplicationService (part 2/2)", () => { }); const runsReplicationService = new RunsReplicationService({ - clickhouse, + clickhouseFactory: new TestReplicationClickhouseFactory(clickhouse), pgConnectionUrl: postgresContainer.getConnectionUri(), serviceName: "runs-replication-multi-event-tx", slotName: "task_runs_to_clickhouse_v1", @@ -522,7 +523,7 @@ describe("RunsReplicationService (part 2/2)", () => { }); const runsReplicationService = new RunsReplicationService({ - clickhouse, + clickhouseFactory: new TestReplicationClickhouseFactory(clickhouse), pgConnectionUrl: postgresContainer.getConnectionUri(), serviceName: "runs-replication-long-tx", slotName: "task_runs_to_clickhouse_v1", @@ -629,7 +630,7 @@ describe("RunsReplicationService (part 2/2)", () => { }); const runsReplicationService = new RunsReplicationService({ - clickhouse, + clickhouseFactory: new TestReplicationClickhouseFactory(clickhouse), pgConnectionUrl: postgresContainer.getConnectionUri(), serviceName: "runs-replication-stress-bulk-insert", slotName: "task_runs_to_clickhouse_v1", @@ -798,7 +799,7 @@ describe("RunsReplicationService (part 2/2)", () => { }); const runsReplicationService = new RunsReplicationService({ - clickhouse, + clickhouseFactory: new TestReplicationClickhouseFactory(clickhouse), pgConnectionUrl: postgresContainer.getConnectionUri(), serviceName: "runs-replication-merge-batch", slotName: "task_runs_to_clickhouse_v1", @@ -923,7 +924,7 @@ describe("RunsReplicationService (part 2/2)", () => { }); const runsReplicationService = new RunsReplicationService({ - clickhouse, + clickhouseFactory: new TestReplicationClickhouseFactory(clickhouse), pgConnectionUrl: postgresContainer.getConnectionUri(), serviceName: "runs-replication-sorting", slotName: "task_runs_to_clickhouse_v1", @@ -1136,7 +1137,7 @@ describe("RunsReplicationService (part 2/2)", () => { }); const runsReplicationService = new RunsReplicationService({ - clickhouse, + clickhouseFactory: new TestReplicationClickhouseFactory(clickhouse), pgConnectionUrl: postgresContainer.getConnectionUri(), serviceName: "runs-replication-exhaustive", slotName: "task_runs_to_clickhouse_v1", diff --git a/apps/webapp/test/utils/replicationUtils.ts b/apps/webapp/test/utils/replicationUtils.ts index ecbf65bacc1..358da0c2cf6 100644 --- a/apps/webapp/test/utils/replicationUtils.ts +++ b/apps/webapp/test/utils/replicationUtils.ts @@ -2,6 +2,7 @@ import { ClickHouse } from "@internal/clickhouse"; import { RedisOptions } from "@internal/redis"; import { PrismaClient } from "~/db.server"; import { RunsReplicationService } from "~/services/runsReplicationService.server"; +import { TestReplicationClickhouseFactory } from "./testReplicationClickhouseFactory"; import { afterEach } from "vitest"; export async function setupClickhouseReplication({ @@ -26,7 +27,7 @@ export async function setupClickhouseReplication({ }); const runsReplicationService = new RunsReplicationService({ - clickhouse, + clickhouseFactory: new TestReplicationClickhouseFactory(clickhouse), pgConnectionUrl: databaseUrl, serviceName: "runs-replication", slotName: "task_runs_to_clickhouse_v1", diff --git a/apps/webapp/test/utils/testReplicationClickhouseFactory.ts b/apps/webapp/test/utils/testReplicationClickhouseFactory.ts new file mode 100644 index 00000000000..4d34fac376c --- /dev/null +++ b/apps/webapp/test/utils/testReplicationClickhouseFactory.ts @@ -0,0 +1,31 @@ +import type { ClickHouse } from "@internal/clickhouse"; +import { + ClickhouseFactory, + type ClientType, +} from "~/services/clickhouse/clickhouseFactory.server"; +import type { OrganizationDataStoresRegistry } from "~/services/dataStores/organizationDataStoresRegistry.server"; + +const testReplicationRegistryStub = { + isLoaded: true, + isReady: Promise.resolve(), + get: () => null, +} as unknown as OrganizationDataStoresRegistry; + +/** + * Routes all `replication` clients to a single test ClickHouse; other client types use the real factory defaults. + */ +export class TestReplicationClickhouseFactory extends ClickhouseFactory { + constructor(private readonly replicationClient: ClickHouse) { + super(testReplicationRegistryStub); + } + + override getClickhouseForOrganizationSync( + organizationId: string, + clientType: ClientType + ): ClickHouse { + if (clientType === "replication") { + return this.replicationClient; + } + return super.getClickhouseForOrganizationSync(organizationId, clientType); + } +} diff --git a/apps/webapp/test/utils/tracing.ts b/apps/webapp/test/utils/tracing.ts index 09500a6c354..ed6119c52b1 100644 --- a/apps/webapp/test/utils/tracing.ts +++ b/apps/webapp/test/utils/tracing.ts @@ -1,6 +1,5 @@ import { NodeTracerProvider } from "@opentelemetry/sdk-trace-node"; import { InMemorySpanExporter, SimpleSpanProcessor } from "@opentelemetry/sdk-trace-base"; -import { trace } from "@opentelemetry/api"; import { MeterProvider, InMemoryMetricExporter, @@ -16,8 +15,9 @@ export function createInMemoryTracing() { }); provider.register(); - // Retrieve the tracer - const tracer = trace.getTracer("test-tracer"); + // Use the provider's tracer so spans hit this exporter even when a global + // NodeTracerProvider was already registered (e.g. via tracer.server import chain). + const tracer = provider.getTracer("test-tracer"); return { exporter, diff --git a/internal-packages/database/prisma/migrations/20260331212308_add_organization_data_stores/migration.sql b/internal-packages/database/prisma/migrations/20260331212308_add_organization_data_stores/migration.sql new file mode 100644 index 00000000000..52b8385539a --- /dev/null +++ b/internal-packages/database/prisma/migrations/20260331212308_add_organization_data_stores/migration.sql @@ -0,0 +1,21 @@ +-- CreateEnum +CREATE TYPE "public"."DataStoreKind" AS ENUM ('CLICKHOUSE'); + +-- CreateTable +CREATE TABLE "public"."OrganizationDataStore" ( + "id" TEXT NOT NULL, + "key" TEXT NOT NULL, + "organizationIds" TEXT[], + "kind" "public"."DataStoreKind" NOT NULL, + "config" JSONB NOT NULL, + "createdAt" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updatedAt" TIMESTAMP(3) NOT NULL, + + CONSTRAINT "OrganizationDataStore_pkey" PRIMARY KEY ("id") +); + +-- CreateIndex +CREATE UNIQUE INDEX "OrganizationDataStore_key_key" ON "public"."OrganizationDataStore"("key"); + +-- CreateIndex +CREATE INDEX "OrganizationDataStore_kind_idx" ON "public"."OrganizationDataStore"("kind"); diff --git a/internal-packages/database/prisma/schema.prisma b/internal-packages/database/prisma/schema.prisma index 7138aeaab0d..0f71994de37 100644 --- a/internal-packages/database/prisma/schema.prisma +++ b/internal-packages/database/prisma/schema.prisma @@ -2917,3 +2917,24 @@ model ErrorGroupState { @@unique([environmentId, taskIdentifier, errorFingerprint]) @@index([environmentId, status]) } + +enum DataStoreKind { + CLICKHOUSE +} + +/// Defines org-scoped data store overrides (e.g. dedicated ClickHouse for HIPAA orgs). +/// Multiple organizations can share a single data store row via organizationIds. +model OrganizationDataStore { + id String @id @default(cuid()) + /// Human-readable unique key (e.g. "hipaa-clickhouse-us-east") + key String @unique + /// Organization IDs that use this data store + organizationIds String[] + kind DataStoreKind + /// Versioned config JSON. Structure is discriminated by the top-level `version` field. + config Json + createdAt DateTime @default(now()) + updatedAt DateTime @updatedAt + + @@index([kind]) +}