[OPIK-8230] [BE] fix: scope the trace delete mutation to the batch's own partitions

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

@@ -1973,22 +1973,57 @@ ORDER BY (workspace_id, project_id, trace_id, id) DESC, last_updated_at DESC
      * project_id}, so no delete is ever project-less - also required once {@code traces} is a Distributed table
      * (OPIK-7455).
      * <p>
-     * {@code <if(partitions)>} adds the table's own weekly partition expression, bound as the set of partitions the
-     * batch's ids resolve to. It is emitted whenever every id in the batch is 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.
+     * The unbounded form: no partition scoping at all, so it is the correct statement whenever
+     * {@link WeeklyPartitions#groupByPartition} cannot derive an exact partition for every id in the batch (a non-UUIDv7
+     * id, or one past the {@code DateTime64} ceiling — see that method's Javadoc). {@link #DELETE_BY_PROJECT_ID_TRACE_ID_PAIRS}
+     * is preferred whenever it can be used; this is the fallback that stays correct on any batch, at the cost of the
+     * mutation visiting every part of the table (OPIK-8230).
      * <p>
-     * No schema flag gates it. {@link WeeklyPartitions} derives a value per {@code id_at} type the mutation may meet
-     * — the legacy 32-bit {@code DateTime} of {@code traces} as well as the {@code DateTime64(0)} of the partitioned
-     * successor — so one rendered statement is correct on both sides of the cutover EXCHANGE, in either direction, with
-     * nothing to flip and nothing to revert on rollback.
+     * 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
+     * as a literal. {@code arrayZip} is a deterministic function, not a subquery - ClickHouse rejects subqueries in
+     * delete mutations. Callers batch to keep each array within the driver's reliable bind size ({@link
+     * com.comet.opik.infrastructure.FilterUtils#ANALYTICS_DELETE_BATCH_SIZE}).
+     */
+    /**
+     * Deletes by the full {@code (workspace_id, project_id, id)} sort key, matching on {@code (project_id, id)} tuples
+     * 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).
+     * <p>
+     * {@code <if(partition)>} adds {@code IN PARTITION <partition>} (OPIK-8230), which scopes the <b>mutation itself</b>
+     * — which parts it is registered against — rather than which rows its {@code WHERE} selects. That distinction is
+     * the entire fix. A ClickHouse mutation selects parts at the <b>partition</b> stage, before the {@code WHERE}
+     * clause is considered at all — a {@code (workspace_id, project_id, id)} predicate, however exact, prunes zero
+     * parts, only rows within whichever parts were already selected. Post-cutover {@code traces_local} is
+     * weekly-partitioned into ~1,900 partitions / ~3,650 parts; a delete matching a handful of rows was registered
+     * against every one of them regardless, at ~19 ms/part fixed overhead, which is what pushed ordinary deletes past
+     * the {@code max_execution_time} ceiling (`Code: 159 TIMEOUT_EXCEEDED`) that pre-cutover deletes (1.5–2.2 s) never
+     * approached. Omitted for the unbounded fallback ({@link TraceDAOImpl#deleteBatch}), whose statement is then
+     * byte-identical to the pre-OPIK-8230 form.
+     * <p>
+     * One statement per partition the batch's ids resolve to ({@link WeeklyPartitions#groupByPartition}), each binding
+     * only the {@code (project_id, id)} pairs whose id belongs to that partition — {@code IN PARTITION} accepts exactly
+     * one partition per statement; ClickHouse's {@code IN PARTITION p1, p2, …} form is for <em>composite</em> partition
+     * keys and rejects a list of distinct partitions outright. Naming a partition that does not exist is a silent
+     * no-op, not an error, which is what makes it safe for a far-future id to be bound into two statements (see
+     * {@link WeeklyPartitions#groupByPartition}'s Javadoc) — the one naming the partition it is not actually in simply
+     * matches nothing there.
      * <p>
-     * Why it matters: a mutation selects parts at the <b>partition</b> 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 <b>3,928 parts / 5.40 TiB</b>. With this predicate the same batch
-     * selects <b>5</b> parts. An {@code id_at} <em>range</em> is not a substitute: on a batch spanning 1996 and 2200 a
-     * range still selected 2,644 parts, where the exact set selected 4.
+     * {@code <partition>} is interpolated by StringTemplate, exactly like {@code <traces_mutation_table>} and
+     * {@code <log_comment>} above it — {@code IN PARTITION {p:UInt32}} is a ClickHouse syntax error, so the value
+     * cannot be bound. Unlike the {@code <log_comment>} interpolation this carries no escaping question: the value is
+     * always a Java {@code long} produced by {@link WeeklyPartitions}, and a {@code long} cannot carry a quote or any
+     * other SQL-meaningful character.
+     * <p>
+     * Kept as a single {@code DELETE FROM} statement rather than a hand-written {@code ALTER TABLE ... UPDATE ... IN
+     * PARTITION}: both compile to the same mutation, but the hand-written form would step outside
+     * {@link #tracesMutationTable()}'s routing guarantee ({@code TraceMutationRoutingArchTest}), outside the project's
+     * one {@code lightweight_deletes_sync} convention, and — decisively — outside {@code lightweight_delete_mode}: a
+     * {@code DELETE FROM} takes the patch-part path when that setting is enabled and the table's data-model columns
+     * allow it; a hand-written {@code ALTER ... UPDATE} does not; so staying on {@code DELETE FROM} means any future
+     * enabling of patch parts benefits this path for free, with no further code change here.
      * <p>
      * 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
@@ -1998,9 +2033,9 @@ ORDER BY (workspace_id, project_id, trace_id, id) DESC, last_updated_at DESC
      */
     private static final String DELETE_BY_PROJECT_ID_TRACE_ID_PAIRS = """
             DELETE FROM <traces_mutation_table>
+            <if(partition)>IN PARTITION <partition><endif>
             WHERE workspace_id = :workspace_id
             AND (project_id, id) IN arrayZip(:project_ids, :trace_ids)
-            <if(partitions)>AND toYYYYMMDD(toDate32(id_at) - toIntervalDay(toDayOfWeek(id_at, 1))) IN :partitions<endif>
             SETTINGS log_comment = '<log_comment>'
             ;
             """;
@@ -2328,7 +2363,7 @@ ORDER BY (workspace_id, project_id, id) DESC, last_updated_at DESC
      * {@code :min_id}'s week and inverted the window, so the fast pass matched nothing and every id fell through to
      * the unbounded pass. It now resolves them.
      * <p>
-     * The fallback remains load-bearing for what {@link com.comet.opik.utils.WeeklyPartitions#of} still cannot derive
+     * The fallback remains load-bearing for what {@link com.comet.opik.utils.WeeklyPartitions#groupByPartition} still cannot derive
      * exactly: an id at or past the end of {@code DateTime64}'s range, where {@code id_at} saturates to
      * {@code 2299-12-31} whatever the real week, so every such id collapses into one partition. Real data contains
      * them, so the bounded query is never a delete's sole resolver.
@@ -3649,38 +3684,103 @@ public Mono<Void> delete(Set<Pair<UUID, UUID>> projectIdTraceIdPairs, @NonNull C
 
         return makeMonoContextAware((userName, workspaceId) -> Flux
                 .fromIterable(Lists.partition(List.copyOf(projectIdTraceIdPairs), ANALYTICS_DELETE_BATCH_SIZE))
-                .concatMap(batch -> {
-                    var template = getSTWithLogComment(DELETE_BY_PROJECT_ID_TRACE_ID_PAIRS, "delete_traces",
-                            workspaceId,
-                            userName, "pairs_size=%s".formatted(batch.size()));
-                    selectTracesMutationTable(template);
-
-                    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 every id in the batch allows it; otherwise emit the
-                    // unbounded form. Needs no schema flag: WeeklyPartitions derives a value per id_at type the
-                    // mutation may meet, so the set is correct on the legacy traces and on the partitioned successor.
-                    var partitions = WeeklyPartitions.of(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);
+                .concatMap(batch -> deleteBatch(batch, workspaceId, userName, connection))
+                .then());
+    }
 
-                    if (partitions.isPresent()) {
-                        statement = statement.bind("partitions", partitions.get().toArray(Long[]::new));
-                    }
+    /**
+     * Deletes one ANALYTICS_DELETE_BATCH_SIZE-sized batch, scoped to its own partitions when every id in it resolves
+     * exactly (OPIK-8230): one {@link #DELETE_BY_PROJECT_ID_TRACE_ID_PAIRS} statement per partition the
+     * batch's ids resolve to ({@link WeeklyPartitions#groupByPartition}), each binding only the pairs whose id belongs
+     * to that partition. Falls back to a single {@link #DELETE_BY_PROJECT_ID_TRACE_ID_PAIRS} statement over the whole
+     * batch either when {@code groupByPartition} cannot derive every id's partition exactly, or when the mutation's
+     * own target table is not actually partitioned — always correct, merely unscoped, and identical to what this
+     * method emitted before OPIK-8230.
+     * <p>
+     * The second condition is load-bearing and easy to miss: {@link WeeklyPartitions} is deliberately schema-agnostic
+     * (see its own Javadoc) and derives a partition value regardless of which table {@code id_at} actually lives in —
+     * including the legacy, UNPARTITIONED {@code traces} that {@link #tracesMutationTable()} still targets before the
+     * sharding-readiness wrap is applied. {@code IN PARTITION <value>} against a table with no partition key is not a
+     * no-op; ClickHouse rejects it outright (Code 248 {@code INVALID_PARTITION_VALUE}, "Wrong number of fields in the
+     * partition expression: 1, must be: 0"). So the scoped path is gated on {@link #tracesMutationTable()}'s
+     * OWN resolved name, not on whether a partition value merely COULD be derived: this method calls the
+     * resolver itself, exactly as {@link #selectTracesMutationTable} does for the template - never the wrap
+     * flag directly ({@code TraceMutationRoutingArchTest} enforces that {@link #tracesDistributedWrapEnabled()}
+     * is called only from {@link #tracesMutationTable()}, so no second place can get the read/mutate split
+     * wrong). Resolved once per batch and reused across every partition-scoped statement it emits: the table
+     * a mutation targets does not vary per partition, only per topology.
+     * <p>
+     * The gate is a NAME check — {@code TRACES_LOCAL_TABLE.equals(table)} — not a runtime read of
+     * {@code traces_local}'s actual engine or partition key. That is an accepted operational precondition, not an
+     * oversight: it relies on the wrap being flipped only after the EXCHANGE that makes {@code traces_local} the
+     * partitioned successor has already run, which is the documented ordering the migration follows and the same
+     * assumption {@link #tracesMutationTable()} already makes for every other mutation in this class (a premature
+     * flip breaks inserts and reads the same way, table-missing rather than Code 248). Verifying the live schema
+     * before every batch would trade a real per-request cost for a misordering that procedure, not this code, is
+     * relied on to prevent.
+     * <p>
+     * Sequential ({@code concatMap}), deliberately not concurrent: bounded concurrency was measured and rejected for
+     * this change (OPIK-8230) — a real gain in the rare many-partition tail, but a knob to hold in reserve rather than
+     * ship alongside the fix itself.
+     */
+    private Mono<Void> deleteBatch(List<Pair<UUID, UUID>> batch, String workspaceId, String userName,
+            Connection connection) {
+        // Resolved once, via the resolver itself - never the flag directly (see the method Javadoc for why).
+        var table = tracesMutationTable();
+        var grouped = TRACES_LOCAL_TABLE.equals(table)
+                ? WeeklyPartitions.groupByPartition(batch.stream().map(Pair::getRight).toList())
+                : Optional.<Map<Long, Set<UUID>>>empty();
+
+        if (grouped.isEmpty()) {
+            return executeDelete(table, batch, null, workspaceId, userName, connection);
+        }
 
-                    var segment = startSegment("traces", "Clickhouse", "delete");
-                    return Mono.from(statement.execute())
-                            .doFinally(_ -> endSegment(segment))
-                            .then();
+        // id -> its pairs, built once per batch rather than rescanning the whole batch once per partition: an id can
+        // map to more than one pair (the same trace id reused across projects, OPIK-7483), so this is a
+        // Collectors.groupingBy, not a plain lookup map.
+        var pairsById = batch.stream().collect(Collectors.groupingBy(Pair::getRight));
+
+        return Flux.fromIterable(grouped.get().entrySet())
+                .concatMap(entry -> {
+                    var partitionPairs = entry.getValue().stream()
+                            .flatMap(id -> pairsById.get(id).stream())
+                            .toList();
+                    return executeDelete(table, partitionPairs, entry.getKey(), workspaceId, userName, connection);
                 })
-                .then());
+                .then();
+    }
+
+    /**
+     * Renders and executes one {@link #DELETE_BY_PROJECT_ID_TRACE_ID_PAIRS} statement — unbounded when
+     * {@code partition} is null, scoped to it otherwise (the weekly partition value the caller has already resolved
+     * every pair in {@code pairs} to belong to). Shares the bind/segment/logging shape in exactly one place rather
+     * than duplicated per call site in {@link #deleteBatch}.
+     * <p>
+     * Takes {@code table} as a parameter rather than resolving it itself via {@link #selectTracesMutationTable}: it
+     * is already resolved once, up front, by {@link #deleteBatch} — the physical table does not vary per statement
+     * within one batch, only per topology — so re-resolving it here per call would be redundant, not incorrect.
+     */
+    private Mono<Void> executeDelete(String table, List<Pair<UUID, UUID>> pairs, Long partition,
+            String workspaceId, String userName, Connection connection) {
+        var template = getSTWithLogComment(DELETE_BY_PROJECT_ID_TRACE_ID_PAIRS, "delete_traces", workspaceId,
+                userName, "pairs_size=%s".formatted(pairs.size()));
+        template.add("traces_mutation_table", table);
+        if (partition != null) {
+            template.add("partition", partition);
+        }
+
+        var projectIds = pairs.stream().map(pair -> pair.getLeft().toString()).toArray(String[]::new);
+        var traceIds = pairs.stream().map(pair -> pair.getRight().toString()).toArray(String[]::new);
+
+        var statement = connection.createStatement(template.render())
+                .bind("workspace_id", workspaceId)
+                .bind("project_ids", projectIds)
+                .bind("trace_ids", traceIds);
+
+        var segment = startSegment("traces", "Clickhouse", "delete");
+        return Mono.from(statement.execute())
+                .doFinally(_ -> endSegment(segment))
+                .then();
     }
 
     /**

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

@@ -537,7 +537,7 @@ public Mono<Void> delete(@NonNull Set<UUID> ids, UUID projectId) {
      * Resolves every owning project for each id: a bounded fast pass, then an unbounded pass over only the ids the
      * bounded one leaves unresolved. Returns id -> owning projects; ids absent from the result have no live row.
      * <p>
-     * The bounded pass's week window can miss a row whose week {@link com.comet.opik.utils.WeeklyPartitions#of}
+     * The bounded pass's week window can miss a row whose week {@link com.comet.opik.utils.WeeklyPartitions#groupByPartition}
      * cannot derive exactly — an id at or past the end of {@code DateTime64}'s range, where {@code id_at} saturates
      * to {@code 2299-12-31} whatever the real week — so the unbounded pass re-resolves the miss set and the bounded
      * query is never a delete's sole resolver. A far-future timestamp short of that ceiling is no longer such a case:
@@ -566,6 +566,20 @@ private Mono<Map<UUID, Set<UUID>>> resolveOwningProjects(Set<UUID> ids) {
                 });
     }
 
+    /**
+     * All-or-nothing over {@code projectIdTraceIdPairs}: {@code TracesDeleted} publishes and the deletion-events
+     * bridge captures only {@code .doOnSuccess}/{@code .then} of {@link TraceDAO#delete}, so any error anywhere in
+     * the batch skips both for the WHOLE batch, not just the pairs that did not complete.
+     * <p>
+     * OPIK-8230 widens the window this can bite in, worth being explicit about here rather than only in that
+     * ticket: {@link TraceDAO#delete} can now emit several sequential statements per batch, one per partition
+     * (via {@code deleteBatch}), so an error partway through leaves the EARLIER statements' rows genuinely,
+     * synchronously deleted while the cascade still does not fire for any pair in the batch - not only for the one
+     * that errored. Restructuring this coupling to be per-statement aware is explicitly out of scope for that
+     * ticket ("this ticket removes the condition that trips it; it does not restructure the coupling") - deletes
+     * are idempotent, so a caller that retries converges, and the ticket's own fix (scoping each mutation so it
+     * completes well inside the statement timeout) is what makes hitting this at all much rarer than it is today.
+     */
     private Mono<Void> delete(Set<Pair<UUID, UUID>> projectIdTraceIdPairs, Connection connection) {
         return Mono.deferContextual(ctx -> {
             String workspaceId = ctx.get(RequestContext.WORKSPACE_ID);