diff --git a/apps/opik-backend/config.yml b/apps/opik-backend/config.yml index bba61f5d02b..7d78715dd5f 100644 --- a/apps/opik-backend/config.yml +++ b/apps/opik-backend/config.yml @@ -155,6 +155,20 @@ databaseAnalyticsDataModel: # MODIFY TTL target `traces_local` only (the Distributed `traces` rejects them); ADD/DROP/MODIFY COLUMN must target # both `traces_local` and `traces`, else reads can't see the column (code 47). tracesDistributedWrapEnabled: ${ANALYTICS_DB_DATA_MODEL_TRACES_DISTRIBUTED_WRAP_ENABLED:-false} + # Default: false + # Description: Enables partition-aware PRUNING of trace deletes - it does NOT create or activate any partitioning. + # With it on, a trace DELETE bounds itself to the weekly partitions its own ids resolve to instead of being planned + # against every part of the table. Turning it on therefore ASSERTS a schema fact rather than causing one: that the + # live mutation target already IS the weekly partitioned successor, id_at as DateTime64(0,'UTC') under + # PARTITION BY toYYYYMMDD(toDate32(id_at) - toIntervalDay(toDayOfWeek(id_at, 1))). Installing that schema is the + # EXCHANGE step of the cutover, never this flag. Purely an optimisation: false keeps the unbounded mutation, which + # is always correct and merely slower. A third flag on purpose - the partitioning appears at the EXCHANGE, and + # neither sibling marks it: + # traceColumnsNonNullable must be rolled out BEFORE the EXCHANGE, tracesDistributedWrapEnabled flips at the wrap, + # which may be deferred long after it. Leave false at deploy time; set true once the EXCHANGE is confirmed, and back + # to false BEFORE a rollback promotes the original `traces` (legacy `traces` has no PARTITION BY and a 32-bit + # DateTime id_at that overflows past 2106, so the predicate would silently match zero rows for a far-future id). + tracesWeeklyPartitionPruningEnabled: ${ANALYTICS_DB_DATA_MODEL_TRACES_WEEKLY_PARTITION_PRUNING_ENABLED:-false} # Description: UUIDv7 ingestion validation. Rejects writes whose `id` embeds a timestamp outside the # window, protecting data quality. diff --git a/apps/opik-backend/data-migrations/traces-local-v2-cutover/README.md b/apps/opik-backend/data-migrations/traces-local-v2-cutover/README.md index bb732b08d2c..83f434ba578 100644 --- a/apps/opik-backend/data-migrations/traces-local-v2-cutover/README.md +++ b/apps/opik-backend/data-migrations/traces-local-v2-cutover/README.md @@ -322,6 +322,30 @@ the `traceColumnsNonNullable` flip"). On rollback, after swapping the Nullable original back, revert the flag to `false` **and** run that repair. +**The `tracesWeeklyPartitionPruningEnabled` flip (optional, and why it goes last).** `databaseAnalyticsDataModel.tracesWeeklyPartitionPruningEnabled` +(env `ANALYTICS_DB_DATA_MODEL_TRACES_WEEKLY_PARTITION_PRUNING_ENABLED`, default `false`) lets a trace `DELETE` bind itself to +the weekly partitions its own ids resolve to (OPIK-6901), instead of being planned against every part of the table — on +prod-test, 12 ids rewrote 3,928 parts / 5.40 TiB without it. It asserts a **schema** fact: that `traces` (or +`traces_local`) is the successor, with `id_at` as `DateTime64(0,'UTC')` under the weekly `PARTITION BY`. + +> **It enables the *pruning*, not the partitioning — the name is deliberate.** Setting it does not create, activate or +> migrate anything; the partitioned schema arrives with the `EXCHANGE` above and nowhere else. So it is never a step that +> *makes* the cutover progress, and setting it early does not bring the partitioning forward — it only starts emitting a +> predicate against whatever table is live, which is the failure below. + +It is a third flag precisely because **neither of the two above marks the `EXCHANGE`**, which is when the partitioning +appears: `traceColumnsNonNullable` must lead it (above), and `tracesDistributedWrapEnabled` flips at the wrap, which may +be deferred long after it (`--skip-wrap` … `--wrap-only`). So gate on this one, not on either of those. + +Unlike its siblings it is **safe to lag and unsafe to lead**: `false` is the previous unbounded delete, always correct +and merely slower, so turn it on at leisure **after** the `EXCHANGE` is confirmed. Turning it on early — while `traces` is +still the original — is the failure mode worth avoiding: the original has **no `PARTITION BY` at all** (nothing to prune) +and declares `id_at` as a 32-bit `DateTime` that overflows past 2106, so a far-future id (the litellm ~2201 rows) is +stored under a wrapped recent timestamp that the derived partition cannot match, and the delete reports success having +matched **zero rows**. For the same reason, a stage B/C **rollback must revert it to `false` — and roll-restart every +instance — before promoting the original**, ahead of the swap rather than after it. Like its siblings it comes from a +startup snapshot of `OpikConfiguration`, so the config change alone changes nothing until each instance restarts. + ## Batching and throttling On a large production table a single week can be enormous, so the backfill does **not** run one INSERT per week. Two @@ -832,6 +856,11 @@ statements, so a failure *between* them needs a restart path: again, so the flip has to be undone in two steps — `rollback.sh` prints both when the stage finishes. The rollback is not complete until they land. +> **If `tracesWeeklyPartitionPruningEnabled` was turned on, revert it *before* the stage runs, not after.** It asserts the +> live table is the partitioned successor, and the restored original is not one, so a stale `true` makes trace deletes +> match zero rows while reporting success — see "The `tracesWeeklyPartitionPruningEnabled` flip". It is the one flag whose +> revert has to lead the swap; the two steps below follow it. + 1. **Revert `traceColumnsNonNullable` to `false` AND roll-restart every backend instance.** The flag is read from a **startup snapshot** of `OpikConfiguration` (bound via `toInstance`), so a config change does **not** take effect until each instance restarts — exactly like the forward rollout before the EXCHANGE. Until the restart completes, the app 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 ceceab54498..0b86a21739e 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 @@ -32,6 +32,7 @@ import com.comet.opik.utils.ErrorUtils; import com.comet.opik.utils.JsonUtils; import com.comet.opik.utils.TruncationUtils; +import com.comet.opik.utils.WeeklyPartitions; import com.comet.opik.utils.template.TemplateUtils; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.node.JsonNodeFactory; @@ -1921,8 +1922,21 @@ ORDER BY (workspace_id, project_id, trace_id, id) DESC, last_updated_at DESC * so a single statement can span several projects (e.g. a reused id resolved to all its owning projects, or a * cross-project batch) instead of one delete per project (OPIK-7483). Every deleted row carries its {@code * project_id}, so no delete is ever project-less - also required once {@code traces} is a Distributed table - * (OPIK-7455). No {@code id_at}/time predicate on purpose, so it still deletes rows whose {@code id_at} is - * untrustworthy (e.g. a wrapped timestamp); correctness here does not depend on {@code id_at}. + * (OPIK-7455). + *

