From 5743d40c521910bd3a1aa4a140074da07f8acd48 Mon Sep 17 00:00:00 2001 From: Thiago Hora Date: Mon, 9 Sep 2024 13:29:53 +0200 Subject: [PATCH 01/10] [NA] Remove creation time from sort key --- .../db-app-analytics/migrations/000001_init_script.sql | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/apps/opik-backend/src/main/resources/liquibase/db-app-analytics/migrations/000001_init_script.sql b/apps/opik-backend/src/main/resources/liquibase/db-app-analytics/migrations/000001_init_script.sql index eb27d66ac2..dd664141b8 100644 --- a/apps/opik-backend/src/main/resources/liquibase/db-app-analytics/migrations/000001_init_script.sql +++ b/apps/opik-backend/src/main/resources/liquibase/db-app-analytics/migrations/000001_init_script.sql @@ -22,7 +22,7 @@ CREATE TABLE IF NOT EXISTS ${ANALYTICS_DB_DATABASE_NAME}.spans created_at DateTime64(9, 'UTC') DEFAULT now64(9), last_updated_at DateTime64(9, 'UTC') DEFAULT now64(9) ) ENGINE = ReplacingMergeTree(last_updated_at) - ORDER BY (workspace_id, project_id, trace_id, parent_span_id, created_at, id); + ORDER BY (workspace_id, project_id, trace_id, parent_span_id, id); CREATE TABLE IF NOT EXISTS ${ANALYTICS_DB_DATABASE_NAME}.traces ( @@ -39,7 +39,7 @@ CREATE TABLE IF NOT EXISTS ${ANALYTICS_DB_DATABASE_NAME}.traces created_at DateTime64(9, 'UTC') DEFAULT now64(9), last_updated_at DateTime64(9, 'UTC') DEFAULT now64(9) ) ENGINE = ReplacingMergeTree(last_updated_at) - ORDER BY (workspace_id, project_id, created_at, id); + ORDER BY (workspace_id, project_id, id); CREATE TABLE IF NOT EXISTS ${ANALYTICS_DB_DATABASE_NAME}.feedback_scores ( @@ -55,7 +55,7 @@ CREATE TABLE IF NOT EXISTS ${ANALYTICS_DB_DATABASE_NAME}.feedback_scores created_at DateTime64(9, 'UTC') DEFAULT now64(9), last_updated_at DateTime64(9, 'UTC') DEFAULT now64(9) ) ENGINE = ReplacingMergeTree(last_updated_at) - ORDER BY (workspace_id, project_id, entity_type, entity_id, created_at, name); + ORDER BY (workspace_id, project_id, entity_type, entity_id, name); CREATE TABLE IF NOT EXISTS ${ANALYTICS_DB_DATABASE_NAME}.dataset_items ( @@ -82,7 +82,7 @@ CREATE TABLE IF NOT EXISTS ${ANALYTICS_DB_DATABASE_NAME}.experiments created_at DateTime64(9, 'UTC') DEFAULT now64(9), last_updated_at DateTime64(9, 'UTC') DEFAULT now64(9) ) ENGINE = ReplacingMergeTree(last_updated_at) - ORDER BY (workspace_id, dataset_id, created_at, id); + ORDER BY (workspace_id, dataset_id, id); CREATE TABLE IF NOT EXISTS ${ANALYTICS_DB_DATABASE_NAME}.experiment_items From 1d26361aa94b19dc15471e5c849a98aad2a11aa2 Mon Sep 17 00:00:00 2001 From: Thiago Hora Date: Tue, 10 Sep 2024 15:13:30 +0200 Subject: [PATCH 02/10] OPIK-75: Add bulk spans creation endpoint --- apps/opik-backend/config.yml | 2 + .../java/com/comet/opik/api/SpanBatch.java | 12 ++ .../api/resources/v1/priv/SpansResource.java | 15 ++ .../com/comet/opik/domain/DatasetItemDAO.java | 7 +- .../comet/opik/domain/ExperimentItemDAO.java | 2 +- .../opik/domain/ExperimentItemService.java | 4 +- .../java/com/comet/opik/domain/SpanDAO.java | 166 +++++++++++++++++- .../com/comet/opik/domain/SpanService.java | 128 ++++++++++++++ .../java/com/comet/opik/domain/TraceDAO.java | 5 +- .../com/comet/opik/domain/TraceService.java | 3 +- .../db/DatabaseAnalyticsModule.java | 1 - .../instrumentation/InstrumentAsyncUtils.java | 7 +- .../infrastructure/redis/LockService.java | 7 + .../redis/RedissonLockService.java | 64 +++++-- .../java/com/comet/opik/utils/JsonUtils.java | 10 ++ .../resources/utils/ClientSupportUtils.java | 4 +- .../utils/ConditionalGZipFilter.java | 42 +++++ .../resources/v1/priv/SpansResourceTest.java | 35 ++++ .../comet/opik/domain/DummyLockService.java | 31 ++++ .../comet/opik/domain/SpanServiceTest.java | 14 +- .../opik/domain/TraceServiceImplTest.java | 14 +- .../src/test/resources/config-test.yml | 3 + 22 files changed, 518 insertions(+), 58 deletions(-) create mode 100644 apps/opik-backend/src/main/java/com/comet/opik/api/SpanBatch.java create mode 100644 apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/ConditionalGZipFilter.java create mode 100644 apps/opik-backend/src/test/java/com/comet/opik/domain/DummyLockService.java diff --git a/apps/opik-backend/config.yml b/apps/opik-backend/config.yml index ad91f3000a..42fee78a19 100644 --- a/apps/opik-backend/config.yml +++ b/apps/opik-backend/config.yml @@ -63,3 +63,5 @@ authentication: server: enableVirtualThreads: ${ENABLE_VIRTUAL_THREADS:-false} + gzip: + enabled: true \ No newline at end of file diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/SpanBatch.java b/apps/opik-backend/src/main/java/com/comet/opik/api/SpanBatch.java new file mode 100644 index 0000000000..f422eb7d74 --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/SpanBatch.java @@ -0,0 +1,12 @@ +package com.comet.opik.api; + +import com.fasterxml.jackson.annotation.JsonView; +import jakarta.validation.Valid; +import jakarta.validation.constraints.NotNull; +import jakarta.validation.constraints.Size; + +import java.util.List; + +public record SpanBatch(@NotNull @Size(min = 1, max = 2000) @JsonView( { + Span.View.Write.class}) @Valid List spans){ +} diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/SpansResource.java b/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/SpansResource.java index 8d7ac2edd9..a1fa71bb6f 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/SpansResource.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/SpansResource.java @@ -5,6 +5,7 @@ import com.comet.opik.api.FeedbackScore; import com.comet.opik.api.FeedbackScoreBatch; import com.comet.opik.api.Span; +import com.comet.opik.api.SpanBatch; import com.comet.opik.api.SpanSearchCriteria; import com.comet.opik.api.SpanUpdate; import com.comet.opik.api.filter.FiltersFactory; @@ -132,6 +133,20 @@ public Response create( return Response.created(uri).build(); } + @POST + @Path("/batch") + @Operation(operationId = "createSpans", summary = "Create spans", description = "Create spans", responses = { + @ApiResponse(responseCode = "204", description = "No Content")}) + public Response createSpans( + @RequestBody(content = @Content(schema = @Schema(implementation = SpanBatch.class))) @JsonView(Span.View.Write.class) @NotNull @Valid SpanBatch spans) { + + spanService.create(spans) + .contextWrite(ctx -> setRequestContext(ctx, requestContext)) + .block(); + + return Response.noContent().build(); + } + @PATCH @Path("{id}") @Operation(operationId = "updateSpan", summary = "Update span by id", description = "Update span by id", responses = { diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetItemDAO.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetItemDAO.java index 73e2d8c790..90b1a3ebac 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetItemDAO.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetItemDAO.java @@ -1,7 +1,6 @@ package com.comet.opik.domain; import com.comet.opik.api.DatasetItem; -import com.comet.opik.api.DatasetItemBatch; import com.comet.opik.api.DatasetItemSearchCriteria; import com.comet.opik.api.DatasetItemSource; import com.comet.opik.api.ExperimentItem; @@ -348,10 +347,12 @@ public Mono save(@NonNull UUID datasetId, @NonNull List items return Mono.empty(); } - return asyncTemplate.nonTransaction(connection -> mapAndInsert(datasetId, items, connection, INSERT_DATASET_ITEM)); + return asyncTemplate + .nonTransaction(connection -> mapAndInsert(datasetId, items, connection, INSERT_DATASET_ITEM)); } - private Mono mapAndInsert(UUID datasetId, List items, Connection connection, String sqlTemplate) { + private Mono mapAndInsert(UUID datasetId, List items, Connection connection, + String sqlTemplate) { List queryItems = getQueryItemPlaceHolder(items.size()); diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/ExperimentItemDAO.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/ExperimentItemDAO.java index 7270c51931..08594f56af 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/domain/ExperimentItemDAO.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/ExperimentItemDAO.java @@ -137,7 +137,7 @@ public Mono insert(@NonNull Set experimentItems) { } return Mono.from(connectionFactory.create()) - .flatMap(connection -> insert(experimentItems, connection)); + .flatMap(connection -> insert(experimentItems, connection)); } private Mono insert(Collection experimentItems, Connection connection) { diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/ExperimentItemService.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/ExperimentItemService.java index 4890197aa3..51e15d26ea 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/domain/ExperimentItemService.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/ExperimentItemService.java @@ -1,6 +1,5 @@ package com.comet.opik.domain; -import com.clickhouse.client.ClickHouseException; import com.comet.opik.api.ExperimentItem; import com.comet.opik.infrastructure.auth.RequestContext; import com.google.common.base.Preconditions; @@ -100,7 +99,8 @@ private Mono validateDatasetItemWorkspace(String workspaceId, Set } return datasetItemDAO.getDatasetItemWorkspace(datasetItemIds) - .map(datasetItemWorkspace -> datasetItemWorkspace.stream().allMatch(datasetItem -> workspaceId.equals(datasetItem.workspaceId()))); + .map(datasetItemWorkspace -> datasetItemWorkspace.stream() + .allMatch(datasetItem -> workspaceId.equals(datasetItem.workspaceId()))); } public Mono get(@NonNull UUID id) { diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/SpanDAO.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/SpanDAO.java index 9d3401eb28..faa0e106bb 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/domain/SpanDAO.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/SpanDAO.java @@ -6,6 +6,7 @@ import com.comet.opik.domain.filter.FilterQueryBuilder; import com.comet.opik.domain.filter.FilterStrategy; import com.comet.opik.utils.JsonUtils; +import com.comet.opik.utils.TemplateUtils; import com.newrelic.api.agent.Segment; import com.newrelic.api.agent.Trace; import io.r2dbc.spi.Connection; @@ -19,12 +20,12 @@ import lombok.extern.slf4j.Slf4j; import org.reactivestreams.Publisher; import org.stringtemplate.v4.ST; -import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import java.time.Instant; import java.util.ArrayList; import java.util.Arrays; +import java.util.Collection; import java.util.List; import java.util.Map; import java.util.Optional; @@ -37,15 +38,62 @@ import static com.comet.opik.domain.AsyncContextUtils.bindWorkspaceIdToFlux; import static com.comet.opik.domain.AsyncContextUtils.bindWorkspaceIdToMono; import static com.comet.opik.domain.FeedbackScoreDAO.EntityType; +import static com.comet.opik.infrastructure.instrumentation.InstrumentAsyncUtils.endSegment; import static com.comet.opik.infrastructure.instrumentation.InstrumentAsyncUtils.startSegment; import static com.comet.opik.utils.AsyncUtils.makeFluxContextAware; import static com.comet.opik.utils.AsyncUtils.makeMonoContextAware; +import static com.comet.opik.utils.TemplateUtils.getQueryItemPlaceHolder; @Singleton @RequiredArgsConstructor(onConstructor_ = @Inject) @Slf4j class SpanDAO { + private static final String PLAIN_INSERT = """ + INSERT INTO spans( + id, + project_id, + workspace_id, + trace_id, + parent_span_id, + name, + type, + start_time, + end_time, + input, + output, + metadata, + tags, + usage, + created_at, + created_by, + last_updated_by + ) VALUES + , + :project_id, + :workspace_id, + :trace_id, + :parent_span_id, + :name, + :type, + parseDateTime64BestEffort(:start_time, 9), + if(:end_time IS NULL, NULL, parseDateTime64BestEffort(:end_time, 9)), + :input, + :output, + :metadata, + :tags, + mapFromArrays(:usage_keys, :usage_values), + parseDateTime64BestEffort(:created_at, 9), + :created_by, + :last_updated_by + ) + , + }> + ; + """; + /** * This query handles the insertion of a new span into the database in two cases: * 1. When the span does not exist in the database. @@ -347,7 +395,7 @@ LEFT JOIN ( * FROM spans - WHERE id = :id + WHERE id IN :ids AND workspace_id = :workspace_id ORDER BY last_updated_at DESC LIMIT 1 @@ -445,6 +493,102 @@ public Mono insert(@NonNull Span span) { .then(); } + @Trace(dispatcher = true) + public Mono batchInsert(@NonNull List spans) { + return Mono.from(connectionFactory.create()) + .flatMapMany(connection -> { + if (spans.isEmpty()) { + return Mono.just(0L); + } + + return insert(spans, connection); + }) + .then(); + } + + private Publisher insert(Collection spans, Connection connection) { + + return makeMonoContextAware((userName, workspaceName, workspaceId) -> { + List queryItems = getQueryItemPlaceHolder(spans.size()); + + var template = new ST(PLAIN_INSERT) + .add("items", queryItems); + + var statement = connection.createStatement(template.render()); + + int i = 0; + for (Span span : spans) { + + statement.bind("id" + i, span.id()) + .bind("project_id" + i, span.projectId()) + .bind("trace_id" + i, span.traceId()) + .bind("name" + i, span.name()) + .bind("type" + i, span.type().toString()) + .bind("start_time" + i, span.startTime().toString()); + + if (span.parentSpanId() != null) { + statement.bind("parent_span_id" + i, span.parentSpanId()); + } else { + statement.bind("parent_span_id" + i, ""); + } + if (span.endTime() != null) { + statement.bind("end_time" + i, span.endTime().toString()); + } else { + statement.bindNull("end_time" + i, String.class); + } + + if (span.input() != null) { + statement.bind("input" + i, span.input().toString()); + } else { + statement.bind("input" + i, ""); + } + if (span.output() != null) { + statement.bind("output" + i, span.output().toString()); + } else { + statement.bind("output" + i, ""); + } + if (span.metadata() != null) { + statement.bind("metadata" + i, span.metadata().toString()); + } else { + statement.bind("metadata" + i, ""); + } + if (span.tags() != null) { + statement.bind("tags" + i, span.tags().toArray(String[]::new)); + } else { + statement.bind("tags" + i, new String[]{}); + } + + if (span.usage() != null) { + Stream.Builder keys = Stream.builder(); + Stream.Builder values = Stream.builder(); + + span.usage().forEach((key, value) -> { + keys.add(key); + values.add(value); + }); + + statement.bind("usage_keys" + i, keys.build().toArray(String[]::new)); + statement.bind("usage_values" + i, values.build().toArray(Integer[]::new)); + } else { + statement.bind("usage_keys" + i, new String[]{}); + statement.bind("usage_values" + i, new Integer[]{}); + } + + statement.bind("created_at" + i, span.createdAt().toString()); + statement.bind("created_by" + i, userName); + statement.bind("last_updated_by" + i, userName); + i++; + } + + statement.bind("workspace_id", workspaceId); + + Segment segment = startSegment("spans", "Clickhouse", "insert_plain"); + + return Mono.from(statement.execute()) + .doFinally(signalType -> endSegment(segment)); + }); + } + private Publisher insert(Span span, Connection connection) { var template = newInsertTemplate(span); var statement = connection.createStatement(template.render()) @@ -617,7 +761,7 @@ public Mono getById(@NonNull UUID id) { private Publisher getById(UUID id, Connection connection) { var statement = connection.createStatement(SELECT_BY_ID) - .bind("id", id); + .bind("ids", new String[]{id.toString()}); Segment segment = startSegment("spans", "Clickhouse", "get_by_id"); @@ -790,4 +934,20 @@ public Mono> getSpanWorkspace(@NonNull Set sp .collectList(); } + public Mono> getByIds(@NonNull List ids) { + + if (ids.isEmpty()) { + return Mono.just(List.of()); + } + + return Mono.from(connectionFactory.create()) + .flatMapMany(connection -> { + var statement = connection.createStatement(SELECT_BY_ID) + .bind("ids", ids.toArray(UUID[]::new)); + + return makeFluxContextAware(bindWorkspaceIdToFlux(statement)); + }) + .flatMap(this::mapToDto) + .collectList(); + } } 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 45cd6536fd..516207ad04 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 @@ -3,6 +3,7 @@ import com.clickhouse.client.ClickHouseException; import com.comet.opik.api.Project; import com.comet.opik.api.Span; +import com.comet.opik.api.SpanBatch; import com.comet.opik.api.SpanSearchCriteria; import com.comet.opik.api.SpanUpdate; import com.comet.opik.api.error.EntityAlreadyExistsException; @@ -18,14 +19,19 @@ import lombok.NonNull; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; +import org.apache.commons.lang3.StringUtils; +import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import reactor.core.scheduler.Schedulers; import java.time.Instant; import java.util.List; +import java.util.Map; import java.util.Objects; import java.util.Set; import java.util.UUID; +import java.util.function.Function; +import java.util.stream.Collectors; import static com.comet.opik.utils.AsyncUtils.makeMonoContextAware; @@ -244,4 +250,126 @@ public Mono validateSpanWorkspace(@NonNull String workspaceId, @NonNull return spanDAO.getSpanWorkspace(spanIds) .map(spanWorkspace -> spanWorkspace.stream().allMatch(span -> workspaceId.equals(span.workspaceId()))); } + + @Trace(dispatcher = true) + public Mono create(SpanBatch batch) { + + if (batch.spans().isEmpty()) { + return Mono.empty(); + } + + List projectNames = batch.spans() + .stream() + .map(Span::projectName) + .distinct() + .toList(); + + Mono> resolveProjects = Flux.fromIterable(projectNames) + .flatMap(this::resolveProject) + .collectList() + .map(projects -> bindSpanToProjectAndId(batch, projects)) + .subscribeOn(Schedulers.boundedElastic()); + + return resolveProjects + .flatMap(spans -> { + List spanIds = spans.stream().map(Span::id).distinct().toList(); + + return lockService.lockAll(spanIds, SPAN_KEY) + .flatMap(locks -> mergeSpans(spans, spanIds) + .flatMap(spanDAO::batchInsert) + .doFinally(signalType -> lockService.unlockAll(locks).subscribe()) + .subscribeOn(Schedulers.boundedElastic())) + .then(); + }); + } + + private Mono> mergeSpans(List spans, List spanIds) { + return spanDAO.getByIds(spanIds) + .map(existingSpans -> { + Map existingSpanMap = existingSpans.stream() + .collect(Collectors.toMap(Span::id, Function.identity())); + + return spans.stream() + .collect(Collectors.groupingBy(Span::id)) + .entrySet() + .stream() + .map(groupedSpans -> { + // START state resolution from existing spans + Span currentSpan = existingSpanMap.getOrDefault(groupedSpans.getKey(), + Span.builder().build()); + + return groupedSpans + .getValue() + .stream() + .reduce(currentSpan, this::mergeSpans); + }).toList(); + }); + } + + private Span mergeSpans(Span currentSpan, Span receivedSpan) { + + if (currentSpan.id() == null) { + return receivedSpan.toBuilder() + .createdAt(getInstant(currentSpan.createdAt(), receivedSpan.createdAt())) + .build(); + } + + return Span.builder() + .id(currentSpan.id()) + .projectId(currentSpan.projectId() != null ? currentSpan.projectId() : receivedSpan.projectId()) + .traceId(currentSpan.traceId() != null ? currentSpan.traceId() : receivedSpan.traceId()) + .parentSpanId( + currentSpan.parentSpanId() != null ? currentSpan.parentSpanId() : receivedSpan.parentSpanId()) + .name(currentSpan.name() != null ? currentSpan.name() : receivedSpan.name()) + .type(currentSpan.type() != null ? currentSpan.type() : receivedSpan.type()) + .startTime(currentSpan.startTime() != null ? currentSpan.startTime() : receivedSpan.startTime()) + .endTime(receivedSpan.endTime() != null ? receivedSpan.endTime() : currentSpan.endTime()) + .input(receivedSpan.input() != null ? receivedSpan.input() : currentSpan.input()) + .output(receivedSpan.output() != null ? receivedSpan.output() : currentSpan.output()) + .metadata(receivedSpan.metadata() != null ? receivedSpan.metadata() : currentSpan.metadata()) + .tags(receivedSpan.tags() != null ? receivedSpan.tags() : currentSpan.tags()) + .usage(receivedSpan.usage() != null ? receivedSpan.usage() : currentSpan.usage()) + .createdAt(getInstant(currentSpan.createdAt(), receivedSpan.createdAt())) + .build(); + } + + private Instant getInstant(Instant current, Instant received) { + if (current == null) { + return received != null ? received : Instant.now(); + } else if (received == null) { + return current; + } else { + return current.isBefore(received) ? current : received; + } + } + + private List bindSpanToProjectAndId(SpanBatch batch, List projects) { + Map projectPerName = projects.stream() + .collect(Collectors.toMap(Project::name, Function.identity())); + + return batch.spans() + .stream() + .map(span -> { + String projectName = WorkspaceUtils.getProjectName(span.projectName()); + Project project = projectPerName.get(projectName); + + if (project == null) { + throw new EntityAlreadyExistsException(new ErrorMessage(List.of("Project not found"))); + } + + UUID id = span.id() == null ? idGenerator.generateId() : span.id(); + IdGenerator.validateVersion(id, SPAN_KEY); + + return span.toBuilder().id(id).projectId(project.id()).build(); + }) + .toList(); + } + + private Mono resolveProject(String projectName) { + if (StringUtils.isEmpty(projectName)) { + return getOrCreateProject(ProjectService.DEFAULT_PROJECT); + } + + return getOrCreateProject(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 9fa1d1d083..ed07fa2736 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 @@ -675,8 +675,9 @@ private void bindSearchCriteria(TraceSearchCriteria traceSearchCriteria, Stateme @Override @com.newrelic.api.agent.Trace(dispatcher = true) - public Mono> getTraceWorkspace(@NonNull Set traceIds, @NonNull Connection connection) { - + public Mono> getTraceWorkspace(@NonNull Set traceIds, + @NonNull Connection connection) { + if (traceIds.isEmpty()) { return Mono.just(List.of()); } 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 43fb93ec8f..3dd3d7c590 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 @@ -249,7 +249,8 @@ public Mono validateTraceWorkspace(@NonNull String workspaceId, @NonNul } return template.nonTransaction(connection -> dao.getTraceWorkspace(traceIds, connection) - .map(traceWorkspace -> traceWorkspace.stream().allMatch(trace -> workspaceId.equals(trace.workspaceId())))); + .map(traceWorkspace -> traceWorkspace.stream() + .allMatch(trace -> workspaceId.equals(trace.workspaceId())))); } } diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/db/DatabaseAnalyticsModule.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/db/DatabaseAnalyticsModule.java index a33c85e85f..ce164a384f 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/db/DatabaseAnalyticsModule.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/db/DatabaseAnalyticsModule.java @@ -7,7 +7,6 @@ import jakarta.inject.Named; import jakarta.inject.Singleton; import ru.vyarus.dropwizard.guice.module.support.DropwizardAwareModule; -import ru.vyarus.dropwizard.guice.module.yaml.bind.Config; public class DatabaseAnalyticsModule extends DropwizardAwareModule { diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/instrumentation/InstrumentAsyncUtils.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/instrumentation/InstrumentAsyncUtils.java index db32cab7e2..0ba6e20572 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/instrumentation/InstrumentAsyncUtils.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/instrumentation/InstrumentAsyncUtils.java @@ -3,12 +3,12 @@ import com.newrelic.api.agent.DatastoreParameters; import com.newrelic.api.agent.NewRelic; import com.newrelic.api.agent.Segment; +import lombok.experimental.UtilityClass; import lombok.extern.slf4j.Slf4j; import reactor.core.scheduler.Schedulers; -import java.lang.reflect.Method; - @Slf4j +@UtilityClass public class InstrumentAsyncUtils { public static Segment startSegment(String segmentName, String product, String operationName) { @@ -30,8 +30,7 @@ public static void endSegment(Segment segment) { Schedulers.boundedElastic().schedule(() -> { try { // End the segment - Method endMethod = segment.getClass().getMethod("end"); - endMethod.invoke(segment); + segment.end(); } catch (Exception e) { log.warn("Failed to end segment", e); } diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/redis/LockService.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/redis/LockService.java index a02bba19d5..e6c542cfea 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/redis/LockService.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/redis/LockService.java @@ -3,10 +3,14 @@ import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; +import java.util.List; import java.util.UUID; public interface LockService { + Mono> lockAll(List spanIds, String spanKey); + Mono unlockAll(List locks); + record Lock(String key) { private static final String KEY_FORMAT = "%s-%s"; @@ -21,6 +25,9 @@ public Lock(String id, String name) { } + record LockRef(Lock lock, String ref) { + } + Mono executeWithLock(Lock lock, Mono action); Flux executeWithLock(Lock lock, Flux action); } diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/redis/RedissonLockService.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/redis/RedissonLockService.java index ded698ca27..322c21cf17 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/redis/RedissonLockService.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/redis/RedissonLockService.java @@ -12,6 +12,9 @@ import reactor.core.scheduler.Schedulers; import java.time.Duration; +import java.util.List; +import java.util.Map; +import java.util.UUID; @RequiredArgsConstructor @Slf4j @@ -21,14 +24,43 @@ class RedissonLockService implements LockService { private final @NonNull DistributedLockConfig distributedLockConfig; @Override - public Mono executeWithLock(Lock lock, Mono action) { + public Mono> lockAll(@NonNull List keys, @NonNull String suffix) { + if (keys.isEmpty()) { + return Mono.empty(); + } - RPermitExpirableSemaphoreReactive semaphore = redisClient.getPermitExpirableSemaphore( - CommonOptions - .name(lock.key()) - .timeout(Duration.ofMillis(distributedLockConfig.getLockTimeoutMS())) - .retryInterval(Duration.ofMillis(10)) - .retryAttempts(distributedLockConfig.getLockTimeoutMS() / 10)); + return Flux.fromIterable(keys) + .map(key -> new Lock(key, suffix)) + .flatMap(lock -> { + RPermitExpirableSemaphoreReactive semaphore = getSemaphore(lock); + log.debug("Trying to lock with {}", lock); + return semaphore.trySetPermits(1).thenReturn(Map.entry(lock, semaphore)); + }) + .flatMap(entry -> entry.getValue().acquire() + .flatMap(locked -> Mono.just(new LockRef(entry.getKey(), locked)))) + .collectList(); + } + + @Override + public Mono unlockAll(@NonNull List locks) { + if (locks.isEmpty()) { + return Mono.empty(); + } + + return Flux.fromIterable(locks) + .flatMap(lock -> { + RPermitExpirableSemaphoreReactive semaphore = getSemaphore(lock.lock()); + log.debug("Trying to unlock with {}", lock); + return semaphore.release(lock.ref()); + }) + .collectList() + .then(); + } + + @Override + public Mono executeWithLock(@NonNull Lock lock, @NonNull Mono action) { + + RPermitExpirableSemaphoreReactive semaphore = getSemaphore(lock); log.debug("Trying to lock with {}", lock); @@ -43,6 +75,15 @@ public Mono executeWithLock(Lock lock, Mono action) { })); } + private RPermitExpirableSemaphoreReactive getSemaphore(Lock lock) { + return redisClient.getPermitExpirableSemaphore( + CommonOptions + .name(lock.key()) + .timeout(Duration.ofMillis(distributedLockConfig.getLockTimeoutMS())) + .retryInterval(Duration.ofMillis(10)) + .retryAttempts(distributedLockConfig.getLockTimeoutMS() / 10)); + } + private Mono runAction(Lock lock, Mono action, String locked) { if (locked != null) { log.debug("Lock {} acquired", lock); @@ -53,13 +94,8 @@ private Mono runAction(Lock lock, Mono action, String locked) { } @Override - public Flux executeWithLock(Lock lock, Flux stream) { - RPermitExpirableSemaphoreReactive semaphore = redisClient.getPermitExpirableSemaphore( - CommonOptions - .name(lock.key()) - .timeout(Duration.ofMillis(distributedLockConfig.getLockTimeoutMS())) - .retryInterval(Duration.ofMillis(10)) - .retryAttempts(distributedLockConfig.getLockTimeoutMS() / 10)); + public Flux executeWithLock(@NonNull Lock lock, @NonNull Flux stream) { + RPermitExpirableSemaphoreReactive semaphore = getSemaphore(lock); return semaphore .trySetPermits(1) diff --git a/apps/opik-backend/src/main/java/com/comet/opik/utils/JsonUtils.java b/apps/opik-backend/src/main/java/com/comet/opik/utils/JsonUtils.java index 71703630d6..34d391f3f2 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/utils/JsonUtils.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/utils/JsonUtils.java @@ -11,6 +11,8 @@ import lombok.NonNull; import lombok.experimental.UtilityClass; +import java.io.ByteArrayOutputStream; +import java.io.IOException; import java.io.UncheckedIOException; import java.math.BigDecimal; @@ -51,4 +53,12 @@ public String writeValueAsString(@NonNull Object value) { throw new UncheckedIOException(exception); } } + + public void writeValueAsString(ByteArrayOutputStream baos, @NonNull Object value) { + try { + MAPPER.writeValue(baos, value); + } catch (IOException e) { + throw new UncheckedIOException(e); + } + } } diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/ClientSupportUtils.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/ClientSupportUtils.java index 11efac3e6e..efbc2fff7d 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/ClientSupportUtils.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/ClientSupportUtils.java @@ -10,8 +10,10 @@ private ClientSupportUtils() { } public static void config(ClientSupport client) { + client.getClient().register(new ConditionalGZipFilter()); + client.getClient().getConfiguration().property(ClientProperties.READ_TIMEOUT, 35_000); - client.getClient().getConfiguration().connectorProvider(new GrizzlyConnectorProvider()); // Required for PATCH: + client.getClient().getConfiguration().connectorProvider(new GrizzlyConnectorProvider()); // Required for PATCH // https://github.com/dropwizard/dropwizard/discussions/6431/ Required for PATCH: } } diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/ConditionalGZipFilter.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/ConditionalGZipFilter.java new file mode 100644 index 0000000000..c6b998fbdf --- /dev/null +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/ConditionalGZipFilter.java @@ -0,0 +1,42 @@ +package com.comet.opik.api.resources.utils; + +import com.comet.opik.utils.JsonUtils; +import jakarta.ws.rs.client.ClientRequestContext; +import jakarta.ws.rs.client.ClientRequestFilter; +import jakarta.ws.rs.core.HttpHeaders; +import jakarta.ws.rs.core.MediaType; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.util.zip.GZIPOutputStream; + +public class ConditionalGZipFilter implements ClientRequestFilter { + + private static final int GZIP_THRESHOLD_500KB = 500 * 1024; + + @Override + public void filter(ClientRequestContext requestContext) throws IOException { + + if (requestContext.hasEntity() && MediaType.APPLICATION_JSON_TYPE.equals(requestContext.getMediaType())) { + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + JsonUtils.writeValueAsString(baos, requestContext.getEntity()); // Serialize the entity to byte array + + byte[] entityBytes = baos.toByteArray(); + int entitySize = entityBytes.length; + + if (entitySize > GZIP_THRESHOLD_500KB) { + ByteArrayOutputStream compressedBaos = new ByteArrayOutputStream(); + try (GZIPOutputStream gzipOutputStream = new GZIPOutputStream(compressedBaos)) { + gzipOutputStream.write(entityBytes); + } + + byte[] compressedEntity = compressedBaos.toByteArray(); + requestContext.setEntity(compressedEntity, null, MediaType.APPLICATION_JSON_TYPE); + requestContext.getHeaders().add(HttpHeaders.CONTENT_ENCODING, "gzip"); + } else { + // Use the original entity if below the threshold + requestContext.setEntity(entityBytes, null, MediaType.APPLICATION_JSON_TYPE); + } + } + } +} \ No newline at end of file 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 39a92f3e5b..1291b9e45d 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 @@ -7,6 +7,7 @@ import com.comet.opik.api.Project; import com.comet.opik.api.ScoreSource; import com.comet.opik.api.Span; +import com.comet.opik.api.SpanBatch; import com.comet.opik.api.SpanUpdate; import com.comet.opik.api.error.ErrorMessage; import com.comet.opik.api.filter.Field; @@ -3189,6 +3190,40 @@ void createWhenTryingToCreateSpanTwice() { } } + @Nested + @DisplayName("Batch:") + @TestInstance(TestInstance.Lifecycle.PER_CLASS) + class BatchInsert { + + @Test + void batch__whenCreateSpans__thenReturnCreated() { + var expectedSpans = IntStream.range(0, 1000) + .mapToObj(i -> podamFactory.manufacturePojo(Span.class).toBuilder() + .projectId(null) + .parentSpanId(null) + .feedbackScores(null) + .build()) + .toList(); + + batchCreateAndAssert(expectedSpans, API_KEY, TEST_WORKSPACE); + } + + } + + private void batchCreateAndAssert(List expectedSpans, String apiKey, String workspaceName) { + + try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) + .path("batch") + .request() + .header(HttpHeaders.AUTHORIZATION, apiKey) + .header(WORKSPACE_HEADER, workspaceName) + .post(Entity.json(new SpanBatch(expectedSpans)))) { + + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204); + assertThat(actualResponse.hasEntity()).isFalse(); + } + } + private Span getAndAssert(Span expectedSpan, String apiKey, String workspaceName) { try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) .path(expectedSpan.id().toString()) diff --git a/apps/opik-backend/src/test/java/com/comet/opik/domain/DummyLockService.java b/apps/opik-backend/src/test/java/com/comet/opik/domain/DummyLockService.java new file mode 100644 index 0000000000..2ee0ffddff --- /dev/null +++ b/apps/opik-backend/src/test/java/com/comet/opik/domain/DummyLockService.java @@ -0,0 +1,31 @@ +package com.comet.opik.domain; + +import com.comet.opik.infrastructure.redis.LockService; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +import java.util.List; +import java.util.UUID; + +public class DummyLockService implements LockService { + + @Override + public Mono> lockAll(List spanIds, String spanKey) { + return Mono.empty(); + } + + @Override + public Mono unlockAll(List locks) { + return Mono.empty(); + } + + @Override + public Mono executeWithLock(LockService.Lock lock, Mono action) { + return action; + } + + @Override + public Flux executeWithLock(LockService.Lock lock, Flux action) { + return action; + } +} \ No newline at end of file diff --git a/apps/opik-backend/src/test/java/com/comet/opik/domain/SpanServiceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/domain/SpanServiceTest.java index 7570c2406c..c5e404b98a 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/domain/SpanServiceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/domain/SpanServiceTest.java @@ -10,7 +10,6 @@ import jakarta.ws.rs.NotFoundException; import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; -import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import uk.co.jemos.podam.api.PodamFactory; @@ -26,18 +25,7 @@ @Disabled class SpanServiceTest { - private static final LockService DUMMY_LOCK_SERVICE = new LockService() { - - @Override - public Mono executeWithLock(Lock lock, Mono action) { - return action; - } - - @Override - public Flux executeWithLock(Lock lock, Flux action) { - return action; - } - }; + private static final LockService DUMMY_LOCK_SERVICE = new DummyLockService(); private final PodamFactory podamFactory = PodamFactoryUtils.newPodamFactory(); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/domain/TraceServiceImplTest.java b/apps/opik-backend/src/test/java/com/comet/opik/domain/TraceServiceImplTest.java index 78d0a40666..2188190fdc 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/domain/TraceServiceImplTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/domain/TraceServiceImplTest.java @@ -19,7 +19,6 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; -import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import uk.co.jemos.podam.api.PodamFactory; import uk.co.jemos.podam.api.PodamFactoryImpl; @@ -40,18 +39,7 @@ @ExtendWith(MockitoExtension.class) class TraceServiceImplTest { - public static final LockService DUMMY_LOCK_SERVICE = new LockService() { - - @Override - public Mono executeWithLock(Lock lock, Mono action) { - return action; - } - - @Override - public Flux executeWithLock(Lock lock, Flux action) { - return action; - } - }; + public static final LockService DUMMY_LOCK_SERVICE = new DummyLockService(); private TraceServiceImpl traceService; diff --git a/apps/opik-backend/src/test/resources/config-test.yml b/apps/opik-backend/src/test/resources/config-test.yml index a77cec3b07..deacc262f4 100644 --- a/apps/opik-backend/src/test/resources/config-test.yml +++ b/apps/opik-backend/src/test/resources/config-test.yml @@ -63,3 +63,6 @@ authentication: server: enableVirtualThreads: ${ENABLE_VIRTUAL_THREADS:-false} + gzip: + enabled: true + From 3b68fd13c487fe32ebcf141c8a0b472f616754c3 Mon Sep 17 00:00:00 2001 From: Thiago Hora Date: Tue, 10 Sep 2024 17:50:15 +0200 Subject: [PATCH 03/10] Add tests --- apps/opik-backend/config.yml | 1 + .../com/comet/opik/domain/SpanService.java | 20 +++- .../infrastructure/DistributedLockConfig.java | 4 + .../redis/RedissonLockService.java | 14 +-- .../resources/v1/priv/SpansResourceTest.java | 95 +++++++++++++++++++ .../src/test/resources/config-test.yml | 1 + 6 files changed, 124 insertions(+), 11 deletions(-) diff --git a/apps/opik-backend/config.yml b/apps/opik-backend/config.yml index 42fee78a19..b22904c242 100644 --- a/apps/opik-backend/config.yml +++ b/apps/opik-backend/config.yml @@ -49,6 +49,7 @@ health: distributedLock: lockTimeoutMS: ${DISTRIBUTED_LOCK_TIME_OUT:-500} + bulkLockTimeoutMS: ${DISTRIBUTED_BULK_LOCK_TIME_OUT:-5000} redis: singleNodeUrl: ${REDIS_URL:-} 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 516207ad04..10ad5124ec 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 @@ -40,10 +40,10 @@ @Slf4j public class SpanService { - public static final String PROJECT_NAME_AND_WORKSPACE_MISMATCH = "Project name and workspace name do not match the existing span"; public static final String PARENT_SPAN_IS_MISMATCH = "parent_span_id does not match the existing span"; public static final String TRACE_ID_MISMATCH = "trace_id does not match the existing span"; public static final String SPAN_KEY = "Span"; + public static final String PROJECT_NAME_MISMATCH = "Project name and workspace name do not match the existing span"; private final @NonNull SpanDAO spanDAO; private final @NonNull ProjectService projectService; @@ -122,7 +122,7 @@ private Mono insertSpan(Span span, Project project, UUID id, Span existing } if (!project.id().equals(existingSpan.projectId())) { - return failWithConflict(PROJECT_NAME_AND_WORKSPACE_MISMATCH); + return failWithConflict(PROJECT_NAME_MISMATCH); } if (!Objects.equals(span.parentSpanId(), existingSpan.parentSpanId())) { @@ -197,7 +197,7 @@ private Mono handleSpanDBError(Throwable ex) { && (ex.getMessage().contains("_CAST(project_id, FixedString(36))") || ex.getMessage() .contains(", CAST(leftPad(workspace_id, 40, '*'), 'FixedString(19)') ::"))) { - return failWithConflict(PROJECT_NAME_AND_WORKSPACE_MISMATCH); + return failWithConflict(PROJECT_NAME_MISMATCH); } if (ex instanceof ClickHouseException @@ -220,7 +220,7 @@ private Mono handleSpanDBError(Throwable ex) { private Mono updateOrFail(SpanUpdate spanUpdate, UUID id, Span existingSpan, Project project) { if (!project.id().equals(existingSpan.projectId())) { - return failWithConflict(PROJECT_NAME_AND_WORKSPACE_MISMATCH); + return failWithConflict(PROJECT_NAME_MISMATCH); } if (!Objects.equals(existingSpan.parentSpanId(), spanUpdate.parentSpanId())) { @@ -314,6 +314,18 @@ private Span mergeSpans(Span currentSpan, Span receivedSpan) { .build(); } + if (!Objects.equals(currentSpan.projectId(), receivedSpan.projectId())) { + throw new EntityAlreadyExistsException(new ErrorMessage(List.of(PROJECT_NAME_MISMATCH))); + } + + if (!Objects.equals(currentSpan.traceId(), receivedSpan.traceId())) { + throw new EntityAlreadyExistsException(new ErrorMessage(List.of(TRACE_ID_MISMATCH))); + } + + if (!Objects.equals(currentSpan.parentSpanId(), receivedSpan.parentSpanId())) { + throw new EntityAlreadyExistsException(new ErrorMessage(List.of(PARENT_SPAN_IS_MISMATCH))); + } + return Span.builder() .id(currentSpan.id()) .projectId(currentSpan.projectId() != null ? currentSpan.projectId() : receivedSpan.projectId()) diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/DistributedLockConfig.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/DistributedLockConfig.java index d57315c266..e34d9d7277 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/DistributedLockConfig.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/DistributedLockConfig.java @@ -12,4 +12,8 @@ public class DistributedLockConfig { @JsonProperty @NotNull private int lockTimeoutMS; + @Valid + @JsonProperty + @NotNull private int bulkLockTimeoutMS; + } diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/redis/RedissonLockService.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/redis/RedissonLockService.java index 322c21cf17..0f5c4b941e 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/redis/RedissonLockService.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/redis/RedissonLockService.java @@ -32,7 +32,7 @@ public Mono> lockAll(@NonNull List keys, @NonNull String suf return Flux.fromIterable(keys) .map(key -> new Lock(key, suffix)) .flatMap(lock -> { - RPermitExpirableSemaphoreReactive semaphore = getSemaphore(lock); + RPermitExpirableSemaphoreReactive semaphore = getSemaphore(lock, distributedLockConfig.getBulkLockTimeoutMS()); log.debug("Trying to lock with {}", lock); return semaphore.trySetPermits(1).thenReturn(Map.entry(lock, semaphore)); }) @@ -49,7 +49,7 @@ public Mono unlockAll(@NonNull List locks) { return Flux.fromIterable(locks) .flatMap(lock -> { - RPermitExpirableSemaphoreReactive semaphore = getSemaphore(lock.lock()); + RPermitExpirableSemaphoreReactive semaphore = getSemaphore(lock.lock(), distributedLockConfig.getBulkLockTimeoutMS()); log.debug("Trying to unlock with {}", lock); return semaphore.release(lock.ref()); }) @@ -60,7 +60,7 @@ public Mono unlockAll(@NonNull List locks) { @Override public Mono executeWithLock(@NonNull Lock lock, @NonNull Mono action) { - RPermitExpirableSemaphoreReactive semaphore = getSemaphore(lock); + RPermitExpirableSemaphoreReactive semaphore = getSemaphore(lock, distributedLockConfig.getLockTimeoutMS()); log.debug("Trying to lock with {}", lock); @@ -75,13 +75,13 @@ public Mono executeWithLock(@NonNull Lock lock, @NonNull Mono action) })); } - private RPermitExpirableSemaphoreReactive getSemaphore(Lock lock) { + private RPermitExpirableSemaphoreReactive getSemaphore(Lock lock, int lockTimeoutMS) { return redisClient.getPermitExpirableSemaphore( CommonOptions .name(lock.key()) - .timeout(Duration.ofMillis(distributedLockConfig.getLockTimeoutMS())) + .timeout(Duration.ofMillis(lockTimeoutMS)) .retryInterval(Duration.ofMillis(10)) - .retryAttempts(distributedLockConfig.getLockTimeoutMS() / 10)); + .retryAttempts(lockTimeoutMS / 10)); } private Mono runAction(Lock lock, Mono action, String locked) { @@ -95,7 +95,7 @@ private Mono runAction(Lock lock, Mono action, String locked) { @Override public Flux executeWithLock(@NonNull Lock lock, @NonNull Flux stream) { - RPermitExpirableSemaphoreReactive semaphore = getSemaphore(lock); + RPermitExpirableSemaphoreReactive semaphore = getSemaphore(lock, distributedLockConfig.getLockTimeoutMS()); return semaphore .trySetPermits(1) 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 1291b9e45d..ec919aab07 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 @@ -24,6 +24,7 @@ import com.comet.opik.api.resources.utils.TestDropwizardAppExtensionUtils; import com.comet.opik.api.resources.utils.WireMockUtils; import com.comet.opik.domain.SpanMapper; +import com.comet.opik.domain.SpanService; import com.comet.opik.domain.SpanType; import com.comet.opik.infrastructure.auth.RequestContext; import com.comet.opik.podam.PodamFactoryUtils; @@ -3208,6 +3209,100 @@ void batch__whenCreateSpans__thenReturnCreated() { batchCreateAndAssert(expectedSpans, API_KEY, TEST_WORKSPACE); } + @Test + void batch__whenCreateSpansAndThenUpdate__thenReturnUpdated() { + var expectedSpans = List.of(podamFactory.manufacturePojo(Span.class).toBuilder() + .projectId(null) + .parentSpanId(null) + .feedbackScores(null) + .build()); + + batchCreateAndAssert(expectedSpans, API_KEY, TEST_WORKSPACE); + + var expectedSpan = expectedSpans.get(0).toBuilder() + .tags(Set.of()) + .endTime(Instant.now()) + .output(JsonUtils.getJsonNodeFromString("{ \"output\": \"data\"}")) + .build(); + + batchCreateAndAssert(List.of(expectedSpan), API_KEY, TEST_WORKSPACE); + + getAndAssert(expectedSpan.toBuilder().tags(null).build(), API_KEY, TEST_WORKSPACE); + } + + @Test + void batch__whenCreateSpansAndThenUpdateInSameRequest__thenReturnUpdated() { + var expectedSpans = List.of(podamFactory.manufacturePojo(Span.class).toBuilder() + .projectId(null) + .parentSpanId(null) + .feedbackScores(null) + .build()); + + var expectedSpan = expectedSpans.getFirst().toBuilder() + .tags(Set.of()) + .endTime(Instant.now()) + .output(JsonUtils.getJsonNodeFromString("{ \"output\": \"data\"}")) + .build(); + + batchCreateAndAssert(List.of(expectedSpans.getFirst(), expectedSpan), API_KEY, TEST_WORKSPACE); + + getAndAssert(expectedSpan.toBuilder().tags(null).build(), API_KEY, TEST_WORKSPACE); + } + + @ParameterizedTest + @MethodSource + void batch__whenCreateSpansWithConflicts__thenReturn409(List expectedSpans, String expectedError) { + + try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) + .path("batch") + .request() + .header(HttpHeaders.AUTHORIZATION, API_KEY) + .header(WORKSPACE_HEADER, TEST_WORKSPACE) + .post(Entity.json(new SpanBatch(expectedSpans)))) { + + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(409); + assertThat(actualResponse.hasEntity()).isTrue(); + assertThat(actualResponse.readEntity(ErrorMessage.class).errors()).contains(expectedError); + } + } + + public Stream batch__whenCreateSpansWithConflicts__thenReturn409() { + Span expectedSpan = podamFactory.manufacturePojo(Span.class).toBuilder() + .projectId(null) + .parentSpanId(null) + .feedbackScores(null) + .build(); + + UUID id = podamFactory.manufacturePojo(UUID.class); + UUID id2 = podamFactory.manufacturePojo(UUID.class); + UUID id3 = podamFactory.manufacturePojo(UUID.class); + + String projectName = UUID.randomUUID().toString(); + String projectName2 = UUID.randomUUID().toString(); + String projectName3 = UUID.randomUUID().toString(); + + return Stream.of( + Arguments.of( + List.of( + expectedSpan.toBuilder().id(id).projectName(projectName).build(), + expectedSpan.toBuilder().id(id).projectName(UUID.randomUUID().toString()).build()), + SpanService.PROJECT_NAME_MISMATCH + ), + Arguments.of( + List.of( + expectedSpan.toBuilder().id(id2).projectName(projectName2).build(), + expectedSpan.toBuilder().id(id2).projectName(projectName2).traceId(podamFactory.manufacturePojo(UUID.class)).build()), + SpanService.TRACE_ID_MISMATCH + ), + Arguments.of( + List.of( + expectedSpan.toBuilder().id(id3).projectName(projectName3).build(), + expectedSpan.toBuilder().id(id3).projectName(projectName3).parentSpanId(podamFactory.manufacturePojo(UUID.class)).build()), + SpanService.PARENT_SPAN_IS_MISMATCH + ) + ); + } + } private void batchCreateAndAssert(List expectedSpans, String apiKey, String workspaceName) { diff --git a/apps/opik-backend/src/test/resources/config-test.yml b/apps/opik-backend/src/test/resources/config-test.yml index deacc262f4..44c1996681 100644 --- a/apps/opik-backend/src/test/resources/config-test.yml +++ b/apps/opik-backend/src/test/resources/config-test.yml @@ -49,6 +49,7 @@ health: distributedLock: lockTimeout: 500 + bulkLockTimeout: 5000 redis: singleNodeUrl: From e03b560e6b7c904a22e775f44a1c429a58dbfb8c Mon Sep 17 00:00:00 2001 From: Thiago Hora Date: Thu, 12 Sep 2024 12:10:28 +0200 Subject: [PATCH 04/10] Add code review suggestions --- .../main/java/com/comet/opik/api/Span.java | 15 +++++ .../api/{SpanBatch.java => SpanBulk.java} | 2 +- .../api/resources/v1/priv/SpansResource.java | 6 +- .../java/com/comet/opik/domain/SpanDAO.java | 65 ++++++------------- .../com/comet/opik/domain/SpanService.java | 24 ++----- .../java/com/comet/opik/utils/JsonUtils.java | 2 +- .../resources/v1/priv/SpansResourceTest.java | 11 ++-- 7 files changed, 51 insertions(+), 74 deletions(-) rename apps/opik-backend/src/main/java/com/comet/opik/api/{SpanBatch.java => SpanBulk.java} (79%) diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/Span.java b/apps/opik-backend/src/main/java/com/comet/opik/api/Span.java index 6071a17f7c..4bb996413c 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/api/Span.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/Span.java @@ -15,6 +15,7 @@ import java.time.Instant; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.Set; import java.util.UUID; @@ -65,4 +66,18 @@ public static class Write { public static class Public { } } + + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + Span span = (Span) o; + return Objects.equals(id, span.id); + } + + @Override + public int hashCode() { + return Objects.hash(id); + } } diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/SpanBatch.java b/apps/opik-backend/src/main/java/com/comet/opik/api/SpanBulk.java similarity index 79% rename from apps/opik-backend/src/main/java/com/comet/opik/api/SpanBatch.java rename to apps/opik-backend/src/main/java/com/comet/opik/api/SpanBulk.java index f422eb7d74..6cc2840470 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/api/SpanBatch.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/SpanBulk.java @@ -7,6 +7,6 @@ import java.util.List; -public record SpanBatch(@NotNull @Size(min = 1, max = 2000) @JsonView( { +public record SpanBulk(@NotNull @Size(min = 1, max = 1000) @JsonView( { Span.View.Write.class}) @Valid List spans){ } diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/SpansResource.java b/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/SpansResource.java index a1fa71bb6f..aec5ba585d 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/SpansResource.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/SpansResource.java @@ -5,7 +5,7 @@ import com.comet.opik.api.FeedbackScore; import com.comet.opik.api.FeedbackScoreBatch; import com.comet.opik.api.Span; -import com.comet.opik.api.SpanBatch; +import com.comet.opik.api.SpanBulk; import com.comet.opik.api.SpanSearchCriteria; import com.comet.opik.api.SpanUpdate; import com.comet.opik.api.filter.FiltersFactory; @@ -134,11 +134,11 @@ public Response create( } @POST - @Path("/batch") + @Path("/bulk") @Operation(operationId = "createSpans", summary = "Create spans", description = "Create spans", responses = { @ApiResponse(responseCode = "204", description = "No Content")}) public Response createSpans( - @RequestBody(content = @Content(schema = @Schema(implementation = SpanBatch.class))) @JsonView(Span.View.Write.class) @NotNull @Valid SpanBatch spans) { + @RequestBody(content = @Content(schema = @Schema(implementation = SpanBulk.class))) @JsonView(Span.View.Write.class) @NotNull @Valid SpanBulk spans) { spanService.create(spans) .contextWrite(ctx -> setRequestContext(ctx, requestContext)) diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/SpanDAO.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/SpanDAO.java index 4f3b075ed5..dfe1b2d707 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/domain/SpanDAO.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/SpanDAO.java @@ -7,6 +7,7 @@ import com.comet.opik.domain.filter.FilterStrategy; import com.comet.opik.utils.JsonUtils; import com.comet.opik.utils.TemplateUtils; +import com.google.common.base.Preconditions; import com.newrelic.api.agent.Segment; import com.newrelic.api.agent.Trace; import io.r2dbc.spi.Connection; @@ -25,7 +26,6 @@ import java.time.Instant; import java.util.ArrayList; import java.util.Arrays; -import java.util.Collection; import java.util.List; import java.util.Map; import java.util.Optional; @@ -49,7 +49,7 @@ @Slf4j class SpanDAO { - private static final String PLAIN_INSERT = """ + private static final String BULK_INSERT = """ INSERT INTO spans( id, project_id, @@ -494,27 +494,25 @@ public Mono insert(@NonNull Span span) { } @Trace(dispatcher = true) - public Mono batchInsert(@NonNull List spans) { - return Mono.from(connectionFactory.create()) - .flatMapMany(connection -> { - if (spans.isEmpty()) { - return Mono.just(0L); - } + public Mono bulkInsert(@NonNull List spans) { - return insert(spans, connection); - }) - .then(); + Preconditions.checkArgument(!spans.isEmpty(), "Spans list must not be empty"); + + return Mono.from(connectionFactory.create()) + .flatMapMany(connection -> insert(spans, connection)) + .flatMap(Result::getRowsUpdated) + .reduce(0L, Long::sum); } - private Publisher insert(Collection spans, Connection connection) { + private Publisher insert(List spans, Connection connection) { return makeMonoContextAware((userName, workspaceName, workspaceId) -> { List queryItems = getQueryItemPlaceHolder(spans.size()); - var template = new ST(PLAIN_INSERT) + var template = new ST(BULK_INSERT) .add("items", queryItems); - var statement = connection.createStatement(template.render()); + Statement statement = connection.createStatement(template.render()); int i = 0; for (Span span : spans) { @@ -524,40 +522,22 @@ private Publisher insert(Collection spans, Connection co .bind("trace_id" + i, span.traceId()) .bind("name" + i, span.name()) .bind("type" + i, span.type().toString()) - .bind("start_time" + i, span.startTime().toString()); + .bind("start_time" + i, span.startTime().toString()) + .bind("parent_span_id" + i, span.parentSpanId() != null ? span.parentSpanId() : "") + .bind("input" + i, span.input() != null ? span.input().toString() : "") + .bind("output" + i, span.output() != null ? span.output().toString() : "") + .bind("metadata" + i, span.metadata() != null ? span.metadata().toString() : "") + .bind("tags" + i, span.tags() != null ? span.tags().toArray(String[]::new) : new String[]{}) + .bind("created_at" + i, span.createdAt().toString()) + .bind("created_by" + i, userName) + .bind("last_updated_by" + i, userName); - if (span.parentSpanId() != null) { - statement.bind("parent_span_id" + i, span.parentSpanId()); - } else { - statement.bind("parent_span_id" + i, ""); - } if (span.endTime() != null) { statement.bind("end_time" + i, span.endTime().toString()); } else { statement.bindNull("end_time" + i, String.class); } - if (span.input() != null) { - statement.bind("input" + i, span.input().toString()); - } else { - statement.bind("input" + i, ""); - } - if (span.output() != null) { - statement.bind("output" + i, span.output().toString()); - } else { - statement.bind("output" + i, ""); - } - if (span.metadata() != null) { - statement.bind("metadata" + i, span.metadata().toString()); - } else { - statement.bind("metadata" + i, ""); - } - if (span.tags() != null) { - statement.bind("tags" + i, span.tags().toArray(String[]::new)); - } else { - statement.bind("tags" + i, new String[]{}); - } - if (span.usage() != null) { Stream.Builder keys = Stream.builder(); Stream.Builder values = Stream.builder(); @@ -574,9 +554,6 @@ private Publisher insert(Collection spans, Connection co statement.bind("usage_values" + i, new Integer[]{}); } - statement.bind("created_at" + i, span.createdAt().toString()); - statement.bind("created_by" + i, userName); - statement.bind("last_updated_by" + i, userName); i++; } 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 10ad5124ec..3e50945944 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 @@ -3,7 +3,7 @@ import com.clickhouse.client.ClickHouseException; import com.comet.opik.api.Project; import com.comet.opik.api.Span; -import com.comet.opik.api.SpanBatch; +import com.comet.opik.api.SpanBulk; import com.comet.opik.api.SpanSearchCriteria; import com.comet.opik.api.SpanUpdate; import com.comet.opik.api.error.EntityAlreadyExistsException; @@ -252,7 +252,7 @@ public Mono validateSpanWorkspace(@NonNull String workspaceId, @NonNull } @Trace(dispatcher = true) - public Mono create(SpanBatch batch) { + public Mono create(SpanBulk batch) { if (batch.spans().isEmpty()) { return Mono.empty(); @@ -276,7 +276,7 @@ public Mono create(SpanBatch batch) { return lockService.lockAll(spanIds, SPAN_KEY) .flatMap(locks -> mergeSpans(spans, spanIds) - .flatMap(spanDAO::batchInsert) + .flatMap(spanDAO::bulkInsert) .doFinally(signalType -> lockService.unlockAll(locks).subscribe()) .subscribeOn(Schedulers.boundedElastic())) .then(); @@ -326,21 +326,7 @@ private Span mergeSpans(Span currentSpan, Span receivedSpan) { throw new EntityAlreadyExistsException(new ErrorMessage(List.of(PARENT_SPAN_IS_MISMATCH))); } - return Span.builder() - .id(currentSpan.id()) - .projectId(currentSpan.projectId() != null ? currentSpan.projectId() : receivedSpan.projectId()) - .traceId(currentSpan.traceId() != null ? currentSpan.traceId() : receivedSpan.traceId()) - .parentSpanId( - currentSpan.parentSpanId() != null ? currentSpan.parentSpanId() : receivedSpan.parentSpanId()) - .name(currentSpan.name() != null ? currentSpan.name() : receivedSpan.name()) - .type(currentSpan.type() != null ? currentSpan.type() : receivedSpan.type()) - .startTime(currentSpan.startTime() != null ? currentSpan.startTime() : receivedSpan.startTime()) - .endTime(receivedSpan.endTime() != null ? receivedSpan.endTime() : currentSpan.endTime()) - .input(receivedSpan.input() != null ? receivedSpan.input() : currentSpan.input()) - .output(receivedSpan.output() != null ? receivedSpan.output() : currentSpan.output()) - .metadata(receivedSpan.metadata() != null ? receivedSpan.metadata() : currentSpan.metadata()) - .tags(receivedSpan.tags() != null ? receivedSpan.tags() : currentSpan.tags()) - .usage(receivedSpan.usage() != null ? receivedSpan.usage() : currentSpan.usage()) + return receivedSpan.toBuilder() .createdAt(getInstant(currentSpan.createdAt(), receivedSpan.createdAt())) .build(); } @@ -355,7 +341,7 @@ private Instant getInstant(Instant current, Instant received) { } } - private List bindSpanToProjectAndId(SpanBatch batch, List projects) { + private List bindSpanToProjectAndId(SpanBulk batch, List projects) { Map projectPerName = projects.stream() .collect(Collectors.toMap(Project::name, Function.identity())); diff --git a/apps/opik-backend/src/main/java/com/comet/opik/utils/JsonUtils.java b/apps/opik-backend/src/main/java/com/comet/opik/utils/JsonUtils.java index 34d391f3f2..cf36565de7 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/utils/JsonUtils.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/utils/JsonUtils.java @@ -54,7 +54,7 @@ public String writeValueAsString(@NonNull Object value) { } } - public void writeValueAsString(ByteArrayOutputStream baos, @NonNull Object value) { + public void writeValueAsString(@NonNull ByteArrayOutputStream baos, @NonNull Object value) { try { MAPPER.writeValue(baos, value); } catch (IOException e) { 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 ec919aab07..aaf4a57daf 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 @@ -7,7 +7,7 @@ import com.comet.opik.api.Project; import com.comet.opik.api.ScoreSource; import com.comet.opik.api.Span; -import com.comet.opik.api.SpanBatch; +import com.comet.opik.api.SpanBulk; import com.comet.opik.api.SpanUpdate; import com.comet.opik.api.error.ErrorMessage; import com.comet.opik.api.filter.Field; @@ -3254,11 +3254,11 @@ void batch__whenCreateSpansAndThenUpdateInSameRequest__thenReturnUpdated() { void batch__whenCreateSpansWithConflicts__thenReturn409(List expectedSpans, String expectedError) { try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) - .path("batch") + .path("bulk") .request() .header(HttpHeaders.AUTHORIZATION, API_KEY) .header(WORKSPACE_HEADER, TEST_WORKSPACE) - .post(Entity.json(new SpanBatch(expectedSpans)))) { + .post(Entity.json(new SpanBulk(expectedSpans)))) { assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(409); assertThat(actualResponse.hasEntity()).isTrue(); @@ -3302,17 +3302,16 @@ public Stream batch__whenCreateSpansWithConflicts__thenReturn409() { ) ); } - } private void batchCreateAndAssert(List expectedSpans, String apiKey, String workspaceName) { try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) - .path("batch") + .path("bulk") .request() .header(HttpHeaders.AUTHORIZATION, apiKey) .header(WORKSPACE_HEADER, workspaceName) - .post(Entity.json(new SpanBatch(expectedSpans)))) { + .post(Entity.json(new SpanBulk(expectedSpans)))) { assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204); assertThat(actualResponse.hasEntity()).isFalse(); From cdc5259cbb460d1db570e8540a7cc62ae65d1cc0 Mon Sep 17 00:00:00 2001 From: Thiago Hora Date: Thu, 12 Sep 2024 16:16:52 +0200 Subject: [PATCH 05/10] PR review feedback --- .../main/java/com/comet/opik/api/Span.java | 1 - .../api/{SpanBulk.java => SpanBatch.java} | 2 +- .../api/resources/v1/priv/SpansResource.java | 18 ++- .../java/com/comet/opik/domain/SpanDAO.java | 7 +- .../com/comet/opik/domain/SpanService.java | 76 +----------- .../redis/RedissonLockService.java | 6 +- .../resources/v1/priv/SpansResourceTest.java | 112 ++++++++---------- 7 files changed, 77 insertions(+), 145 deletions(-) rename apps/opik-backend/src/main/java/com/comet/opik/api/{SpanBulk.java => SpanBatch.java} (79%) diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/Span.java b/apps/opik-backend/src/main/java/com/comet/opik/api/Span.java index 4bb996413c..d351406ac2 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/api/Span.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/Span.java @@ -67,7 +67,6 @@ public static class Public { } } - @Override public boolean equals(Object o) { if (this == o) return true; diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/SpanBulk.java b/apps/opik-backend/src/main/java/com/comet/opik/api/SpanBatch.java similarity index 79% rename from apps/opik-backend/src/main/java/com/comet/opik/api/SpanBulk.java rename to apps/opik-backend/src/main/java/com/comet/opik/api/SpanBatch.java index 6cc2840470..57435ebf26 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/api/SpanBulk.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/SpanBatch.java @@ -7,6 +7,6 @@ import java.util.List; -public record SpanBulk(@NotNull @Size(min = 1, max = 1000) @JsonView( { +public record SpanBatch(@NotNull @Size(min = 1, max = 1000) @JsonView( { Span.View.Write.class}) @Valid List spans){ } diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/SpansResource.java b/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/SpansResource.java index aec5ba585d..e337363f76 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/SpansResource.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/SpansResource.java @@ -5,7 +5,7 @@ import com.comet.opik.api.FeedbackScore; import com.comet.opik.api.FeedbackScoreBatch; import com.comet.opik.api.Span; -import com.comet.opik.api.SpanBulk; +import com.comet.opik.api.SpanBatch; import com.comet.opik.api.SpanSearchCriteria; import com.comet.opik.api.SpanUpdate; import com.comet.opik.api.filter.FiltersFactory; @@ -28,6 +28,7 @@ import jakarta.validation.Valid; import jakarta.validation.constraints.Min; import jakarta.validation.constraints.NotNull; +import jakarta.ws.rs.ClientErrorException; import jakarta.ws.rs.Consumes; import jakarta.ws.rs.DELETE; import jakarta.ws.rs.DefaultValue; @@ -48,6 +49,7 @@ import lombok.extern.slf4j.Slf4j; import java.util.UUID; +import java.util.stream.Collectors; import static com.comet.opik.utils.AsyncUtils.setRequestContext; import static com.comet.opik.utils.ValidationUtils.validateProjectNameAndProjectId; @@ -134,11 +136,21 @@ public Response create( } @POST - @Path("/bulk") + @Path("/batch") @Operation(operationId = "createSpans", summary = "Create spans", description = "Create spans", responses = { @ApiResponse(responseCode = "204", description = "No Content")}) public Response createSpans( - @RequestBody(content = @Content(schema = @Schema(implementation = SpanBulk.class))) @JsonView(Span.View.Write.class) @NotNull @Valid SpanBulk spans) { + @RequestBody(content = @Content(schema = @Schema(implementation = SpanBatch.class))) @JsonView(Span.View.Write.class) @NotNull @Valid SpanBatch spans) { + + spans.spans() + .stream() + .filter(span -> span.id() != null) // Filter out spans with null IDs + .collect(Collectors.groupingBy(Span::id)) + .forEach((id, spanGroup) -> { + if (spanGroup.size() > 1) { + throw new ClientErrorException("Duplicate span id '%s'".formatted(id), 422); + } + }); spanService.create(spans) .contextWrite(ctx -> setRequestContext(ctx, requestContext)) diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/SpanDAO.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/SpanDAO.java index dfe1b2d707..0141e175e5 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/domain/SpanDAO.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/SpanDAO.java @@ -65,7 +65,6 @@ INSERT INTO spans( metadata, tags, usage, - created_at, created_by, last_updated_by ) VALUES @@ -85,7 +84,6 @@ INSERT INTO spans( :metadata, :tags, mapFromArrays(:usage_keys, :usage_values), - parseDateTime64BestEffort(:created_at, 9), :created_by, :last_updated_by ) @@ -494,7 +492,7 @@ public Mono insert(@NonNull Span span) { } @Trace(dispatcher = true) - public Mono bulkInsert(@NonNull List spans) { + public Mono batchInsert(@NonNull List spans) { Preconditions.checkArgument(!spans.isEmpty(), "Spans list must not be empty"); @@ -528,7 +526,6 @@ private Publisher insert(List spans, Connection connecti .bind("output" + i, span.output() != null ? span.output().toString() : "") .bind("metadata" + i, span.metadata() != null ? span.metadata().toString() : "") .bind("tags" + i, span.tags() != null ? span.tags().toArray(String[]::new) : new String[]{}) - .bind("created_at" + i, span.createdAt().toString()) .bind("created_by" + i, userName) .bind("last_updated_by" + i, userName); @@ -559,7 +556,7 @@ private Publisher insert(List spans, Connection connecti statement.bind("workspace_id", workspaceId); - Segment segment = startSegment("spans", "Clickhouse", "insert_plain"); + Segment segment = startSegment("spans", "Clickhouse", "batch_insert"); return Mono.from(statement.execute()) .doFinally(signalType -> endSegment(segment)); 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 3e50945944..b7f8406f50 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 @@ -3,7 +3,7 @@ import com.clickhouse.client.ClickHouseException; import com.comet.opik.api.Project; import com.comet.opik.api.Span; -import com.comet.opik.api.SpanBulk; +import com.comet.opik.api.SpanBatch; import com.comet.opik.api.SpanSearchCriteria; import com.comet.opik.api.SpanUpdate; import com.comet.opik.api.error.EntityAlreadyExistsException; @@ -252,7 +252,7 @@ public Mono validateSpanWorkspace(@NonNull String workspaceId, @NonNull } @Trace(dispatcher = true) - public Mono create(SpanBulk batch) { + public Mono create(@NonNull SpanBatch batch) { if (batch.spans().isEmpty()) { return Mono.empty(); @@ -271,77 +271,11 @@ public Mono create(SpanBulk batch) { .subscribeOn(Schedulers.boundedElastic()); return resolveProjects - .flatMap(spans -> { - List spanIds = spans.stream().map(Span::id).distinct().toList(); - - return lockService.lockAll(spanIds, SPAN_KEY) - .flatMap(locks -> mergeSpans(spans, spanIds) - .flatMap(spanDAO::bulkInsert) - .doFinally(signalType -> lockService.unlockAll(locks).subscribe()) - .subscribeOn(Schedulers.boundedElastic())) - .then(); - }); + .flatMap(spanDAO::batchInsert) + .then(); } - private Mono> mergeSpans(List spans, List spanIds) { - return spanDAO.getByIds(spanIds) - .map(existingSpans -> { - Map existingSpanMap = existingSpans.stream() - .collect(Collectors.toMap(Span::id, Function.identity())); - - return spans.stream() - .collect(Collectors.groupingBy(Span::id)) - .entrySet() - .stream() - .map(groupedSpans -> { - // START state resolution from existing spans - Span currentSpan = existingSpanMap.getOrDefault(groupedSpans.getKey(), - Span.builder().build()); - - return groupedSpans - .getValue() - .stream() - .reduce(currentSpan, this::mergeSpans); - }).toList(); - }); - } - - private Span mergeSpans(Span currentSpan, Span receivedSpan) { - - if (currentSpan.id() == null) { - return receivedSpan.toBuilder() - .createdAt(getInstant(currentSpan.createdAt(), receivedSpan.createdAt())) - .build(); - } - - if (!Objects.equals(currentSpan.projectId(), receivedSpan.projectId())) { - throw new EntityAlreadyExistsException(new ErrorMessage(List.of(PROJECT_NAME_MISMATCH))); - } - - if (!Objects.equals(currentSpan.traceId(), receivedSpan.traceId())) { - throw new EntityAlreadyExistsException(new ErrorMessage(List.of(TRACE_ID_MISMATCH))); - } - - if (!Objects.equals(currentSpan.parentSpanId(), receivedSpan.parentSpanId())) { - throw new EntityAlreadyExistsException(new ErrorMessage(List.of(PARENT_SPAN_IS_MISMATCH))); - } - - return receivedSpan.toBuilder() - .createdAt(getInstant(currentSpan.createdAt(), receivedSpan.createdAt())) - .build(); - } - - private Instant getInstant(Instant current, Instant received) { - if (current == null) { - return received != null ? received : Instant.now(); - } else if (received == null) { - return current; - } else { - return current.isBefore(received) ? current : received; - } - } - - private List bindSpanToProjectAndId(SpanBulk batch, List projects) { + private List bindSpanToProjectAndId(SpanBatch batch, List projects) { Map projectPerName = projects.stream() .collect(Collectors.toMap(Project::name, Function.identity())); diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/redis/RedissonLockService.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/redis/RedissonLockService.java index 0f5c4b941e..97be6da029 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/redis/RedissonLockService.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/redis/RedissonLockService.java @@ -32,7 +32,8 @@ public Mono> lockAll(@NonNull List keys, @NonNull String suf return Flux.fromIterable(keys) .map(key -> new Lock(key, suffix)) .flatMap(lock -> { - RPermitExpirableSemaphoreReactive semaphore = getSemaphore(lock, distributedLockConfig.getBulkLockTimeoutMS()); + RPermitExpirableSemaphoreReactive semaphore = getSemaphore(lock, + distributedLockConfig.getBulkLockTimeoutMS()); log.debug("Trying to lock with {}", lock); return semaphore.trySetPermits(1).thenReturn(Map.entry(lock, semaphore)); }) @@ -49,7 +50,8 @@ public Mono unlockAll(@NonNull List locks) { return Flux.fromIterable(locks) .flatMap(lock -> { - RPermitExpirableSemaphoreReactive semaphore = getSemaphore(lock.lock(), distributedLockConfig.getBulkLockTimeoutMS()); + RPermitExpirableSemaphoreReactive semaphore = getSemaphore(lock.lock(), + distributedLockConfig.getBulkLockTimeoutMS()); log.debug("Trying to unlock with {}", lock); return semaphore.release(lock.ref()); }) 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 aaf4a57daf..9aade09878 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 @@ -7,7 +7,7 @@ import com.comet.opik.api.Project; import com.comet.opik.api.ScoreSource; import com.comet.opik.api.Span; -import com.comet.opik.api.SpanBulk; +import com.comet.opik.api.SpanBatch; import com.comet.opik.api.SpanUpdate; import com.comet.opik.api.error.ErrorMessage; import com.comet.opik.api.filter.Field; @@ -24,7 +24,6 @@ import com.comet.opik.api.resources.utils.TestDropwizardAppExtensionUtils; import com.comet.opik.api.resources.utils.WireMockUtils; import com.comet.opik.domain.SpanMapper; -import com.comet.opik.domain.SpanService; import com.comet.opik.domain.SpanType; import com.comet.opik.infrastructure.auth.RequestContext; import com.comet.opik.podam.PodamFactoryUtils; @@ -3197,7 +3196,7 @@ void createWhenTryingToCreateSpanTwice() { class BatchInsert { @Test - void batch__whenCreateSpans__thenReturnCreated() { + void batch__whenCreateSpans__thenReturnNoContent() { var expectedSpans = IntStream.range(0, 1000) .mapToObj(i -> podamFactory.manufacturePojo(Span.class).toBuilder() .projectId(null) @@ -3210,28 +3209,7 @@ void batch__whenCreateSpans__thenReturnCreated() { } @Test - void batch__whenCreateSpansAndThenUpdate__thenReturnUpdated() { - var expectedSpans = List.of(podamFactory.manufacturePojo(Span.class).toBuilder() - .projectId(null) - .parentSpanId(null) - .feedbackScores(null) - .build()); - - batchCreateAndAssert(expectedSpans, API_KEY, TEST_WORKSPACE); - - var expectedSpan = expectedSpans.get(0).toBuilder() - .tags(Set.of()) - .endTime(Instant.now()) - .output(JsonUtils.getJsonNodeFromString("{ \"output\": \"data\"}")) - .build(); - - batchCreateAndAssert(List.of(expectedSpan), API_KEY, TEST_WORKSPACE); - - getAndAssert(expectedSpan.toBuilder().tags(null).build(), API_KEY, TEST_WORKSPACE); - } - - @Test - void batch__whenCreateSpansAndThenUpdateInSameRequest__thenReturnUpdated() { + void batch__whenSendingMultipleSpansWithSameId__thenReturn422() { var expectedSpans = List.of(podamFactory.manufacturePojo(Span.class).toBuilder() .projectId(null) .parentSpanId(null) @@ -3244,74 +3222,84 @@ void batch__whenCreateSpansAndThenUpdateInSameRequest__thenReturnUpdated() { .output(JsonUtils.getJsonNodeFromString("{ \"output\": \"data\"}")) .build(); - batchCreateAndAssert(List.of(expectedSpans.getFirst(), expectedSpan), API_KEY, TEST_WORKSPACE); + List expectedSpans1 = List.of(expectedSpans.getFirst(), expectedSpan); + + try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) + .path("batch") + .request() + .header(HttpHeaders.AUTHORIZATION, API_KEY) + .header(WORKSPACE_HEADER, TEST_WORKSPACE) + .post(Entity.json(new SpanBatch(expectedSpans1)))) { - getAndAssert(expectedSpan.toBuilder().tags(null).build(), API_KEY, TEST_WORKSPACE); + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(422); + assertThat(actualResponse.hasEntity()).isTrue(); + + var errorMessage = actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class); + assertThat(errorMessage.getMessage()).isEqualTo("Duplicate span id '%s'".formatted(expectedSpan.id())); + } } @ParameterizedTest @MethodSource - void batch__whenCreateSpansWithConflicts__thenReturn409(List expectedSpans, String expectedError) { + void batch__whenBatchIsInvalid__thenReturn422(List spans, String errorMessage) { try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) - .path("bulk") + .path("batch") .request() .header(HttpHeaders.AUTHORIZATION, API_KEY) .header(WORKSPACE_HEADER, TEST_WORKSPACE) - .post(Entity.json(new SpanBulk(expectedSpans)))) { + .post(Entity.json(new SpanBatch(spans)))) { - assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(409); + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(422); assertThat(actualResponse.hasEntity()).isTrue(); - assertThat(actualResponse.readEntity(ErrorMessage.class).errors()).contains(expectedError); + + var responseBody = actualResponse.readEntity(ErrorMessage.class); + assertThat(responseBody.errors()).contains(errorMessage); } } - public Stream batch__whenCreateSpansWithConflicts__thenReturn409() { - Span expectedSpan = podamFactory.manufacturePojo(Span.class).toBuilder() + Stream batch__whenBatchIsInvalid__thenReturn422() { + return Stream.of( + Arguments.of(List.of(), "spans size must be between 1 and 1000"), + Arguments.of(IntStream.range(0, 1001) + .mapToObj(i -> podamFactory.manufacturePojo(Span.class).toBuilder() + .projectId(null) + .parentSpanId(null) + .feedbackScores(null) + .build()) + .toList(), "spans size must be between 1 and 1000")); + } + + @Test + void batch__whenSendingMultipleSpansWithNoId__thenReturnNoContent() { + var newSpan = podamFactory.manufacturePojo(Span.class).toBuilder() .projectId(null) + .id(null) .parentSpanId(null) .feedbackScores(null) .build(); - UUID id = podamFactory.manufacturePojo(UUID.class); - UUID id2 = podamFactory.manufacturePojo(UUID.class); - UUID id3 = podamFactory.manufacturePojo(UUID.class); + var expectedSpan = newSpan.toBuilder() + .tags(Set.of()) + .endTime(Instant.now()) + .output(JsonUtils.getJsonNodeFromString("{ \"output\": \"data\"}")) + .build(); - String projectName = UUID.randomUUID().toString(); - String projectName2 = UUID.randomUUID().toString(); - String projectName3 = UUID.randomUUID().toString(); + List expectedSpans = List.of(newSpan, expectedSpan); - return Stream.of( - Arguments.of( - List.of( - expectedSpan.toBuilder().id(id).projectName(projectName).build(), - expectedSpan.toBuilder().id(id).projectName(UUID.randomUUID().toString()).build()), - SpanService.PROJECT_NAME_MISMATCH - ), - Arguments.of( - List.of( - expectedSpan.toBuilder().id(id2).projectName(projectName2).build(), - expectedSpan.toBuilder().id(id2).projectName(projectName2).traceId(podamFactory.manufacturePojo(UUID.class)).build()), - SpanService.TRACE_ID_MISMATCH - ), - Arguments.of( - List.of( - expectedSpan.toBuilder().id(id3).projectName(projectName3).build(), - expectedSpan.toBuilder().id(id3).projectName(projectName3).parentSpanId(podamFactory.manufacturePojo(UUID.class)).build()), - SpanService.PARENT_SPAN_IS_MISMATCH - ) - ); + batchCreateAndAssert(expectedSpans, API_KEY, TEST_WORKSPACE); } + } private void batchCreateAndAssert(List expectedSpans, String apiKey, String workspaceName) { try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) - .path("bulk") + .path("batch") .request() .header(HttpHeaders.AUTHORIZATION, apiKey) .header(WORKSPACE_HEADER, workspaceName) - .post(Entity.json(new SpanBulk(expectedSpans)))) { + .post(Entity.json(new SpanBatch(expectedSpans)))) { assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204); assertThat(actualResponse.hasEntity()).isFalse(); From 91589e7fa902e7813d440ce22d3ac01cb74a6248 Mon Sep 17 00:00:00 2001 From: Thiago Hora Date: Thu, 12 Sep 2024 17:14:47 +0200 Subject: [PATCH 06/10] OPIK-75 Batch traces creation endpoint --- .../java/com/comet/opik/api/TraceBatch.java | 12 ++ .../api/resources/v1/priv/TraceResource.java | 27 +++++ .../java/com/comet/opik/domain/TraceDAO.java | 97 +++++++++++++++ .../com/comet/opik/domain/TraceService.java | 62 ++++++++++ .../resources/v1/priv/TracesResourceTest.java | 114 ++++++++++++++++++ 5 files changed, 312 insertions(+) create mode 100644 apps/opik-backend/src/main/java/com/comet/opik/api/TraceBatch.java diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/TraceBatch.java b/apps/opik-backend/src/main/java/com/comet/opik/api/TraceBatch.java new file mode 100644 index 0000000000..ac5c164940 --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/TraceBatch.java @@ -0,0 +1,12 @@ +package com.comet.opik.api; + +import com.fasterxml.jackson.annotation.JsonView; +import jakarta.validation.Valid; +import jakarta.validation.constraints.NotNull; +import jakarta.validation.constraints.Size; + +import java.util.List; + +public record TraceBatch(@NotNull @Size(min = 1, max = 1000) @JsonView( { + Trace.View.Write.class}) @Valid List traces){ +} diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/TraceResource.java b/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/TraceResource.java index d5843cac98..bd1c77aca1 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/TraceResource.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/TraceResource.java @@ -5,6 +5,7 @@ import com.comet.opik.api.FeedbackScore; import com.comet.opik.api.FeedbackScoreBatch; import com.comet.opik.api.Trace; +import com.comet.opik.api.TraceBatch; import com.comet.opik.api.TraceSearchCriteria; import com.comet.opik.api.TraceUpdate; import com.comet.opik.api.filter.FiltersFactory; @@ -25,6 +26,7 @@ import jakarta.validation.Valid; import jakarta.validation.constraints.Min; import jakarta.validation.constraints.NotNull; +import jakarta.ws.rs.ClientErrorException; import jakarta.ws.rs.Consumes; import jakarta.ws.rs.DELETE; import jakarta.ws.rs.DefaultValue; @@ -45,6 +47,7 @@ import lombok.extern.slf4j.Slf4j; import java.util.UUID; +import java.util.stream.Collectors; import static com.comet.opik.utils.AsyncUtils.setRequestContext; import static com.comet.opik.utils.ValidationUtils.validateProjectNameAndProjectId; @@ -115,6 +118,30 @@ public Response create( return Response.created(uri).build(); } + @POST + @Path("/batch") + @Operation(operationId = "createTraces", summary = "Create traces", description = "Create traces", responses = { + @ApiResponse(responseCode = "204", description = "No Content")}) + public Response createSpans( + @RequestBody(content = @Content(schema = @Schema(implementation = TraceBatch.class))) @JsonView(Trace.View.Write.class) @NotNull @Valid TraceBatch traces) { + + traces.traces() + .stream() + .filter(trace -> trace.id() != null) // Filter out spans with null IDs + .collect(Collectors.groupingBy(Trace::id)) + .forEach((id, traceGroup) -> { + if (traceGroup.size() > 1) { + throw new ClientErrorException("Duplicate trace id '%s'".formatted(id), 422); + } + }); + + service.create(traces) + .contextWrite(ctx -> setRequestContext(ctx, requestContext)) + .block(); + + return Response.noContent().build(); + } + @PATCH @Path("{id}") @Operation(operationId = "updateTrace", summary = "Update trace by id", description = "Update trace by id", responses = { 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 256daa3ad3..5a054f3245 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 @@ -6,6 +6,9 @@ import com.comet.opik.domain.filter.FilterQueryBuilder; import com.comet.opik.domain.filter.FilterStrategy; import com.comet.opik.utils.JsonUtils; +import com.comet.opik.utils.TemplateUtils; +import com.fasterxml.jackson.databind.JsonNode; +import com.google.common.base.Preconditions; import com.google.inject.ImplementedBy; import com.newrelic.api.agent.Segment; import io.r2dbc.spi.Connection; @@ -38,6 +41,7 @@ import static com.comet.opik.infrastructure.instrumentation.InstrumentAsyncUtils.startSegment; import static com.comet.opik.utils.AsyncUtils.makeFluxContextAware; import static com.comet.opik.utils.AsyncUtils.makeMonoContextAware; +import static com.comet.opik.utils.TemplateUtils.getQueryItemPlaceHolder; @ImplementedBy(TraceDAOImpl.class) interface TraceDAO { @@ -57,6 +61,7 @@ Mono partialInsert(UUID projectId, TraceUpdate traceUpdate, UUID traceId, Mono> getTraceWorkspace(Set traceIds, Connection connection); + Mono batchInsert(List traces, Connection connection); } @Slf4j @@ -64,6 +69,41 @@ Mono partialInsert(UUID projectId, TraceUpdate traceUpdate, UUID traceId, @RequiredArgsConstructor(onConstructor_ = @Inject) class TraceDAOImpl implements TraceDAO { + private static final String BATCH_INSERT = """ + INSERT INTO traces( + id, + project_id, + workspace_id, + name, + start_time, + end_time, + input, + output, + metadata, + tags, + created_by, + last_updated_by + ) VALUES + , + :project_id, + :workspace_id, + :name, + parseDateTime64BestEffort(:start_time, 9), + if(:end_time IS NULL, NULL, parseDateTime64BestEffort(:end_time, 9)), + :input, + :output, + :metadata, + :tags, + :user_name, + :user_name + ) + , + }> + ; + """; + /** * This query handles the insertion of a new trace into the database in two cases: * 1. When the trace does not exist in the database. @@ -695,4 +735,61 @@ public Mono> getTraceWorkspace( .collectList(); } + @Override + public Mono batchInsert(@NonNull List traces, @NonNull Connection connection) { + + Preconditions.checkArgument(!traces.isEmpty(), "traces must not be empty"); + + return Mono.from(insert(traces, connection)) + .flatMapMany(Result::getRowsUpdated) + .reduce(0L, Long::sum); + + } + + private Publisher insert(List traces, Connection connection) { + + return makeMonoContextAware((userName, workspaceName, workspaceId) -> { + List queryItems = getQueryItemPlaceHolder(traces.size()); + + var template = new ST(BATCH_INSERT) + .add("items", queryItems); + + Statement statement = connection.createStatement(template.render()); + + int i = 0; + for (Trace trace : traces) { + + statement.bind("id" + i, trace.id()) + .bind("project_id" + i, trace.projectId()) + .bind("name" + i, trace.name()) + .bind("start_time" + i, trace.startTime().toString()) + .bind("input" + i, getOrDefault(trace.input())) + .bind("output" + i, getOrDefault(trace.output())) + .bind("metadata" + i, getOrDefault(trace.metadata())) + .bind("tags" + i, trace.tags() != null ? trace.tags().toArray(String[]::new) : new String[]{}); + + if (trace.endTime() != null) { + statement.bind("end_time" + i, trace.endTime().toString()); + } else { + statement.bindNull("end_time" + i, String.class); + } + + i++; + } + + statement + .bind("workspace_id", workspaceId) + .bind("user_name", userName); + + Segment segment = startSegment("traces", "Clickhouse", "batch_insert"); + + return Mono.from(statement.execute()) + .doFinally(signalType -> endSegment(segment)); + }); + } + + private String getOrDefault(JsonNode value) { + return value != null ? value.toString() : ""; + } + } 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 3dd3d7c590..a0c2b1e409 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 @@ -3,6 +3,7 @@ import com.clickhouse.client.ClickHouseException; import com.comet.opik.api.Project; import com.comet.opik.api.Trace; +import com.comet.opik.api.TraceBatch; import com.comet.opik.api.TraceSearchCriteria; import com.comet.opik.api.TraceUpdate; import com.comet.opik.api.error.EntityAlreadyExistsException; @@ -20,13 +21,18 @@ import jakarta.ws.rs.core.Response; import lombok.NonNull; import lombok.RequiredArgsConstructor; +import org.apache.commons.lang3.StringUtils; +import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import reactor.core.scheduler.Schedulers; import java.time.Instant; import java.util.List; +import java.util.Map; import java.util.Set; import java.util.UUID; +import java.util.function.Function; +import java.util.stream.Collectors; import static com.comet.opik.domain.FeedbackScoreDAO.EntityType; @@ -35,6 +41,8 @@ public interface TraceService { Mono create(Trace trace); + Mono create(TraceBatch batch); + Mono update(TraceUpdate trace, UUID id); Mono get(UUID id); @@ -77,6 +85,60 @@ public Mono create(@NonNull Trace trace) { Mono.defer(() -> insertTrace(trace, project, id)))); } + @com.newrelic.api.agent.Trace(dispatcher = true) + public Mono create(TraceBatch batch) { + + if (batch.traces().isEmpty()) { + return Mono.empty(); + } + + List projectNames = batch.traces() + .stream() + .map(Trace::projectName) + .distinct() + .toList(); + + Mono> resolveProjects = Flux.fromIterable(projectNames) + .flatMap(this::resolveProject) + .collectList() + .map(projects -> bindTraceToProjectAndId(batch, projects)) + .subscribeOn(Schedulers.boundedElastic()); + + return resolveProjects + .flatMap(traces -> template.nonTransaction(connection -> dao.batchInsert(traces, connection))) + .then(); + } + + private List bindTraceToProjectAndId(TraceBatch batch, List projects) { + Map projectPerName = projects.stream() + .collect(Collectors.toMap(Project::name, Function.identity())); + + return batch.traces() + .stream() + .map(trace -> { + String projectName = WorkspaceUtils.getProjectName(trace.projectName()); + Project project = projectPerName.get(projectName); + + if (project == null) { + throw new EntityAlreadyExistsException(new ErrorMessage(List.of("Project not found"))); + } + + UUID id = trace.id() == null ? idGenerator.generateId() : trace.id(); + IdGenerator.validateVersion(id, TRACE_KEY); + + return trace.toBuilder().id(id).projectId(project.id()).build(); + }) + .toList(); + } + + private Mono resolveProject(String projectName) { + if (StringUtils.isEmpty(projectName)) { + return getOrCreateProject(ProjectService.DEFAULT_PROJECT); + } + + return getOrCreateProject(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)) 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 74687bfee5..f811ff3c73 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 @@ -7,6 +7,7 @@ import com.comet.opik.api.Project; import com.comet.opik.api.ScoreSource; import com.comet.opik.api.Trace; +import com.comet.opik.api.TraceBatch; import com.comet.opik.api.TraceUpdate; import com.comet.opik.api.error.ErrorMessage; import com.comet.opik.api.filter.Filter; @@ -3035,6 +3036,119 @@ void create__whenProjectNameIsNull__thenAcceptAndUseDefaultProject() { } + @Nested + @DisplayName("Batch:") + @TestInstance(TestInstance.Lifecycle.PER_CLASS) + class BatchInsert { + + @Test + void batch__whenCreateTraces__thenReturnNoContent() { + var expectedTraces = IntStream.range(0, 1000) + .mapToObj(i -> factory.manufacturePojo(Trace.class).toBuilder() + .projectId(null) + .endTime(null) + .feedbackScores(null) + .build()) + .toList(); + + batchCreateAndAssert(expectedTraces, API_KEY, TEST_WORKSPACE); + } + + @Test + void batch__whenSendingMultipleTracesWithSameId__thenReturn422() { + var trace = factory.manufacturePojo(Trace.class).toBuilder() + .projectId(null) + .feedbackScores(null) + .build(); + + var expectedTrace = trace.toBuilder() + .tags(Set.of()) + .endTime(Instant.now()) + .output(JsonUtils.getJsonNodeFromString("{ \"output\": \"data\"}")) + .build(); + + List traces = List.of(trace, expectedTrace); + + try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) + .path("batch") + .request() + .header(HttpHeaders.AUTHORIZATION, API_KEY) + .header(WORKSPACE_HEADER, TEST_WORKSPACE) + .post(Entity.json(new TraceBatch(traces)))) { + + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(422); + assertThat(actualResponse.hasEntity()).isTrue(); + + var errorMessage = actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class); + assertThat(errorMessage.getMessage()).isEqualTo("Duplicate trace id '%s'".formatted(trace.id())); + } + } + + @ParameterizedTest + @MethodSource + void batch__whenBatchIsInvalid__thenReturn422(List traces, String errorMessage) { + + try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) + .path("batch") + .request() + .header(HttpHeaders.AUTHORIZATION, API_KEY) + .header(WORKSPACE_HEADER, TEST_WORKSPACE) + .post(Entity.json(new TraceBatch(traces)))) { + + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(422); + assertThat(actualResponse.hasEntity()).isTrue(); + + var responseBody = actualResponse.readEntity(ErrorMessage.class); + assertThat(responseBody.errors()).contains(errorMessage); + } + } + + Stream batch__whenBatchIsInvalid__thenReturn422() { + return Stream.of( + Arguments.of(List.of(), "traces size must be between 1 and 1000"), + Arguments.of(IntStream.range(0, 1001) + .mapToObj(i -> factory.manufacturePojo(Trace.class).toBuilder() + .projectId(null) + .feedbackScores(null) + .build()) + .toList(), "traces size must be between 1 and 1000")); + } + + @Test + void batch__whenSendingMultipleTracesWithNoId__thenReturnNoContent() { + var newTrace = factory.manufacturePojo(Trace.class).toBuilder() + .projectId(null) + .id(null) + .feedbackScores(null) + .build(); + + var expectedTrace = newTrace.toBuilder() + .tags(Set.of()) + .endTime(Instant.now()) + .output(JsonUtils.getJsonNodeFromString("{ \"output\": \"data\"}")) + .build(); + + List expectedTraces = List.of(newTrace, expectedTrace); + + batchCreateAndAssert(expectedTraces, API_KEY, TEST_WORKSPACE); + } + + private void batchCreateAndAssert(List traces, String apiKey, String workspaceName) { + + try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) + .path("batch") + .request() + .header(HttpHeaders.AUTHORIZATION, apiKey) + .header(WORKSPACE_HEADER, workspaceName) + .post(Entity.json(new TraceBatch(traces)))) { + + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204); + assertThat(actualResponse.hasEntity()).isFalse(); + } + } + + } + @Nested @DisplayName("Delete:") @TestInstance(TestInstance.Lifecycle.PER_CLASS) From 59c9e5bed164702f4ed4df4f0a5e0cb28ef29e9f Mon Sep 17 00:00:00 2001 From: Thiago Hora Date: Fri, 13 Sep 2024 10:44:28 +0200 Subject: [PATCH 07/10] Remove remaining dead code --- apps/opik-backend/config.yml | 1 - .../main/java/com/comet/opik/api/Span.java | 13 ------- .../com/comet/opik/domain/SpanService.java | 16 ++------ .../infrastructure/DistributedLockConfig.java | 5 --- .../infrastructure/redis/LockService.java | 7 ---- .../redis/RedissonLockService.java | 39 ------------------- .../resources/v1/priv/SpansResourceTest.java | 2 + .../comet/opik/domain/DummyLockService.java | 13 ------- 8 files changed, 5 insertions(+), 91 deletions(-) diff --git a/apps/opik-backend/config.yml b/apps/opik-backend/config.yml index d707513892..ec31499b02 100644 --- a/apps/opik-backend/config.yml +++ b/apps/opik-backend/config.yml @@ -49,7 +49,6 @@ health: distributedLock: lockTimeoutMS: ${DISTRIBUTED_LOCK_TIME_OUT:-500} - bulkLockTimeoutMS: ${DISTRIBUTED_BULK_LOCK_TIME_OUT:-5000} redis: singleNodeUrl: ${REDIS_URL:-} diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/Span.java b/apps/opik-backend/src/main/java/com/comet/opik/api/Span.java index d351406ac2..edc0a55e6a 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/api/Span.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/Span.java @@ -15,7 +15,6 @@ import java.time.Instant; import java.util.List; import java.util.Map; -import java.util.Objects; import java.util.Set; import java.util.UUID; @@ -67,16 +66,4 @@ public static class Public { } } - @Override - public boolean equals(Object o) { - if (this == o) return true; - if (o == null || getClass() != o.getClass()) return false; - Span span = (Span) o; - return Objects.equals(id, span.id); - } - - @Override - public int hashCode() { - return Objects.hash(id); - } } 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 b7f8406f50..b3860f7e5a 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 @@ -12,6 +12,7 @@ import com.comet.opik.infrastructure.auth.RequestContext; import com.comet.opik.infrastructure.redis.LockService; import com.comet.opik.utils.WorkspaceUtils; +import com.google.common.base.Preconditions; import com.newrelic.api.agent.Trace; import jakarta.inject.Inject; import jakarta.inject.Singleton; @@ -19,7 +20,6 @@ import lombok.NonNull; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; -import org.apache.commons.lang3.StringUtils; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import reactor.core.scheduler.Schedulers; @@ -254,9 +254,7 @@ public Mono validateSpanWorkspace(@NonNull String workspaceId, @NonNull @Trace(dispatcher = true) public Mono create(@NonNull SpanBatch batch) { - if (batch.spans().isEmpty()) { - return Mono.empty(); - } + Preconditions.checkArgument(!batch.spans().isEmpty(), "Batch spans must not be empty"); List projectNames = batch.spans() .stream() @@ -285,10 +283,6 @@ private List bindSpanToProjectAndId(SpanBatch batch, List project String projectName = WorkspaceUtils.getProjectName(span.projectName()); Project project = projectPerName.get(projectName); - if (project == null) { - throw new EntityAlreadyExistsException(new ErrorMessage(List.of("Project not found"))); - } - UUID id = span.id() == null ? idGenerator.generateId() : span.id(); IdGenerator.validateVersion(id, SPAN_KEY); @@ -298,10 +292,6 @@ private List bindSpanToProjectAndId(SpanBatch batch, List project } private Mono resolveProject(String projectName) { - if (StringUtils.isEmpty(projectName)) { - return getOrCreateProject(ProjectService.DEFAULT_PROJECT); - } - - return getOrCreateProject(projectName); + return getOrCreateProject(WorkspaceUtils.getProjectName(projectName)); } } diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/DistributedLockConfig.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/DistributedLockConfig.java index e34d9d7277..d6b713d8c2 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/DistributedLockConfig.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/DistributedLockConfig.java @@ -11,9 +11,4 @@ public class DistributedLockConfig { @Valid @JsonProperty @NotNull private int lockTimeoutMS; - - @Valid - @JsonProperty - @NotNull private int bulkLockTimeoutMS; - } diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/redis/LockService.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/redis/LockService.java index e6c542cfea..a02bba19d5 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/redis/LockService.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/redis/LockService.java @@ -3,14 +3,10 @@ import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; -import java.util.List; import java.util.UUID; public interface LockService { - Mono> lockAll(List spanIds, String spanKey); - Mono unlockAll(List locks); - record Lock(String key) { private static final String KEY_FORMAT = "%s-%s"; @@ -25,9 +21,6 @@ public Lock(String id, String name) { } - record LockRef(Lock lock, String ref) { - } - Mono executeWithLock(Lock lock, Mono action); Flux executeWithLock(Lock lock, Flux action); } diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/redis/RedissonLockService.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/redis/RedissonLockService.java index 97be6da029..5ef02eac86 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/redis/RedissonLockService.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/redis/RedissonLockService.java @@ -12,9 +12,6 @@ import reactor.core.scheduler.Schedulers; import java.time.Duration; -import java.util.List; -import java.util.Map; -import java.util.UUID; @RequiredArgsConstructor @Slf4j @@ -23,42 +20,6 @@ class RedissonLockService implements LockService { private final @NonNull RedissonReactiveClient redisClient; private final @NonNull DistributedLockConfig distributedLockConfig; - @Override - public Mono> lockAll(@NonNull List keys, @NonNull String suffix) { - if (keys.isEmpty()) { - return Mono.empty(); - } - - return Flux.fromIterable(keys) - .map(key -> new Lock(key, suffix)) - .flatMap(lock -> { - RPermitExpirableSemaphoreReactive semaphore = getSemaphore(lock, - distributedLockConfig.getBulkLockTimeoutMS()); - log.debug("Trying to lock with {}", lock); - return semaphore.trySetPermits(1).thenReturn(Map.entry(lock, semaphore)); - }) - .flatMap(entry -> entry.getValue().acquire() - .flatMap(locked -> Mono.just(new LockRef(entry.getKey(), locked)))) - .collectList(); - } - - @Override - public Mono unlockAll(@NonNull List locks) { - if (locks.isEmpty()) { - return Mono.empty(); - } - - return Flux.fromIterable(locks) - .flatMap(lock -> { - RPermitExpirableSemaphoreReactive semaphore = getSemaphore(lock.lock(), - distributedLockConfig.getBulkLockTimeoutMS()); - log.debug("Trying to unlock with {}", lock); - return semaphore.release(lock.ref()); - }) - .collectList() - .then(); - } - @Override public Mono executeWithLock(@NonNull Lock lock, @NonNull Mono action) { 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 9aade09878..4c48b98619 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 @@ -3206,6 +3206,8 @@ void batch__whenCreateSpans__thenReturnNoContent() { .toList(); batchCreateAndAssert(expectedSpans, API_KEY, TEST_WORKSPACE); + + expectedSpans.forEach(expectedSpan -> getAndAssert(expectedSpan, API_KEY, TEST_WORKSPACE)); } @Test diff --git a/apps/opik-backend/src/test/java/com/comet/opik/domain/DummyLockService.java b/apps/opik-backend/src/test/java/com/comet/opik/domain/DummyLockService.java index 2ee0ffddff..faf96f9979 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/domain/DummyLockService.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/domain/DummyLockService.java @@ -4,21 +4,8 @@ import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; -import java.util.List; -import java.util.UUID; - public class DummyLockService implements LockService { - @Override - public Mono> lockAll(List spanIds, String spanKey) { - return Mono.empty(); - } - - @Override - public Mono unlockAll(List locks) { - return Mono.empty(); - } - @Override public Mono executeWithLock(LockService.Lock lock, Mono action) { return action; From be2be343f637dc8fc75d53caca4f87fbe23a4d30 Mon Sep 17 00:00:00 2001 From: Thiago Hora Date: Fri, 13 Sep 2024 10:59:24 +0200 Subject: [PATCH 08/10] Fix conflicts --- .../com/comet/opik/domain/TraceService.java | 23 +++++-------------- .../resources/v1/priv/TracesResourceTest.java | 6 +++++ 2 files changed, 12 insertions(+), 17 deletions(-) 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 a0c2b1e409..77002bf5ca 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 @@ -14,6 +14,7 @@ import com.comet.opik.infrastructure.redis.LockService; import com.comet.opik.utils.AsyncUtils; import com.comet.opik.utils.WorkspaceUtils; +import com.google.common.base.Preconditions; import com.google.inject.ImplementedBy; import jakarta.inject.Inject; import jakarta.inject.Singleton; @@ -21,7 +22,6 @@ import jakarta.ws.rs.core.Response; import lombok.NonNull; import lombok.RequiredArgsConstructor; -import org.apache.commons.lang3.StringUtils; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import reactor.core.scheduler.Schedulers; @@ -41,7 +41,7 @@ public interface TraceService { Mono create(Trace trace); - Mono create(TraceBatch batch); + Mono create(TraceBatch batch); Mono update(TraceUpdate trace, UUID id); @@ -86,11 +86,9 @@ public Mono create(@NonNull Trace trace) { } @com.newrelic.api.agent.Trace(dispatcher = true) - public Mono create(TraceBatch batch) { + public Mono create(TraceBatch batch) { - if (batch.traces().isEmpty()) { - return Mono.empty(); - } + Preconditions.checkArgument(!batch.traces().isEmpty(), "Batch traces cannot be empty"); List projectNames = batch.traces() .stream() @@ -105,8 +103,7 @@ public Mono create(TraceBatch batch) { .subscribeOn(Schedulers.boundedElastic()); return resolveProjects - .flatMap(traces -> template.nonTransaction(connection -> dao.batchInsert(traces, connection))) - .then(); + .flatMap(traces -> template.nonTransaction(connection -> dao.batchInsert(traces, connection))); } private List bindTraceToProjectAndId(TraceBatch batch, List projects) { @@ -119,10 +116,6 @@ private List bindTraceToProjectAndId(TraceBatch batch, List proj String projectName = WorkspaceUtils.getProjectName(trace.projectName()); Project project = projectPerName.get(projectName); - if (project == null) { - throw new EntityAlreadyExistsException(new ErrorMessage(List.of("Project not found"))); - } - UUID id = trace.id() == null ? idGenerator.generateId() : trace.id(); IdGenerator.validateVersion(id, TRACE_KEY); @@ -132,11 +125,7 @@ private List bindTraceToProjectAndId(TraceBatch batch, List proj } private Mono resolveProject(String projectName) { - if (StringUtils.isEmpty(projectName)) { - return getOrCreateProject(ProjectService.DEFAULT_PROJECT); - } - - return getOrCreateProject(projectName); + return getOrCreateProject(WorkspaceUtils.getProjectName(projectName)); } private Mono insertTrace(Trace newTrace, Project project, UUID id) { 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 f811ff3c73..97da4e36ff 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 @@ -3052,6 +3052,12 @@ void batch__whenCreateTraces__thenReturnNoContent() { .toList(); batchCreateAndAssert(expectedTraces, API_KEY, TEST_WORKSPACE); + + expectedTraces.forEach(trace -> { + UUID projectId = getProjectId(client, trace.projectName(), TEST_WORKSPACE, API_KEY); + + getAndAssert(trace, trace.id(), projectId, Instant.now(), API_KEY, TEST_WORKSPACE); + }); } @Test From 7c0eab24dae51ad09b904a0bb504ee4ddc9ef507 Mon Sep 17 00:00:00 2001 From: Thiago Hora Date: Fri, 13 Sep 2024 12:57:49 +0200 Subject: [PATCH 09/10] Fix test issues --- .../opik/domain/FeedbackScoreService.java | 7 - .../opik/api/resources/utils/TestUtils.java | 14 ++ .../v1/priv/DatasetsResourceTest.java | 10 +- .../v1/priv/ExperimentsResourceTest.java | 4 +- .../priv/FeedbackDefinitionResourceTest.java | 8 +- .../v1/priv/ProjectsResourceTest.java | 16 +- .../resources/v1/priv/SpansResourceTest.java | 3 +- .../resources/v1/priv/TracesResourceTest.java | 228 ++++++++++-------- 8 files changed, 154 insertions(+), 136 deletions(-) create mode 100644 apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/TestUtils.java diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/FeedbackScoreService.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/FeedbackScoreService.java index dae9a80750..6cb1084da8 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/domain/FeedbackScoreService.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/FeedbackScoreService.java @@ -37,7 +37,6 @@ @ImplementedBy(FeedbackScoreServiceImpl.class) public interface FeedbackScoreService { - Flux getScores(EntityType entityType, UUID entityId); Mono scoreTrace(UUID traceId, FeedbackScore score); Mono scoreSpan(UUID spanId, FeedbackScore score); @@ -66,12 +65,6 @@ class FeedbackScoreServiceImpl implements FeedbackScoreService { record ProjectDto(Project project, List scores) { } - @Override - public Flux getScores(@NonNull EntityType entityType, @NonNull UUID entityId) { - return asyncTemplate.nonTransaction(connection -> dao.getScores(entityType, List.of(entityId), connection)) - .flatMapIterable(entityIdToFeedbackScoresMap -> entityIdToFeedbackScoresMap.get(entityId)); - } - @Override public Mono scoreTrace(@NonNull UUID traceId, @NonNull FeedbackScore score) { return lockService.executeWithLock( diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/TestUtils.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/TestUtils.java new file mode 100644 index 0000000000..99d5fd1d33 --- /dev/null +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/TestUtils.java @@ -0,0 +1,14 @@ +package com.comet.opik.api.resources.utils; + +import lombok.experimental.UtilityClass; + +import java.net.URI; +import java.util.UUID; + +@UtilityClass +public class TestUtils { + + public static UUID getIdFromLocation(URI location) { + return UUID.fromString(location.getPath().substring(location.getPath().lastIndexOf('/') + 1)); + } +} diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/DatasetsResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/DatasetsResourceTest.java index 79e6cb9a51..537d8ccc94 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/DatasetsResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/DatasetsResourceTest.java @@ -24,6 +24,7 @@ import com.comet.opik.api.resources.utils.MySQLContainerUtils; import com.comet.opik.api.resources.utils.RedisContainerUtils; import com.comet.opik.api.resources.utils.TestDropwizardAppExtensionUtils; +import com.comet.opik.api.resources.utils.TestUtils; import com.comet.opik.api.resources.utils.WireMockUtils; import com.comet.opik.domain.FeedbackScoreMapper; import com.comet.opik.podam.PodamFactoryUtils; @@ -195,8 +196,7 @@ private UUID createAndAssert(Dataset dataset, String apiKey, String workspaceNam assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(201); assertThat(actualResponse.hasEntity()).isFalse(); - var id = UUID.fromString(actualResponse.getHeaderString("Location") - .substring(actualResponse.getHeaderString("Location").lastIndexOf('/') + 1)); + var id = TestUtils.getIdFromLocation(actualResponse.getLocation()); assertThat(id).isNotNull(); assertThat(id.version()).isEqualTo(7); @@ -2532,8 +2532,7 @@ private UUID createTrace(Trace trace, String apiKey, String workspaceName) { .post(Entity.entity(trace, MediaType.APPLICATION_JSON_TYPE))) { assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(201); - return UUID.fromString(actualResponse.getHeaderString("Location") - .substring(actualResponse.getHeaderString("Location").lastIndexOf('/') + 1)); + return TestUtils.getIdFromLocation(actualResponse.getLocation()); } } @@ -2545,8 +2544,7 @@ private UUID createSpan(Span span, String apiKey, String workspaceName) { .post(Entity.entity(span, MediaType.APPLICATION_JSON_TYPE))) { assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(201); - return UUID.fromString(actualResponse.getHeaderString("Location") - .substring(actualResponse.getHeaderString("Location").lastIndexOf('/') + 1)); + return TestUtils.getIdFromLocation(actualResponse.getLocation()); } } diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ExperimentsResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ExperimentsResourceTest.java index 3880b5eb96..2beeb2df2b 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ExperimentsResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ExperimentsResourceTest.java @@ -18,6 +18,7 @@ import com.comet.opik.api.resources.utils.MySQLContainerUtils; import com.comet.opik.api.resources.utils.RedisContainerUtils; import com.comet.opik.api.resources.utils.TestDropwizardAppExtensionUtils; +import com.comet.opik.api.resources.utils.TestUtils; import com.comet.opik.api.resources.utils.WireMockUtils; import com.comet.opik.domain.FeedbackScoreMapper; import com.comet.opik.podam.PodamFactoryUtils; @@ -1536,8 +1537,7 @@ private UUID createAndAssert(Experiment expectedExperiment, String apiKey, Strin assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(201); - var path = actualResponse.getLocation().getPath(); - var actualId = UUID.fromString(path.substring(path.lastIndexOf('/') + 1)); + var actualId = TestUtils.getIdFromLocation(actualResponse.getLocation()); assertThat(actualResponse.hasEntity()).isFalse(); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/FeedbackDefinitionResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/FeedbackDefinitionResourceTest.java index 27f543ce41..5390a9f49c 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/FeedbackDefinitionResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/FeedbackDefinitionResourceTest.java @@ -8,6 +8,7 @@ import com.comet.opik.api.resources.utils.MySQLContainerUtils; import com.comet.opik.api.resources.utils.RedisContainerUtils; import com.comet.opik.api.resources.utils.TestDropwizardAppExtensionUtils; +import com.comet.opik.api.resources.utils.TestUtils; import com.comet.opik.api.resources.utils.WireMockUtils; import com.comet.opik.podam.PodamFactoryUtils; import com.fasterxml.uuid.Generators; @@ -135,8 +136,7 @@ private UUID create(final FeedbackDefinition feedback, String apiKey, String assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(201); - return UUID.fromString(actualResponse.getLocation().getPath() - .substring(actualResponse.getLocation().getPath().lastIndexOf('/') + 1)); + return TestUtils.getIdFromLocation(actualResponse.getLocation()); } } @@ -848,9 +848,7 @@ void create() { assertThat(actualResponse.hasEntity()).isFalse(); assertThat(actualResponse.getHeaderString("Location")).matches(Pattern.compile(URL_PATTERN)); - id = UUID.fromString(actualResponse.getLocation().getPath() - .substring(actualResponse.getLocation().getPath().lastIndexOf('/') + 1)); - + id = TestUtils.getIdFromLocation(actualResponse.getLocation()); } var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ProjectsResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ProjectsResourceTest.java index 117fabcd7d..ea8b879d8a 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ProjectsResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/ProjectsResourceTest.java @@ -9,6 +9,7 @@ import com.comet.opik.api.resources.utils.MySQLContainerUtils; import com.comet.opik.api.resources.utils.RedisContainerUtils; import com.comet.opik.api.resources.utils.TestDropwizardAppExtensionUtils; +import com.comet.opik.api.resources.utils.TestUtils; import com.comet.opik.api.resources.utils.WireMockUtils; import com.comet.opik.podam.PodamFactoryUtils; import com.github.tomakehurst.wiremock.client.WireMock; @@ -134,8 +135,7 @@ private UUID createProject(Project project, String apiKey, String workspaceName) assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(201); - return UUID.fromString(actualResponse.getLocation().getPath() - .substring(actualResponse.getLocation().getPath().lastIndexOf('/') + 1)); + return TestUtils.getIdFromLocation(actualResponse.getLocation()); } } @@ -820,8 +820,7 @@ void create() { assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(201); assertThat(actualResponse.hasEntity()).isFalse(); - id = UUID.fromString(actualResponse.getHeaderString("Location") - .substring(actualResponse.getHeaderString("Location").lastIndexOf('/') + 1)); + id = TestUtils.getIdFromLocation(actualResponse.getLocation()); } assertProject(project.toBuilder().id(id) @@ -848,8 +847,7 @@ void create__whenWorkspaceNameIsSpecified__thenAcceptTheRequest() { assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(201); assertThat(actualResponse.hasEntity()).isFalse(); - id = UUID.fromString(actualResponse.getHeaderString("Location") - .substring(actualResponse.getHeaderString("Location").lastIndexOf('/') + 1)); + id = TestUtils.getIdFromLocation(actualResponse.getLocation()); } @@ -932,8 +930,7 @@ void create__whenProjectsHaveSameNameButDifferentWorkspace__thenAcceptTheRequest assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(201); assertThat(actualResponse.hasEntity()).isFalse(); - id = UUID.fromString(actualResponse.getHeaderString("Location") - .substring(actualResponse.getHeaderString("Location").lastIndexOf('/') + 1)); + id = TestUtils.getIdFromLocation(actualResponse.getLocation()); } var project2 = project1.toBuilder() @@ -950,8 +947,7 @@ void create__whenProjectsHaveSameNameButDifferentWorkspace__thenAcceptTheRequest assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(201); assertThat(actualResponse.hasEntity()).isFalse(); - id2 = UUID.fromString(actualResponse.getHeaderString("Location") - .substring(actualResponse.getHeaderString("Location").lastIndexOf('/') + 1)); + id2 = TestUtils.getIdFromLocation(actualResponse.getLocation()); } assertProject(project1.toBuilder().id(id).build()); 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 4c48b98619..78ed1332e4 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 @@ -22,6 +22,7 @@ import com.comet.opik.api.resources.utils.MySQLContainerUtils; import com.comet.opik.api.resources.utils.RedisContainerUtils; import com.comet.opik.api.resources.utils.TestDropwizardAppExtensionUtils; +import com.comet.opik.api.resources.utils.TestUtils; import com.comet.opik.api.resources.utils.WireMockUtils; import com.comet.opik.domain.SpanMapper; import com.comet.opik.domain.SpanType; @@ -3092,7 +3093,7 @@ private UUID createAndAssert(Span expectedSpan, String apiKey, String workspaceN if (expectedSpan.id() != null) { expectedSpanId = expectedSpan.id(); } else { - expectedSpanId = UUID.fromString(actualHeaderString.substring(actualHeaderString.lastIndexOf('/') + 1)); + expectedSpanId = TestUtils.getIdFromLocation(actualResponse.getLocation()); } assertThat(actualHeaderString).isEqualTo(URL_TEMPLATE.formatted(baseURI) 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 97da4e36ff..aa5e155592 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 @@ -22,6 +22,7 @@ import com.comet.opik.api.resources.utils.MySQLContainerUtils; import com.comet.opik.api.resources.utils.RedisContainerUtils; import com.comet.opik.api.resources.utils.TestDropwizardAppExtensionUtils; +import com.comet.opik.api.resources.utils.TestUtils; import com.comet.opik.api.resources.utils.WireMockUtils; import com.comet.opik.infrastructure.auth.RequestContext; import com.comet.opik.podam.PodamFactoryUtils; @@ -163,7 +164,7 @@ void tearDownAll() { wireMock.server().stop(); } - private UUID getProjectId(ClientSupport client, String projectName, String workspaceName, String apiKey) { + private UUID getProjectId(String projectName, String workspaceName, String apiKey) { return client.target("%s/v1/private/projects".formatted(baseURI)) .queryParam("name", projectName) .request() @@ -178,6 +179,19 @@ private UUID getProjectId(ClientSupport client, String projectName, String works .id(); } + private UUID createProject(String projectName, String workspaceName, String apiKey) { + try (Response response = client.target("%s/v1/private/projects".formatted(baseURI)) + .queryParam("name", projectName) + .request() + .header(HttpHeaders.AUTHORIZATION, apiKey) + .header(WORKSPACE_HEADER, workspaceName) + .post(Entity.json(Project.builder().name(projectName).build()))) { + + assertThat(response.getStatusInfo().getStatusCode()).isEqualTo(201); + return TestUtils.getIdFromLocation(response.getLocation()); + } + } + @Nested @DisplayName("Api Key Authentication:") @TestInstance(TestInstance.Lifecycle.PER_CLASS) @@ -838,7 +852,7 @@ void getByProjectName__whenProjectIdIsNotEmpty__thenReturnTracesByProjectId() { .feedbackScores(null) .build(), apiKey, workspaceName); - UUID projectId = getProjectId(client, projectName, workspaceName, apiKey); + UUID projectId = getProjectId(projectName, workspaceName, apiKey); var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) .queryParam("workspace_name", workspaceName) @@ -2636,93 +2650,96 @@ void getByProjectName__whenFilterInvalidValueOrKeyForFieldType__thenReturn400(Fi var actualError = actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class); assertThat(actualError).isEqualTo(expectedError); } + } - private void getAndAssertPage(String workspaceName, String projectName, List filters, - List traces, - List expectedTraces, List unexpectedTraces, String apiKey) { - int page = 1; - int size = traces.size() + expectedTraces.size() + unexpectedTraces.size(); - getAndAssertPage(page, size, projectName, filters, expectedTraces, unexpectedTraces, - workspaceName, apiKey); - } + private void getAndAssertPage(String workspaceName, String projectName, List filters, + List traces, + List expectedTraces, List unexpectedTraces, String apiKey) { + int page = 1; + int size = traces.size() + expectedTraces.size() + unexpectedTraces.size(); + getAndAssertPage(page, size, projectName, filters, expectedTraces, unexpectedTraces, + workspaceName, apiKey); + } - private void getAndAssertPage(int page, int size, String projectName, List filters, - List expectedTraces, List unexpectedTraces, String workspaceName, String apiKey) { - var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) - .queryParam("page", page) - .queryParam("size", size) - .queryParam("project_name", projectName) - .queryParam("filters", toURLEncodedQueryParam(filters)) - .request() - .header(HttpHeaders.AUTHORIZATION, apiKey) - .header(WORKSPACE_HEADER, workspaceName) - .get(); + private void getAndAssertPage(int page, int size, String projectName, List filters, + List expectedTraces, List unexpectedTraces, String workspaceName, String apiKey) { + var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) + .queryParam("page", page) + .queryParam("size", size) + .queryParam("project_name", projectName) + .queryParam("filters", toURLEncodedQueryParam(filters)) + .request() + .header(HttpHeaders.AUTHORIZATION, apiKey) + .header(WORKSPACE_HEADER, workspaceName) + .get(); - assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200); + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200); - var actualPage = actualResponse.readEntity(Trace.TracePage.class); - var actualTraces = actualPage.content(); + var actualPage = actualResponse.readEntity(Trace.TracePage.class); + var actualTraces = actualPage.content(); - assertThat(actualPage.page()).isEqualTo(page); - assertThat(actualPage.size()).isEqualTo(expectedTraces.size()); - assertThat(actualPage.total()).isEqualTo(expectedTraces.size()); - assertThat(actualTraces) - .usingRecursiveFieldByFieldElementComparatorIgnoringFields(IGNORED_FIELDS_LIST) - .containsExactlyElementsOf(expectedTraces); - assertIgnoredFields(actualTraces, expectedTraces); + assertThat(actualPage.page()).isEqualTo(page); + assertThat(actualPage.size()).isEqualTo(expectedTraces.size()); + assertThat(actualPage.total()).isEqualTo(expectedTraces.size()); + assertThat(actualTraces) + .usingRecursiveFieldByFieldElementComparatorIgnoringFields(IGNORED_FIELDS_LIST) + .containsExactlyElementsOf(expectedTraces); + assertIgnoredFields(actualTraces, expectedTraces); + + if (!unexpectedTraces.isEmpty()) { assertThat(actualTraces) .usingRecursiveFieldByFieldElementComparatorIgnoringFields(IGNORED_FIELDS_LIST) .doesNotContainAnyElementsOf(unexpectedTraces); } + } - private String toURLEncodedQueryParam(List filters) { - return URLEncoder.encode(JsonUtils.writeValueAsString(filters), StandardCharsets.UTF_8); - } + private String toURLEncodedQueryParam(List filters) { + return URLEncoder.encode(JsonUtils.writeValueAsString(filters), StandardCharsets.UTF_8); + } - private void assertIgnoredFields(List actualTraces, List expectedTraces) { - for (int i = 0; i < actualTraces.size(); i++) { - var actualTrace = actualTraces.get(i); - var expectedTrace = expectedTraces.get(i); - var expectedFeedbackScores = expectedTrace.feedbackScores() == null - ? null - : expectedTrace.feedbackScores().reversed(); - assertThat(actualTrace.projectId()).isNotNull(); - assertThat(actualTrace.projectName()).isNull(); - assertThat(actualTrace.createdAt()).isAfter(expectedTrace.createdAt()); - assertThat(actualTrace.lastUpdatedAt()).isAfter(expectedTrace.lastUpdatedAt()); - assertThat(actualTrace.lastUpdatedBy()).isEqualTo(USER); - assertThat(actualTrace.lastUpdatedBy()).isEqualTo(USER); - assertThat(actualTrace.feedbackScores()) - .usingRecursiveComparison( - RecursiveComparisonConfiguration.builder() - .withComparatorForType(BigDecimal::compareTo, BigDecimal.class) - .withIgnoredFields(IGNORED_FIELDS) - .build()) - .isEqualTo(expectedFeedbackScores); - - if (expectedTrace.feedbackScores() != null) { - actualTrace.feedbackScores().forEach(feedbackScore -> { - assertThat(feedbackScore.createdAt()).isAfter(expectedTrace.createdAt()); - assertThat(feedbackScore.lastUpdatedAt()).isAfter(expectedTrace.createdAt()); - assertThat(feedbackScore.lastUpdatedBy()).isEqualTo(USER); - assertThat(feedbackScore.lastUpdatedBy()).isEqualTo(USER); - }); - } + private void assertIgnoredFields(List actualTraces, List expectedTraces) { + for (int i = 0; i < actualTraces.size(); i++) { + var actualTrace = actualTraces.get(i); + var expectedTrace = expectedTraces.get(i); + var expectedFeedbackScores = expectedTrace.feedbackScores() == null + ? null + : expectedTrace.feedbackScores().reversed(); + assertThat(actualTrace.projectId()).isNotNull(); + assertThat(actualTrace.projectName()).isNull(); + assertThat(actualTrace.createdAt()).isAfter(expectedTrace.createdAt()); + assertThat(actualTrace.lastUpdatedAt()).isAfter(expectedTrace.lastUpdatedAt()); + assertThat(actualTrace.lastUpdatedBy()).isEqualTo(USER); + assertThat(actualTrace.lastUpdatedBy()).isEqualTo(USER); + assertThat(actualTrace.feedbackScores()) + .usingRecursiveComparison( + RecursiveComparisonConfiguration.builder() + .withComparatorForType(BigDecimal::compareTo, BigDecimal.class) + .withIgnoredFields(IGNORED_FIELDS) + .build()) + .isEqualTo(expectedFeedbackScores); + + if (expectedTrace.feedbackScores() != null) { + actualTrace.feedbackScores().forEach(feedbackScore -> { + assertThat(feedbackScore.createdAt()).isAfter(expectedTrace.createdAt()); + assertThat(feedbackScore.lastUpdatedAt()).isAfter(expectedTrace.createdAt()); + assertThat(feedbackScore.lastUpdatedBy()).isEqualTo(USER); + assertThat(feedbackScore.lastUpdatedBy()).isEqualTo(USER); + }); } } + } - private List updateFeedbackScore(List feedbackScores, int index, double val) { - feedbackScores.set(index, feedbackScores.get(index).toBuilder() - .value(BigDecimal.valueOf(val)) - .build()); - return feedbackScores; - } + private List updateFeedbackScore(List feedbackScores, int index, double val) { + feedbackScores.set(index, feedbackScores.get(index).toBuilder() + .value(BigDecimal.valueOf(val)) + .build()); + return feedbackScores; + } - private List updateFeedbackScore( - List destination, List source, int index) { - destination.set(index, source.get(index).toBuilder().build()); - return destination; - } + private List updateFeedbackScore( + List destination, List source, int index) { + destination.set(index, source.get(index).toBuilder().build()); + return destination; } @Nested @@ -2809,8 +2826,7 @@ private UUID create(Trace trace, String apiKey, String workspaceName) { assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(201); - return UUID.fromString(actualResponse.getHeaderString("Location") - .substring(actualResponse.getHeaderString("Location").lastIndexOf('/') + 1)); + return TestUtils.getIdFromLocation(actualResponse.getLocation()); } } @@ -2894,7 +2910,7 @@ void create() { assertThat(actualResponse.getHeaderString("Location")).matches(Pattern.compile(URL_PATTERN)); } - UUID projectId = getProjectId(client, trace.projectName(), TEST_WORKSPACE, API_KEY); + UUID projectId = getProjectId(trace.projectName(), TEST_WORKSPACE, API_KEY); getAndAssert(trace, id, projectId, now, API_KEY, TEST_WORKSPACE); } @@ -2922,8 +2938,8 @@ void create__whenCreatingTracesWithDifferentWorkspacesNames__thenReturnCreatedTr var createdTrace2 = Instant.now(); UUID id2 = TracesResourceTest.this.create(trace2, API_KEY, TEST_WORKSPACE); - UUID projectId1 = getProjectId(client, DEFAULT_PROJECT, TEST_WORKSPACE, API_KEY); - UUID projectId2 = getProjectId(client, projectName, TEST_WORKSPACE, API_KEY); + UUID projectId1 = getProjectId(DEFAULT_PROJECT, TEST_WORKSPACE, API_KEY); + UUID projectId2 = getProjectId(projectName, TEST_WORKSPACE, API_KEY); getAndAssert(trace1, id1, projectId1, createdTrace1, API_KEY, TEST_WORKSPACE); getAndAssert(trace2, id2, projectId2, createdTrace2, API_KEY, TEST_WORKSPACE); @@ -2954,10 +2970,9 @@ void create__whenIdComesFromClient__thenAcceptAndUseId() { assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(201); - String actualId = actualResponse.getLocation().toString() - .substring(actualResponse.getLocation().toString().lastIndexOf('/') + 1); + UUID actualId = TestUtils.getIdFromLocation(actualResponse.getLocation()); - assertThat(UUID.fromString(actualId)).isEqualTo(traceId); + assertThat(actualId).isEqualTo(traceId); } } @@ -3028,7 +3043,7 @@ void create__whenProjectNameIsNull__thenAcceptAndUseDefaultProject() { var actualResponse = getById(id, TEST_WORKSPACE, API_KEY); assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200); - UUID projectId = getProjectId(client, DEFAULT_PROJECT, TEST_WORKSPACE, API_KEY); + UUID projectId = getProjectId(DEFAULT_PROJECT, TEST_WORKSPACE, API_KEY); var actualEntity = actualResponse.readEntity(Trace.class); assertThat(actualEntity.projectId()).isEqualTo(projectId); @@ -3043,9 +3058,15 @@ class BatchInsert { @Test void batch__whenCreateTraces__thenReturnNoContent() { + + String projectName = UUID.randomUUID().toString(); + + UUID projectId = createProject(projectName, TEST_WORKSPACE, API_KEY); + var expectedTraces = IntStream.range(0, 1000) .mapToObj(i -> factory.manufacturePojo(Trace.class).toBuilder() - .projectId(null) + .projectName(projectName) + .projectId(projectId) .endTime(null) .feedbackScores(null) .build()) @@ -3053,11 +3074,8 @@ void batch__whenCreateTraces__thenReturnNoContent() { batchCreateAndAssert(expectedTraces, API_KEY, TEST_WORKSPACE); - expectedTraces.forEach(trace -> { - UUID projectId = getProjectId(client, trace.projectName(), TEST_WORKSPACE, API_KEY); - - getAndAssert(trace, trace.id(), projectId, Instant.now(), API_KEY, TEST_WORKSPACE); - }); + getAndAssertPage(TEST_WORKSPACE, projectName, List.of(), List.of(), expectedTraces.reversed(), List.of(), + API_KEY); } @Test @@ -3277,7 +3295,7 @@ void when__traceDoesNotExist__thenReturnCreateIt() { assertThat(actualEntity.metadata()).isEqualTo(traceUpdate.metadata()); assertThat(actualEntity.tags()).isEqualTo(traceUpdate.tags()); - UUID projectId = getProjectId(client, traceUpdate.projectName(), TEST_WORKSPACE, API_KEY); + UUID projectId = getProjectId(traceUpdate.projectName(), TEST_WORKSPACE, API_KEY); assertThat(actualEntity.name()).isEmpty(); assertThat(actualEntity.startTime()).isEqualTo(Instant.EPOCH); @@ -3495,7 +3513,7 @@ void update__whenTagsIsEmpty__thenAcceptUpdate() { runPatchAndAssertStatus(id, traceUpdate, API_KEY, TEST_WORKSPACE); - UUID projectId = getProjectId(client, trace.projectName(), TEST_WORKSPACE, API_KEY); + UUID projectId = getProjectId(trace.projectName(), TEST_WORKSPACE, API_KEY); Trace actualTrace = getAndAssert(trace, id, projectId, trace.createdAt().minusMillis(1), API_KEY, TEST_WORKSPACE); @@ -3516,7 +3534,7 @@ void update__whenMetadataIsEmpty__thenAcceptUpdate() { runPatchAndAssertStatus(id, traceUpdate, API_KEY, TEST_WORKSPACE); - UUID projectId = getProjectId(client, trace.projectName(), TEST_WORKSPACE, API_KEY); + UUID projectId = getProjectId(trace.projectName(), TEST_WORKSPACE, API_KEY); Trace actualTrace = getAndAssert(trace.toBuilder().metadata(metadata).build(), id, projectId, trace.createdAt().minusMillis(1), API_KEY, TEST_WORKSPACE); @@ -3537,7 +3555,7 @@ void update__whenInputIsEmpty__thenAcceptUpdate() { runPatchAndAssertStatus(id, traceUpdate, API_KEY, TEST_WORKSPACE); - UUID projectId = getProjectId(client, trace.projectName(), TEST_WORKSPACE, API_KEY); + UUID projectId = getProjectId(trace.projectName(), TEST_WORKSPACE, API_KEY); Trace actualTrace = getAndAssert(trace.toBuilder().input(input).build(), id, projectId, trace.createdAt().minusMillis(1), API_KEY, TEST_WORKSPACE); @@ -3558,7 +3576,7 @@ void update__whenOutputIsEmpty__thenAcceptUpdate() { runPatchAndAssertStatus(id, traceUpdate, API_KEY, TEST_WORKSPACE); - UUID projectId = getProjectId(client, trace.projectName(), TEST_WORKSPACE, API_KEY); + UUID projectId = getProjectId(trace.projectName(), TEST_WORKSPACE, API_KEY); Trace actualTrace = getAndAssert(trace.toBuilder().output(output).build(), id, projectId, trace.createdAt().minusMillis(1), API_KEY, TEST_WORKSPACE); @@ -3570,7 +3588,7 @@ void update__whenOutputIsEmpty__thenAcceptUpdate() { @DisplayName("when updating using projectId, then accept update") void update__whenUpdatingUsingProjectId__thenAcceptUpdate() { - var projectId = getProjectId(client, trace.projectName(), TEST_WORKSPACE, API_KEY); + var projectId = getProjectId(trace.projectName(), TEST_WORKSPACE, API_KEY); var traceUpdate = factory.manufacturePojo(TraceUpdate.class).toBuilder() .projectId(projectId) @@ -3698,7 +3716,7 @@ void feedback__whenFeedbackWithoutCategoryNameOrReason__thenReturnNoContent() { create(id, score, TEST_WORKSPACE, API_KEY); - UUID projectId = getProjectId(client, trace.projectName(), TEST_WORKSPACE, API_KEY); + UUID projectId = getProjectId(trace.projectName(), TEST_WORKSPACE, API_KEY); var actualEntity = getAndAssert(trace, id, projectId, now, API_KEY, TEST_WORKSPACE); @@ -3735,7 +3753,7 @@ void feedback__whenFeedbackWithCategoryNameOrReason__thenReturnNoContent() { create(id, score, TEST_WORKSPACE, API_KEY); - UUID projectId = getProjectId(client, trace.projectName(), TEST_WORKSPACE, API_KEY); + UUID projectId = getProjectId(trace.projectName(), TEST_WORKSPACE, API_KEY); Trace actualEntity = getAndAssert(trace, id, projectId, now, API_KEY, TEST_WORKSPACE); @@ -3773,7 +3791,7 @@ void feedback__whenOverridingFeedbackValue__thenReturnNoContent() { FeedbackScore newScore = score.toBuilder().value(BigDecimal.valueOf(2)).build(); create(id, newScore, TEST_WORKSPACE, API_KEY); - UUID projectId = getProjectId(client, trace.projectName(), TEST_WORKSPACE, API_KEY); + UUID projectId = getProjectId(trace.projectName(), TEST_WORKSPACE, API_KEY); var actualEntity = getAndAssert(trace, id, projectId, now, API_KEY, TEST_WORKSPACE); assertThat(actualEntity.feedbackScores()).hasSize(1); @@ -3966,8 +3984,8 @@ void feedback() { assertThat(actualResponse.hasEntity()).isFalse(); } - UUID projectId = getProjectId(client, trace.projectName(), TEST_WORKSPACE, API_KEY); - UUID projectId2 = getProjectId(client, trace2.projectName(), TEST_WORKSPACE, API_KEY); + UUID projectId = getProjectId(trace.projectName(), TEST_WORKSPACE, API_KEY); + UUID projectId2 = getProjectId(trace2.projectName(), TEST_WORKSPACE, API_KEY); var actualTrace1 = getAndAssert(trace, id, projectId, now, API_KEY, TEST_WORKSPACE); var actualTrace2 = getAndAssert(trace2, id2, projectId2, now, API_KEY, TEST_WORKSPACE); @@ -4037,8 +4055,8 @@ void feedback__whenWorkspaceIsSpecified__thenReturnNoContent() { assertThat(actualResponse.hasEntity()).isFalse(); } - UUID projectId = getProjectId(client, DEFAULT_PROJECT, workspaceName, apiKey); - UUID projectId2 = getProjectId(client, projectName, workspaceName, apiKey); + UUID projectId = getProjectId(DEFAULT_PROJECT, workspaceName, apiKey); + UUID projectId2 = getProjectId(projectName, workspaceName, apiKey); var actualTrace1 = getAndAssert(expectedTrace1, id, projectId, now, apiKey, workspaceName); var actualTrace2 = getAndAssert(expectedTrace2, id2, projectId2, now, apiKey, workspaceName); @@ -4107,7 +4125,7 @@ void feedback__whenFeedbackWithoutCategoryNameOrReason__thenReturnNoContent() { assertThat(actualResponse.hasEntity()).isFalse(); } - UUID projectId = getProjectId(client, trace.projectName(), TEST_WORKSPACE, API_KEY); + UUID projectId = getProjectId(trace.projectName(), TEST_WORKSPACE, API_KEY); var actualEntity = getAndAssert(trace, id, projectId, now, API_KEY, TEST_WORKSPACE); @@ -4158,7 +4176,7 @@ void feedback__whenFeedbackWithCategoryNameOrReason__thenReturnNoContent() { } var actualEntity = getAndAssert(expectedTrace, id, - getProjectId(client, expectedTrace.projectName(), TEST_WORKSPACE, API_KEY), now, API_KEY, + getProjectId(expectedTrace.projectName(), TEST_WORKSPACE, API_KEY), now, API_KEY, TEST_WORKSPACE); assertThat(actualEntity.feedbackScores()).hasSize(1); @@ -4220,7 +4238,7 @@ void feedback__whenOverridingFeedbackValue__thenReturnNoContent() { assertThat(actualResponse.hasEntity()).isFalse(); } - UUID projectId = getProjectId(client, trace.projectName(), TEST_WORKSPACE, API_KEY); + UUID projectId = getProjectId(trace.projectName(), TEST_WORKSPACE, API_KEY); var actualEntity = getAndAssert(trace, id, projectId, now, API_KEY, TEST_WORKSPACE); From bd67126e5079a18d6f541d5bca77375158c670c9 Mon Sep 17 00:00:00 2001 From: Thiago Hora Date: Fri, 13 Sep 2024 13:16:32 +0200 Subject: [PATCH 10/10] Fix tests --- .../utils/ConditionalGZipFilter.java | 1 - .../resources/v1/priv/SpansResourceTest.java | 245 ++++++++++-------- 2 files changed, 133 insertions(+), 113 deletions(-) diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/ConditionalGZipFilter.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/ConditionalGZipFilter.java index 24ce718593..8df64f0cfb 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/ConditionalGZipFilter.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/ConditionalGZipFilter.java @@ -1,6 +1,5 @@ package com.comet.opik.api.resources.utils; - import com.comet.opik.utils.JsonUtils; import jakarta.ws.rs.client.ClientRequestContext; import jakarta.ws.rs.client.ClientRequestFilter; 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 4c48b98619..60db8e83ce 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 @@ -161,7 +161,7 @@ private static void mockTargetWorkspace(String apiKey, String workspaceName, Str AuthTestUtils.mockTargetWorkspace(wireMock.server(), apiKey, workspaceName, workspaceId, USER); } - private UUID getProjectId(ClientSupport client, String projectName, String workspaceName, String apiKey) { + private UUID getProjectId(String projectName, String workspaceName, String apiKey) { return client.target("%s/v1/private/projects".formatted(baseURI)) .queryParam("name", projectName) .request() @@ -176,6 +176,19 @@ private UUID getProjectId(ClientSupport client, String projectName, String works .id(); } + private UUID createProject(String projectName, String workspaceName, String apiKey) { + try (Response response = client.target("%s/v1/private/projects".formatted(baseURI)) + .request() + .header(HttpHeaders.AUTHORIZATION, apiKey) + .header(WORKSPACE_HEADER, workspaceName) + .post(Entity.json(Project.builder().name(projectName).build()))) { + + assertThat(response.getStatusInfo().getStatusCode()).isEqualTo(201); + return UUID.fromString( + response.getLocation().getPath().substring(response.getLocation().getPath().lastIndexOf('/') + 1)); + } + } + @Nested @DisplayName("Api Key Authentication:") @TestInstance(TestInstance.Lifecycle.PER_CLASS) @@ -2959,122 +2972,124 @@ void getByProjectName__whenFilterInvalidValueOrKeyForFieldType__thenReturn400(Fi assertThat(actualError).isEqualTo(expectedError); } - private void getAndAssertPage( - String workspaceName, - String projectName, - List filters, - List spans, - List expectedSpans, - List unexpectedSpans, String apiKey) { - int page = 1; - int size = spans.size() + expectedSpans.size() + unexpectedSpans.size(); - getAndAssertPage( - workspaceName, - projectName, - null, - null, - null, - filters, - page, - size, - expectedSpans, - expectedSpans.size(), - unexpectedSpans, apiKey); + private List updateFeedbackScore(List feedbackScores, int index, double val) { + feedbackScores.set(index, feedbackScores.get(index).toBuilder() + .value(BigDecimal.valueOf(val)) + .build()); + return feedbackScores; } - private void getAndAssertPage( - String workspaceName, - String projectName, - UUID projectId, - UUID traceId, - SpanType type, - List filters, - int page, - int size, - List expectedSpans, - int expectedTotal, - List unexpectedSpans, String apiKey) { - try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) - .queryParam("page", page) - .queryParam("size", size) - .queryParam("project_name", projectName) - .queryParam("project_id", projectId) - .queryParam("trace_id", traceId) - .queryParam("type", type) - .queryParam("filters", toURLEncodedQueryParam(filters)) - .request() - .header(HttpHeaders.AUTHORIZATION, apiKey) - .header(WORKSPACE_HEADER, workspaceName) - .get()) { - var actualPage = actualResponse.readEntity(Span.SpanPage.class); - var actualSpans = actualPage.content(); + private List updateFeedbackScore( + List destination, List source, int index) { + destination.set(index, source.get(index).toBuilder().build()); + return destination; + } + } + + private void getAndAssertPage( + String workspaceName, + String projectName, + List filters, + List spans, + List expectedSpans, + List unexpectedSpans, String apiKey) { + int page = 1; + int size = spans.size() + expectedSpans.size() + unexpectedSpans.size(); + getAndAssertPage( + workspaceName, + projectName, + null, + null, + null, + filters, + page, + size, + expectedSpans, + expectedSpans.size(), + unexpectedSpans, apiKey); + } + + private void getAndAssertPage( + String workspaceName, + String projectName, + UUID projectId, + UUID traceId, + SpanType type, + List filters, + int page, + int size, + List expectedSpans, + int expectedTotal, + List unexpectedSpans, String apiKey) { + try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) + .queryParam("page", page) + .queryParam("size", size) + .queryParam("project_name", projectName) + .queryParam("project_id", projectId) + .queryParam("trace_id", traceId) + .queryParam("type", type) + .queryParam("filters", toURLEncodedQueryParam(filters)) + .request() + .header(HttpHeaders.AUTHORIZATION, apiKey) + .header(WORKSPACE_HEADER, workspaceName) + .get()) { + var actualPage = actualResponse.readEntity(Span.SpanPage.class); + var actualSpans = actualPage.content(); - assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200); + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200); - assertThat(actualPage.page()).isEqualTo(page); - assertThat(actualPage.size()).isEqualTo(expectedSpans.size()); - assertThat(actualPage.total()).isEqualTo(expectedTotal); + assertThat(actualPage.page()).isEqualTo(page); + assertThat(actualPage.size()).isEqualTo(expectedSpans.size()); + assertThat(actualPage.total()).isEqualTo(expectedTotal); - assertThat(actualSpans.size()).isEqualTo(expectedSpans.size()); - assertThat(actualSpans) - .usingRecursiveFieldByFieldElementComparatorIgnoringFields(IGNORED_FIELDS) - .containsExactlyElementsOf(expectedSpans); - assertIgnoredFields(actualSpans, expectedSpans); + assertThat(actualSpans.size()).isEqualTo(expectedSpans.size()); + assertThat(actualSpans) + .usingRecursiveFieldByFieldElementComparatorIgnoringFields(IGNORED_FIELDS) + .containsExactlyElementsOf(expectedSpans); + assertIgnoredFields(actualSpans, expectedSpans); + if (!unexpectedSpans.isEmpty()) { assertThat(actualSpans) .usingRecursiveFieldByFieldElementComparatorIgnoringFields(IGNORED_FIELDS) .doesNotContainAnyElementsOf(unexpectedSpans); } } + } - private String toURLEncodedQueryParam(List filters) { - return CollectionUtils.isEmpty(filters) - ? null - : URLEncoder.encode(JsonUtils.writeValueAsString(filters), StandardCharsets.UTF_8); - } + private String toURLEncodedQueryParam(List filters) { + return CollectionUtils.isEmpty(filters) + ? null + : URLEncoder.encode(JsonUtils.writeValueAsString(filters), StandardCharsets.UTF_8); + } - private void assertIgnoredFields(List actualSpans, List expectedSpans) { - for (int i = 0; i < actualSpans.size(); i++) { - var actualSpan = actualSpans.get(i); - var expectedSpan = expectedSpans.get(i); - var expectedFeedbackScores = expectedSpan.feedbackScores() == null - ? null - : expectedSpan.feedbackScores().reversed(); - assertThat(actualSpan.projectId()).isNotNull(); - assertThat(actualSpan.projectName()).isNull(); - assertThat(actualSpan.createdAt()).isAfter(expectedSpan.createdAt()); - assertThat(actualSpan.lastUpdatedAt()).isAfter(expectedSpan.lastUpdatedAt()); - assertThat(actualSpan.feedbackScores()) - .usingRecursiveComparison( - RecursiveComparisonConfiguration.builder() - .withComparatorForType(BigDecimal::compareTo, BigDecimal.class) - .withIgnoredFields(IGNORED_FIELDS_SCORES) - .build()) - .isEqualTo(expectedFeedbackScores); - - if (actualSpan.feedbackScores() != null) { - actualSpan.feedbackScores().forEach(feedbackScore -> { - assertThat(feedbackScore.createdAt()).isAfter(expectedSpan.createdAt()); - assertThat(feedbackScore.lastUpdatedAt()).isAfter(expectedSpan.lastUpdatedAt()); - assertThat(feedbackScore.createdBy()).isEqualTo(USER); - assertThat(feedbackScore.lastUpdatedBy()).isEqualTo(USER); - }); - } + private void assertIgnoredFields(List actualSpans, List expectedSpans) { + for (int i = 0; i < actualSpans.size(); i++) { + var actualSpan = actualSpans.get(i); + var expectedSpan = expectedSpans.get(i); + var expectedFeedbackScores = expectedSpan.feedbackScores() == null + ? null + : expectedSpan.feedbackScores().reversed(); + assertThat(actualSpan.projectId()).isNotNull(); + assertThat(actualSpan.projectName()).isNull(); + assertThat(actualSpan.createdAt()).isAfter(expectedSpan.createdAt()); + assertThat(actualSpan.lastUpdatedAt()).isAfter(expectedSpan.lastUpdatedAt()); + assertThat(actualSpan.feedbackScores()) + .usingRecursiveComparison( + RecursiveComparisonConfiguration.builder() + .withComparatorForType(BigDecimal::compareTo, BigDecimal.class) + .withIgnoredFields(IGNORED_FIELDS_SCORES) + .build()) + .isEqualTo(expectedFeedbackScores); + + if (actualSpan.feedbackScores() != null) { + actualSpan.feedbackScores().forEach(feedbackScore -> { + assertThat(feedbackScore.createdAt()).isAfter(expectedSpan.createdAt()); + assertThat(feedbackScore.lastUpdatedAt()).isAfter(expectedSpan.lastUpdatedAt()); + assertThat(feedbackScore.createdBy()).isEqualTo(USER); + assertThat(feedbackScore.lastUpdatedBy()).isEqualTo(USER); + }); } } - - private List updateFeedbackScore(List feedbackScores, int index, double val) { - feedbackScores.set(index, feedbackScores.get(index).toBuilder() - .value(BigDecimal.valueOf(val)) - .build()); - return feedbackScores; - } - - private List updateFeedbackScore( - List destination, List source, int index) { - destination.set(index, source.get(index).toBuilder().build()); - return destination; - } } private UUID createAndAssert(Span expectedSpan, String apiKey, String workspaceName) { @@ -3197,9 +3212,14 @@ class BatchInsert { @Test void batch__whenCreateSpans__thenReturnNoContent() { + + String projectName = UUID.randomUUID().toString(); + UUID projectId = createProject(projectName, TEST_WORKSPACE, API_KEY); + var expectedSpans = IntStream.range(0, 1000) .mapToObj(i -> podamFactory.manufacturePojo(Span.class).toBuilder() - .projectId(null) + .projectId(projectId) + .projectName(projectName) .parentSpanId(null) .feedbackScores(null) .build()) @@ -3207,7 +3227,8 @@ void batch__whenCreateSpans__thenReturnNoContent() { batchCreateAndAssert(expectedSpans, API_KEY, TEST_WORKSPACE); - expectedSpans.forEach(expectedSpan -> getAndAssert(expectedSpan, API_KEY, TEST_WORKSPACE)); + getAndAssertPage(TEST_WORKSPACE, projectName, List.of(), List.of(), expectedSpans.reversed(), List.of(), + API_KEY); } @Test @@ -3508,7 +3529,7 @@ void when__spanDoesNotExist__thenReturnCreateIt() { var actualResponse = getById(id, TEST_WORKSPACE, API_KEY); - var projectId = getProjectId(client, spanUpdate.projectName(), TEST_WORKSPACE, API_KEY); + var projectId = getProjectId(spanUpdate.projectName(), TEST_WORKSPACE, API_KEY); var actualEntity = actualResponse.readEntity(Span.class); assertThat(actualEntity.id()).isEqualTo(id); @@ -3552,7 +3573,7 @@ void when__spanUpdateAndInsertAreProcessedOutOfOther__thenReturnSpan() { var actualResponse = getById(id, TEST_WORKSPACE, API_KEY); - var projectId = getProjectId(client, spanUpdate.projectName(), TEST_WORKSPACE, API_KEY); + var projectId = getProjectId(spanUpdate.projectName(), TEST_WORKSPACE, API_KEY); var actualEntity = actualResponse.readEntity(Span.class); assertThat(actualEntity.id()).isEqualTo(id); @@ -3760,7 +3781,7 @@ void when__multipleSpanUpdateAndInsertAreProcessedOutOfOtherAndConcurrent__thenR var actualEntity = actualResponse.readEntity(Span.class); assertThat(actualEntity.id()).isEqualTo(id); - var projectId = getProjectId(client, projectName, TEST_WORKSPACE, API_KEY); + var projectId = getProjectId(projectName, TEST_WORKSPACE, API_KEY); assertThat(actualEntity.projectId()).isEqualTo(projectId); assertThat(actualEntity.traceId()).isEqualTo(spanUpdate1.traceId()); @@ -3799,7 +3820,7 @@ void update__whenTagsIsEmpty__thenAcceptUpdate() { runPatchAndAssertStatus(expectedSpan.id(), spanUpdate, API_KEY, TEST_WORKSPACE); - UUID projectId = getProjectId(client, spanUpdate.projectName(), TEST_WORKSPACE, API_KEY); + UUID projectId = getProjectId(spanUpdate.projectName(), TEST_WORKSPACE, API_KEY); Span updatedSpan = expectedSpan.toBuilder() .tags(spanUpdate.tags()) @@ -3832,7 +3853,7 @@ void update__whenMetadataIsEmpty__thenAcceptUpdate() { runPatchAndAssertStatus(expectedSpan.id(), spanUpdate, API_KEY, TEST_WORKSPACE); - UUID projectId = getProjectId(client, spanUpdate.projectName(), TEST_WORKSPACE, API_KEY); + UUID projectId = getProjectId(spanUpdate.projectName(), TEST_WORKSPACE, API_KEY); Span updatedSpan = expectedSpan.toBuilder() .metadata(metadata) @@ -3865,7 +3886,7 @@ void update__whenInputIsEmpty__thenAcceptUpdate() { runPatchAndAssertStatus(expectedSpan.id(), spanUpdate, API_KEY, TEST_WORKSPACE); - UUID projectId = getProjectId(client, spanUpdate.projectName(), TEST_WORKSPACE, API_KEY); + UUID projectId = getProjectId(spanUpdate.projectName(), TEST_WORKSPACE, API_KEY); Span updatedSpan = expectedSpan.toBuilder() .input(input) @@ -3897,7 +3918,7 @@ void update__whenOutputIsEmpty__thenAcceptUpdate() { runPatchAndAssertStatus(expectedSpan.id(), spanUpdate, API_KEY, TEST_WORKSPACE); - UUID projectId = getProjectId(client, spanUpdate.projectName(), TEST_WORKSPACE, API_KEY); + UUID projectId = getProjectId(spanUpdate.projectName(), TEST_WORKSPACE, API_KEY); Span updatedSpan = expectedSpan.toBuilder() .output(output) @@ -3919,7 +3940,7 @@ void update__whenUpdatingUsingProjectId__thenAcceptUpdate() { createAndAssert(expectedSpan, API_KEY, TEST_WORKSPACE); - var projectId = getProjectId(client, expectedSpan.projectName(), TEST_WORKSPACE, API_KEY); + var projectId = getProjectId(expectedSpan.projectName(), TEST_WORKSPACE, API_KEY); var spanUpdate = podamFactory.manufacturePojo(SpanUpdate.class).toBuilder() .traceId(expectedSpan.traceId())