Skip to content

Commit 0d3b353

Browse files
thiagohorathiagohora
committed
Merge branch 'thiaghora/OPIK-7402-uuid-validation-audit-mode' of https://github.com/comet-ml/opik into thiaghora/OPIK-7402-uuid-validation-audit-mode
2 parents c45d938 + 3ef46ad commit 0d3b353

48 files changed

Lines changed: 2629 additions & 94 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎apps/opik-backend/src/main/java/com/comet/opik/domain/AlertService.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -555,6 +555,8 @@ private Alert prepareAlert(Alert alert, String userName, String workspaceId) {
555555

556556
UUID id = alert.id() == null ? idGenerator.generateId() : alert.id();
557557
IdGenerator.validateVersion(id, "Alert");
558+
// projectId is persisted without an existence check here, so enforce v7 to avoid storing an orphan v4.
559+
idGenerator.validateIdNotInFutureIfPresent(alert.projectId(), "project");
558560

559561
UUID webhookId = alert.webhook().id() == null ? idGenerator.generateId() : alert.webhook().id();
560562
IdGenerator.validateVersion(webhookId, "Webhook");

‎apps/opik-backend/src/main/java/com/comet/opik/domain/AnnotationQueueService.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,10 @@ public Mono<Long> addItems(@NonNull UUID queueId, @NonNull Set<UUID> itemIds) {
134134
return Mono.just(0L);
135135
}
136136

137+
// Queue items reference trace/thread ids (v7 by construction); enforce so the referenced-id
138+
// policy is uniform. Past allowed — queues commonly collect older traces/threads.
139+
itemIds.forEach(itemId -> idGenerator.validateIdNotInFuture(itemId, "AnnotationQueue item"));
140+
137141
return annotationQueueDAO.findQueueInfoById(queueId)
138142
.switchIfEmpty(Mono.error(createNotFoundError(queueId)))
139143
.flatMap(queue -> annotationQueueDAO.addItems(queueId, itemIds, queue.projectId()))
@@ -257,6 +261,8 @@ private Mono<AnnotationQueue.AnnotationQueuePage> enhancePageWithProjectNames(
257261
private AnnotationQueue prepareAnnotationQueue(AnnotationQueue annotationQueue) {
258262
UUID id = annotationQueue.id() == null ? idGenerator.generateId() : annotationQueue.id();
259263
IdGenerator.validateVersion(id, "AnnotationQueue");
264+
// projectId is persisted without an existence check here, so enforce v7 to avoid storing an orphan v4.
265+
idGenerator.validateIdNotInFutureIfPresent(annotationQueue.projectId(), "project");
260266

261267
log.debug("Preparing annotation queue with id '{}', name '{}', project '{}'",
262268
id, annotationQueue.name(), annotationQueue.projectId());

‎apps/opik-backend/src/main/java/com/comet/opik/domain/AssertionResultService.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ class AssertionResultServiceImpl implements AssertionResultService {
4343
private final @NonNull AssertionResultDAO assertionResultDAO;
4444
private final @NonNull ProjectService projectService;
4545
private final @NonNull EventBus eventBus;
46+
private final @NonNull IdGenerator idGenerator;
4647

4748
@Override
4849
public Mono<Long> insertBatch(@NonNull EntityType entityType,
@@ -63,7 +64,7 @@ public Mono<Void> saveBatch(@NonNull EntityType entityType,
6364
}
6465

6566
// Validate up front so a bad id fails fast and independently of project-name normalisation.
66-
assertionResults.forEach(item -> IdGenerator.validateVersion(item.entityId(), entityType.getType()));
67+
assertionResults.forEach(item -> idGenerator.validateIdNotInFuture(item.entityId(), entityType.getType()));
6768

6869
return Mono.deferContextual(ctx -> {
6970
String workspaceId = ctx.get(RequestContext.WORKSPACE_ID);

‎apps/opik-backend/src/main/java/com/comet/opik/domain/CommentService.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ class CommentServiceImpl implements CommentService {
5353

5454
@Override
5555
public Mono<UUID> create(@NonNull UUID entityId, @NonNull Comment comment, CommentDAO.EntityType entityType) {
56+
idGenerator.validateIdNotInFuture(entityId, entityType.getType());
5657
UUID id = idGenerator.generateId();
5758
var monoProjectId = resolveProjectId(entityType, entityId);
5859

‎apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetItemService.java‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -187,6 +187,8 @@ public Mono<Void> createFromTraces(
187187

188188
log.info("Creating dataset items from '{}' traces for dataset '{}'", traceIds.size(), datasetId);
189189

190+
traceIds.forEach(traceId -> idGenerator.validateIdNotInFuture(traceId, "dataset_item trace"));
191+
190192
// Verify dataset exists
191193
return Mono.deferContextual(ctx -> {
192194
String workspaceId = ctx.get(RequestContext.WORKSPACE_ID);
@@ -240,6 +242,8 @@ public Mono<Void> createFromSpans(
240242

241243
log.info("Creating dataset items from '{}' spans for dataset '{}'", spanIds.size(), datasetId);
242244

245+
spanIds.forEach(spanId -> idGenerator.validateIdNotInFuture(spanId, "dataset_item span"));
246+
243247
// Verify dataset exists
244248
return Mono.deferContextual(ctx -> {
245249
String workspaceId = ctx.get(RequestContext.WORKSPACE_ID);
@@ -382,6 +386,7 @@ private Mono<DatasetItem> authorizeItem(Mono<DatasetItem> itemMono) {
382386
@Override
383387
@WithSpan
384388
public Mono<Void> patch(@NonNull UUID id, @NonNull DatasetItem item) {
389+
validateReferencedTraceAndSpan(item);
385390
return Mono.deferContextual(ctx -> {
386391
String workspaceId = ctx.get(RequestContext.WORKSPACE_ID);
387392
String userName = ctx.get(RequestContext.USER_NAME);
@@ -903,11 +908,23 @@ private List<DatasetItem> addIdIfAbsent(DatasetItemBatch batch) {
903908
.stream()
904909
.map(item -> {
905910
IdGenerator.validateVersion(item.id(), "dataset_item");
911+
validateReferencedTraceAndSpan(item);
906912
return item;
907913
})
908914
.toList();
909915
}
910916

917+
// The dataset_item's referenced trace_id / span_id must be a time-ordered UUIDv7 (past allowed:
918+
// items are commonly linked to older traces/spans). Reuses the shared referenced-id policy.
919+
private void validateReferencedTraceAndSpan(DatasetItem item) {
920+
if (item.traceId() != null) {
921+
idGenerator.validateIdNotInFuture(item.traceId(), "dataset_item trace");
922+
}
923+
if (item.spanId() != null) {
924+
idGenerator.validateIdNotInFuture(item.spanId(), "dataset_item span");
925+
}
926+
}
927+
911928
private <T> Mono<T> failWithConflict(String message) {
912929
log.info(message);
913930
return Mono.error(new IdentifierMismatchException(new ErrorMessage(List.of(message))));

‎apps/opik-backend/src/main/java/com/comet/opik/domain/ExperimentService.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -495,6 +495,8 @@ private Set<UUID> getPromptVersionIds(Experiment experiment) {
495495
public Mono<UUID> create(@NonNull Experiment experiment) {
496496
var id = experiment.id() == null ? idGenerator.generateId() : experiment.id();
497497
IdGenerator.validateVersion(id, "Experiment");
498+
// optimizationId is stored without an existence check, so enforce v7 to avoid storing an orphan v4.
499+
idGenerator.validateIdNotInFutureIfPresent(experiment.optimizationId(), "optimization");
498500
var name = StringUtils.getIfBlank(experiment.name(), nameGenerator::generateName);
499501
return resolveProjectId(experiment)
500502
.flatMap(resolvedExperiment -> datasetService

‎apps/opik-backend/src/main/java/com/comet/opik/domain/FeedbackScoreService.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,7 @@ class FeedbackScoreServiceImpl implements FeedbackScoreService {
8989
private final @NonNull TraceThreadService traceThreadService;
9090
private final @NonNull Provider<RequestContext> requestContext;
9191
private final @NonNull EventBus eventBus;
92+
private final @NonNull IdGenerator idGenerator;
9293

9394
@Builder(toBuilder = true)
9495
record ProjectDto<T extends FeedbackScoreItem>(Project project, List<T> scores) {
@@ -100,6 +101,8 @@ public Mono<Void> scoreTrace(@NonNull UUID traceId, @NonNull FeedbackScore score
100101
String workspaceId = ctx.get(RequestContext.WORKSPACE_ID);
101102
String userName = ctx.get(RequestContext.USER_NAME);
102103

104+
idGenerator.validateIdNotInFuture(traceId, EntityType.TRACE.getType());
105+
idGenerator.validateIdNotInFutureIfPresent(score.sourceQueueId(), "annotation queue");
103106
return traceDAO.getProjectIdFromTrace(traceId)
104107
.switchIfEmpty(Mono.error(failWithNotFound("Trace", traceId)))
105108
.flatMap(projectId -> getAuthor()
@@ -118,6 +121,8 @@ public Mono<Void> scoreSpan(@NonNull UUID spanId, @NonNull FeedbackScore score)
118121
String workspaceId = ctx.get(RequestContext.WORKSPACE_ID);
119122
String userName = ctx.get(RequestContext.USER_NAME);
120123

124+
idGenerator.validateIdNotInFuture(spanId, EntityType.SPAN.getType());
125+
idGenerator.validateIdNotInFutureIfPresent(score.sourceQueueId(), "annotation queue");
121126
return spanDAO.getProjectIdFromSpan(spanId)
122127
.switchIfEmpty(Mono.error(failWithNotFound("Span", spanId)))
123128
.flatMap(projectId -> getAuthor()
@@ -173,7 +178,8 @@ private Mono<Void> processScoreBatch(EntityType entityType, List<FeedbackScoreBa
173178
Map<String, List<FeedbackScoreItem>> scoresPerProject = scores
174179
.stream()
175180
.map(score -> {
176-
IdGenerator.validateVersion(score.id(), entityType.getType()); // validate span/trace id
181+
idGenerator.validateIdNotInFuture(score.id(), entityType.getType()); // validate span/trace id
182+
idGenerator.validateIdNotInFutureIfPresent(score.sourceQueueId(), "annotation queue");
177183

178184
return score.toBuilder()
179185
.projectName(WorkspaceUtils.getProjectName(score.projectName()))

‎apps/opik-backend/src/main/java/com/comet/opik/domain/GuardrailsService.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,8 @@ public Mono<Void> addTraceGuardrails(List<Guardrail> guardrails) {
5151
.stream()
5252
.map(guardrail -> {
5353
UUID id = idGenerator.generateId();
54-
IdGenerator.validateVersion(guardrail.entityId(), entityType.getType()); // validate trace id
54+
idGenerator.validateIdNotInFuture(guardrail.entityId(), entityType.getType());
55+
idGenerator.validateIdNotInFuture(guardrail.secondaryId(), "guardrail secondary");
5556

5657
return guardrail.toBuilder()
5758
.id(id)

‎apps/opik-backend/src/main/java/com/comet/opik/domain/IdGenerator.java‎

Lines changed: 36 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -38,12 +38,27 @@ public interface IdGenerator {
3838
Mono<UUID> validateIdAsync(UUID id, String resource);
3939

4040
/**
41-
* Validates an ingested {@code id} on the update path: it must be a version 7 UUID
42-
* ({@link #validateVersion(UUID, String)}) and must not embed a timestamp far in the future (which
43-
* would corrupt the partition layout). Unlike {@link #validateId}, old ids are allowed, because
44-
* updating a long-lived entity (e.g. created months ago) is a legitimate operation.
41+
* Validates an {@code id} that may legitimately point at an entity created in the past: it must be a
42+
* version 7 UUID ({@link #validateVersion(UUID, String)}) and must not embed a timestamp far in the
43+
* future (which would corrupt the partition layout / retention id-range). Unlike {@link #validateId},
44+
* old ids are allowed.
45+
*
46+
* <p>Used both on the update path (updating a long-lived entity created months ago is legitimate) and
47+
* for referenced/foreign ids on ingest (e.g. a span's {@code traceId}: retention orders spans by the
48+
* {@code trace_id} range assuming it is a time-ordered UUIDv7, and late spans on old traces are common,
49+
* so old is fine but non-v7 or future-dated must be rejected).
4550
*/
46-
Mono<UUID> validateIdForUpdateAsync(UUID id, String resource);
51+
void validateIdNotInFuture(UUID id, String resource);
52+
53+
Mono<UUID> validateIdNotInFutureAsync(UUID id, String resource);
54+
55+
/**
56+
* Null-safe variant of {@link #validateIdNotInFuture} for optional referenced ids (e.g. an optional
57+
* {@code projectId} that may be resolved by name instead). No-op when {@code id} is null.
58+
*/
59+
void validateIdNotInFutureIfPresent(UUID id, String resource);
60+
61+
Mono<UUID> validateIdNotInFutureIfPresentAsync(UUID id, String resource);
4762

4863
static Mono<UUID> validateVersionAsync(@NonNull UUID id, String resource) {
4964
return Mono.fromCallable(() -> {
@@ -95,15 +110,17 @@ public Mono<UUID> validateIdAsync(@NonNull UUID id, String resource) {
95110
}));
96111
}
97112

98-
private void validateIdForUpdate(UUID id, String resource, String workspaceId) {
113+
@Override
114+
115+
public void validateIdForUpdate(@NonNull UUID id, String resource, String workspaceId) {
99116
IdGenerator.validateVersion(id, resource);
100117
uuidV7TimestampValidator.validateNotInFuture(id, resource, workspaceId);
101118
}
102119

103120
@Override
104121
public Mono<UUID> validateIdForUpdateAsync(@NonNull UUID id, String resource) {
105122
return Mono.deferContextual(ctx -> Mono.fromCallable(() -> {
106-
validateIdForUpdate(id, resource, workspaceId(ctx));
123+
validateIdNotInFuture(id, resource, workspaceId(ctx));
107124
return id;
108125
}));
109126
}
@@ -116,4 +133,16 @@ public Mono<UUID> validateIdForUpdateAsync(@NonNull UUID id, String resource) {
116133
private static String workspaceId(ContextView ctx) {
117134
return ctx.getOrDefault(RequestContext.WORKSPACE_ID, ErrorMetricsResolver.UNKNOWN);
118135
}
136+
137+
@Override
138+
public void validateIdNotInFutureIfPresent(UUID id, String resource) {
139+
if (id != null) {
140+
validateIdNotInFuture(id, resource);
141+
}
142+
}
143+
144+
@Override
145+
public Mono<UUID> validateIdNotInFutureIfPresentAsync(UUID id, String resource) {
146+
return id == null ? Mono.empty() : validateIdNotInFutureAsync(id, resource);
147+
}
119148
}

‎apps/opik-backend/src/main/java/com/comet/opik/domain/SpanService.java‎

Lines changed: 27 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,8 @@ public class SpanService {
6262
public static final String PARENT_SPAN_IS_MISMATCH = "parent_span_id does not match the existing span";
6363
public static final String TRACE_ID_MISMATCH = "trace_id does not match the existing span";
6464
public static final String SPAN_KEY = "Span";
65+
public static final String SPAN_TRACE_KEY = "Span trace";
66+
public static final String SPAN_PARENT_KEY = "Span parent";
6567
public static final String PROJECT_AND_WORKSPACE_NAME_MISMATCH = "Project name and workspace name do not match the existing span";
6668

6769
private final @NonNull SpanDAO spanDAO;
@@ -159,6 +161,7 @@ public Mono<UUID> create(@NonNull Span span) {
159161
var projectName = WorkspaceUtils.getProjectName(span.projectName());
160162
return idGenerator
161163
.validateIdAsync(id, SPAN_KEY)
164+
.then(Mono.fromRunnable(() -> validateSpanReferences(span.traceId(), span.parentSpanId())))
162165
.then(projectService.getOrCreate(projectName))
163166
.flatMap(project -> lockService.executeWithLock(
164167
new LockService.Lock(id, SPAN_KEY),
@@ -220,7 +223,9 @@ public Mono<Void> update(@NonNull UUID id, @NonNull SpanUpdate spanUpdate) {
220223
String userName = ctx.get(RequestContext.USER_NAME);
221224

222225
return idGenerator
223-
.validateIdForUpdateAsync(id, SPAN_KEY)
226+
.validateIdNotInFutureAsync(id, SPAN_KEY)
227+
.then(Mono.fromRunnable(
228+
() -> validateSpanReferences(spanUpdate.traceId(), spanUpdate.parentSpanId())))
224229
.then(Mono.defer(() -> getProjectById(spanUpdate)
225230
.switchIfEmpty(Mono.defer(() -> projectService.getOrCreate(projectName)))
226231
.subscribeOn(Schedulers.boundedElastic()))
@@ -247,7 +252,10 @@ public Mono<Void> batchUpdate(@NonNull SpanBatchUpdate batchUpdate) {
247252
String workspaceId = ctx.get(RequestContext.WORKSPACE_ID);
248253
String userName = ctx.get(RequestContext.USER_NAME);
249254

250-
return spanDAO.bulkUpdate(batchUpdate.ids(), batchUpdate.update(), mergeTags)
255+
return Mono
256+
.fromRunnable(() -> validateSpanReferences(batchUpdate.update().traceId(),
257+
batchUpdate.update().parentSpanId()))
258+
.then(spanDAO.bulkUpdate(batchUpdate.ids(), batchUpdate.update(), mergeTags))
251259
.onErrorResume(TagOperations::mapTagLimitError)
252260
.doOnSuccess(__ -> {
253261
log.info("Completed batch update for '{}' spans", batchUpdate.ids().size());
@@ -367,6 +375,15 @@ public Mono<Long> create(@NonNull SpanBatch batch) {
367375

368376
List<Span> dedupedSpans = dedupSpans(batch.spans());
369377

378+
// Fail fast on invalid ids BEFORE any side effect below (auto-stripped attachment deletion, project
379+
// creation), so a rejected batch never mutates state.
380+
dedupedSpans.forEach(span -> {
381+
if (span.id() != null) {
382+
idGenerator.validateId(span.id(), SPAN_KEY);
383+
}
384+
validateSpanReferences(span.traceId(), span.parentSpanId());
385+
});
386+
370387
List<String> projectNames = dedupedSpans
371388
.stream()
372389
.map(Span::projectName)
@@ -438,6 +455,13 @@ private List<Span> dedupSpans(List<Span> initialSpans) {
438455
return result;
439456
}
440457

458+
// Shared span reference-id policy: the trace (required) and parent (optional) must be time-ordered
459+
// UUIDv7, past allowed. Used by every span write path so the rules can't drift between them.
460+
private void validateSpanReferences(UUID traceId, UUID parentSpanId) {
461+
idGenerator.validateIdNotInFuture(traceId, SPAN_TRACE_KEY);
462+
idGenerator.validateIdNotInFutureIfPresent(parentSpanId, SPAN_PARENT_KEY);
463+
}
464+
441465
private List<Span> bindSpanToProjectAndId(List<Span> spans, List<Project> projects, String workspaceId) {
442466
Map<String, Project> projectPerName = projects.stream()
443467
.collect(Collectors.toMap(
@@ -459,6 +483,7 @@ private List<Span> bindSpanToProjectAndId(List<Span> spans, List<Project> projec
459483
}
460484

461485
UUID id = span.id() == null ? idGenerator.generateId() : span.id();
486+
462487
idGenerator.validateId(id, SPAN_KEY, workspaceId);
463488

464489
return span.toBuilder().id(id).projectId(project.id()).build();

0 commit comments

Comments
 (0)