+ * {@code } adds the table's own weekly partition expression, bound as the exact set of partitions + * the batch's ids resolve to. Both conditions must hold for it to be emitted: the live table must be the weekly + * partitioned successor ({@link #tracesWeeklyPartitionPruningEnabled()}), and every id in the batch must be one whose + * partition can be derived exactly ({@link WeeklyPartitions#of}). Otherwise the predicate is omitted and the + * statement is byte-identical to the previous unbounded form. That is what preserves the original guarantee — a + * row whose {@code id_at} cannot be trusted is still deleted, because no id in such a batch is used to derive a + * partition. + *

+ * Why it matters: a mutation selects parts at the partition stage, where the (workspace_id, project_id, id) + * predicate prunes nothing, so deleting a handful of rows rewrote every part of the table. Measured on prod-test + * (271.6 M rows, 3,928 parts): 12 ids rewrote 3,928 parts / 5.40 TiB. With this predicate the same batch + * selects 5 parts. An {@code id_at} range is not a substitute: on a batch spanning 1996 and 2200 a + * range still selected 2,644 parts, where the exact set selected 4. *

* The pairs are bound (never inlined) as two positional string arrays and zipped back into {@code (project_id, id)} * tuples with {@code arrayZip}, so the query text is constant regardless of batch size and no value reaches the SQL @@ -1934,6 +1948,7 @@ ORDER BY (workspace_id, project_id, trace_id, id) DESC, last_updated_at DESC DELETE FROM traces_localtraces WHERE workspace_id = :workspace_id AND (project_id, id) IN arrayZip(:project_ids, :trace_ids) + AND toYYYYMMDD(toDate32(id_at) - toIntervalDay(toDayOfWeek(id_at, 1))) IN :partitions SETTINGS log_comment = '' ; """; @@ -3356,6 +3371,49 @@ private void selectTracesMutationTable(ST template) { } } + /** + * Whether a trace mutation may prune to the partitions its ids resolve to ({@link WeeklyPartitions}). The flag + * enables the pruning, never the partitioning: it asserts that the live mutation target already is the + * weekly partitioned successor — {@code id_at} as {@code DateTime64(0, 'UTC')} under + * {@code PARTITION BY toYYYYMMDD(toDate32(id_at) - toIntervalDay(toDayOfWeek(id_at, 1)))} — which the cutover's + * EXCHANGE installs, not this flag. + *

+ * This has to be its own flag because neither existing schema flag marks the EXCHANGE, which is the moment the + * partitioning appears, and both are wrong in a different direction: + *

+ * Emitting it against the legacy table is not merely unhelpful, it is wrong: legacy {@code traces} has no + * {@code PARTITION BY} at all (one {@code all} partition, so nothing to prune) and declares {@code id_at} as a + * 32-bit {@code DateTime} that overflows past 2106, so a far-future id — the litellm ~2201 ids, real + * customer-facing rows — is stored under a wrapped recent timestamp that the derived partition cannot match. The + * delete would then match zero rows and report success. + *

+ * Only ever {@code true} while {@code traces} really is that successor, so unlike its siblings it is safe to lag: + * {@code false} is the always-correct unbounded behaviour, and only {@code true} asserts something about the + * schema. Turn it on once the EXCHANGE is confirmed, and back off before a rollback stage B/C promotes the + * original. + */ + private boolean tracesWeeklyPartitionPruningEnabled() { + return configuration.getDatabaseAnalyticsDataModel().tracesWeeklyPartitionPruningEnabled(); + } + + /** + * The partitions a delete batch may bound itself to, or empty to leave the mutation unbounded. Empty whenever the + * live table is not the partitioned successor, ahead of asking {@link WeeklyPartitions} at all — the derivation is + * only meaningful against a table that partitions on it. + */ + private Optional> weeklyPartitionsFor(Collection ids) { + return tracesWeeklyPartitionPruningEnabled() ? WeeklyPartitions.of(ids) : Optional.empty(); + } + /** * Binds input, output, metadata, and their slim versions (input_slim, output_slim) to a statement. * Centralizes the JSON conversion and binding logic for consistency across single and batch inserts. @@ -3529,11 +3587,22 @@ public Mono delete(Set> projectIdTraceIdPairs, @NonNull C var projectIds = batch.stream().map(pair -> pair.getLeft().toString()).toArray(String[]::new); var traceIds = batch.stream().map(pair -> pair.getRight().toString()).toArray(String[]::new); + // Prune to the batch's own partitions when the schema and every id in the batch allow it; + // otherwise emit the unbounded form. + var partitions = weeklyPartitionsFor(batch.stream().map(Pair::getRight).toList()); + // Flag only, exactly like distributed_wrap: the values reach ClickHouse via the bind below, + // never through the template, so the rendered SQL is constant regardless of batch contents. + partitions.ifPresent(_ -> template.add("partitions", true)); + var statement = connection.createStatement(template.render()) .bind("workspace_id", workspaceId) .bind("project_ids", projectIds) .bind("trace_ids", traceIds); + if (partitions.isPresent()) { + statement = statement.bind("partitions", partitions.get().toArray(Long[]::new)); + } + var segment = startSegment("traces", "Clickhouse", "delete"); return Mono.from(statement.execute()) .doFinally(_ -> endSegment(segment)) diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/DatabaseAnalyticsDataModelConfig.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/DatabaseAnalyticsDataModelConfig.java index 090ad3028c3..1c27e719b15 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/DatabaseAnalyticsDataModelConfig.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/DatabaseAnalyticsDataModelConfig.java @@ -50,6 +50,26 @@ * {@code MODIFY COLUMN} must be applied to both {@code traces_local} and the {@code Distributed} {@code traces} * (the wrapper accepts them as metadata-only, and targeting only {@code traces_local} leaves the wrapper without the * column, so reads fail with code 47).

+ * + *

{@code tracesWeeklyPartitionPruningEnabled}: enables partition-aware pruning of trace deletes — a trace + * {@code DELETE} bounds itself to the weekly partitions its own ids resolve to instead of being planned against every + * part of the table (OPIK-6901). It does not create or activate any partitioning; installing the partitioned + * schema is the EXCHANGE step of the cutover. Turning it on therefore asserts a schema fact rather than causing + * one: that the live mutation target already is the weekly partitioned successor — {@code id_at} as + * {@code DateTime64(0, 'UTC')} under + * {@code PARTITION BY toYYYYMMDD(toDate32(id_at) - toIntervalDay(toDayOfWeek(id_at, 1)))}. Purely an optimisation: + * {@code false} keeps the unbounded mutation, which is always correct and merely slower, and only {@code true} asserts + * anything about the schema.

+ * + *

It is deliberately a third flag rather than a reuse of the two above, because the partitioning appears at the + * EXCHANGE and neither of them marks that moment. {@code traceColumnsNonNullable} must be rolled out + * before the EXCHANGE (a rolling restart cannot be atomic with a metadata swap), and + * {@code tracesDistributedWrapEnabled} flips at the wrap, a separate step that may be deferred long after it — so one + * flag would be true too early and the other true too late. Emitting the predicate too early is the harmful direction: + * legacy {@code traces} has no {@code PARTITION BY} at all and declares {@code id_at} as a 32-bit {@code DateTime} that + * overflows past 2106, so a far-future id is stored under a wrapped timestamp the derived partition cannot match and + * the delete would silently affect zero rows. Left {@code false} at deploy time; set {@code true} once the EXCHANGE is + * confirmed on the target, and back to {@code false} before a rollback promotes the original {@code traces}.

*/ @Builder(toBuilder = true) public record DatabaseAnalyticsDataModelConfig( @@ -58,5 +78,6 @@ public record DatabaseAnalyticsDataModelConfig( boolean traceDeletionEventsCaptureEnabled, boolean spanDeletionEventsCaptureEnabled, @Min(1) @Max(2_000) int deletionEventsInsertBatchSize, - boolean tracesDistributedWrapEnabled) { + boolean tracesDistributedWrapEnabled, + boolean tracesWeeklyPartitionPruningEnabled) { } diff --git a/apps/opik-backend/src/main/java/com/comet/opik/utils/WeeklyPartitions.java b/apps/opik-backend/src/main/java/com/comet/opik/utils/WeeklyPartitions.java new file mode 100644 index 00000000000..193d00bc841 --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/utils/WeeklyPartitions.java @@ -0,0 +1,106 @@ +package com.comet.opik.utils; + +import lombok.NonNull; +import lombok.experimental.UtilityClass; + +import java.time.DayOfWeek; +import java.time.Instant; +import java.time.LocalDate; +import java.time.ZoneOffset; +import java.time.temporal.TemporalAdjusters; +import java.util.Collection; +import java.util.HashSet; +import java.util.Optional; +import java.util.Set; +import java.util.UUID; + +/** + * Single derivation point for the {@code id_at} weekly partition values a batch of ids resolves to, so a mutation can + * name its own partitions instead of being planned against every part of the table. + * + *

Mirrors the partition expression of {@code traces_local_v2} / {@code spans_local_v2} exactly — + * {@code PARTITION BY toYYYYMMDD(toDate32(id_at) - toIntervalDay(toDayOfWeek(id_at, 1)))}, where {@code id_at} is + * {@code DateTime64(0, 'UTC') MATERIALIZED UUIDv7ToDateTime(toUUID(id))} — i.e. the Monday of the id's UTC week as + * {@code yyyyMMdd}. Both tables share the expression, so it is derived here once rather than per DAO.

+ * + *

Why the caller must treat an empty result as "no predicate", never as "no partitions". The whole value of + * the derivation is that a partition the batch resolves to is the only place its rows can be; a value that is + * merely close is a silently skipped delete, not a slower one. So this returns a set only when every id in the batch + * is one it can derive exactly, and empty otherwise — leaving the caller to emit its unbounded form, which is always + * correct and merely slower. The two rejections are:

+ * + * + *

Far-future ids within the range are supported on purpose and are the common case of the two: a UUIDv7 + * minted with a bad clock (litellm BerriAI/litellm#31294 + * mints ~2201) has a bogus but self-consistent {@code id_at}, so it lives in the far-future partition this computes and + * stays deletable and prunable. Verified on prod-test: 0 partition mismatches across 11.23 M far-future rows.

+ */ +@UtilityClass +public class WeeklyPartitions { + + /** + * First instant {@code id_at} cannot represent. {@code DateTime64} spans + * {@code [1900-01-01 00:00:00, 2299-12-31 23:59:59.99999999]} and saturates rather than wrapping or throwing, and + * it saturates twice over before {@code toDate32} is reached: {@code UUIDv7ToDateTime} already returns + * {@code DateTime64(3)}, and the column is {@code DateTime64(0)}. So an id at or past this instant is stored as + * {@code 2299-12-31 23:59:59} and partitions as {@code 22991225} (that Sunday's Monday) whatever its real week — + * observable on prod-test, whose far-future rows top out at exactly {@code 2299-12-31}. + */ + private static final long ID_AT_CEILING = LocalDate.of(2300, 1, 1) + .atStartOfDay() + .toInstant(ZoneOffset.UTC) + .toEpochMilli(); + + /** + * The weekly partition values the ids resolve to, or empty if the batch contains an id whose partition cannot be + * derived exactly (see the class javadoc) — in which case the caller must omit its partition predicate entirely. + * An empty batch yields empty for the same reason: there is nothing to bound the mutation to. + *

+ * A {@code null} batch throws rather than reading as empty, which is the one place this class is deliberately + * intolerant. Empty is a documented answer — "this batch cannot be pruned, emit the unbounded form" — and a + * caller that lost its batch would receive that answer, silently issue a correct-but-unbounded mutation, and never + * learn it had a bug. A null collection here is a programming error, not a data condition; a null element + * is the data condition, and that keeps returning empty. + * + * @throws NullPointerException if {@code ids} is null. + */ + public static Optional> of(@NonNull Collection ids) { + var partitions = new HashSet(); + + for (UUID id : ids) { + if (id == null || id.version() != 7) { + return Optional.empty(); + } + // UUIDv7: the high 48 bits are the unix epoch in milliseconds. `>>> 16` reads them unsigned, so the value + // is in [0, 2^48) — never negative, and never large enough for Instant.ofEpochMilli to overflow. That is + // also why only the ceiling is checked: the floor of the range is 1900-01-01, the smallest id_at any + // UUIDv7 can carry is the epoch, and even its Monday (1969-12-29) is comfortably inside Date32. + long epochMilli = id.getMostSignificantBits() >>> 16; + if (epochMilli >= ID_AT_CEILING) { + return Optional.empty(); + } + LocalDate monday = Instant.ofEpochMilli(epochMilli) + .atZone(ZoneOffset.UTC) + .toLocalDate() + .with(TemporalAdjusters.previousOrSame(DayOfWeek.MONDAY)); + partitions.add(monday.getYear() * 10000L + monday.getMonthValue() * 100L + monday.getDayOfMonth()); + } + + // Set.copyOf, not the working HashSet: what escapes here decides which partitions a DELETE mutation touches, so + // a caller holding a mutable reference could narrow the set after it was derived and turn a correct delete into + // a silent no-op. Immutable by default per apps/opik-backend/AGENTS.md, and the accumulator stays local. + return partitions.isEmpty() ? Optional.empty() : Optional.of(Set.copyOf(partitions)); + } +} diff --git a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/TracesPartitionPruningMutationTest.java b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/TracesPartitionPruningMutationTest.java new file mode 100644 index 00000000000..98285e7d2e9 --- /dev/null +++ b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/TracesPartitionPruningMutationTest.java @@ -0,0 +1,1120 @@ +package com.comet.opik.infrastructure; + +import com.comet.opik.api.Trace; +import com.comet.opik.api.resources.utils.ClickHouseContainerUtils; +import com.comet.opik.api.resources.utils.ClientSupportUtils; +import com.comet.opik.api.resources.utils.MigrationUtils; +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.TestDropwizardAppExtensionUtils.AppContextConfig; +import com.comet.opik.api.resources.utils.TestDropwizardAppExtensionUtils.CustomConfig; +import com.comet.opik.api.resources.utils.TestUtils; +import com.comet.opik.api.resources.utils.WireMockUtils; +import com.comet.opik.api.resources.utils.resources.TraceResourceClient; +import com.comet.opik.domain.IdGenerator; +import com.comet.opik.domain.TestIdGeneratorFactory; +import com.comet.opik.domain.TraceDAO; +import com.comet.opik.extensions.DropwizardAppExtensionProvider; +import com.comet.opik.extensions.RegisterApp; +import com.comet.opik.infrastructure.auth.RequestContext; +import com.comet.opik.infrastructure.db.TransactionTemplateAsync; +import com.comet.opik.podam.PodamFactoryUtils; +import com.comet.opik.utils.JsonUtils; +import com.comet.opik.utils.WeeklyPartitions; +import com.comet.opik.utils.template.TemplateUtils; +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.databind.JsonNode; +import com.redis.testcontainers.RedisContainer; +import io.r2dbc.spi.Statement; +import lombok.Builder; +import org.apache.commons.lang3.RandomStringUtils; +import org.apache.commons.lang3.tuple.Pair; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; +import org.junit.jupiter.api.extension.ExtendWith; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; +import org.testcontainers.clickhouse.ClickHouseContainer; +import org.testcontainers.containers.GenericContainer; +import org.testcontainers.containers.Network; +import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; +import ru.vyarus.dropwizard.guice.test.ClientSupport; +import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; +import uk.co.jemos.podam.api.PodamFactory; + +import java.time.LocalDate; +import java.time.ZoneOffset; +import java.util.ArrayList; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.UUID; +import java.util.function.Consumer; +import java.util.regex.Pattern; +import java.util.stream.Collectors; +import java.util.stream.Stream; + +import static com.comet.opik.api.resources.utils.AuthTestUtils.mockTargetWorkspace; +import static com.comet.opik.infrastructure.FilterUtils.ANALYTICS_DELETE_BATCH_SIZE; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.junit.jupiter.params.provider.Arguments.arguments; + +/** + * Exercises the partition pruning of {@code DELETE_BY_PROJECT_ID_TRACE_ID_PAIRS} against the post-EXCHANGE topology, + * where {@code traces} is the weekly partitioned successor (OPIK-6901). {@code WeeklyPartitionsTest} covers the + * derivation itself; what only a real ClickHouse can show is the half between it and the mutation — that the template + * renders the predicate exactly when it should, that the derived {@code Long[]} binds to {@code IN :partitions}, and + * that the row still goes away either way. + * + *

Everything is asserted through the DAO's own delete. This suite writes no query that re-implements or + * re-evaluates the partition expression, and it asserts nothing about the EXCHANGE — the EXCHANGE is only setup (see + * below). Its own SQL is three statements, all plain binds: seed a row, count a row, read a statement back from + * {@code system.query_log}. Correctness is established the way production would feel it — the rows go away. If + * the DAO's predicate resolved to any partition other than the one ClickHouse filed a row under, the mutation would + * select the wrong parts and that row would survive, so + * {@link #deleteClearsEveryEraAndBindsExactlyThosePartitions} passing is the agreement between the migration's + * {@code PARTITION BY} as installed, the DAO's predicate, and {@link WeeklyPartitions#of}. + * + *

Each test then pairs that with the SQL ClickHouse actually received, because rows alone cannot see pruning + * silently stop — a delete that stopped bounding itself is still correct, just slow, and that is the regression this + * change exists to prevent. Read back by {@code log_comment} plus the test's own trace id, and checked as an + * exact bound partition set: a superset would keep every delete correct while handing back the whole benefit. + * + *

The eras in {@link #deleteClearsEveryEraAndBindsExactlyThosePartitions} are load-bearing, not variety. + * {@code toMonday} agrees with the {@code Date32} expression across the ordinary calendar and diverges only far-future + * or at the epoch, so a recent-only batch would accept the very expression migration 000114 was written to escape. The + * 2200 row also covers the {@code DateTime64} half of what the flag asserts: against a 32-bit {@code id_at} it would be + * stored under a wrapped recent timestamp, the derived partition would not match, and it would survive. + * + *

Two internal touches, on the pattern of {@code TracesDistributedWrapMutationTest}: the EXCHANGE has no public API, + * so {@link #beforeAll} runs it in raw SQL identical to the swap block of {@code 000003_exchange_and_wrap.sql}; and the + * ingestion path rejects a non-v7, backdated or far-future {@code id} by design ({@code IdGenerator.validateId}), so + * those rows are seeded raw and those batches are handed to {@link TraceDAO#delete} directly — the only way to reach + * that arm. + * + *

Why the topology setup is here at all — it is setup, never the subject. Nothing asserts anything about the + * EXCHANGE or the wrap themselves; {@code TracesLocalV2CutoverTest} owns that. They are required because the DAO names + * its target table: after the Liquibase migrations the live {@code traces} is still the legacy table — no + * {@code PARTITION BY} at all and a 32-bit {@code DateTime} {@code id_at} — while the partitioned successor exists only + * as the empty {@code traces_local_v2}, which no DAO query can reach. Without installing it these tests would run + * against the one table where this predicate must never be emitted, and would pass while proving nothing: the predicate + * is harmless against an unpartitioned table for recent ids. Hand-authoring a partitioned {@code traces} in the test + * instead would duplicate migration 000114 and reintroduce exactly the drift this suite exists to detect. Both steps are + * idempotent, so the suite keeps working once the cutover migration lands and they become no-ops — and that is + * asserted rather than assumed, by {@link PruningEnabled#topologySetupIsANoOpOnceTheEstateProvidesIt}. Two states + * are therefore covered: with the swap, which every other test here needs today, and without it, + * which is what the estate will look like once the migrations create {@code traces_local} partitioned with the + * {@code Distributed} {@code traces} over it and this suite's {@code EXCHANGE} stops being needed at all. + * + *

Topology: the post-cutover end state, with both schema flags on. The wrap is applied and + * {@code tracesDistributedWrapEnabled} set, so the DAO's mutations reach the data the way production routes them — + * {@code DELETE FROM traces_local}, chosen by the configuration switch that governs it, not by a table this suite + * renamed under the DAO. {@code traces} is the {@code Distributed} wrapper that reads and inserts flow through, which + * is why the endpoint-created row and the raw-seeded ones land in the same place. + * {@link #distributedTracesRejectsDirectMutation} keeps that claim honest: had the wrap not taken effect, + * {@code traces} would still be a {@code MergeTree} and every pruned delete here would have run against it. + * + *

This supersedes an earlier note in this file that both flags on at once was untested and would need its own suite. + * It does not: the wrap flag is how the DAO is pointed at the data, so enabling it costs two setup statements and + * covers the combination the fleet actually ends up in. What is no longer covered here is the transient + * post-EXCHANGE/pre-wrap window — the pruning predicate is identical in both, since {@code traces_local} is the same + * physical table under a different name, and the flag's own javadoc records that it must hold in that window. + * + *

Dedicated, non-reused ClickHouse and ZooKeeper containers are required because the setup destructively renames the + * live {@code traces} table — the EXCHANGE swaps it, and the wrap then renames it to {@code traces_local} — so a reused + * container would corrupt other suites and reruns. + */ +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +class TracesPartitionPruningMutationTest { + + private static final String API_KEY = "apiKey-" + UUID.randomUUID(); + private static final String WORKSPACE_NAME = "workspace-" + RandomStringUtils.secure().nextAlphanumeric(32); + private static final String WORKSPACE_ID = UUID.randomUUID().toString(); + private static final String USER = "user-" + RandomStringUtils.secure().nextAlphanumeric(32); + + /** + * The partition-key fragment the DAO's template emits. Only ever compared against the statement read back + * from {@code system.query_log} — never spliced into a query this suite runs. Verbatim comparison is safe because + * {@code query_log} stores the query as submitted, so both sides are the DAO's own template text. + */ + private static final String PARTITION_PREDICATE = "toYYYYMMDD(toDate32(id_at) - toIntervalDay(toDayOfWeek(id_at, 1)))"; + + // The suite's whole SQL surface, per .agents/skills/opik-backend/SKILL.md "SQL Query Construction": one text block + // per query, every varying value a :placeholder. There are no StringTemplate fragments and no interpolation at all, + // because nothing here re-implements the DAO's predicate - PARTITION_PREDICATE is only ever compared against the + // statement the DAO emitted, never spliced into a query of ours. + // + // The EXCHANGE/RENAME pair in installPartitionedSuccessorUnderTraces() stays as inline literals on purpose: they + // are single literals built by no Java string operation, and they are kept byte-identical to + // 000003_exchange_and_wrap.sql by eye, so they belong at the call site next to the javadoc that says so - as + // TracesDistributedWrapMutationTest does. + private static final String PARTITION_KEY_OF_TABLE = """ + SELECT partition_key + FROM system.tables + WHERE database = currentDatabase() + AND name = :table + """; + + private static final String TABLE_COUNT = """ + SELECT toString(count()) + FROM system.tables + WHERE database = currentDatabase() + AND name = :table + """; + + private static final String TABLE_ENGINE_FULL = """ + SELECT engine_full + FROM system.tables + WHERE database = currentDatabase() + AND name = 'traces' + """; + + private static final String INSERT_RAW_TRACE = """ + INSERT INTO traces_local (workspace_id, project_id, id) + VALUES (:workspace_id, :project_id, :id) + """; + + private static final String LAST_TRACE_DELETE = """ + SELECT query + FROM system.query_log + WHERE log_comment LIKE 'delete_traces:%' + AND type = 'QueryFinish' + AND query LIKE concat('%', :trace_id, '%') + ORDER BY event_time_microseconds DESC + LIMIT 1 + """; + + /** + * The newest {@code delete_traces} statement carrying exactly {@code pairs_size} pairs. A request larger than + * {@link com.comet.opik.infrastructure.FilterUtils#ANALYTICS_DELETE_BATCH_SIZE} is chunked by the DAO into one + * statement per chunk, and pruning is derived per chunk — so a test that reads only one statement cannot see + * the second chunk at all. + */ + private static final String DELETE_BY_PAIR_COUNT = """ + SELECT query + FROM system.query_log + WHERE log_comment LIKE concat('delete_traces:%pairs_size=', :pairs_size) + AND type = 'QueryFinish' + ORDER BY event_time_microseconds DESC + LIMIT 1 + """; + + /** + * Scoped by the same key {@code TraceDAO.delete} matches on — {@code (workspace_id, project_id, id)}. An oracle + * narrower than the delete would answer a different question: a row for the same id in another project would keep + * the count at {@code 1} after a successful delete, failing a test whose subject worked. + */ + private static final String LIVE_ROW_COUNT = """ + SELECT toString(uniqExact(id)) + FROM traces_local + WHERE workspace_id = :workspace_id + AND project_id = :project_id + AND id = :id + """; + + /** + * The {@code IN} clause the DAO emitted, captured to end of line: the predicate sits on its own line in the + * template with {@code SETTINGS log_comment} on the next, so the line boundary delimits it exactly. Read this way + * rather than by matching a bracket style, because the driver's rendering of a {@code Long[]} is its own choice — + * what matters is which partition values are in the clause, not how it punctuates them. + */ + private static final Pattern EMITTED_IN_CLAUSE = Pattern.compile( + Pattern.quote(PARTITION_PREDICATE) + "\\s+IN\\s+([^\\n]*)"); + + /** + * The emitted statement's shape, so its {@code WHERE} clause can be lifted verbatim and re-asked as a + * {@code SELECT}: {@code EXPLAIN} does not accept a mutation. Captures the target table too, since the DAO picks + * {@code traces} or {@code traces_local} depending on the wrap flag. + */ + private static final Pattern DELETE_SHAPE = Pattern.compile( + "DELETE\\s+FROM\\s+(\\S+)\\s+(WHERE\\b.*?)\\s+SETTINGS\\b", Pattern.DOTALL); + + /** {@code EXPLAIN} index entries that reflect partition-level part selection. */ + private static final Set PARTITION_INDEX_TYPES = Set.of("MinMax", "Partition"); + + /** + * Asks the planner how many parts the DAO's partition predicate selects. Declared once as a text block, per + * {@code .agents/skills/opik-backend/SKILL.md}: the table {@code } and the predicate + * {@code } are fragments and go through {@link TemplateUtils#newST}, the partition + * values are values and are bound — nothing is spliced with {@code .formatted(...)}. + *

+ * The predicate fragment is {@link #PARTITION_PREDICATE}, and using the constant does not re-author what is under + * test: every caller has already asserted the emitted statement contains that exact text, so the constant + * is pinned to the DAO's own SQL by assertion rather than by string surgery on it. The partition values come from + * the emitted statement too, parsed by {@link #boundPartitionsOf} and bound here. + *

+ * The DAO's {@code workspace_id} and {@code (project_id, id)} predicates are deliberately not reproduced. They are + * sort-key filters, not partition filters, so they cannot change partition selection — and leaving them out makes + * the unbounded case a full scan, which is the conservative direction for an assertion that the fallback prunes + * nothing. + */ + private static final String EXPLAIN_SELECTED_PARTS = """ + EXPLAIN indexes = 1, json = 1 + SELECT id + FROM

+ WHERE IN :partitions + """; + + /** + * The wrap block of {@code 000003_exchange_and_wrap.sql}. The database name is a fragment (an identifier + * inside a function argument, not a bindable value), so it goes through {@link TemplateUtils#newST} rather than + * {@code .formatted(...)} — the sibling suites still splice it, but the rule says not to add new ones. The + * {@code {cluster}} macros are ClickHouse's own and pass through StringTemplate untouched, which uses {@code <>}. + */ + private static final String CREATE_DISTRIBUTED_WRAPPER = """ + CREATE TABLE traces_dist ON CLUSTER '{cluster}' AS traces + ENGINE = Distributed('{cluster}', '', 'traces_local', sipHash64(project_id)) + """; + + /** A weekly partition value as it appears in SQL — {@code yyyyMMdd}, so always eight digits. */ + private static final Pattern PARTITION_VALUE = Pattern.compile("\\d{8}"); + + /** + * One id per era the derivation has to get right, with the partition each resolves to. Not interchangeable samples: + * {@code toMonday} agrees with the {@code Date32} expression across the ordinary calendar and diverges only + * far-future or at the epoch, so a recent-only batch would accept the very expression migration 000114 was written + * to escape. The 2200 id is the litellm shape and the one that makes the assertion bite. + */ + private static final IdGenerator ID_GENERATOR = TestIdGeneratorFactory.create(); + + /** + * The weekly partitions this suite works in, each named by its Monday — which is also the partition name, since the + * key is that Monday as {@code yyyyMMdd}. Fixed rather than {@code now}-derived, for the reason + * {@code TracesLocalV2PartitioningTest} gives for its own anchor: the partition math stays deterministic and cannot + * drift across a week boundary mid-suite. + *

+ * Ids are minted mid-week ({@link #idInWeekOf}), so the assertions exercise the map back to Monday rather + * than identity. The three eras are not interchangeable samples: {@code toMonday} agrees with the {@code Date32} + * expression across the ordinary calendar and diverges only far-future or at the epoch, so a recent-only batch would + * accept the very expression migration 000114 was written to escape. The 2199 row is the litellm shape and the one + * that makes it bite; it also covers the {@code DateTime64} half of what the flag asserts, since a 32-bit + * {@code id_at} would store it under a wrapped recent timestamp and the derived partition would miss it. + */ + private static final List ERA_MONDAYS = List.of( + LocalDate.of(1996, 2, 5), + LocalDate.of(2025, 3, 3), + LocalDate.of(2199, 12, 30)); + + /** {@link UUID#randomUUID()} is a v4 by definition: no embedded timestamp to derive a partition from. */ + private static final UUID NON_V7_ID = UUID.randomUUID(); + + /** + * A UUIDv7 minted one second past the first instant {@code DateTime64} cannot represent, so its {@code id_at} + * saturates to the ceiling and the honest week is not the partition the row lands in. Same rejection as + * {@link #NON_V7_ID}, different cause — see {@code WeeklyPartitions}. Minted rather than written out, so the + * boundary it sits past is visible. + */ + private static final UUID OUT_OF_RANGE_ID = ID_GENERATOR + .getTimeOrderedEpoch(LocalDate.of(2300, 1, 1).atStartOfDay().toInstant(ZoneOffset.UTC).toEpochMilli()); + + // Dedicated, non-reused ClickHouse + ZooKeeper on their own network: the EXCHANGE destructively swaps `traces`, so a + // shared/reused container would corrupt other suites and reruns. Redis/MySQL are only read, so the shared ones are + // fine. + private final Network network = Network.newNetwork(); + private final GenericContainer zookeeperContainer = ClickHouseContainerUtils.newZookeeperContainer(false, + network); + private final ClickHouseContainer clickHouseContainer = ClickHouseContainerUtils + .newClickHouseContainer(false, network, zookeeperContainer); + private final RedisContainer redisContainer = RedisContainerUtils.newRedisContainer(); + private final MySQLContainer mysqlContainer = MySQLContainerUtils.newMySQLContainer(); + + private final WireMockUtils.WireMockRuntime wireMock; + + private final PodamFactory factory = PodamFactoryUtils.newPodamFactory(); + + /** + * Runs the topology setup and this suite's own raw reads and seeds straight against the container, with no app in + * the way — the idiom the sibling partition suites use. It has to be app-independent: the topology must be + * installed once, before either nested app boots, and both nested classes then read through the same handle. + *

+ * Never use this for anything the DAO executes — see {@link #appTemplate}. It carries none of the production + * {@code queryParameters}, so a statement the real connection would accept can fail here for reasons production + * would never hit. + */ + private final TransactionTemplateAsync template; + + { + Startables.deepStart(redisContainer, mysqlContainer, clickHouseContainer, zookeeperContainer) + .join(); + wireMock = WireMockUtils.startWireMock(); + MigrationUtils.runMysqlDbMigration(mysqlContainer); + MigrationUtils.runClickhouseDbMigration(clickHouseContainer); + template = TransactionTemplateAsync.create(ClickHouseContainerUtils + .newDatabaseAnalyticsFactory(clickHouseContainer, ClickHouseContainerUtils.DATABASE_NAME) + .build()); + ensurePartitionedSuccessorUnderTraces(); + ensureDistributedWrap(); + } + + private TraceResourceClient traceResourceClient; + private TraceDAO traceDAO; + + /** + * The app's own connection handle, used for anything the DAO executes. It is not interchangeable with + * {@link #template}: the app builds its factory from {@code config-test.yml}, which carries the production + * {@code queryParameters} — including {@code max_query_size=100000000}. The container-derived handle sets none of + * them, so a full-size delete chunk (10,000 pairs inline to ~762 KiB of SQL) dies on ClickHouse's 256 KiB default. + * Routing the DAO through the container handle silently ran every delete in this suite on non-production settings. + */ + private TransactionTemplateAsync appTemplate; + + /** + * One app per flag state, identical in every other respect. {@code tracesWeeklyPartitionPruningEnabled} is the only + * thing that varies between the two nested classes, which is what lets them be read as an A/B: same schema, same + * topology, same fixtures, one flag. + *

+ * The other two flags are fixed on, as production runs them post-cutover — the successor's {@code end_time}/ + * {@code ttft} are non-nullable sentinel columns, and the wrap is what points the DAO's mutations at + * {@code traces_local}. + */ + private TestDropwizardAppExtension newApp(boolean pruningEnabled) { + return TestDropwizardAppExtensionUtils.newTestDropwizardAppExtension( + AppContextConfig.builder() + .jdbcUrl(mysqlContainer.getJdbcUrl()) + .databaseAnalyticsFactory(ClickHouseContainerUtils.newDatabaseAnalyticsFactory( + clickHouseContainer, ClickHouseContainerUtils.DATABASE_NAME)) + .redisUrl(redisContainer.getRedisURI()) + .runtimeInfo(wireMock.runtimeInfo()) + .customConfigs(List.of( + new CustomConfig("databaseAnalyticsDataModel.traceColumnsNonNullable", "true"), + new CustomConfig("databaseAnalyticsDataModel.tracesDistributedWrapEnabled", "true"), + new CustomConfig("databaseAnalyticsDataModel.tracesWeeklyPartitionPruningEnabled", + String.valueOf(pruningEnabled)))) + .build()); + } + + /** + * Wires the currently-running nested class's app in. Safe on shared outer fields because JUnit runs nested + * containers one at a time — this tree configures no parallel execution — so only one app is live at a time. + */ + private void bindApp(ClientSupport clientSupport, TraceDAO traceDAO, TransactionTemplateAsync appTemplate) { + var baseUrl = TestUtils.getBaseUrl(clientSupport); + ClientSupportUtils.config(clientSupport); + mockTargetWorkspace(wireMock.server(), API_KEY, WORKSPACE_NAME, WORKSPACE_ID, USER); + this.traceResourceClient = new TraceResourceClient(clientSupport, baseUrl); + this.traceDAO = traceDAO; + this.appTemplate = appTemplate; + } + + @AfterAll + void afterAll() { + wireMock.server().stop(); + clickHouseContainer.stop(); + zookeeperContainer.stop(); + network.close(); + } + + /** A UUIDv7 in the given week, minted mid-week so the partition assertion exercises the map back to Monday. */ + private static UUID idInWeekOf(LocalDate monday) { + return ID_GENERATOR.generateId(monday.plusDays(2).atTime(12, 0).toInstant(ZoneOffset.UTC)); + } + + /** + * The partition name for a week, which is its Monday as {@code yyyyMMdd}. Formatting a Monday the test already + * names — not re-deriving "the Monday of an arbitrary date", which is the part under test. + */ + private static long partitionNameOf(LocalDate monday) { + return monday.getYear() * 10000L + monday.getMonthValue() * 100L + monday.getDayOfMonth(); + } + + /** + * The partition a single id resolves to. Derived through {@code WeeklyPartitions} on purpose: this suite is about + * the predicate reaching ClickHouse, and the derivation's own expected values are pinned against real ClickHouse + * output in {@code WeeklyPartitionsTest}, so restating them here would only duplicate that. Where the expectation + * needs to be readable on its own — the era batch — the partitions are stated as literals instead. + */ + private static long partitionOf(UUID id) { + return WeeklyPartitions.of(List.of(id)).orElseThrow().iterator().next(); + } + + /** + * The partition values actually bound into the emitted {@code IN} clause, so a test can assert the set is exact + * rather than merely inclusive. The driver substitutes bound values into the query text client-side, which is why + * they are readable here at all; the two assertions below are what make a change in that behaviour say so plainly + * instead of quietly turning every set assertion into a tautology on an empty set. + */ + private static Set boundPartitionsOf(String sql) { + var clause = EMITTED_IN_CLAUSE.matcher(sql); + assertThat(clause.find()) + .as("the delete SQL carries the partition predicate followed by an IN clause:%n%s", sql) + .isTrue(); + var bound = PARTITION_VALUE.matcher(clause.group(1)).results() + .map(match -> Long.parseLong(match.group())) + .collect(Collectors.toUnmodifiableSet()); + assertThat(bound) + .as("the IN clause carries inlined partition values — got '%s'", clause.group(1)) + .isNotEmpty(); + return bound; + } + + /** Invokes the DAO under a workspace/user context, as {@code TraceService} does for the live delete path. */ + private void delete(Set> projectIdTraceIdPairs) { + appTemplate.nonTransaction(connection -> traceDAO.delete(projectIdTraceIdPairs, connection)) + .contextWrite(ctx -> ctx + .put(RequestContext.WORKSPACE_ID, WORKSPACE_ID) + .put(RequestContext.USER_NAME, USER)) + .block(); + } + + /** + * The SQL of the trace delete that carried {@code traceId}, as ClickHouse received it. Two filters, because either + * alone is ambiguous: {@code log_comment} narrows to this one template ({@code TraceDAO} stamps every statement + * {@code :::

}, and {@code delete_traces} names it alone), and the id narrows + * to this test's statement. + *

+ * The id filter is what makes the lookup deterministic. Every test in this class deletes under the same workspace, + * so ordering by {@code event_time_microseconds} alone would hand back a neighbouring test's delete whenever this + * one had not reached {@code query_log} yet, or whenever two landed in the same microsecond — and that flake could + * pass for the wrong reason, since a neighbour's statement has the same shape. Each test's target id is freshly + * minted, so it identifies the statement exactly; the {@code ORDER BY}/{@code LIMIT} now only picks the newest + * attempt when surefire retries a test. + *

+ * Filtering on the id rather than on a marker injected into {@code details}: the DAO owns {@code log_comment} and + * puts {@code pairs_size} there, and the id is already in the statement text, so this needs no production change to + * serve a test. + */ + private String lastTraceDeleteSql(UUID traceId) { + execute("SYSTEM FLUSH LOGS", _ -> { + }); + var sql = queryOneString(LAST_TRACE_DELETE, statement -> statement.bind("trace_id", traceId.toString())); + assertThat(sql) + .as("query_log holds a delete_traces statement mentioning id '%s'", traceId) + .isNotNull(); + return sql; + } + + /** + * The newest {@code delete_traces} statement that carried exactly {@code pairsSize} pairs. Chunks are identified by + * their pair count rather than by a contained id, because the point is to inspect a specific chunk of one + * request — and the DAO stamps each chunk's size into its {@code log_comment}. + *

+ * The returned text is truncated for large statements. ClickHouse caps {@code query_log.query} at + * {@code log_queries_cut_to_length} (100,000 bytes by default), and a full 10,000-pair chunk inlines to ~762 KiB, + * so anything at the tail of such a statement — the partition predicate included — is simply absent. Only assert on + * the text of statements small enough to be recorded whole. + */ + private String deleteSqlForChunkOf(int pairsSize) { + execute("SYSTEM FLUSH LOGS", _ -> { + }); + var sql = queryOneString(DELETE_BY_PAIR_COUNT, + statement -> statement.bind("pairs_size", String.valueOf(pairsSize))); + assertThat(sql) + .as("query_log holds a delete_traces statement with pairs_size=%s", pairsSize) + .isNotBlank(); + return sql; + } + + /** {@code "1"} while a live (non-lightweight-deleted) row exists for the id, {@code "0"} once it is gone. */ + private String liveRowCount(UUID projectId, UUID id) { + return queryOneString(LIVE_ROW_COUNT, statement -> statement + .bind("project_id", projectId.toString()) + .bind("workspace_id", WORKSPACE_ID) + .bind("id", id.toString())); + } + + /** + * Setup, not a test: puts the partitioned successor under the name the DAO deletes from, so the pruning assertions + * are made against a table that actually has weekly partitions. + *

+ * Idempotent on purpose, because the estate this runs against is going to change. Today {@code traces} is + * the legacy table and the successor is the empty {@code traces_local_v2}, so the two cutover statements are needed. + * Once the cutover migration lands, {@code traces} is the partitioned successor and {@code traces_local_v2} + * is gone — at which point this is a no-op and the suite keeps working unchanged, instead of dying in + * {@link #beforeAll} on an {@code EXCHANGE} against a table that no longer exists. If neither state holds it fails + * with that said plainly, rather than surfacing as a bare "table not found". + *

+ * The EXCHANGE (000003 exchange block): puts the successor under {@code traces} and the original under + * {@code traces_local_v2}, then a RENAME parks the original as {@code traces_pre_cutover_backup}. The wrap is + * deliberately not applied — it is a separate, deferrable step, and the flag under test must hold on its own between + * the two (which is why it is not the wrap flag). Kept identical to the cutover SQL by eye, as + * {@code TracesLocalV2CutoverTest.exchangeTables} and the wrap suite do. + */ + private void ensurePartitionedSuccessorUnderTraces() { + if (tableExists("traces_local") || partitionKeyOf("traces").contains("id_at")) { + return; // Already installed, or the cutover migration has landed. + } + assertThat(tableExists("traces_local_v2")) + .as("neither `traces_local` nor a partitioned `traces` nor `traces_local_v2` is present (partition key" + + " of `traces`: '%s') - this suite needs one of those states to install the successor from", + partitionKeyOf("traces")) + .isTrue(); + execute("EXCHANGE TABLES traces AND traces_local_v2 ON CLUSTER '{cluster}'", _ -> { + }); + execute("RENAME TABLE traces_local_v2 TO traces_pre_cutover_backup ON CLUSTER '{cluster}'", _ -> { + }); + } + + /** + * Setup, not a test: applies the sharding-readiness wrap, so the DAO's mutations reach the data through the + * configuration switch that governs them in production ({@code tracesDistributedWrapEnabled}) rather than through + * a table this suite renamed under it. After this, {@code traces} is a {@code Distributed} wrapper and + * {@code traces_local} holds the partitioned data — the post-cutover end state. + *

+ * Kept identical to the wrap block of {@code 000003_exchange_and_wrap.sql}, as + * {@code TracesDistributedWrapMutationTest.applyDistributedWrap} and + * {@code TracesLocalV2CutoverTest.wrapInDistributed} do: build the wrapper under a temp name first, then one atomic + * multi-target {@code RENAME} rotates the data to {@code traces_local} and the wrapper into {@code traces}, so + * {@code traces} is never absent. Re-entrant: it returns early when the wrap is already applied, and + * clears a wrapper stranded by an interrupted run before rebuilding it. + */ + private void ensureDistributedWrap() { + // Accepting any Distributed table would let a wrapper pointing at another database or another local table + // block the rebuild and silently route reads and inserts elsewhere - so the check is on the definition, not the + // engine name. Matched on the two parts that decide where rows actually go (the database and the local target) + // rather than the whole string, which ClickHouse re-prints and which would make this brittle about spacing. + var engineFull = Optional.ofNullable(queryOneString(TABLE_ENGINE_FULL, _ -> { + })).orElse(""); + if (engineFull.startsWith("Distributed")) { + assertThat(engineFull) + .as("`traces` is already Distributed but not over this database's traces_local: %s", engineFull) + .contains("'" + ClickHouseContainerUtils.DATABASE_NAME + "'") + .contains("'traces_local'"); + return; + } + // Clear a wrapper left behind by a run that died between the CREATE and the RENAME. It holds no data - a + // Distributed table is a routing definition - so dropping it is safe, and without this the CREATE below fails + // on a duplicate name and buries the real state. Same reset TracesLocalV2CutoverTest performs. + execute("DROP TABLE IF EXISTS traces_dist ON CLUSTER '{cluster}' SYNC", _ -> { + }); + execute(TemplateUtils.newST(CREATE_DISTRIBUTED_WRAPPER) + .add("database", ClickHouseContainerUtils.DATABASE_NAME) + .render(), _ -> { + }); + execute(""" + RENAME TABLE + traces TO traces_local, + traces_dist TO traces + ON CLUSTER '{cluster}' + """, _ -> { + }); + } + + private boolean tableExists(String table) { + return "1".equals(queryOneString(TABLE_COUNT, statement -> statement.bind("table", table))); + } + + /** + * The table's partition-key expression, or {@code ""} when there is no such table — never {@code null}. A missing + * row is a legitimate state here (a half-applied wrap can leave {@code traces} renamed away), and the setup guard + * has to be able to report that rather than die dereferencing it, which is what would bury the diagnostic. + */ + private String partitionKeyOf(String table) { + return Optional + .ofNullable(queryOneString(PARTITION_KEY_OF_TABLE, statement -> statement.bind("table", table))) + .orElse(""); + } + + /** + * A trace with every trace-table column populated. Only the span-derived aggregates podam would otherwise + * fabricate ({@code feedbackScores}, {@code usage}) are nulled, since they are not columns of the {@code traces} + * table and would only add noise. The generated {@code projectName} is fresh per trace, so a test that wants two + * traces in one project has to say so; that keeps one test's rows out of another's reads. + */ + private Trace.TraceBuilder newTrace() { + return factory.manufacturePojo(Trace.class).toBuilder() + .feedbackScores(null) + .usage(null); + } + + private UUID projectIdOf(Trace trace) { + return traceResourceClient + .getTraces(trace.projectName(), null, API_KEY, WORKSPACE_NAME, List.of(), List.of(), 100, Map.of()) + .content().stream() + .filter(found -> found.id().equals(trace.id())) + .map(Trace::projectId) + .findFirst() + .orElseThrow(); + } + + private List traceIdsOf(String projectName) { + return traceResourceClient + .getTraces(projectName, null, API_KEY, WORKSPACE_NAME, List.of(), List.of(), 100, Map.of()) + .content().stream() + .map(Trace::id) + .toList(); + } + + /** + * Seeds one row through the table's real column definitions, in the caller's project. Only the three columns + * without a {@code DEFAULT} are supplied; {@code id_at} is {@code MATERIALIZED}, so ClickHouse derives it from + * {@code id} exactly as it does for a row the ingestion path wrote - which is the point, since restating the + * {@code id_at} expression here would create the very drift this suite exists to detect. Raw SQL because ingestion + * rejects a backdated, far-future or non-v7 {@code id} by design ({@code IdGenerator.validateId}). + */ + private void insertRawTrace(UUID projectId, UUID id) { + execute(INSERT_RAW_TRACE, + statement -> statement + .bind("workspace_id", WORKSPACE_ID) + .bind("project_id", projectId.toString()) + .bind("id", id.toString())); + } + + /** + * Parts the planner selects for the DAO's own statement, or empty when it reports no partition index at all. + *

+ * {@code SELECT id} rather than {@code count()}, so no trivial-count optimisation can answer from metadata without + * selecting parts at all. + *

+ * Only {@code MinMax} and {@code Partition} entries are considered, and {@code PrimaryKey} is deliberately + * excluded: the DAO's {@code WHERE} also filters {@code workspace_id} and {@code (project_id, id)}, which are the + * sort key, so {@code PrimaryKey} prunes parts for the unbounded statement too — counting it would make the + * fallback look pruned and destroy the discrimination this test rests on. Across the entries that do qualify it + * takes the smallest selected count and the largest initial count, so it does not depend on which of the two + * reports the pruning. + *

+ * Empty is a meaningful answer rather than a failure: the {@code Indexes} block carries a partition entry only when + * the query filters on the partition key, so its absence is exactly what the fallback should produce. + */ + private Optional partsSelectedBy(String daoDeleteSql) { + var shape = DELETE_SHAPE.matcher(daoDeleteSql); + assertThat(shape.find()) + .as("the emitted statement has the expected DELETE shape:%n%s", daoDeleteSql) + .isTrue(); + Set bound = EMITTED_IN_CLAUSE.matcher(daoDeleteSql).find() + ? boundPartitionsOf(daoDeleteSql) + : Set.of(); + + var explainSql = TemplateUtils.newST(EXPLAIN_SELECTED_PARTS) + .add("table", shape.group(1)); + if (!bound.isEmpty()) { + explainSql.add("partition_expression", PARTITION_PREDICATE); + } + var sql = explainSql.render(); + + var explainRows = template.stream(connection -> { + var statement = connection.createStatement(sql); + if (!bound.isEmpty()) { + statement.bind("partitions", bound.toArray(Long[]::new)); + } + return Flux.from(statement.execute()) + .flatMap(result -> result.map((row, _) -> row.get("explain", String.class))); + }) + .collectList() + .block(); + var explain = String.join("\n", explainRows); + + var indexes = JsonUtils.getJsonNodeFromString(explain).findValue("Indexes"); + if (indexes == null) { + return Optional.empty(); + } + SelectedParts partition = null; + for (JsonNode index : indexes) { + if (!PARTITION_INDEX_TYPES.contains(index.path("Type").asText()) || !index.has("Selected Parts")) { + continue; + } + var entry = JsonUtils.treeToValue(index, SelectedParts.class); + partition = partition == null + ? entry + : partition.toBuilder() + .selected(Math.min(partition.selected(), entry.selected())) + .total(Math.max(partition.total(), entry.total())) + .build(); + } + return Optional.ofNullable(partition); + } + + /** + * First column of the first row, as a string. Every read in this suite is a single scalar, so this is the only + * mapper needed; values go in as binds. + */ + private String queryOneString(String sql, Consumer binder) { + return template.nonTransaction(connection -> { + var statement = connection.createStatement(sql); + binder.accept(statement); + return Mono.from(statement.execute()) + .flatMap(result -> Mono.from(result.map((row, _) -> row.get(0, String.class)))); + }).block(); + } + + private void execute(String sql, Consumer binder) { + template.nonTransaction(connection -> { + var statement = connection.createStatement(sql); + binder.accept(statement); + return Mono.from(statement.execute()).flatMap(result -> Mono.from(result.getRowsUpdated())); + }).block(); + } + + /** + * The part counts {@code EXPLAIN indexes = 1, json = 1} reports for one index entry: how many parts the query + * started from, and how many survived pruning. + */ + @Builder(toBuilder = true) + @JsonIgnoreProperties(ignoreUnknown = true) + private record SelectedParts( + @JsonProperty("Selected Parts") int selected, + @JsonProperty("Initial Parts") int total) { + } + + /** + * The flag on: the pruning this change exists to add. Every assertion here would also hold with the feature + * deleted except the ones about the emitted SQL and the planner — which is exactly why those exist, and why + * {@link PruningDisabled} runs the same fixtures with the flag off. Read as a pair, the two classes pin the flag as + * the cause: remove the pruning and this class fails; remove the flag gate and the other one does. + */ + @Nested + @TestInstance(TestInstance.Lifecycle.PER_CLASS) + @ExtendWith(DropwizardAppExtensionProvider.class) + class PruningEnabled { + + @RegisterApp + private final TestDropwizardAppExtension app = newApp(true); + + @BeforeAll + void beforeAll(ClientSupport clientSupport, TraceDAO traceDAO, TransactionTemplateAsync appTemplate) { + bindApp(clientSupport, traceDAO, appTemplate); + } + + @Test + @DisplayName("traces is a mutation-rejecting Distributed wrapper, so these deletes ran on traces_local") + void distributedTracesRejectsDirectMutation() { + // Keeps the both-flags claim from being vacuous. If the wrap had not taken effect, `traces` would still be a + // MergeTree, every pruned delete in this suite would have run against it, and nothing here would say so. + // Asserting the specific rejection - not merely that something threw - is what proves `traces` is Distributed, + // so a green delete could only have reached `traces_local`. + assertThatThrownBy(() -> execute("DELETE FROM traces WHERE workspace_id = :workspace_id", + statement -> statement.bind("workspace_id", WORKSPACE_ID))) + .hasMessageContaining("DELETE query is not supported for table"); + } + + @Test + @DisplayName("an all-UUIDv7 delete prunes to the batch's own partitions and removes the target row") + void allUuidV7DeletePrunesAndRemovesTheTargetRow() { + var target = newTrace().build(); + // Same project, so one read shows both: the pruned delete must take the target and leave this one. + var bystander = newTrace().projectName(target.projectName()).build(); + traceResourceClient.createTrace(target, API_KEY, WORKSPACE_NAME); + traceResourceClient.createTrace(bystander, API_KEY, WORKSPACE_NAME); + assertThat(traceIdsOf(target.projectName())).contains(target.id(), bystander.id()); + + // The live user path, end to end. + traceResourceClient.deleteTrace(target.id(), WORKSPACE_NAME, API_KEY); + + assertThat(traceIdsOf(target.projectName())) + .as("only the target row is gone") + .doesNotContain(target.id()) + .contains(bystander.id()); + var sql = lastTraceDeleteSql(target.id()); + assertThat(sql).as("the mutation carries the partition predicate").contains(PARTITION_PREDICATE); + assertThat(boundPartitionsOf(sql)) + .as("bounded to exactly the target's own partition, nothing wider") + .containsExactly(partitionOf(target.id())); + } + + @Test + @DisplayName("the DAO's own delete clears every era and binds exactly those partitions") + void deleteClearsEveryEraAndBindsExactlyThosePartitions() { + // The three-way agreement - the migration's PARTITION BY as installed, the DAO's predicate, and + // WeeklyPartitions.of - asserted through the DAO's own delete rather than by re-evaluating the expression in + // test SQL. If the predicate resolved to any partition other than the one ClickHouse filed a row under, the + // mutation would select the wrong parts and that row would SURVIVE. So "every row is gone" IS the agreement. + // + // Ids are minted from the named Mondays rather than written out, so both sides of the assertion are the same + // arithmetic a reader can check, and every era in one batch is also the multi-value Long[] bind. Seeded raw and + // in a minted project: the ingestion window is 24h, so no endpoint can create a 1996 or 2199 row, and the + // project id is a real UUIDv7 straight from IdGenerator - which is how the sibling partition suites get one. + var projectId = ID_GENERATOR.generateId(); + var ids = ERA_MONDAYS.stream().map(TracesPartitionPruningMutationTest::idInWeekOf).toList(); + ids.forEach(id -> insertRawTrace(projectId, id)); + assertThat(ids.stream().map(id -> liveRowCount(projectId, id))) + .as("every era is seeded").containsOnly("1"); + + delete(ids.stream().map(id -> Pair.of(projectId, id)).collect(Collectors.toUnmodifiableSet())); + + assertThat(ids.stream().map(id -> liveRowCount(projectId, id))) + .as("every era's row is gone, so the predicate named the partition each was actually filed under") + .containsOnly("0"); + var sql = lastTraceDeleteSql(ids.getFirst()); + assertThat(sql).as("the mutation carries the partition predicate").contains(PARTITION_PREDICATE); + assertThat(boundPartitionsOf(sql)) + .as("exactly the partitions the batch resolves to, not a range across two centuries") + .containsExactlyInAnyOrderElementsOf(ERA_MONDAYS.stream() + .map(TracesPartitionPruningMutationTest::partitionNameOf) + .collect(Collectors.toUnmodifiableSet())); + } + + @Test + @DisplayName("the planner actually prunes, and the fallback provably does not") + void pruningReachesThePlannerAndTheFallbackDoesNot() { + // Correctness and pruning are different claims, and this is the only test that makes the second one. Deletes + // were already correct before OPIK-6901 - what the change buys is parts touched (3,928/3,928 -> 5/3,928 on + // prod-test), so a suite that cannot see pruning stop does not test what this change exists to do. + // + // The regression it guards is specific: a migration rewrites the partition expression to something semantically + // identical but textually different, the planner stops recognising the DAO's predicate as the partition key, + // pruning silently stops - and values still agree, so every row is still deleted and every other assertion in + // this suite stays green. That is the property the removed AST pin covered; this asks the planner directly + // instead of inferring it from text. + // + // EXPLAIN does not accept a mutation, so the WHERE clause is lifted verbatim out of the DAO's own emitted + // DELETE and put behind a SELECT - predicate and bound partition values included. Only the verb changes; the + // statement being explained is still the DAO's. Same instrument and record shape as + // TracesLocalV2PartitioningTest. + // One row per era, so the table holds several partitions for the planner to prune between. + var projectId = ID_GENERATOR.generateId(); + var ids = ERA_MONDAYS.stream().map(TracesPartitionPruningMutationTest::idInWeekOf).toList(); + ids.forEach(id -> insertRawTrace(projectId, id)); + + // Bounded: one derivable id, so the predicate names a single one of those partitions. + delete(Set.of(Pair.of(projectId, ids.getFirst()))); + var bounded = partsSelectedBy(lastTraceDeleteSql(ids.getFirst())) + .orElseThrow(() -> new AssertionError( + "EXPLAIN reported no partition index for the bounded delete")); + + // Unbounded: a non-v7 id in the batch, so no predicate at all. Its partner is a different era, so the query_log + // lookup finds this statement rather than the one above. + var partner = ids.get(1); + delete(Set.of(Pair.of(projectId, partner), Pair.of(projectId, NON_V7_ID))); + var unbounded = partsSelectedBy(lastTraceDeleteSql(partner)); + + assertThat(bounded.selected()) + .as("the bounded delete selects fewer parts than the table holds: %s", bounded) + .isLessThan(bounded.total()); + // Shown to discriminate, or it proves nothing - the same trap as `.contains(partition)` and + // `doesNotContain("toDayOfWeek")`. The fallback must not prune: either the planner reports no partition index at + // all, because nothing filters on the key, or it reports every part still selected. + assertThat(unbounded.map(parts -> parts.selected() == parts.total()).orElse(true)) + .as("the fallback prunes nothing: %s", unbounded) + .isTrue(); + } + + @Test + @DisplayName("the topology setup is a no-op once the estate provides it - the path it takes post-cutover") + void topologySetupIsANoOpOnceTheEstateProvidesIt() { + // Today this suite installs the topology itself, so both setup steps take their INSTALL path and their + // early returns are dead code. Once the cutover migration lands, the migrations will provide + // `traces_local` partitioned with the Distributed `traces` over it: the EXCHANGE and the wrap are no + // longer needed, and that early return becomes the ONLY path either step takes. Nothing would exercise it + // until the day it becomes load-bearing, which is the wrong day to find out it was wrong. + // + // Re-running the setup against the topology it already installed IS that shape - `traces_local` exists and + // `traces` is Distributed over it, which is what the migration will hand us. So this covers the second of + // the two states: with the swap (every other test here) and without it (this one). + // A row that already exists BEFORE the setup runs, created through the ingestion path so the check reads it + // back the way production does. Metadata on its own cannot see data loss: a table recreated from the same + // DDL reports the same engine_full and partition_key while being empty, so a step that rebuilt the topology + // rather than skipping it would satisfy every metadata assertion here. The surviving row is the assertion + // that distinguishes "skipped" from "rebuilt", and reading it through the wrapper also shows routing is + // intact rather than merely that the wrapper exists. + var existing = newTrace().build(); + traceResourceClient.createTrace(existing, API_KEY, WORKSPACE_NAME); + assertThat(traceIdsOf(existing.projectName())) + .as("the pre-existing row is readable before the setup re-runs") + .contains(existing.id()); + + // Not a tautology either: if either step failed to early-return it would THROW, not quietly repeat itself. + // The EXCHANGE needs `traces_local_v2`, which the install renamed to `traces_pre_cutover_backup`; and the + // wrap ends in a RENAME onto `traces_local`, which by now exists. So a non-idempotent step fails loudly. + var tracesEngineBefore = queryOneString(TABLE_ENGINE_FULL, _ -> { + }); + var localPartitionKeyBefore = partitionKeyOf("traces_local"); + + ensurePartitionedSuccessorUnderTraces(); + ensureDistributedWrap(); + + assertThat(traceIdsOf(existing.projectName())) + .as("the row that existed before the setup is still there, and still routed through the wrapper") + .contains(existing.id()); + assertThat(queryOneString(TABLE_ENGINE_FULL, _ -> { + })) + .as("re-running the setup left the Distributed wrapper untouched") + .isEqualTo(tracesEngineBefore); + assertThat(partitionKeyOf("traces_local")) + .as("and left the partitioned table's key untouched") + .isEqualTo(localPartitionKeyBefore); + + // Not just survivable - still testing what it claims. A pruned delete on the untouched topology, so a + // no-op setup cannot quietly leave the suite asserting against something that is no longer partitioned. + var projectId = ID_GENERATOR.generateId(); + var id = idInWeekOf(ERA_MONDAYS.getFirst()); + insertRawTrace(projectId, id); + + delete(Set.of(Pair.of(projectId, id))); + + assertThat(liveRowCount(projectId, id)) + .as("the row is still deleted after a no-op setup") + .isEqualTo("0"); + assertThat(lastTraceDeleteSql(id)) + .as("and the delete is still pruned") + .contains(PARTITION_PREDICATE); + } + + @Test + @DisplayName("a request spanning two chunks prunes each chunk on its own") + void requestSpanningTwoChunksPrunesEachChunkIndependently() { + // The DAO chunks a request at ANALYTICS_DELETE_BATCH_SIZE and derives partitions PER CHUNK, inside the + // concatMap - so "all-or-nothing" is a per-statement guarantee, not a per-request one. Every other test + // here passes a handful of pairs, which is one chunk, so none of them can see that. + // + // The non-v7 id goes in the FIRST chunk and the derivable ids in the second, which is the only arrangement + // that can actually discriminate. Chunks are sized [BATCH_SIZE, remainder], so the first is always full and + // never inspectable: query_log truncates at log_queries_cut_to_length (100,000 bytes here) and a + // 10,000-pair statement inlines to ~762 KiB, so its tail - where the predicate sits - is not recorded. + // Only the remainder chunk is small enough to read back, so the assertion that matters has to live there. + // + // That makes the test bite on the refactor it exists to catch. Hoisting weeklyPartitionsFor out of the + // lambda, so the whole request is derived once, would let the non-v7 id in chunk one strip pruning from + // chunk two as well - and chunk two's predicate is exactly what is asserted below. Removing the pruning + // outright fails the same assertion. With the arrangement reversed, both regressions passed. + var projectId = ID_GENERATOR.generateId(); + var firstChunkRow = idInWeekOf(ERA_MONDAYS.getFirst()); + var secondChunkRow = idInWeekOf(ERA_MONDAYS.getLast()); + insertRawTrace(projectId, firstChunkRow); + insertRawTrace(projectId, secondChunkRow); + + // Chunk one: a real row, the non-v7 id, and filler up to exactly ANALYTICS_DELETE_BATCH_SIZE. Filler ids + // match no row - a delete does not need its ids to exist, and the chunk boundary is what is under test. + var ordered = new ArrayList(); + ordered.add(firstChunkRow); + ordered.add(NON_V7_ID); + while (ordered.size() < ANALYTICS_DELETE_BATCH_SIZE) { + ordered.add(idInWeekOf(ERA_MONDAYS.getFirst())); + } + // Chunk two: the remainder, all derivable, in two different weeks so the bound set is exact rather than + // trivially a single value. + var secondChunkCompanion = idInWeekOf(ERA_MONDAYS.get(1)); + ordered.add(secondChunkRow); + ordered.add(secondChunkCompanion); + + delete(ordered.stream().map(id -> Pair.of(projectId, id)).collect(Collectors.toCollection( + LinkedHashSet::new))); + + assertThat(liveRowCount(projectId, firstChunkRow)) + .as("the row in the chunk that fell back to unbounded is deleted") + .isEqualTo("0"); + assertThat(liveRowCount(projectId, secondChunkRow)) + .as("and so is the row in the chunk that pruned") + .isEqualTo("0"); + + // The full chunk's statement is only checked to exist - that is what shows the request was split at all. + // Nothing about its text can be asserted, for the truncation reason above. + deleteSqlForChunkOf(ANALYTICS_DELETE_BATCH_SIZE); + + var secondChunkSql = deleteSqlForChunkOf(2); + assertThat(secondChunkSql) + .as("the all-derivable chunk prunes even though an earlier chunk could not") + .contains(PARTITION_PREDICATE); + assertThat(boundPartitionsOf(secondChunkSql)) + .as("bounded to exactly its own two weeks") + .containsExactlyInAnyOrder(partitionNameOf(ERA_MONDAYS.getLast()), + partitionNameOf(ERA_MONDAYS.get(1))); + } + + @ParameterizedTest + @MethodSource + @DisplayName("an id with no derivable partition disables pruning for the batch, and the delete still lands") + void underivableIdDisablesPruning(String cause, UUID underivableId) { + // The fallback that preserves the pre-OPIK-6901 guarantee, as the original javadoc stated it: a row whose id_at + // cannot be trusted is STILL DELETED. That is a claim about the underivable row ITSELF, so it gets a real row + // here - seeded raw, since ingestion rejects both id shapes by design. Passing it as an id matching nothing + // would let an implementation that quietly drops underivable ids from the batch pass, which is the very bug the + // all-or-nothing rule exists to prevent. + var target = newTrace().build(); + traceResourceClient.createTrace(target, API_KEY, WORKSPACE_NAME); + var projectId = projectIdOf(target); + insertRawTrace(projectId, underivableId); + assertThat(liveRowCount(projectId, underivableId)) + .as("the %s row is seeded before the delete", cause).isEqualTo("1"); + + delete(Set.of(Pair.of(projectId, target.id()), Pair.of(projectId, underivableId))); + + assertThat(liveRowCount(projectId, underivableId)) + .as("the %s row is itself deleted, not skipped", cause) + .isEqualTo("0"); + assertThat(traceIdsOf(target.projectName())) + .as("and the derivable row batched alongside the %s id goes too", cause) + .doesNotContain(target.id()); + + // Asserted as the absence of ANY id_at predicate, not just of this PR's expression. A regression that narrowed + // the mutation with toMonday(id_at), an id_at range, or any other partition predicate would skip exactly the + // rows this fallback exists to reach, and rejecting one function name would not see it. The unbounded template + // mentions id_at nowhere at all, so that is the whole check. + var sql = lastTraceDeleteSql(target.id()); + assertThat(sql) + .as("the unbounded form for a %s batch carries no id_at predicate of any kind", cause) + .doesNotContain("id_at"); + assertThat(EMITTED_IN_CLAUSE.matcher(sql).find()) + .as("and no partition IN clause: %s", sql) + .isFalse(); + } + + private static Stream underivableIdDisablesPruning() { + return Stream.of( + arguments("non-v7", NON_V7_ID), + arguments("beyond-2299", OUT_OF_RANGE_ID)); + } + } + + /** + * The flag off, against the same post-cutover topology and the same fixtures — so the only difference + * from {@link PruningEnabled} is the flag itself. That is what makes the pair a control rather than two unrelated + * suites: the sibling {@code TracesPruningDisabledMutationTest} also runs with pruning off, but against the legacy + * table, so it varies the schema at the same time and cannot attribute anything to the flag alone. + *

+ * Both assertions here are the inverse of one in {@link PruningEnabled}, on identical data: no partition predicate + * is emitted, and the planner selects every part. Delete the flag gate so pruning always happens and these fail; + * delete the pruning and the other class fails. Correctness is unaffected either way, which is the point — the rows + * go away in both, so only these assertions can tell the two states apart. + */ + @Nested + @TestInstance(TestInstance.Lifecycle.PER_CLASS) + @ExtendWith(DropwizardAppExtensionProvider.class) + class PruningDisabled { + + @RegisterApp + private final TestDropwizardAppExtension app = newApp(false); + + @BeforeAll + void beforeAll(ClientSupport clientSupport, TraceDAO traceDAO, TransactionTemplateAsync appTemplate) { + bindApp(clientSupport, traceDAO, appTemplate); + } + + @Test + @DisplayName("the same batch still clears every era, and emits no partition predicate") + void deleteStillClearsEveryEraWithoutPruning() { + var projectId = ID_GENERATOR.generateId(); + var ids = ERA_MONDAYS.stream().map(TracesPartitionPruningMutationTest::idInWeekOf).toList(); + ids.forEach(id -> insertRawTrace(projectId, id)); + + delete(ids.stream().map(id -> Pair.of(projectId, id)).collect(Collectors.toUnmodifiableSet())); + + assertThat(ids.stream().map(id -> liveRowCount(projectId, id))) + .as("the delete is still correct with the flag off - that is what makes it an optimisation") + .containsOnly("0"); + var sql = lastTraceDeleteSql(ids.getFirst()); + assertThat(sql) + .as("and carries no id_at narrowing of any kind") + .doesNotContain("id_at"); + assertThat(EMITTED_IN_CLAUSE.matcher(sql).find()) + .as("nor a partition IN clause: %s", sql) + .isFalse(); + } + + @Test + @DisplayName("the planner selects every part - the enabled class's assertion, inverted on the same data") + void plannerPrunesNothingWithoutPruning() { + var projectId = ID_GENERATOR.generateId(); + var ids = ERA_MONDAYS.stream().map(TracesPartitionPruningMutationTest::idInWeekOf).toList(); + ids.forEach(id -> insertRawTrace(projectId, id)); + + delete(Set.of(Pair.of(projectId, ids.getFirst()))); + + var parts = partsSelectedBy(lastTraceDeleteSql(ids.getFirst())); + assertThat(parts.map(selected -> selected.selected() == selected.total()).orElse(true)) + .as("no partition pruning with the flag off: %s", parts) + .isTrue(); + } + } +} diff --git a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/TracesPruningDisabledMutationTest.java b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/TracesPruningDisabledMutationTest.java new file mode 100644 index 00000000000..736e2c50726 --- /dev/null +++ b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/TracesPruningDisabledMutationTest.java @@ -0,0 +1,287 @@ +package com.comet.opik.infrastructure; + +import com.comet.opik.api.Trace; +import com.comet.opik.api.resources.utils.ClickHouseContainerUtils; +import com.comet.opik.api.resources.utils.ClientSupportUtils; +import com.comet.opik.api.resources.utils.MigrationUtils; +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.TestDropwizardAppExtensionUtils.AppContextConfig; +import com.comet.opik.api.resources.utils.TestUtils; +import com.comet.opik.api.resources.utils.WireMockUtils; +import com.comet.opik.api.resources.utils.resources.TraceResourceClient; +import com.comet.opik.domain.TraceDAO; +import com.comet.opik.extensions.DropwizardAppExtensionProvider; +import com.comet.opik.extensions.RegisterApp; +import com.comet.opik.infrastructure.auth.RequestContext; +import com.comet.opik.infrastructure.db.TransactionTemplateAsync; +import com.comet.opik.podam.PodamFactoryUtils; +import com.redis.testcontainers.RedisContainer; +import io.r2dbc.spi.Statement; +import org.apache.commons.lang3.RandomStringUtils; +import org.apache.commons.lang3.tuple.Pair; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; +import org.junit.jupiter.api.extension.ExtendWith; +import org.testcontainers.clickhouse.ClickHouseContainer; +import org.testcontainers.containers.GenericContainer; +import org.testcontainers.lifecycle.Startables; +import org.testcontainers.mysql.MySQLContainer; +import reactor.core.publisher.Mono; +import ru.vyarus.dropwizard.guice.test.ClientSupport; +import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension; +import uk.co.jemos.podam.api.PodamFactory; + +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.UUID; +import java.util.function.Consumer; + +import static com.comet.opik.api.resources.utils.AuthTestUtils.mockTargetWorkspace; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assumptions.assumeThat; + +/** + * The other side of {@code TracesPartitionPruningMutationTest}: the trace delete with + * {@code databaseAnalyticsDataModel.tracesWeeklyPartitionPruningEnabled} off — its default — against the + * legacy {@code traces}. That is not an exotic topology; it is the state every deployment is in today, and the + * state a stage B/C rollback returns to. + * + *

What it guards is a false-flag regression: the partition predicate emitted while the flag is off, or any + * other {@code id_at} narrowing creeping into the unbounded form. Legacy {@code traces} has no {@code PARTITION BY} at + * all and declares {@code id_at} as a 32-bit {@code DateTime} that overflows past 2106 (migrations 000001 and 000091), + * so a predicate here prunes nothing and can silently exclude rows. + * + *

Nothing else catches it. Every other trace-delete suite runs in exactly this state and asserts only that rows go + * away — which they do for a recent id, because the legacy {@code id_at} is accurate for one. The damage shows + * only on a far-future id (litellm BerriAI/litellm#31294 + * mints ~2201): the 32-bit column stores a wrapped recent timestamp, a derived partition cannot match it, and the + * delete reports success having matched zero rows. No existing test has such a row, because ingestion rejects + * far-future ids by design ({@code IdGenerator.validateId}) — so this seeds one in raw SQL and deletes it through + * {@link TraceDAO#delete}, which is also the only way to reach that arm. + * + *

Shared, reusable containers: unlike the post-EXCHANGE suite this changes no topology, so there is nothing here + * that could corrupt another suite or a rerun. + */ +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +@ExtendWith(DropwizardAppExtensionProvider.class) +class TracesPruningDisabledMutationTest { + + private static final String API_KEY = "apiKey-" + UUID.randomUUID(); + private static final String WORKSPACE_NAME = "workspace-" + RandomStringUtils.secure().nextAlphanumeric(32); + private static final String WORKSPACE_ID = UUID.randomUUID().toString(); + private static final String USER = "user-" + RandomStringUtils.secure().nextAlphanumeric(32); + + /** A UUIDv7 carrying id_at 2200-01-01 — the litellm shape, and the id the legacy 32-bit column wraps. */ + private static final UUID FAR_FUTURE_ID = UUID.fromString("0699eb8a-59dd-7215-8000-03b8d2a8d5e2"); + + private static final String INSERT_RAW_TRACE = """ + INSERT INTO traces (workspace_id, project_id, id) + VALUES (:workspace_id, :project_id, :id) + """; + + /** + * Scoped by the same key {@code TraceDAO.delete} matches on — {@code (workspace_id, project_id, id)}. An oracle + * narrower than the delete answers a different question, and this suite runs on shared containers where other + * suites' rows are present, so the scoping is what keeps the count about this test's row. + */ + private static final String LIVE_ROW_COUNT = """ + SELECT toString(uniqExact(id)) + FROM traces + WHERE workspace_id = :workspace_id + AND project_id = :project_id + AND id = :id + """; + + private static final String LAST_TRACE_DELETE = """ + SELECT query + FROM system.query_log + WHERE log_comment LIKE 'delete_traces:%' + AND type = 'QueryFinish' + AND query LIKE concat('%', :trace_id, '%') + ORDER BY event_time_microseconds DESC + LIMIT 1 + """; + + private static final String LEGACY_SCHEMA = """ + SELECT concat( + (SELECT partition_key FROM system.tables + WHERE database = currentDatabase() AND name = 'traces'), + '|', + (SELECT type FROM system.columns + WHERE database = currentDatabase() AND table = 'traces' AND name = 'id_at')) + """; + + private final RedisContainer redisContainer = RedisContainerUtils.newRedisContainer(); + private final MySQLContainer mysqlContainer = MySQLContainerUtils.newMySQLContainer(); + private final GenericContainer zookeeperContainer = ClickHouseContainerUtils.newZookeeperContainer(); + private final ClickHouseContainer clickHouseContainer = ClickHouseContainerUtils + .newClickHouseContainer(zookeeperContainer); + + private final WireMockUtils.WireMockRuntime wireMock; + + private final PodamFactory factory = PodamFactoryUtils.newPodamFactory(); + + @RegisterApp + private final TestDropwizardAppExtension app; + + { + Startables.deepStart(redisContainer, mysqlContainer, clickHouseContainer, zookeeperContainer).join(); + wireMock = WireMockUtils.startWireMock(); + var databaseAnalyticsFactory = ClickHouseContainerUtils.newDatabaseAnalyticsFactory( + clickHouseContainer, ClickHouseContainerUtils.DATABASE_NAME); + MigrationUtils.runMysqlDbMigration(mysqlContainer); + MigrationUtils.runClickhouseDbMigration(clickHouseContainer); + // No customConfigs on purpose: the defaults ARE the state under test. + app = TestDropwizardAppExtensionUtils.newTestDropwizardAppExtension( + AppContextConfig.builder() + .jdbcUrl(mysqlContainer.getJdbcUrl()) + .databaseAnalyticsFactory(databaseAnalyticsFactory) + .redisUrl(redisContainer.getRedisURI()) + .runtimeInfo(wireMock.runtimeInfo()) + .build()); + } + + private TraceResourceClient traceResourceClient; + private TransactionTemplateAsync template; + private TraceDAO traceDAO; + + @BeforeAll + void beforeAll(ClientSupport clientSupport, TransactionTemplateAsync template, TraceDAO traceDAO) { + var baseUrl = TestUtils.getBaseUrl(clientSupport); + ClientSupportUtils.config(clientSupport); + mockTargetWorkspace(wireMock.server(), API_KEY, WORKSPACE_NAME, WORKSPACE_ID, USER); + traceResourceClient = new TraceResourceClient(clientSupport, baseUrl); + this.template = template; + this.traceDAO = traceDAO; + + // This whole suite is about the pre-cutover / post-rollback state. Once the cutover migration lands there is no + // legacy `traces` left to test - it becomes the partitioned successor with a DateTime64 id_at - and the + // far-future wrapping hazard this guards stops existing. Skip with the reason stated rather than fail on a + // premise that has legitimately gone away; JUnit reports the skip and its description in the test report. + assumeThat(legacySchema()) + .as("pre-cutover estate: legacy `traces`, unpartitioned with a 32-bit DateTime id_at") + .isEqualTo("|DateTime('UTC')"); + } + + @AfterAll + void afterAll() { + wireMock.server().stop(); + } + + @Test + @DisplayName("with pruning off, a far-future row on the legacy table is still deleted, unbounded") + void farFutureRowOnLegacyTableIsStillDeleted() { + // The project and its id come from the real ingestion path - create a trace through the endpoint, then read the + // project id back off it - so the delete runs against a project that exists and a genuine UUIDv7 project id, + // not a fabricated one. + var seedTrace = factory.manufacturePojo(Trace.class).toBuilder() + .feedbackScores(null) + .usage(null) + .build(); + traceResourceClient.createTrace(seedTrace, API_KEY, WORKSPACE_NAME); + var projectId = projectIdOf(seedTrace); + assertThat(projectId.version()).as("the project id is a real UUIDv7, as the backend mints them").isEqualTo(7); + + // Only the far-future row is raw, because ingestion rejects it by design (24h window). A recent id would prove + // nothing here anyway: the legacy id_at is accurate for one, so even a wrongly-emitted predicate would match it + // and the row would still go. + insertRawTrace(projectId, FAR_FUTURE_ID); + assertThat(liveRowCount(projectId, FAR_FUTURE_ID)).as("the far-future row is seeded").isEqualTo("1"); + + delete(Set.of(Pair.of(projectId, FAR_FUTURE_ID))); + + assertThat(liveRowCount(projectId, FAR_FUTURE_ID)) + .as("it is deleted - a partition predicate here would have matched nothing and reported success") + .isEqualTo("0"); + + var sql = lastTraceDeleteSql(FAR_FUTURE_ID); + assertThat(sql) + .as("no partition predicate is emitted while the flag is off") + .doesNotContain("toYYYYMMDD", "toDayOfWeek", "toIntervalDay"); + assertThat(sql) + .as("and no id_at narrowing of any kind - toMonday, a range, or otherwise") + .doesNotContain("id_at"); + } + + /** Invokes the DAO under a workspace/user context, as {@code TraceService} does for the live delete path. */ + private void delete(Set> projectIdTraceIdPairs) { + template.nonTransaction(connection -> traceDAO.delete(projectIdTraceIdPairs, connection)) + .contextWrite(ctx -> ctx + .put(RequestContext.WORKSPACE_ID, WORKSPACE_ID) + .put(RequestContext.USER_NAME, USER)) + .block(); + } + + /** + * Seeds one row through the table's real column definitions. Only the identity columns are supplied; {@code id_at} + * is {@code MATERIALIZED}, so ClickHouse derives it — and, on this table, wraps it — exactly as it would for a row + * the ingestion path wrote. + */ + private void insertRawTrace(UUID projectId, UUID id) { + execute(INSERT_RAW_TRACE, statement -> statement + .bind("workspace_id", WORKSPACE_ID) + .bind("project_id", projectId.toString()) + .bind("id", id.toString())); + } + + /** The real project id the ingestion path minted, read back off the created trace. */ + private UUID projectIdOf(Trace trace) { + return traceResourceClient + .getTraces(trace.projectName(), null, API_KEY, WORKSPACE_NAME, List.of(), List.of(), 100, Map.of()) + .content().stream() + .filter(found -> found.id().equals(trace.id())) + .map(Trace::projectId) + .findFirst() + .orElseThrow(); + } + + /** {@code "1"} while a live (non-lightweight-deleted) row exists for the id, {@code "0"} once it is gone. */ + private String liveRowCount(UUID projectId, UUID id) { + return queryOneString(LIVE_ROW_COUNT, statement -> statement + .bind("workspace_id", WORKSPACE_ID) + .bind("project_id", projectId.toString()) + .bind("id", id.toString())); + } + + private String legacySchema() { + return queryOneString(LEGACY_SCHEMA, _ -> { + }); + } + + /** + * The SQL of the trace delete that carried {@code traceId}. Filtered by {@code log_comment} and the id, so a + * neighbouring suite's delete in this shared container cannot be picked up instead — the id narrows it to this test. + */ + private String lastTraceDeleteSql(UUID traceId) { + execute("SYSTEM FLUSH LOGS", _ -> { + }); + var sql = queryOneString(LAST_TRACE_DELETE, statement -> statement.bind("trace_id", traceId.toString())); + assertThat(sql) + .as("query_log holds a delete_traces statement mentioning id '%s'", traceId) + .isNotNull(); + return sql; + } + + private String queryOneString(String sql, Consumer binder) { + return template.nonTransaction(connection -> { + var statement = connection.createStatement(sql); + binder.accept(statement); + return Mono.from(statement.execute()) + .flatMap(result -> Mono.from(result.map((row, _) -> row.get(0, String.class)))); + }).block(); + } + + private void execute(String sql, Consumer binder) { + template.nonTransaction(connection -> { + var statement = connection.createStatement(sql); + binder.accept(statement); + return Mono.from(statement.execute()).flatMap(result -> Mono.from(result.getRowsUpdated())); + }).block(); + } +} diff --git a/apps/opik-backend/src/test/java/com/comet/opik/utils/WeeklyPartitionsTest.java b/apps/opik-backend/src/test/java/com/comet/opik/utils/WeeklyPartitionsTest.java new file mode 100644 index 00000000000..30f3ea12582 --- /dev/null +++ b/apps/opik-backend/src/test/java/com/comet/opik/utils/WeeklyPartitionsTest.java @@ -0,0 +1,179 @@ +package com.comet.opik.utils; + +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +import java.util.List; +import java.util.Set; +import java.util.UUID; +import java.util.stream.Stream; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.junit.jupiter.params.provider.Arguments.arguments; + +/** + * Covers {@link WeeklyPartitions#of}, which derives the weekly partition values a delete batch resolves to so the + * mutation can prune instead of rewriting every part. + *

+ * The in-range expected values are not hand-computed: each is what ClickHouse itself returned for + * {@code toYYYYMMDD(toDate32(id_at) - toIntervalDay(toDayOfWeek(id_at, 1)))} on prod-test for that id. If the table's + * partition expression ever changes, these assertions are what should fail. + *

+ * The range-boundary cases are the exception and say so: the ids are constructed rather than observed, because the point + * of each is an {@code id_at} value the column cannot store, which is exactly what no real row has. + */ +class WeeklyPartitionsTest { + + @ParameterizedTest(name = "{0}") + @MethodSource + @DisplayName("matches the partition ClickHouse computed") + void matchesThePartitionClickHouseComputed(String era, UUID id, long expectedPartition) { + assertThat(WeeklyPartitions.of(List.of(id))).contains(Set.of(expectedPartition)); + } + + /** + * One row per era the derivation has to get right, each pinned to the value ClickHouse itself returned for that id + * on prod-test — not to a hand-computed Monday. The eras are not interchangeable samples: an expression that is + * correct for the ordinary calendar and wrong at the extremes is exactly the {@code toMonday} trap migration 000114 + * was written to escape, so dropping any of the three would weaken the pin rather than tidy it. + */ + private static Stream matchesThePartitionClickHouseComputed() { + return Stream.of( + // The ordinary case: id_at 2026-08-19 (a Wednesday) -> Monday 2026-08-17. + arguments("ordinary id", UUID.fromString("01a01a75-76de-785e-ae84-8870ed5e6db3"), 20260817L), + // Long before Opik existed but well after the Unix epoch, and well inside Date32's 1900 floor: an id + // this old prunes like any other. id_at 1996-02-09 -> Monday 1996-02-05. + arguments("id from 1996", UUID.fromString("00bfd451-fa93-7c10-9923-88a219a974c8"), 19960205L), + // Far-future ids are supported, not excluded: a bogus-but-self-consistent timestamp, id_at 2200-01-01 -> + // Monday 2199-12-30. 4.1% of rows on prod-test look like this; they must still be deletable and prune. + arguments("far-future id", UUID.fromString("0699eb8a-59dd-7215-8000-03b8d2a8d5e2"), 21991230L)); + } + + @Test + @DisplayName("a scattered batch yields the exact set, not a range") + void scatteredBatch() { + // 1996 and 2026 in one batch. An id_at RANGE over this span selected 2,644 of 3,928 parts on prod-test; + // the exact set selected 4. This is the reason the predicate is a set. + assertThat(WeeklyPartitions.of(List.of( + UUID.fromString("01a01a75-76de-785e-ae84-8870ed5e6db3"), + UUID.fromString("00bfd451-fa93-7c10-9923-88a219a974c8")))) + .contains(Set.of(20260817L, 19960205L)); + } + + @Test + @DisplayName("a non-v7 id disables pruning for the whole batch") + void nonV7DisablesPruning() { + // All-or-nothing on purpose: deriving a partition from a non-v7 id reads whatever sits in the timestamp + // field, and a wrong partition is a SILENTLY skipped delete. Omitting the predicate keeps the old behaviour. + assertThat(WeeklyPartitions.of(List.of( + UUID.fromString("01a01a75-76de-785e-ae84-8870ed5e6db3"), + UUID.fromString("9f527bac-527a-4f92-8875-0fa8af8e4f22")))) // v4 + .isEmpty(); + } + + @Test + @DisplayName("a single non-v7 id yields no partitions") + void singleNonV7() { + assertThat(WeeklyPartitions.of(List.of(UUID.fromString("9f527bac-527a-4f92-8875-0fa8af8e4f22")))) + .isEmpty(); + } + + @Test + @DisplayName("duplicate ids in the same week collapse to one partition") + void duplicatesCollapse() { + assertThat(WeeklyPartitions.of(List.of( + UUID.fromString("01a01a75-76de-785e-ae84-8870ed5e6db3"), + UUID.fromString("01a01a75-6f8e-7f22-9279-ee4f7ca7810d"), + UUID.fromString("01a01a75-609d-7935-8d22-2dd8dfeb2454")))) + .contains(Set.of(20260817L)); + } + + @Test + @DisplayName("an empty batch yields no partitions") + void emptyBatch() { + assertThat(WeeklyPartitions.of(List.of())).isEmpty(); + } + + @Test + @DisplayName("the returned set is immutable, so a delete's partitions cannot be narrowed after derivation") + void returnedSetIsImmutable() { + // Not a general hygiene assertion: this set IS the partition list a DELETE binds, so a caller that removed an + // entry would turn a correct delete into one that matches nothing and reports success. + var partitions = WeeklyPartitions.of(List.of( + UUID.fromString("01a01a75-76de-785e-ae84-8870ed5e6db3"), + UUID.fromString("00bfd451-fa93-7c10-9923-88a219a974c8"))) + .orElseThrow(); + + assertThatThrownBy(() -> partitions.remove(20260817L)) + .isInstanceOf(UnsupportedOperationException.class); + assertThat(partitions).containsExactlyInAnyOrder(20260817L, 19960205L); + } + + @Test + @DisplayName("a null batch throws rather than reading as an unprunable one") + void nullBatchThrows() { + // The one intolerant case, and deliberately not folded into the empty result above: empty is a documented + // answer ("emit the unbounded form"), so a caller that lost its batch would get a valid-looking answer, issue a + // correct-but-unbounded mutation, and never learn it had a bug. Matches every other collection-taking method in + // this package, all of which are @NonNull. + assertThatThrownBy(() -> WeeklyPartitions.of(null)) + .isInstanceOf(NullPointerException.class) + .hasMessageContaining("ids"); + } + + @Test + @DisplayName("the last id_at the column can store still prunes") + void lastRepresentableIdStillPrunes() { + // id_at 2299-12-31T23:59:59.999 — the last instant DateTime64 represents, so nothing saturates and the two + // sides agree: DateTime64(0) truncates to 23:59:59, whose Date32 is 2299-12-31 (a Sunday) -> Monday 2299-12-25. + // The ceiling check must be exclusive at exactly this point, hence a case sitting on it. + assertThat(WeeklyPartitions.of(List.of(UUID.fromString("0978a65f-77ff-7abc-8000-000000000001")))) + .contains(Set.of(22991225L)); + } + + @Test + @DisplayName("an id one millisecond past the id_at ceiling disables pruning") + void firstUnrepresentableIdDisablesPruning() { + // id_at 2300-01-01T00:00:00 — one ms past the previous case and outside DateTime64. ClickHouse saturates it to + // 2299-12-31 23:59:59 and files the row under 22991225, so the honest week this would compute (2300-01-01 is + // itself a Monday, giving 23000101) is a partition the row is NOT in. Pruning off rather than clamped: matching + // ClickHouse here would mean reproducing its saturation semantics, for ids that should not exist. + assertThat(WeeklyPartitions.of(List.of(UUID.fromString("0978a65f-7800-7abc-8000-000000000001")))) + .isEmpty(); + } + + @Test + @DisplayName("the largest timestamp a UUIDv7 can carry disables pruning, and does not throw") + void largestUuidV7TimestampDisablesPruning() { + // All 48 timestamp bits set: 10889-08-02, the furthest future any UUIDv7 can encode. Read unsigned it is still + // only ~2.8e14 ms, far inside Instant's range — so the guard is what excludes it, not an exception, and there + // is no input on which this can throw. + assertThat(WeeklyPartitions.of(List.of(UUID.fromString("ffffffff-ffff-7abc-8000-000000000001")))) + .isEmpty(); + } + + @Test + @DisplayName("one out-of-range id disables pruning for the whole batch") + void outOfRangeIdDisablesPruningForTheWholeBatch() { + // Same all-or-nothing rule as a non-v7 id, for the same reason: a set derived from the rest of the batch is a + // set this row is not in. + assertThat(WeeklyPartitions.of(List.of( + UUID.fromString("01a01a75-76de-785e-ae84-8870ed5e6db3"), + UUID.fromString("ffffffff-ffff-7abc-8000-000000000001")))) + .isEmpty(); + } + + @Test + @DisplayName("the earliest id a UUIDv7 can carry is inside Date32, so there is no floor to guard") + void earliestUuidV7TimestampIsInRange() { + // All 48 timestamp bits clear: id_at 1970-01-01, the earliest any UUIDv7 can encode (the field is unsigned). + // Its Monday, 1969-12-29, is 70 years above Date32's 1900 floor, so a below-1900 id_at is unreachable by + // construction rather than merely untested — which is why `of` guards only the ceiling. + assertThat(WeeklyPartitions.of(List.of(UUID.fromString("00000000-0000-7abc-8000-000000000001")))) + .contains(Set.of(19691229L)); + } +} diff --git a/apps/opik-backend/src/test/resources/config-test.yml b/apps/opik-backend/src/test/resources/config-test.yml index 3891da65200..114caa3463e 100644 --- a/apps/opik-backend/src/test/resources/config-test.yml +++ b/apps/opik-backend/src/test/resources/config-test.yml @@ -128,6 +128,20 @@ databaseAnalyticsDataModel: # MODIFY TTL target `traces_local` only (the Distributed `traces` rejects them); ADD/DROP/MODIFY COLUMN must target # both `traces_local` and `traces`, else reads can't see the column (code 47). tracesDistributedWrapEnabled: false + # Default: false + # Description: Enables partition-aware PRUNING of trace deletes - it does NOT create or activate any partitioning. + # With it on, a trace DELETE bounds itself to the weekly partitions its own ids resolve to instead of being planned + # against every part of the table. Turning it on therefore ASSERTS a schema fact rather than causing one: that the + # live mutation target already IS the weekly partitioned successor, id_at as DateTime64(0,'UTC') under + # PARTITION BY toYYYYMMDD(toDate32(id_at) - toIntervalDay(toDayOfWeek(id_at, 1))). Installing that schema is the + # EXCHANGE step of the cutover, never this flag. Purely an optimisation: false keeps the unbounded mutation, which + # is always correct and merely slower. A third flag on purpose - the partitioning appears at the EXCHANGE, and + # neither sibling marks it: + # traceColumnsNonNullable must be rolled out BEFORE the EXCHANGE, tracesDistributedWrapEnabled flips at the wrap, + # which may be deferred long after it. Leave false at deploy time; set true once the EXCHANGE is confirmed, and back + # to false BEFORE a rollback promotes the original `traces` (legacy `traces` has no PARTITION BY and a 32-bit + # DateTime id_at that overflows past 2106, so the predicate would silently match zero rows for a far-future id). + tracesWeeklyPartitionPruningEnabled: false # Description: UUIDv7 ingestion validation. Rejects writes whose `id` embeds a timestamp outside the # window, protecting data quality. diff --git a/deployment/docker-compose/docker-compose.yaml b/deployment/docker-compose/docker-compose.yaml index 4470ccfa05f..4346af90387 100644 --- a/deployment/docker-compose/docker-compose.yaml +++ b/deployment/docker-compose/docker-compose.yaml @@ -188,6 +188,11 @@ services: # Flip TRACES_DISTRIBUTED_WRAP_ENABLED to true in lockstep with applying the Distributed wrap: it routes trace # delete/retention mutations to traces_local, since a Distributed traces rejects mutations. ANALYTICS_DB_DATA_MODEL_TRACES_DISTRIBUTED_WRAP_ENABLED: ${ANALYTICS_DB_DATA_MODEL_TRACES_DISTRIBUTED_WRAP_ENABLED:-false} + # Flip TRACES_WEEKLY_PARTITION_PRUNING_ENABLED to true only once the EXCHANGE has put the weekly partitioned + # successor under `traces`: it lets a trace DELETE prune to the partitions its own ids resolve to. It enables the + # PRUNING, not the partitioning - setting it does not create or activate anything. Purely an optimisation - false + # keeps the always-correct unbounded mutation. + ANALYTICS_DB_DATA_MODEL_TRACES_WEEKLY_PARTITION_PRUNING_ENABLED: ${ANALYTICS_DB_DATA_MODEL_TRACES_WEEKLY_PARTITION_PRUNING_ENABLED:-false} # Readiness probes for the Hyperscale topology, both off for the single-node Compose stack: the cluster probe # checks the Distributed cluster definition is visible, the cold-storage probe checks the cold_s3 tiered disk. ANALYTICS_DB_CLUSTER_HEALTH_CHECK_ENABLED: ${ANALYTICS_DB_CLUSTER_HEALTH_CHECK_ENABLED:-false} diff --git a/deployment/helm_chart/opik/README.md b/deployment/helm_chart/opik/README.md index 10daecbac12..0f5fea5267b 100644 --- a/deployment/helm_chart/opik/README.md +++ b/deployment/helm_chart/opik/README.md @@ -477,6 +477,7 @@ Call opik api on http://localhost:5173/api | databaseAnalyticsDataModel.traceColumnsNonNullable | bool | `false` | | | databaseAnalyticsDataModel.traceDeletionEventsCaptureEnabled | bool | `false` | | | databaseAnalyticsDataModel.tracesDistributedWrapEnabled | bool | `false` | | +| databaseAnalyticsDataModel.tracesWeeklyPartitionPruningEnabled | bool | `false` | | | demoDataJob.enabled | bool | `true` | | | fullnameOverride | string | `""` | | | global.argocd | bool | `false` | | diff --git a/deployment/helm_chart/opik/templates/configmap-backend.yaml b/deployment/helm_chart/opik/templates/configmap-backend.yaml index d80c4a8e53b..4e2d3491e40 100644 --- a/deployment/helm_chart/opik/templates/configmap-backend.yaml +++ b/deployment/helm_chart/opik/templates/configmap-backend.yaml @@ -50,6 +50,7 @@ data: "ANALYTICS_DB_DATA_MODEL_SPAN_DELETION_EVENTS_CAPTURE_ENABLED" $dm.spanDeletionEventsCaptureEnabled "ANALYTICS_DB_DATA_MODEL_DELETION_EVENTS_INSERT_BATCH_SIZE" $dm.deletionEventsInsertBatchSize "ANALYTICS_DB_DATA_MODEL_TRACES_DISTRIBUTED_WRAP_ENABLED" $dm.tracesDistributedWrapEnabled + "ANALYTICS_DB_DATA_MODEL_TRACES_WEEKLY_PARTITION_PRUNING_ENABLED" $dm.tracesWeeklyPartitionPruningEnabled "PARTITION_METRICS_ENABLED" $pm.enabled "PARTITION_METRICS_INTERVAL" $pm.interval "PARTITION_METRICS_LWD_TABLES" $pm.lwdTables diff --git a/deployment/helm_chart/opik/tests/configmap_env_test.yaml b/deployment/helm_chart/opik/tests/configmap_env_test.yaml index c4a1d9e68e1..458a24bfd91 100644 --- a/deployment/helm_chart/opik/tests/configmap_env_test.yaml +++ b/deployment/helm_chart/opik/tests/configmap_env_test.yaml @@ -73,6 +73,9 @@ tests: - equal: path: data.ANALYTICS_DB_DATA_MODEL_TRACES_DISTRIBUTED_WRAP_ENABLED value: "false" + - equal: + path: data.ANALYTICS_DB_DATA_MODEL_TRACES_WEEKLY_PARTITION_PRUNING_ENABLED + value: "false" - equal: path: data.PARTITION_METRICS_ENABLED value: "false" diff --git a/deployment/helm_chart/opik/values.yaml b/deployment/helm_chart/opik/values.yaml index 9393b1121df..73b32acf8e3 100644 --- a/deployment/helm_chart/opik/values.yaml +++ b/deployment/helm_chart/opik/values.yaml @@ -740,6 +740,10 @@ databaseAnalyticsDataModel: # Turn true in lockstep with applying the Distributed wrap, so trace delete/retention mutations target # traces_local (a Distributed traces rejects mutations). Leave false while traces is still a MergeTree. tracesDistributedWrapEnabled: false + # Turn true only once the EXCHANGE has put the weekly partitioned successor under `traces`, so a trace DELETE + # can prune to the partitions its own ids resolve to. Enables the PRUNING, not the partitioning - setting it creates + # nothing. Purely an optimisation; false keeps the unbounded delete. + tracesWeeklyPartitionPruningEnabled: false # ClickHouse partition-health observability (OPIK-6904). Polls system.parts plus the # lightweight-delete mask and publishes opik.clickhouse.partition.* gauges; a distributed lock keeps