@mastra/duckdb 1.2.0 → 1.3.0-alpha.1

package/dist/{observability-YJBOVLPV.js → observability-7LL4LLXY.js}

@@ -1,6 +1,6 @@
 import { DuckDBConnection } from './chunk-37GBWD4M.js';
 import { MastraError, ErrorCategory, ErrorDomain } from '@mastra/core/error';
-import { ObservabilityStorage, createStorageErrorId, toTraceSpans } from '@mastra/core/storage';
+import { BRANCH_SPAN_TYPES, ObservabilityStorage, createStorageErrorId, toTraceSpans, METRIC_DISTINCT_COLUMNS } from '@mastra/core/storage';
 import { EntityType } from '@mastra/core/observability';
 import { parseFieldKey } from '@mastra/core/utils';
 
@@ -1186,7 +1186,16 @@ async function listLogs(db, args) {
     logs
   };
 }
-function getAggregationSql2(aggregation, measure = "value") {
+function resolveDistinctColumnSql(distinctColumn) {
+  if (!distinctColumn) {
+    throw new Error(`count_distinct aggregation requires a 'distinctColumn' argument`);
+  }
+  if (!METRIC_DISTINCT_COLUMNS.includes(distinctColumn)) {
+    throw new Error(`Invalid distinctColumn: ${distinctColumn}`);
+  }
+  return parseFieldKey(distinctColumn);
+}
+function getAggregationSql2(aggregation, measure = "value", distinctColumn) {
   switch (aggregation) {
     case "sum":
       return `SUM(${measure})`;
@@ -1198,6 +1207,9 @@ function getAggregationSql2(aggregation, measure = "value") {
       return `MAX(${measure})`;
     case "count":
       return `CAST(COUNT(${measure}) AS DOUBLE)`;
+    case "count_distinct": {
+      return `CAST(approx_count_distinct(${resolveDistinctColumnSql(distinctColumn)}) AS DOUBLE)`;
+    }
     case "last":
       return `arg_max(${measure}, timestamp)`;
     default:
@@ -1443,7 +1455,7 @@ async function listMetrics(db, args) {
   };
 }
 async function getMetricAggregate(db, args) {
-  const aggSql = getAggregationSql2(args.aggregation);
+  const aggSql = getAggregationSql2(args.aggregation, "value", args.distinctColumn);
   const { clause: nameClause, params: nameParams } = buildMetricNameFilter(args.name);
   const { clause: filterClause, params: filterParams } = buildWhereClause(
     args.filters
@@ -1527,7 +1539,7 @@ async function getMetricAggregate(db, args) {
   return { value, estimatedCost: costSummary.estimatedCost, costUnit: costSummary.costUnit };
 }
 async function getMetricBreakdown(db, args) {
-  const aggSql = getAggregationSql2(args.aggregation);
+  const aggSql = getAggregationSql2(args.aggregation, "value", args.distinctColumn);
   const { clause: nameClause, params: nameParams } = buildMetricNameFilter(args.name);
   const { clause: filterClause, params: filterParams } = buildWhereClause(
     args.filters
@@ -1541,8 +1553,11 @@ async function getMetricBreakdown(db, args) {
   const resolvedGroupBy = resolveGroupBy(args.groupBy);
   const selectGroupBy = resolvedGroupBy.map((entry) => entry.selectSql).join(", ");
   const groupByCols = resolvedGroupBy.map((entry) => entry.groupSql).join(", ");
-  const sql = `SELECT ${selectGroupBy}, ${aggSql} AS value, ${getCostSummarySelect()} FROM metric_events ${whereClause} GROUP BY ${groupByCols} ORDER BY value DESC`;
-  const rows = await db.query(sql, allParams);
+  const orderDirection = args.orderDirection === "ASC" ? "ASC" : "DESC";
+  const limitClause = typeof args.limit === "number" ? `LIMIT ?` : "";
+  const limitParams = typeof args.limit === "number" ? [args.limit] : [];
+  const sql = `SELECT ${selectGroupBy}, ${aggSql} AS value, ${getCostSummarySelect()} FROM metric_events ${whereClause} GROUP BY ${groupByCols} ORDER BY value ${orderDirection} ${limitClause}`;
+  const rows = await db.query(sql, [...allParams, ...limitParams]);
   const groups = rows.map((row) => {
     const dimensions = {};
     for (const entry of resolvedGroupBy) {
@@ -1560,7 +1575,7 @@ async function getMetricBreakdown(db, args) {
   return { groups };
 }
 async function getMetricTimeSeries(db, args) {
-  const aggSql = getAggregationSql2(args.aggregation);
+  const aggSql = getAggregationSql2(args.aggregation, "value", args.distinctColumn);
   const intervalSql = getIntervalSql2(args.interval);
   const { clause: nameClause, params: nameParams } = buildMetricNameFilter(args.name);
   const { clause: filterClause, params: filterParams } = buildWhereClause(
@@ -2144,6 +2159,12 @@ async function listScores(db, args) {
     scores: rows.map((row) => rowToScoreRecord(row))
   };
 }
+async function getScoreById(db, scoreId) {
+  const rows = await db.query(`SELECT * FROM score_events WHERE scoreId = ? LIMIT 1`, [
+    scoreId
+  ]);
+  return rows[0] ? rowToScoreRecord(rows[0]) : null;
+}
 async function getScoreAggregate(db, args) {
   const aggSql = getAggregationSql3(args.aggregation);
   const { clause, params } = buildScoreWhereClause(args);
@@ -2413,11 +2434,91 @@ function rowToSpanRecord(row) {
     updatedAt: null
   };
 }
-function buildHasChildErrorClause(hasChildError) {
+function buildHasChildErrorClause(hasChildError, rootAlias) {
   if (hasChildError === void 0) return "";
-  const base = `SELECT 1 FROM reconstructed_spans c WHERE c.traceId = root_spans.traceId AND c.spanId != root_spans.spanId AND c.error IS NOT NULL`;
+  const base = `SELECT 1 FROM span_events c WHERE c.traceId = ${rootAlias}.traceId AND c.spanId != ${rootAlias}.spanId AND c.error IS NOT NULL`;
   return hasChildError ? `EXISTS (${base})` : `NOT EXISTS (${base})`;
 }
+var PREFILTER_KEYS = /* @__PURE__ */ new Set([
+  "traceId",
+  "spanId",
+  "parentSpanId",
+  "name",
+  "spanType",
+  "source",
+  "entityType",
+  "entityId",
+  "entityName",
+  "entityVersionId",
+  "experimentId",
+  "userId",
+  "organizationId",
+  "resourceId",
+  "runId",
+  "sessionId",
+  "threadId",
+  "requestId",
+  "environment",
+  "serviceName"
+]);
+var SAFE_PREFILTER_ORDER_FIELDS = /* @__PURE__ */ new Set(["startedAt"]);
+function intersectTimestampRange(existing, incoming) {
+  if (!existing) return { ...incoming };
+  const merged = { ...existing };
+  if (incoming.start !== void 0) {
+    if (merged.start === void 0 || incoming.start.getTime() > merged.start.getTime()) {
+      merged.start = incoming.start;
+      merged.startExclusive = incoming.startExclusive;
+    } else if (incoming.start.getTime() === merged.start.getTime()) {
+      merged.startExclusive = (merged.startExclusive ?? false) || (incoming.startExclusive ?? false);
+    }
+  }
+  if (incoming.end !== void 0) {
+    if (merged.end === void 0 || incoming.end.getTime() < merged.end.getTime()) {
+      merged.end = incoming.end;
+      merged.endExclusive = incoming.endExclusive;
+    } else if (incoming.end.getTime() === merged.end.getTime()) {
+      merged.endExclusive = (merged.endExclusive ?? false) || (incoming.endExclusive ?? false);
+    }
+  }
+  return merged;
+}
+function partitionAnchorFilters(filters) {
+  const prefilter = {};
+  const postAgg = {};
+  let hasChildError;
+  for (const [key, value] of Object.entries(filters)) {
+    if (value === void 0 || value === null) continue;
+    if (key === "hasChildError") {
+      if (typeof value === "boolean") hasChildError = value;
+      continue;
+    }
+    if (key === "startedAt") {
+      prefilter.timestamp = intersectTimestampRange(
+        prefilter.timestamp,
+        value
+      );
+      continue;
+    }
+    if (key === "endedAt") {
+      postAgg.endedAt = value;
+      const dateRange = value;
+      if (dateRange?.end) {
+        prefilter.timestamp = intersectTimestampRange(prefilter.timestamp, {
+          end: dateRange.end,
+          endExclusive: dateRange.endExclusive
+        });
+      }
+      continue;
+    }
+    if (PREFILTER_KEYS.has(key)) {
+      prefilter[key] = value;
+      continue;
+    }
+    postAgg[key] = value;
+  }
+  return { prefilter, postAgg, hasChildError };
+}
 function toValuesTuple(row) {
   return [
     v(row.eventType),
@@ -2596,47 +2697,211 @@ async function listTraces(db, args) {
   const page = Number(args.pagination?.page ?? 0);
   const perPage = Number(args.pagination?.perPage ?? 10);
   const orderBy = { field: args.orderBy?.field ?? "startedAt", direction: args.orderBy?.direction ?? "DESC" };
-  const { clause: filterClause, params: filterParams } = buildWhereClause(filters);
-  const orderByClause = buildOrderByClause(orderBy);
-  const { clause: paginationClause, params: paginationParams } = buildPaginationClause({ page, perPage });
-  const filterParts = [];
-  if (filterClause) filterParts.push(filterClause.replace(/^WHERE\s+/i, ""));
-  const hasChildError = typeof filters.hasChildError === "boolean" ? filters.hasChildError : void 0;
-  const childErrorClause = buildHasChildErrorClause(hasChildError);
-  if (childErrorClause) filterParts.push(childErrorClause);
-  const combinedFilterClause = filterParts.length > 0 ? `WHERE ${filterParts.join(" AND ")}` : "";
-  const cteSql = `
-    WITH reconstructed_spans AS (
+  const { prefilter, postAgg, hasChildError } = partitionAnchorFilters(filters);
+  const { clause: prefilterClause, params: prefilterParams } = buildWhereClause(prefilter);
+  const prefilterParts = [`eventType = 'start'`, `parentSpanId IS NULL`];
+  if (prefilterClause) prefilterParts.push(prefilterClause.replace(/^WHERE\s+/i, ""));
+  const prefilterWhere = `WHERE ${prefilterParts.join(" AND ")}`;
+  const outerAlias = "outer_root";
+  const orderDir = orderBy.direction.toUpperCase();
+  if (orderDir !== "ASC" && orderDir !== "DESC") {
+    throw new Error(`Invalid sort direction: ${orderBy.direction}`);
+  }
+  const canOrderInPrefilter = SAFE_PREFILTER_ORDER_FIELDS.has(orderBy.field);
+  const hasPostAggFilters = Object.keys(postAgg).length > 0 || hasChildError !== void 0 || !canOrderInPrefilter;
+  if (!hasPostAggFilters) {
+    const prefilterOrderBy = `ORDER BY timestamp ${orderDir}`;
+    const offset = page * perPage;
+    const countSql2 = `
+      SELECT COUNT(*) as total
+      FROM span_events AS ${outerAlias}
+      ${prefilterWhere}
+    `;
+    const countResult2 = await db.query(countSql2, prefilterParams);
+    const total2 = Number(countResult2[0]?.total ?? 0);
+    const pageSql = `
+      WITH page_roots AS (
+        SELECT traceId, spanId
+        FROM span_events AS ${outerAlias}
+        ${prefilterWhere}
+        ${prefilterOrderBy}
+        LIMIT ? OFFSET ?
+      )
       ${SPAN_RECONSTRUCT_SELECT}
+      WHERE (traceId, spanId) IN (SELECT traceId, spanId FROM page_roots)
       GROUP BY traceId, spanId
+      ${buildOrderByClause(orderBy)}
+    `;
+    const rows2 = await db.query(pageSql, [...prefilterParams, perPage, offset]);
+    const spans2 = rows2.map((row) => rowToSpanRecord(row));
+    return {
+      pagination: { total: total2, page, perPage, hasMore: (page + 1) * perPage < total2 },
+      spans: toTraceSpans(spans2)
+    };
+  }
+  const { clause: postAggClause, params: postAggParams } = buildWhereClause(postAgg);
+  const postAggParts = [];
+  if (postAggClause) postAggParts.push(postAggClause.replace(/^WHERE\s+/i, ""));
+  const childErrorClause = buildHasChildErrorClause(hasChildError, "root_spans");
+  if (childErrorClause) postAggParts.push(childErrorClause);
+  const postAggWhere = postAggParts.length > 0 ? `WHERE ${postAggParts.join(" AND ")}` : "";
+  const cteSql = `
+    WITH candidate_roots AS (
+      SELECT traceId, spanId
+      FROM span_events AS ${outerAlias}
+      ${prefilterWhere}
     ),
     root_spans AS (
-      SELECT * FROM reconstructed_spans
-      WHERE parentSpanId IS NULL
+      ${SPAN_RECONSTRUCT_SELECT}
+      WHERE (traceId, spanId) IN (SELECT traceId, spanId FROM candidate_roots)
+      GROUP BY traceId, spanId
     )
   `;
+  const orderByClause = buildOrderByClause(orderBy);
+  const { clause: paginationClause, params: paginationParams } = buildPaginationClause({ page, perPage });
   const countSql = `
     ${cteSql}
-    SELECT COUNT(*) as total FROM root_spans ${combinedFilterClause}
+    SELECT COUNT(*) as total FROM root_spans ${postAggWhere}
   `;
-  const countResult = await db.query(countSql, filterParams);
+  const countResult = await db.query(countSql, [...prefilterParams, ...postAggParams]);
   const total = Number(countResult[0]?.total ?? 0);
   const dataSql = `
     ${cteSql}
-    SELECT * FROM root_spans ${combinedFilterClause} ${orderByClause} ${paginationClause}
+    SELECT * FROM root_spans ${postAggWhere} ${orderByClause} ${paginationClause}
   `;
-  const rows = await db.query(dataSql, [...filterParams, ...paginationParams]);
+  const rows = await db.query(dataSql, [...prefilterParams, ...postAggParams, ...paginationParams]);
   const spans = rows.map((row) => rowToSpanRecord(row));
   return {
-    pagination: {
-      total,
-      page,
-      perPage,
-      hasMore: (page + 1) * perPage < total
-    },
+    pagination: { total, page, perPage, hasMore: (page + 1) * perPage < total },
     spans: toTraceSpans(spans)
   };
 }
+var BRANCH_SPAN_TYPE_PLACEHOLDERS = BRANCH_SPAN_TYPES.map(() => "?").join(", ");
+async function getSpans(db, args) {
+  if (args.spanIds.length === 0) {
+    return { traceId: args.traceId, spans: [] };
+  }
+  const placeholders = args.spanIds.map(() => "?").join(", ");
+  const rows = await db.query(
+    `${SPAN_RECONSTRUCT_SELECT}
+     WHERE traceId = ? AND spanId IN (${placeholders})
+     GROUP BY traceId, spanId`,
+    [args.traceId, ...args.spanIds]
+  );
+  return {
+    traceId: args.traceId,
+    spans: rows.map((row) => rowToSpanRecord(row))
+  };
+}
+async function listBranches(db, args) {
+  const filters = args.filters ?? {};
+  const page = Number(args.pagination?.page ?? 0);
+  const perPage = Number(args.pagination?.perPage ?? 10);
+  const orderBy = { field: args.orderBy?.field ?? "startedAt", direction: args.orderBy?.direction ?? "DESC" };
+  const userSpanType = filters.spanType;
+  if (typeof userSpanType === "string" && !BRANCH_SPAN_TYPES.includes(userSpanType)) {
+    return {
+      pagination: { total: 0, page, perPage, hasMore: false },
+      branches: []
+    };
+  }
+  const { spanType: _spanType, ...rest } = filters;
+  const { prefilter, postAgg} = partitionAnchorFilters(rest);
+  const { clause: prefilterClause, params: prefilterFilterParams } = buildWhereClause(prefilter);
+  const prefilterParts = [`eventType = 'start'`];
+  let spanTypeParams;
+  if (typeof userSpanType === "string") {
+    prefilterParts.push(`spanType = ?`);
+    spanTypeParams = [userSpanType];
+  } else {
+    prefilterParts.push(`spanType IN (${BRANCH_SPAN_TYPE_PLACEHOLDERS})`);
+    spanTypeParams = [...BRANCH_SPAN_TYPES];
+  }
+  if (prefilterClause) prefilterParts.push(prefilterClause.replace(/^WHERE\s+/i, ""));
+  const prefilterWhere = `WHERE ${prefilterParts.join(" AND ")}`;
+  const prefilterParams = [...spanTypeParams, ...prefilterFilterParams];
+  const outerAlias = "outer_anchor";
+  const orderDir = orderBy.direction.toUpperCase();
+  if (orderDir !== "ASC" && orderDir !== "DESC") {
+    throw new Error(`Invalid sort direction: ${orderBy.direction}`);
+  }
+  const canOrderInPrefilter = SAFE_PREFILTER_ORDER_FIELDS.has(orderBy.field);
+  const hasPostAggFilters = Object.keys(postAgg).length > 0 || !canOrderInPrefilter;
+  if (!hasPostAggFilters) {
+    const prefilterOrderBy = `ORDER BY timestamp ${orderDir}`;
+    const offset = page * perPage;
+    const countSql2 = `
+      SELECT COUNT(*) as total
+      FROM span_events AS ${outerAlias}
+      ${prefilterWhere}
+    `;
+    const countResult2 = await db.query(countSql2, prefilterParams);
+    const total2 = Number(countResult2[0]?.total ?? 0);
+    if (total2 === 0) {
+      return {
+        pagination: { total: 0, page, perPage, hasMore: false },
+        branches: []
+      };
+    }
+    const pageSql = `
+      WITH page_anchors AS (
+        SELECT traceId, spanId
+        FROM span_events AS ${outerAlias}
+        ${prefilterWhere}
+        ${prefilterOrderBy}
+        LIMIT ? OFFSET ?
+      )
+      ${SPAN_RECONSTRUCT_SELECT}
+      WHERE (traceId, spanId) IN (SELECT traceId, spanId FROM page_anchors)
+      GROUP BY traceId, spanId
+      ${buildOrderByClause(orderBy)}
+    `;
+    const rows2 = await db.query(pageSql, [...prefilterParams, perPage, offset]);
+    const spans2 = rows2.map((row) => rowToSpanRecord(row));
+    return {
+      pagination: { total: total2, page, perPage, hasMore: (page + 1) * perPage < total2 },
+      branches: toTraceSpans(spans2)
+    };
+  }
+  const { clause: postAggClause, params: postAggParams } = buildWhereClause(postAgg);
+  const postAggWhere = postAggClause ? postAggClause : "";
+  const orderByClause = buildOrderByClause(orderBy);
+  const { clause: paginationClause, params: paginationParams } = buildPaginationClause({ page, perPage });
+  const cteSql = `
+    WITH candidate_anchors AS (
+      SELECT traceId, spanId
+      FROM span_events AS ${outerAlias}
+      ${prefilterWhere}
+    ),
+    branch_anchors AS (
+      ${SPAN_RECONSTRUCT_SELECT}
+      WHERE (traceId, spanId) IN (SELECT traceId, spanId FROM candidate_anchors)
+      GROUP BY traceId, spanId
+    )
+  `;
+  const countSql = `
+    ${cteSql}
+    SELECT COUNT(*) as total FROM branch_anchors ${postAggWhere}
+  `;
+  const countResult = await db.query(countSql, [...prefilterParams, ...postAggParams]);
+  const total = Number(countResult[0]?.total ?? 0);
+  if (total === 0) {
+    return {
+      pagination: { total: 0, page, perPage, hasMore: false },
+      branches: []
+    };
+  }
+  const dataSql = `
+    ${cteSql}
+    SELECT * FROM branch_anchors ${postAggWhere} ${orderByClause} ${paginationClause}
+  `;
+  const rows = await db.query(dataSql, [...prefilterParams, ...postAggParams, ...paginationParams]);
+  const spans = rows.map((row) => rowToSpanRecord(row));
+  return {
+    pagination: { total, page, perPage, hasMore: (page + 1) * perPage < total },
+    branches: toTraceSpans(spans)
+  };
+}
 
 // src/storage/domains/observability/index.ts
 function buildSignalMigrationRequiredMessage(args) {
@@ -2739,6 +3004,9 @@ var ObservabilityStorageDuckDB = class extends ObservabilityStorage {
   async getSpan(args) {
     return getSpan(this.db, args);
   }
+  async getSpans(args) {
+    return getSpans(this.db, args);
+  }
   async getRootSpan(args) {
     return getRootSpan(this.db, args);
   }
@@ -2751,6 +3019,9 @@ var ObservabilityStorageDuckDB = class extends ObservabilityStorage {
   async listTraces(args) {
     return listTraces(this.db, args);
   }
+  async listBranches(args) {
+    return listBranches(this.db, args);
+  }
   // Logs
   async batchCreateLogs(args) {
     return batchCreateLogs(this.db, args);
@@ -2813,6 +3084,9 @@ var ObservabilityStorageDuckDB = class extends ObservabilityStorage {
   async listScores(args) {
     return listScores(this.db, args);
   }
+  async getScoreById(scoreId) {
+    return getScoreById(this.db, scoreId);
+  }
   async getScoreAggregate(args) {
     return getScoreAggregate(this.db, args);
   }
@@ -2850,5 +3124,5 @@ var ObservabilityStorageDuckDB = class extends ObservabilityStorage {
 };
 
 export { ObservabilityStorageDuckDB };
-//# sourceMappingURL=observability-YJBOVLPV.js.map
-//# sourceMappingURL=observability-YJBOVLPV.js.map
+//# sourceMappingURL=observability-7LL4LLXY.js.map
+//# sourceMappingURL=observability-7LL4LLXY.js.map