Rework the API/SDK so we can get nice spans

Author: matt-aitken

apps/webapp/app/components/runs/v3/RunIcon.tsx

@@ -4,6 +4,7 @@ import {
   InformationCircleIcon,
   RectangleStackIcon,
   Squares2X2Icon,
+  TableCellsIcon,
   TagIcon,
 } from "@heroicons/react/20/solid";
 import { AttemptIcon } from "~/assets/icons/AttemptIcon";
@@ -81,6 +82,8 @@ export function RunIcon({ name, className, spanName }: TaskIconProps) {
       return <WaitpointTokenIcon className={cn(className, "text-sky-500")} />;
     case "function":
       return <FunctionIcon className={cn(className, "text-text-dimmed")} />;
+    case "query":
+      return <TableCellsIcon className={cn(className, "text-query")} />;
     //log levels
     case "debug":
     case "log":
@@ -110,7 +113,7 @@ export function RunIcon({ name, className, spanName }: TaskIconProps) {
     case "task-hook-catchError":
       return <FunctionIcon className={cn(className, "text-error")} />;
     case "streams":
-        return <StreamsIcon className={cn(className, "text-text-dimmed")} />;
+      return <StreamsIcon className={cn(className, "text-text-dimmed")} />;
   }
 
   return <InformationCircleIcon className={cn(className, "text-text-dimmed")} />;

apps/webapp/app/routes/api.v1.query.ts

@@ -58,17 +58,15 @@ const { action, loader } = createActionApiRoute(
     if (format === "csv") {
       const csv = rowsToCSV(result.rows, result.columns);
 
-      return new Response(csv, {
-        status: 200,
-        headers: {
-          "Content-Type": "text/csv; charset=utf-8",
-          "Content-Disposition": "attachment; filename=query-results.csv",
-        },
+      return json({
+        format: "csv",
+        results: csv,
       });
     }
 
     return json({
-      rows: result.rows,
+      format: "json",
+      results: result.rows,
     });
   }
 );

packages/core/src/v3/apiClient/index.ts

@@ -1419,7 +1419,7 @@ export class ApiClient {
       format?: "json" | "csv";
     },
     requestOptions?: ZodFetchOptions
-  ): Promise<QueryExecuteResponseBody | QueryExecuteCSVResponseBody> {
+  ): Promise<QueryExecuteResponseBody> {
     const body = {
       query,
       scope: options?.scope ?? "environment",
@@ -1430,37 +1430,9 @@ export class ApiClient {
     };
 
     const format = options?.format ?? "json";
-
-    if (format === "csv") {
-      // For CSV, we get a text response
-      const response = await fetch(`${this.baseUrl}/api/v1/query`, {
-        method: "POST",
-        headers: this.#getHeaders(false),
-        body: JSON.stringify(body),
-      });
-
-      if (!response.ok) {
-        const errText = await response.text().catch((e) => (e as Error).message);
-        let errJSON: Object | undefined;
-        try {
-          errJSON = JSON.parse(errText) as Object;
-        } catch {
-          // ignore
-        }
-        const errMessage = errJSON ? undefined : errText;
-        const responseHeaders = Object.fromEntries(response.headers.entries());
-
-        throw ApiError.generate(response.status, errJSON, errMessage, responseHeaders);
-      }
-
-      return await response.text();
-    }
-
     // For JSON, use zodfetch
     return zodfetch(
-      z.object({
-        rows: z.array(z.record(z.any())),
-      }),
+      QueryExecuteResponseBody,
       `${this.baseUrl}/api/v1/query`,
       {
         method: "POST",

packages/core/src/v3/schemas/query.ts

@@ -1,4 +1,4 @@
-import { z } from "zod";
+import { TypeOf, z } from "zod";
 
 /**
  * Request body schema for executing a query
@@ -17,13 +17,25 @@ export type QueryExecuteRequestBody = z.infer<typeof QueryExecuteRequestBody>;
 /**
  * Response body schema for JSON format queries
  */
-export const QueryExecuteResponseBody = z.object({
-  rows: z.array(z.record(z.any())),
+export const QueryExecuteJSONResponseBody = z.object({
+  format: z.literal("json"),
+  results: z.array(z.record(z.any())),
 });
 
-export type QueryExecuteResponseBody = z.infer<typeof QueryExecuteResponseBody>;
+export type QueryExecuteJSONResponseBody = z.infer<typeof QueryExecuteResponseBody>;
 
 /**
- * Response body type for CSV format queries (returns a string)
+ * Response body type for CSV format queries
  */
-export type QueryExecuteCSVResponseBody = string;
+export const QueryExecuteCSVResponseBody = z.object({
+  format: z.literal("json"),
+  results: z.string(),
+});
+
+export type QueryExecuteCSVResponseBody = z.infer<typeof QueryExecuteCSVResponseBody>;
+
+export const QueryExecuteResponseBody = z.discriminatedUnion("format", [
+  QueryExecuteJSONResponseBody,
+  QueryExecuteCSVResponseBody,
+]);
+export type QueryExecuteResponseBody = z.infer<typeof QueryExecuteResponseBody>;

packages/trigger-sdk/src/v3/query.ts

@@ -81,14 +81,14 @@ function execute<TRow extends Record<string, any> = Record<string, any>>(
  * @example
  * ```typescript
  * // Basic query with defaults (environment scope, json format)
- * const result = await query.execute("SELECT * FROM runs LIMIT 10");
+ * const result = await query.execute("SELECT run_id, status FROM runs LIMIT 10");
  * console.log(result.format); // "json"
  * console.log(result.results); // Array<Record<string, any>>
  *
  * // Type-safe query with row type
  * type RunRow = { id: string; status: string; duration: number };
  * const typedResult = await query.execute<RunRow>(
- *   "SELECT id, status, duration FROM runs LIMIT 10"
+ *   "SELECT run_id, status, triggered_at FROM runs LIMIT 10"
  * );
  * typedResult.results.forEach(row => {
  *   console.log(row.id, row.status); // Fully typed!
@@ -105,14 +105,14 @@ function execute<TRow extends Record<string, any> = Record<string, any>>(
  * // Query with custom period
  * const lastMonth = await query.execute(
  *   "SELECT COUNT(*) as count FROM runs",
- *   { period: "30d" }
+ *   { period: "3d" }
  * );
  * console.log(lastMonth.results[0].count); // Type-safe access
  *
  * // Export as CSV - automatically narrowed!
  * const csvResult = await query.execute(
  *   "SELECT * FROM runs",
- *   { format: "csv", period: "7d" }
+ *   { format: "csv", period: "1d" }
  * );
  * console.log(csvResult.format); // "csv"
  * const lines = csvResult.results.split('\n'); // ✓ results is string
@@ -133,26 +133,27 @@ function execute<TRow extends Record<string, any> = Record<string, any>>(
 ): Promise<{ format: "json"; results: Array<TRow> } | { format: "csv"; results: string }> {
   const apiClient = apiClientManager.clientOrThrow();
 
+  const format = options?.format ?? "json";
+
   const $requestOptions = mergeRequestOptions(
     {
       tracer,
       name: "query.execute()",
-      icon: "sparkles",
+      icon: "query",
       attributes: {
         scope: options?.scope ?? "environment",
         format: options?.format ?? "json",
+        query: tsql,
+        period: options?.period,
+        from: options?.from,
+        to: options?.to,
       },
     },
     requestOptions
   );
 
-  const format = options?.format ?? "json";
-
   return apiClient.executeQuery(tsql, options, $requestOptions).then((response) => {
-    if (typeof response === "string") {
-      return { format: "csv" as const, results: response };
-    }
-    return { format: "json" as const, results: response.rows };
+    return response;
   }) as Promise<{ format: "json"; results: Array<TRow> } | { format: "csv"; results: string }>;
 }
 

references/hello-world/src/trigger/query.ts

@@ -4,8 +4,8 @@ import { logger, query, task } from "@trigger.dev/sdk";
 type RunRow = {
   id: string;
   status: string;
-  created_at: string;
-  duration: number;
+  triggered_at: string;
+  total_duration: number;
 };
 
 // Simple query example - just the query string, all defaults
@@ -25,7 +25,7 @@ export const simpleQueryTask = task({
 
     // Type-safe query with explicit row type
     const typedResult = await query.execute<RunRow>(
-      "SELECT id, status, created_at, duration FROM runs LIMIT 10"
+      "SELECT run_id, status, triggered_at, total_duration FROM runs LIMIT 10"
     );
 
     logger.info("Query results (typed)", {
@@ -39,7 +39,7 @@ export const simpleQueryTask = task({
       logger.info(`Run ${index + 1}`, {
         id: row.id, // TypeScript knows this is a string
         status: row.status, // TypeScript knows this is a string
-        duration: row.duration, // TypeScript knows this is a number
+        total_duration: row.total_duration, // TypeScript knows this is a number
       });
     });
 
@@ -65,7 +65,7 @@ export const fullJsonQueryTask = task({
       `SELECT
         status,
         COUNT(*) as count,
-        AVG(duration) as avg_duration
+        AVG(total_duration) as avg_duration
       FROM runs
       WHERE status IN ('COMPLETED', 'FAILED')
       GROUP BY status`,
@@ -104,7 +104,7 @@ export const csvQueryTask = task({
 
     // Query with CSV format - automatically typed as discriminated union!
     const result = await query.execute(
-      "SELECT id, status, created_at, duration FROM runs LIMIT 100",
+      "SELECT run_id, status, triggered_at, total_duration FROM runs LIMIT 10",
       {
         scope: "project", // Query all environments in the project
         period: "7d", // Last 7 days
@@ -116,22 +116,12 @@ export const csvQueryTask = task({
     logger.info("CSV query completed", {
       format: result.format,
       dataLength: result.results.length,
-      preview: result.results.substring(0, 200), // Show first 200 chars
-    });
-
-    // Count the number of rows (lines - 1 for header)
-    const lines = result.results.split("\n");
-    const rowCount = lines.length - 1;
-
-    logger.info("CSV stats", {
-      totalRows: rowCount,
-      headerLine: lines[0],
+      results: result.results,
     });
 
     return {
       format: result.format,
       csv: result.results,
-      rowCount,
     };
   },
 });