301 lines
9.1 KiB
TypeScript
301 lines
9.1 KiB
TypeScript
import { env } from "~/env.server";
|
|
import { eventRepository } from "./eventRepository.server";
|
|
import { type IEventRepository, type TraceEventOptions } from "./eventRepository.types";
|
|
import { prisma } from "~/db.server";
|
|
import { runStore } from "../runStore.server";
|
|
import { controlPlaneResolver } from "~/v3/runOpsMigration/controlPlaneResolver.server";
|
|
import { logger } from "~/services/logger.server";
|
|
import { FEATURE_FLAG } from "../featureFlags";
|
|
import { flag } from "../featureFlags.server";
|
|
import { getTaskEventStore } from "../taskEventStore.server";
|
|
import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactoryInstance.server";
|
|
|
|
export const EVENT_STORE_TYPES = {
|
|
POSTGRES: "postgres",
|
|
CLICKHOUSE: "clickhouse",
|
|
CLICKHOUSE_V2: "clickhouse_v2",
|
|
} as const;
|
|
|
|
export type EventStoreType = (typeof EVENT_STORE_TYPES)[keyof typeof EVENT_STORE_TYPES];
|
|
/**
|
|
* Async variant of {@link resolveEventRepositoryForStore}. Awaits the factory's
|
|
* registry readiness before returning the ClickHouse event repository; for
|
|
* non-ClickHouse stores (e.g. the "taskEvent" DB default for Postgres-backed
|
|
* runs) it returns the Prisma event repository without ever touching the
|
|
* factory — so the factory never needs to know about Postgres.
|
|
*/
|
|
export async function getEventRepositoryForStore(
|
|
store: string,
|
|
organizationId: string
|
|
): Promise<IEventRepository> {
|
|
if (store !== EVENT_STORE_TYPES.CLICKHOUSE && store !== EVENT_STORE_TYPES.CLICKHOUSE_V2) {
|
|
return eventRepository;
|
|
}
|
|
const { repository } = await clickhouseFactory.getEventRepositoryForOrganization(
|
|
store,
|
|
organizationId
|
|
);
|
|
return repository;
|
|
}
|
|
|
|
export async function getConfiguredEventRepository(
|
|
organizationId: string
|
|
): Promise<{ repository: IEventRepository; store: EventStoreType }> {
|
|
const organization = await prisma.organization.findFirst({
|
|
select: {
|
|
id: true,
|
|
featureFlags: true,
|
|
},
|
|
where: {
|
|
id: organizationId,
|
|
},
|
|
});
|
|
|
|
if (!organization) {
|
|
throw new Error("Organization not found when configuring event repository");
|
|
}
|
|
|
|
// resolveTaskEventRepositoryFlag checks:
|
|
// 1. organization.featureFlags (highest priority)
|
|
// 2. global feature flags (via flags() function)
|
|
// 3. env.EVENT_REPOSITORY_DEFAULT_STORE (fallback)
|
|
const taskEventStore = await resolveTaskEventRepositoryFlag(
|
|
(organization.featureFlags as Record<string, unknown> | null) ?? undefined
|
|
);
|
|
|
|
if (taskEventStore === EVENT_STORE_TYPES.CLICKHOUSE_V2) {
|
|
const { repository: resolvedRepository } =
|
|
await clickhouseFactory.getEventRepositoryForOrganization(taskEventStore, organizationId);
|
|
return { repository: resolvedRepository, store: EVENT_STORE_TYPES.CLICKHOUSE_V2 };
|
|
}
|
|
|
|
if (taskEventStore === EVENT_STORE_TYPES.CLICKHOUSE) {
|
|
const { repository: resolvedRepository } =
|
|
await clickhouseFactory.getEventRepositoryForOrganization(taskEventStore, organizationId);
|
|
return { repository: resolvedRepository, store: EVENT_STORE_TYPES.CLICKHOUSE };
|
|
}
|
|
|
|
return { repository: eventRepository, store: EVENT_STORE_TYPES.POSTGRES };
|
|
}
|
|
|
|
export async function getEventRepository(
|
|
organizationId: string,
|
|
featureFlags: Record<string, unknown> | undefined,
|
|
parentStore: string | undefined
|
|
): Promise<{ repository: IEventRepository; store: string }> {
|
|
const taskEventStore = parentStore ?? (await resolveTaskEventRepositoryFlag(featureFlags));
|
|
|
|
// Non-ClickHouse stores (e.g. the "taskEvent" DB default for Postgres-backed
|
|
// runs, or the legacy "postgres" value) resolve to the Prisma event repo.
|
|
if (
|
|
taskEventStore !== EVENT_STORE_TYPES.CLICKHOUSE &&
|
|
taskEventStore !== EVENT_STORE_TYPES.CLICKHOUSE_V2
|
|
) {
|
|
return { repository: eventRepository, store: getTaskEventStore() };
|
|
}
|
|
|
|
const { repository: resolvedRepository } =
|
|
await clickhouseFactory.getEventRepositoryForOrganization(taskEventStore, organizationId);
|
|
|
|
switch (taskEventStore) {
|
|
case EVENT_STORE_TYPES.CLICKHOUSE_V2: {
|
|
return { repository: resolvedRepository, store: EVENT_STORE_TYPES.CLICKHOUSE_V2 };
|
|
}
|
|
case EVENT_STORE_TYPES.CLICKHOUSE: {
|
|
return { repository: resolvedRepository, store: EVENT_STORE_TYPES.CLICKHOUSE };
|
|
}
|
|
default: {
|
|
return { repository: eventRepository, store: getTaskEventStore() };
|
|
}
|
|
}
|
|
}
|
|
|
|
export async function getV3EventRepository(
|
|
organizationId: string,
|
|
parentStore: string | undefined
|
|
): Promise<{ repository: IEventRepository; store: string }> {
|
|
if (typeof parentStore === "string") {
|
|
// Support legacy Postgres store for self-hosters and runs persisted with a
|
|
// non-ClickHouse store — fall back to the Prisma-based event repository.
|
|
if (
|
|
parentStore !== EVENT_STORE_TYPES.CLICKHOUSE &&
|
|
parentStore !== EVENT_STORE_TYPES.CLICKHOUSE_V2
|
|
) {
|
|
return { repository: eventRepository, store: parentStore };
|
|
}
|
|
|
|
const { repository: resolvedRepository } =
|
|
await clickhouseFactory.getEventRepositoryForOrganization(parentStore, organizationId);
|
|
return { repository: resolvedRepository, store: parentStore };
|
|
}
|
|
|
|
if (env.EVENT_REPOSITORY_DEFAULT_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") {
|
|
const { repository: resolvedRepository } =
|
|
await clickhouseFactory.getEventRepositoryForOrganization("clickhouse", organizationId);
|
|
return { repository: resolvedRepository, store: "clickhouse" };
|
|
} else {
|
|
return { repository: eventRepository, store: getTaskEventStore() };
|
|
}
|
|
}
|
|
|
|
async function resolveTaskEventRepositoryFlag(
|
|
featureFlags: Record<string, unknown> | undefined
|
|
): Promise<"clickhouse" | "clickhouse_v2" | "postgres"> {
|
|
const flagResult = await flag({
|
|
key: FEATURE_FLAG.taskEventRepository,
|
|
defaultValue: env.EVENT_REPOSITORY_DEFAULT_STORE,
|
|
overrides: featureFlags,
|
|
});
|
|
|
|
if (flagResult === "clickhouse_v2") {
|
|
return "clickhouse_v2";
|
|
}
|
|
|
|
if (flagResult === "clickhouse") {
|
|
return "clickhouse";
|
|
}
|
|
|
|
return flagResult;
|
|
}
|
|
|
|
export async function recordRunDebugLog(
|
|
runId: string,
|
|
message: string,
|
|
options: Omit<TraceEventOptions, "environment" | "taskSlug" | "startTime"> & {
|
|
duration?: number;
|
|
parentId?: string;
|
|
startTime?: Date;
|
|
}
|
|
): Promise<
|
|
| {
|
|
success: true;
|
|
}
|
|
| {
|
|
success: false;
|
|
code: "RUN_NOT_FOUND" | "FAILED_TO_RECORD_EVENT";
|
|
error?: unknown;
|
|
}
|
|
> {
|
|
if (env.EVENT_REPOSITORY_DEBUG_LOGS_DISABLED) {
|
|
// drop debug events silently
|
|
return {
|
|
success: true,
|
|
};
|
|
}
|
|
|
|
return recordRunEvent(runId, message, {
|
|
...options,
|
|
attributes: {
|
|
...options?.attributes,
|
|
isDebug: true,
|
|
},
|
|
});
|
|
}
|
|
|
|
async function recordRunEvent(
|
|
runId: string,
|
|
message: string,
|
|
options: Omit<TraceEventOptions, "environment" | "taskSlug" | "startTime"> & {
|
|
duration?: number;
|
|
parentId?: string;
|
|
startTime?: Date;
|
|
}
|
|
): Promise<
|
|
| {
|
|
success: true;
|
|
}
|
|
| {
|
|
success: false;
|
|
code: "RUN_NOT_FOUND" | "FAILED_TO_RECORD_EVENT";
|
|
error?: unknown;
|
|
}
|
|
> {
|
|
try {
|
|
const foundRun = await findRunForEventCreation(runId);
|
|
|
|
if (!foundRun) {
|
|
logger.error("Failed to find run for event creation", { runId });
|
|
return {
|
|
success: false,
|
|
code: "RUN_NOT_FOUND",
|
|
};
|
|
}
|
|
|
|
const $eventRepository = await getEventRepositoryForStore(
|
|
foundRun.taskEventStore,
|
|
foundRun.runtimeEnvironment.organizationId
|
|
);
|
|
|
|
const { attributes, startTime, ...optionsRest } = options;
|
|
|
|
await $eventRepository.recordEvent(message, {
|
|
environment: foundRun.runtimeEnvironment,
|
|
taskSlug: foundRun.taskIdentifier,
|
|
context: foundRun.traceContext as Record<string, string | undefined>,
|
|
attributes: {
|
|
runId: foundRun.friendlyId,
|
|
...attributes,
|
|
},
|
|
startTime: BigInt((startTime?.getTime() ?? Date.now()) * 1_000_000),
|
|
...optionsRest,
|
|
});
|
|
|
|
return {
|
|
success: true,
|
|
};
|
|
} catch (error) {
|
|
logger.error("Failed to record event for run", {
|
|
error: error instanceof Error ? error.message : error,
|
|
runId,
|
|
});
|
|
|
|
return {
|
|
success: false,
|
|
code: "FAILED_TO_RECORD_EVENT",
|
|
error,
|
|
};
|
|
}
|
|
}
|
|
|
|
async function findRunForEventCreation(runId: string) {
|
|
const foundRun = await runStore.findRun(
|
|
{
|
|
id: runId,
|
|
},
|
|
{
|
|
select: {
|
|
friendlyId: true,
|
|
taskIdentifier: true,
|
|
traceContext: true,
|
|
taskEventStore: true,
|
|
runtimeEnvironmentId: true,
|
|
},
|
|
},
|
|
prisma
|
|
);
|
|
|
|
if (!foundRun) {
|
|
return null;
|
|
}
|
|
|
|
const environment = await controlPlaneResolver.resolveAuthenticatedEnv(
|
|
foundRun.runtimeEnvironmentId
|
|
);
|
|
|
|
if (!environment) {
|
|
// Run exists but its environment could not be resolved (e.g. a lagging replica
|
|
// under split); distinguish this from a genuinely missing run.
|
|
logger.warn("Run found but environment unresolved for event creation", {
|
|
runId,
|
|
runtimeEnvironmentId: foundRun.runtimeEnvironmentId,
|
|
});
|
|
return null;
|
|
}
|
|
|
|
return { ...foundRun, runtimeEnvironment: environment };
|
|
}
|