WIP

Author: andresgutgon

.husky/pre-commit

@@ -1,6 +1,6 @@
 #!/usr/bin/env sh
 set -e
 
-# pnpm format
-# pnpm check:fix
-# pnpm knip
+pnpm format
+pnpm check:fix
+pnpm knip

packages/domain/datasets/src/ports/dataset-row-repository.ts

@@ -1,8 +1,13 @@
-import type { DatasetId, DatasetRowId, OrganizationId, RepositoryError } from "@domain/shared"
+import type { DatasetId, DatasetRowId, OrganizationId, RepositoryError, TraceId } from "@domain/shared"
 import { type Effect, ServiceMap } from "effect"
 import type { DatasetRow, RowFieldValue, RowNotFoundError } from "../entities/dataset-row.ts"
 
 export interface DatasetRowRepositoryShape {
+  findExistingTraceIds(args: {
+    readonly organizationId: OrganizationId
+    readonly datasetId: DatasetId
+    readonly traceIds: readonly TraceId[]
+  }): Effect.Effect<ReadonlySet<TraceId>, RepositoryError>
   insertBatch(args: {
     readonly organizationId: OrganizationId
     readonly datasetId: DatasetId

packages/domain/datasets/src/use-cases/add-traces-to-dataset.ts

@@ -3,6 +3,7 @@ import type { TraceDetail } from "@domain/spans"
 import { TraceRepository } from "@domain/spans"
 import { Effect } from "effect"
 import { DatasetRepository } from "../ports/dataset-repository.ts"
+import { DatasetRowRepository } from "../ports/dataset-row-repository.ts"
 import { createDataset } from "./create-dataset.ts"
 import { insertRows } from "./insert-rows.ts"
 
@@ -24,15 +25,14 @@ function mapTraceToRow(t: TraceDetail) {
   }
 }
 
-function fetchAndMapTraces(args: {
+function fetchTraces(args: {
   readonly organizationId: OrganizationId
   readonly projectId: ProjectId
   readonly traceIds: readonly TraceId[]
 }) {
   return Effect.gen(function* () {
     const repo = yield* TraceRepository
-    const traces = yield* repo.findByTraceIds(args)
-    return traces.map(mapTraceToRow)
+    return yield* repo.findByTraceIds(args)
   })
 }
 
@@ -43,7 +43,19 @@ export function addTracesToDataset(args: {
   readonly traceIds: readonly TraceId[]
 }) {
   return Effect.gen(function* () {
-    const rows = yield* fetchAndMapTraces(args)
+    const rowRepo = yield* DatasetRowRepository
+
+    const existingTraceIds = yield* rowRepo.findExistingTraceIds({
+      organizationId: args.organizationId,
+      datasetId: args.datasetId,
+      traceIds: args.traceIds,
+    })
+
+    const newTraceIds = args.traceIds.filter((id) => !existingTraceIds.has(id))
+    if (newTraceIds.length === 0) return { versionId: "" as const, version: 0, rowIds: [] as string[] }
+
+    const traces = yield* fetchTraces({ ...args, traceIds: newTraceIds })
+    const rows = traces.map(mapTraceToRow)
     if (rows.length === 0) return { versionId: "" as const, version: 0, rowIds: [] as string[] }
 
     return yield* insertRows({
@@ -71,7 +83,8 @@ export function createDatasetFromTraces(args: {
     })
 
     const populateDataset = Effect.gen(function* () {
-      const rows = yield* fetchAndMapTraces(args)
+      const traces = yield* fetchTraces(args)
+      const rows = traces.map(mapTraceToRow)
       if (rows.length === 0) {
         return { datasetId: dataset.id, versionId: "" as const, version: 0, rowIds: [] as string[] }
       }

packages/platform/db-clickhouse/src/repositories/dataset-row-repository.ts

@@ -1,7 +1,7 @@
 import type { ClickHouseClient } from "@clickhouse/client"
 import type { DatasetRow } from "@domain/datasets"
 import { DatasetRowRepository, RowNotFoundError } from "@domain/datasets"
-import { ChSqlClient, type ChSqlClientShape, DatasetId, DatasetRowId } from "@domain/shared"
+import { ChSqlClient, type ChSqlClientShape, DatasetId, DatasetRowId, TraceId } from "@domain/shared"
 import { safeParseJson, safeStringifyJson } from "@repo/utils"
 import { Effect, Layer } from "effect"
 
@@ -47,6 +47,35 @@ export const DatasetRowRepositoryLive = Layer.effect(
     const chSqlClient = (yield* ChSqlClient) as ChSqlClientShape<ClickHouseClient>
 
     return {
+      findExistingTraceIds: (args) =>
+        chSqlClient.query(async (client) => {
+          if (args.traceIds.length === 0) return new Set<TraceId>()
+
+          const result = await client
+            .query({
+              query: `
+                SELECT DISTINCT JSONExtractString(
+                  argMax(metadata, xact_id), 'traceId'
+                ) AS trace_id
+                FROM dataset_rows
+                WHERE organization_id = {organizationId:String}
+                  AND dataset_id = {datasetId:String}
+                GROUP BY row_id
+                HAVING argMax(_object_delete, xact_id) = false
+                  AND trace_id IN ({traceIds:Array(String)})
+              `,
+              query_params: {
+                organizationId: args.organizationId as string,
+                datasetId: args.datasetId as string,
+                traceIds: Array.from(args.traceIds) as string[],
+              },
+              format: "JSONEachRow",
+            })
+            .then((r) => r.json<{ trace_id: string }>())
+
+          return new Set(result.map((r) => TraceId(r.trace_id)))
+        }),
+
       insertBatch: (args) =>
         chSqlClient.query(async (client) => {
           const values = args.rows.map((row) => ({