Record each query

2025-12-20 20:57:58 +00:00
parent 558b94eca9
commit 402ea21d75
5 changed files with 209 additions and 43 deletions
@@ -521,6 +521,7 @@ const EnvironmentSchema = z
    PROD_USAGE_HEARTBEAT_INTERVAL_MS: z.coerce.number().int().optional(),

    CENTS_PER_RUN: z.coerce.number().default(0),
+    CENTS_PER_QUERY_BYTE_SECOND: z.coerce.number().default(0),

    EVENT_LOOP_MONITOR_ENABLED: z.string().default("1"),
    RESOURCE_MONITOR_ENABLED: z.string().default("0"),
@@ -1,5 +1,5 @@
 import { ArrowDownTrayIcon, ClipboardIcon, LightBulbIcon } from "@heroicons/react/20/solid";
-import type { FieldMappings, OutputColumnMetadata } from "@internal/clickhouse";
+import type { OutputColumnMetadata } from "@internal/clickhouse";
 import type { ColumnSchema } from "@internal/tsql";
 import { Form, useNavigation } from "@remix-run/react";
 import {
@@ -39,10 +39,9 @@ import { Switch } from "~/components/primitives/Switch";
 import { useEnvironment } from "~/hooks/useEnvironment";
 import { useOrganization } from "~/hooks/useOrganizations";
 import { useProject } from "~/hooks/useProject";
-import { prisma } from "~/db.server";
 import { findProjectBySlug } from "~/models/project.server";
 import { findEnvironmentBySlug } from "~/models/runtimeEnvironment.server";
-import { executeQuery } from "~/services/queryService.server";
+import { executeQuery, type QueryScope } from "~/services/queryService.server";
 import { requireUser } from "~/services/session.server";
 import { downloadFile, rowsToCSV, rowsToJSON } from "~/utils/dataExport";
 import { EnvironmentParamSchema } from "~/utils/pathBuilder";
@@ -54,8 +53,6 @@ const scopeOptions = [
  { value: "organization", label: "Organization" },
 ] as const;

-type QueryScope = (typeof scopeOptions)[number]["value"];
-
 export const loader = async ({ request, params }: LoaderFunctionArgs) => {
  const user = await requireUser(request);

@@ -143,39 +140,6 @@ export const action = async ({ request, params }: ActionFunctionArgs) => {

  const { query, scope } = parsed.data;

-  // Build tenant IDs based on scope
-  const tenantOptions: {
-    organizationId: string;
-    projectId?: string;
-    environmentId?: string;
-  } = {
-    organizationId: project.organizationId,
-  };
-
-  if (scope === "project" || scope === "environment") {
-    tenantOptions.projectId = project.id;
-  }
-
-  if (scope === "environment") {
-    tenantOptions.environmentId = environment.id;
-  }
-
-  // Build field mappings for project_ref → project_id and environment_id → slug translation
-  const projects = await prisma.project.findMany({
-    where: { organizationId: project.organizationId },
-    select: { id: true, externalRef: true },
-  });
-
-  const environments = await prisma.runtimeEnvironment.findMany({
-    where: { project: { organizationId: project.organizationId } },
-    select: { id: true, slug: true },
-  });
-
-  const fieldMappings: FieldMappings = {
-    project: Object.fromEntries(projects.map((p) => [p.id, p.externalRef])),
-    environment: Object.fromEntries(environments.map((e) => [e.id, e.slug])),
-  };
-
  try {
    const [error, result] = await executeQuery({
      name: "query-page",
@@ -183,8 +147,14 @@ export const action = async ({ request, params }: ActionFunctionArgs) => {
      schema: z.record(z.any()),
      tableSchema: querySchemas,
      transformValues: true,
-      fieldMappings,
-      ...tenantOptions,
+      scope,
+      organizationId: project.organizationId,
+      projectId: project.id,
+      environmentId: environment.id,
+      history: {
+        source: "DASHBOARD",
+        userId: user.id,
+      },
    });

    if (error) {
@@ -1,15 +1,116 @@
-import { executeTSQL, type ExecuteTSQLOptions, type TSQLQueryResult } from "@internal/clickhouse";
+import {
+  executeTSQL,
+  type ExecuteTSQLOptions,
+  type FieldMappings,
+  type TSQLQueryResult,
+} from "@internal/clickhouse";
+import type { CustomerQuerySource } from "@trigger.dev/database";
 import type { TableSchema } from "@internal/tsql";
 import { type z } from "zod";
+import { prisma } from "~/db.server";
+import { env } from "~/env.server";
 import { clickhouseClient } from "./clickhouseInstance.server";

 export type { TableSchema, TSQLQueryResult };

+export type QueryScope = "organization" | "project" | "environment";
+
+const scopeToEnum = {
+  organization: "ORGANIZATION",
+  project: "PROJECT",
+  environment: "ENVIRONMENT",
+} as const;
+
+export type ExecuteQueryOptions<TOut extends z.ZodSchema> = Omit<
+  ExecuteTSQLOptions<TOut>,
+  "tableSchema" | "organizationId" | "projectId" | "environmentId" | "fieldMappings"
+> & {
+  tableSchema: TableSchema[];
+  /** The scope of the query - determines tenant isolation */
+  scope: QueryScope;
+  /** Organization ID (required) */
+  organizationId: string;
+  /** Project ID (required for project/environment scope) */
+  projectId: string;
+  /** Environment ID (required for environment scope) */
+  environmentId: string;
+  /** History options for saving query to billing/audit */
+  history?: {
+    /** Where the query originated from */
+    source: CustomerQuerySource;
+    /** User ID (optional, null for API calls) */
+    userId?: string | null;
+  };
+};
+
 /**
 * Execute a TSQL query against ClickHouse with tenant isolation
+ * Handles building tenant options, field mappings, and optionally saves to history
 */
 export async function executeQuery<TOut extends z.ZodSchema>(
-  options: Omit<ExecuteTSQLOptions<TOut>, "tableSchema"> & { tableSchema: TableSchema[] }
+  options: ExecuteQueryOptions<TOut>
 ): Promise<TSQLQueryResult<z.output<TOut>>> {
-  return executeTSQL(clickhouseClient.reader, options);
+  const { scope, organizationId, projectId, environmentId, history, ...baseOptions } = options;
+
+  // Build tenant IDs based on scope
+  const tenantOptions: {
+    organizationId: string;
+    projectId?: string;
+    environmentId?: string;
+  } = {
+    organizationId,
+  };
+
+  if (scope === "project" || scope === "environment") {
+    tenantOptions.projectId = projectId;
+  }
+
+  if (scope === "environment") {
+    tenantOptions.environmentId = environmentId;
+  }
+
+  // Build field mappings for project_ref → project_id and environment_id → slug translation
+  const projects = await prisma.project.findMany({
+    where: { organizationId },
+    select: { id: true, externalRef: true },
+  });
+
+  const environments = await prisma.runtimeEnvironment.findMany({
+    where: { project: { organizationId } },
+    select: { id: true, slug: true },
+  });
+
+  const fieldMappings: FieldMappings = {
+    project: Object.fromEntries(projects.map((p) => [p.id, p.externalRef])),
+    environment: Object.fromEntries(environments.map((e) => [e.id, e.slug])),
+  };
+
+  const result = await executeTSQL(clickhouseClient.reader, {
+    ...baseOptions,
+    ...tenantOptions,
+    fieldMappings,
+  });
+
+  // If query succeeded and history options provided, save to history
+  if (result[0] === null && history) {
+    const stats = result[1].stats;
+    const byteSeconds = parseFloat(stats.byte_seconds);
+    const costInCents = byteSeconds * env.CENTS_PER_QUERY_BYTE_SECOND;
+
+    await prisma.customerQuery.create({
+      data: {
+        query: options.query,
+        scope: scopeToEnum[scope],
+        stats: { ...stats },
+        costInCents,
+        source: history.source,
+        organizationId,
+        projectId: scope === "project" || scope === "environment" ? projectId : null,
+        environmentId: scope === "environment" ? environmentId : null,
+        userId: history.userId ?? null,
+      },
+    });
+  }
+
+  return result;
 }
@@ -0,0 +1,40 @@
+-- CreateEnum
+CREATE TYPE "public"."CustomerQuerySource" AS ENUM ('DASHBOARD', 'API');
+
+-- CreateEnum
+CREATE TYPE "public"."CustomerQueryScope" AS ENUM ('ORGANIZATION', 'PROJECT', 'ENVIRONMENT');
+
+-- CreateTable
+CREATE TABLE
+    "public"."CustomerQuery" (
+        "id" TEXT NOT NULL,
+        "query" TEXT NOT NULL,
+        "scope" "public"."CustomerQueryScope" NOT NULL,
+        "stats" JSONB NOT NULL,
+        "costInCents" DOUBLE PRECISION NOT NULL DEFAULT 0,
+        "source" "public"."CustomerQuerySource" NOT NULL DEFAULT 'DASHBOARD',
+        "organizationId" TEXT NOT NULL,
+        "projectId" TEXT,
+        "environmentId" TEXT,
+        "userId" TEXT,
+        "createdAt" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+        CONSTRAINT "CustomerQuery_pkey" PRIMARY KEY ("id")
+    );
+
+-- CreateIndex
+CREATE INDEX "CustomerQuery_organizationId_createdAt_idx" ON "public"."CustomerQuery" ("organizationId", "createdAt" DESC);
+
+-- CreateIndex
+CREATE INDEX "CustomerQuery_createdAt_idx" ON "public"."CustomerQuery" ("createdAt");
+
+-- AddForeignKey
+ALTER TABLE "public"."CustomerQuery" ADD CONSTRAINT "CustomerQuery_organizationId_fkey" FOREIGN KEY ("organizationId") REFERENCES "public"."Organization" ("id") ON DELETE CASCADE ON UPDATE CASCADE;
+
+-- AddForeignKey
+ALTER TABLE "public"."CustomerQuery" ADD CONSTRAINT "CustomerQuery_projectId_fkey" FOREIGN KEY ("projectId") REFERENCES "public"."Project" ("id") ON DELETE CASCADE ON UPDATE CASCADE;
+
+-- AddForeignKey
+ALTER TABLE "public"."CustomerQuery" ADD CONSTRAINT "CustomerQuery_environmentId_fkey" FOREIGN KEY ("environmentId") REFERENCES "public"."RuntimeEnvironment" ("id") ON DELETE CASCADE ON UPDATE CASCADE;
+
+-- AddForeignKey
+ALTER TABLE "public"."CustomerQuery" ADD CONSTRAINT "CustomerQuery_userId_fkey" FOREIGN KEY ("userId") REFERENCES "public"."User" ("id") ON DELETE SET NULL ON UPDATE CASCADE;
@@ -59,6 +59,7 @@ model User {
  deployments          WorkerDeployment[]
  backupCodes          MfaBackupCode[]
  bulkActions          BulkActionGroup[]
+  customerQueries      CustomerQuery[]
 }

 model MfaBackupCode {
@@ -218,6 +219,7 @@ model Organization {
  workerGroups             WorkerInstanceGroup[]
  workerInstances          WorkerInstance[]
  githubAppInstallations   GithubAppInstallation[]
+  customerQueries          CustomerQuery[]
 }

 model OrgMember {
@@ -335,6 +337,7 @@ model RuntimeEnvironment {
  workerInstances      WorkerInstance[]
  waitpointTags        WaitpointTag[]
  BulkActionGroup      BulkActionGroup[]
+  customerQueries      CustomerQuery[]

  @@unique([projectId, slug, orgMemberId])
  @@unique([projectId, shortcode])
@@ -399,6 +402,7 @@ model Project {
  taskRunCheckpoints        TaskRunCheckpoint[]
  waitpointTags             WaitpointTag[]
  connectedGithubRepository ConnectedGithubRepository?
+  customerQueries           CustomerQuery[]

  buildSettings Json?
 }
@@ -2383,3 +2387,53 @@ model ConnectedGithubRepository {
  @@unique([projectId])
  @@index([repositoryId])
 }
+
+enum CustomerQuerySource {
+  DASHBOARD
+  API
+}
+
+enum CustomerQueryScope {
+  ORGANIZATION
+  PROJECT
+  ENVIRONMENT
+}
+
+model CustomerQuery {
+  id String @id @default(cuid())
+
+  /// The TSQL query text that was executed
+  query String
+
+  /// The scope of the query (determines which tenant IDs were used)
+  scope CustomerQueryScope
+
+  /// Query execution statistics from ClickHouse
+  stats Json
+
+  /// Cost of the query in cents (for Stripe metering)
+  costInCents Float @default(0)
+
+  /// Where the query originated from
+  source CustomerQuerySource @default(DASHBOARD)
+
+  organization   Organization @relation(fields: [organizationId], references: [id], onDelete: Cascade, onUpdate: Cascade)
+  organizationId String
+
+  project   Project? @relation(fields: [projectId], references: [id], onDelete: Cascade, onUpdate: Cascade)
+  projectId String?
+
+  environment   RuntimeEnvironment? @relation(fields: [environmentId], references: [id], onDelete: Cascade, onUpdate: Cascade)
+  environmentId String?
+
+  /// Optional user who executed the query (null for API calls)
+  user   User?   @relation(fields: [userId], references: [id], onDelete: SetNull, onUpdate: Cascade)
+  userId String?
+
+  createdAt DateTime @default(now())
+
+  /// Fast lookup for history menu (most recent 20 per org)
+  @@index([organizationId, createdAt(sort: Desc)])
+  /// For Stripe metering job - find unprocessed queries
+  @@index([createdAt])
+}