diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/SpanService.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/SpanService.java index 315f20eee9..648708b59d 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/domain/SpanService.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/SpanService.java @@ -262,11 +262,12 @@ public Mono create(@NonNull SpanBatch batch) { List projectNames = batch.spans() .stream() .map(Span::projectName) + .map(WorkspaceUtils::getProjectName) .distinct() .toList(); Mono> resolveProjects = Flux.fromIterable(projectNames) - .flatMap(this::resolveProject) + .flatMap(this::getOrCreateProject) .collectList() .map(projects -> bindSpanToProjectAndId(batch, projects)) .subscribeOn(Schedulers.boundedElastic()); @@ -293,7 +294,4 @@ private List bindSpanToProjectAndId(SpanBatch batch, List project .toList(); } - private Mono resolveProject(String projectName) { - return getOrCreateProject(WorkspaceUtils.getProjectName(projectName)); - } } diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/TraceDAO.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/TraceDAO.java index 38994473b3..275a5f8f9e 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/domain/TraceDAO.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/TraceDAO.java @@ -1,7 +1,6 @@ package com.comet.opik.domain; import com.comet.opik.api.Trace; -import com.comet.opik.api.TraceCountResponse; import com.comet.opik.api.TraceSearchCriteria; import com.comet.opik.api.TraceUpdate; import com.comet.opik.domain.filter.FilterQueryBuilder; @@ -36,6 +35,7 @@ import java.util.stream.Collectors; import static com.comet.opik.api.Trace.TracePage; +import static com.comet.opik.api.TraceCountResponse.WorkspaceTraceCount; import static com.comet.opik.domain.AsyncContextUtils.bindUserNameAndWorkspaceContext; import static com.comet.opik.domain.AsyncContextUtils.bindWorkspaceIdToFlux; import static com.comet.opik.domain.AsyncContextUtils.bindWorkspaceIdToMono; @@ -67,7 +67,7 @@ interface TraceDAO { Mono batchInsert(List traces, Connection connection); - Flux countTracesPerWorkspace(Connection connection); + Flux countTracesPerWorkspace(Connection connection); } @Slf4j @@ -877,12 +877,12 @@ private String getOrDefault(JsonNode value) { } @com.newrelic.api.agent.Trace(dispatcher = true) - public Flux countTracesPerWorkspace(Connection connection) { + public Flux countTracesPerWorkspace(Connection connection) { var statement = connection.createStatement(TRACE_COUNT_BY_WORKSPACE_ID); return Mono.from(statement.execute()) - .flatMapMany(result -> result.map((row, rowMetadata) -> TraceCountResponse.WorkspaceTraceCount.builder() + .flatMapMany(result -> result.map((row, rowMetadata) -> WorkspaceTraceCount.builder() .workspace(row.get("workspace_id", String.class)) .traceCount(row.get("trace_count", Integer.class)).build())); } diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/TraceService.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/TraceService.java index 1b2ab35dab..ac03ca7b34 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/domain/TraceService.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/TraceService.java @@ -101,11 +101,12 @@ public Mono create(TraceBatch batch) { List projectNames = batch.traces() .stream() .map(Trace::projectName) + .map(WorkspaceUtils::getProjectName) .distinct() .toList(); Mono> resolveProjects = Flux.fromIterable(projectNames) - .flatMap(this::resolveProject) + .flatMap(this::getOrCreateProject) .collectList() .map(projects -> bindTraceToProjectAndId(batch, projects)) .subscribeOn(Schedulers.boundedElastic()); @@ -132,10 +133,6 @@ private List bindTraceToProjectAndId(TraceBatch batch, List proj .toList(); } - private Mono resolveProject(String projectName) { - return getOrCreateProject(WorkspaceUtils.getProjectName(projectName)); - } - private Mono insertTrace(Trace newTrace, Project project, UUID id) { //TODO: refactor to implement proper conflict resolution return template.nonTransaction(connection -> dao.findById(id, connection)) @@ -327,6 +324,7 @@ public Mono validateTraceWorkspace(@NonNull String workspaceId, @NonNul } @Override + @com.newrelic.api.agent.Trace(dispatcher = true) public Mono countTracesPerWorkspace() { return template.stream(dao::countTracesPerWorkspace) .collectList() diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/ClickHouseContainerUtils.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/ClickHouseContainerUtils.java index 122dc0eed6..9282783960 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/ClickHouseContainerUtils.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/ClickHouseContainerUtils.java @@ -12,10 +12,14 @@ public class ClickHouseContainerUtils { public static final String DATABASE_NAME_VARIABLE = "ANALYTICS_DB_DATABASE_NAME"; public static ClickHouseContainer newClickHouseContainer() { + return newClickHouseContainer(true); + } + + public static ClickHouseContainer newClickHouseContainer(boolean reusable) { // TODO: Use non-deprecated ClickHouseContainer: https://github.com/comet-ml/opik/issues/58 return new ClickHouseContainer( DockerImageName.parse("clickhouse/clickhouse-server:24.3.8.13-alpine")) - .withReuse(true); + .withReuse(reusable); } public static DatabaseAnalyticsFactory newDatabaseAnalyticsFactory(ClickHouseContainer clickHouseContainer, diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/internal/UsageResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/internal/UsageResourceTest.java index 78e22c95ec..864d7b462c 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/internal/UsageResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/internal/UsageResourceTest.java @@ -55,7 +55,7 @@ public class UsageResourceTest { private static final MySQLContainer MYSQL_CONTAINER = MySQLContainerUtils.newMySQLContainer(); - private static final ClickHouseContainer CLICK_HOUSE_CONTAINER = ClickHouseContainerUtils.newClickHouseContainer(); + private static final ClickHouseContainer CLICK_HOUSE_CONTAINER = ClickHouseContainerUtils.newClickHouseContainer(false); @RegisterExtension private static final TestDropwizardAppExtension app; @@ -134,21 +134,21 @@ void tracesCountForWorkspace() { // Setup second workspace with traces, but leave created_at date set to today, so traces do not end up in the pool var workspaceNameForToday = UUID.randomUUID().toString(); var workspaceIdForToday = UUID.randomUUID().toString(); - setupTracesForWorkspace(workspaceNameForToday, workspaceIdForToday, okApikey); + var apikey = UUID.randomUUID().toString(); + + setupTracesForWorkspace(workspaceNameForToday, workspaceIdForToday, apikey); try (var actualResponse = client.target(USAGE_RESOURCE_URL_TEMPLATE.formatted(baseURI)) .path("/workspace-trace-counts") .request() - .header(HttpHeaders.AUTHORIZATION, okApikey) - .header(WORKSPACE_HEADER, workspaceName) .get()) { assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200); assertThat(actualResponse.hasEntity()).isTrue(); var response = actualResponse.readEntity(TraceCountResponse.class); - assertThat(response.workspacesTracesCount().size()).isEqualTo(1); - assertThat(response.workspacesTracesCount().get(0)) + assertThat(response.workspacesTracesCount()).hasSize(1); + assertThat(response.workspacesTracesCount().getFirst()) .isEqualTo(new TraceCountResponse.WorkspaceTraceCount(workspaceId, tracesCount)); } } diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/SpansResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/SpansResourceTest.java index ecc0b11279..2d5640ad40 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/SpansResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/SpansResourceTest.java @@ -3231,6 +3231,30 @@ void batch__whenCreateSpans__thenReturnNoContent() { API_KEY); } + @Test + void batch__whenSpansProjectNameIsNull__thenUserDefaultProjectAndReturnNoContent() { + + String apiKey = UUID.randomUUID().toString(); + String workspaceName = UUID.randomUUID().toString(); + String workspaceId = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var expectedSpans = PodamFactoryUtils.manufacturePojoList(podamFactory, Span.class).stream() + .map(trace -> trace.toBuilder() + .projectName(null) + .endTime(null) + .usage(null) + .feedbackScores(null) + .build()) + .toList(); + + batchCreateAndAssert(expectedSpans, apiKey, workspaceName); + + getAndAssertPage(workspaceName, DEFAULT_PROJECT, List.of(), List.of(), expectedSpans.reversed(), List.of(), + apiKey); + } + @Test void batch__whenSendingMultipleSpansWithSameId__thenReturn422() { var expectedSpans = List.of(podamFactory.manufacturePojo(Span.class).toBuilder() diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/TracesResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/TracesResourceTest.java index 4a7d72f5d9..48d12ebc55 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/TracesResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/TracesResourceTest.java @@ -3434,12 +3434,11 @@ void batch__whenCreateTraces__thenReturnNoContent() { var projectName = UUID.randomUUID().toString(); - var projectId = createProject(projectName, TEST_WORKSPACE, API_KEY); + createProject(projectName, TEST_WORKSPACE, API_KEY); var expectedTraces = IntStream.range(0, 1000) .mapToObj(i -> factory.manufacturePojo(Trace.class).toBuilder() .projectName(projectName) - .projectId(projectId) .endTime(null) .usage(null) .feedbackScores(null) @@ -3452,6 +3451,30 @@ void batch__whenCreateTraces__thenReturnNoContent() { API_KEY); } + @Test + void batch__whenTraceProjectNameIsNull__thenUserDefaultProjectAndReturnNoContent() { + + String apiKey = UUID.randomUUID().toString(); + String workspaceName = UUID.randomUUID().toString(); + String workspaceId = UUID.randomUUID().toString(); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var expectedTraces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class).stream() + .map(trace -> trace.toBuilder() + .projectName(null) + .endTime(null) + .usage(null) + .feedbackScores(null) + .build()) + .toList(); + + batchCreateTracesAndAssert(expectedTraces, apiKey, workspaceName); + + getAndAssertPage(workspaceName, DEFAULT_PROJECT, List.of(), List.of(), expectedTraces.reversed(), List.of(), + apiKey); + } + @Test void batch__whenSendingMultipleTracesWithSameId__thenReturn422() { var trace = factory.manufacturePojo(Trace.class).toBuilder()