comet-ml
diff --git a/‎apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/MigrationUtils.java‎
Lines changed: 80 additions & 0 deletions b/‎apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/MigrationUtils.java‎
Lines changed: 80 additions & 0 deletions
diff --git a/‎apps/opik-backend/src/test/java/com/comet/opik/db/TableSchema.java‎
Lines changed: 187 additions & 0 deletions b/‎apps/opik-backend/src/test/java/com/comet/opik/db/TableSchema.java‎
Lines changed: 187 additions & 0 deletions
@@ -1,6 +1,9 @@
 package com.comet.opik.api.resources.utils;
 
+import liquibase.Contexts;
+import liquibase.LabelExpression;
 import liquibase.Liquibase;
+import liquibase.changelog.ChangeSet;
 import liquibase.database.DatabaseConnection;
 import liquibase.database.DatabaseFactory;
 import liquibase.database.jvm.JdbcConnection;
@@ -14,6 +17,7 @@
 import ru.yandex.clickhouse.ClickHouseConnectionImpl;
 
 import java.sql.SQLException;
+import java.util.List;
 import java.util.Map;
 
 @UtilityClass
@@ -48,6 +52,82 @@ public static void runClickhouseDbMigration(ClickHouseContainer container) {
         }
     }
 
+    /**
+     * Applies the ClickHouse changelog only up to and including the changesets of {@code migrationFileName}, leaving
+     * every later migration unrun so a caller can transform the schema mid-changelog and then resume with
+     * {@link #runClickhouseDbMigration(ClickHouseContainer)}.
+     * <p>
+     * This exists for the post-cutover topology gate: the cutover's {@code EXCHANGE} + {@code Distributed} wrap is
+     * produced by the operator runbook rather than by Liquibase, so the only way to run the later migrations against
+     * the topology they will really meet in production is to stop the changelog at the shadow-table migration, splice
+     * the transform in, and carry on. The cut is expressed as a migration <i>file name</i> rather than a changeset
+     * count so appending migrations never silently moves it.
+     *
+     * @param migrationFileName the migration file the apply stops after, e.g.
+     *            {@code 000114_recreate_traces_local_v2_id_at_datetime64.sql}
+     */
+    public static void runClickhouseDbMigrationThrough(ClickHouseContainer container, String migrationFileName) {
+        try (var connection = container.createConnection("")) {
+            DatabaseConnection dbConnection = new JdbcConnection(
+                    new ClickHouseConnectionImpl(connection.getMetaData().getURL()));
+            var database = DatabaseFactory.getInstance().findCorrectDatabaseImplementation(dbConnection);
+            try (var liquibase = new Liquibase(CLICKHOUSE_CHANGELOG_FILE, new ClassLoaderResourceAccessor(),
+                    database)) {
+                ClickHouseContainerUtils.migrationParameters().forEach(liquibase::setChangeLogParameter);
+                liquibase.update(countChangeSetsThrough(liquibase, migrationFileName), new Contexts(),
+                        new LabelExpression());
+            }
+        } catch (SQLException e) {
+            throw new RuntimeException("Failed to run ClickHouse DB migration", e);
+        } catch (LiquibaseException e) {
+            throw new UnexpectedLiquibaseException(e);
+        }
+    }
+
+    /**
+     * Identifiers of the ClickHouse changesets the database has <b>not</b> applied. Empty means the changelog is fully
+     * applied — which is the assertion a topology gate needs after resuming a spliced apply, because a changeset the
+     * extension skipped (a precondition evaluating to {@code MARK_RAN} is recorded as run; an unsupported statement is
+     * not) would otherwise leave the schema short without anything throwing.
+     */
+    public static List<String> unrunClickhouseChangeSetIds(ClickHouseContainer container) {
+        try (var connection = container.createConnection("")) {
+            DatabaseConnection dbConnection = new JdbcConnection(
+                    new ClickHouseConnectionImpl(connection.getMetaData().getURL()));
+            var database = DatabaseFactory.getInstance().findCorrectDatabaseImplementation(dbConnection);
+            try (var liquibase = new Liquibase(CLICKHOUSE_CHANGELOG_FILE, new ClassLoaderResourceAccessor(),
+                    database)) {
+                ClickHouseContainerUtils.migrationParameters().forEach(liquibase::setChangeLogParameter);
+                return liquibase.listUnrunChangeSets(new Contexts(), new LabelExpression())
+                        .stream()
+                        .map(ChangeSet::getId)
+                        .toList();
+            }
+        } catch (SQLException e) {
+            throw new RuntimeException("Failed to list unrun ClickHouse changesets", e);
+        } catch (LiquibaseException e) {
+            throw new UnexpectedLiquibaseException(e);
+        }
+    }
+
+    /**
+     * Number of changesets from the start of the changelog through the last one declared in {@code migrationFileName}.
+     * The changelog is a single {@code includeAll}, so this is the file's position in lexicographic order; counting the
+     * parsed changesets rather than the files keeps it right for a migration that declares more than one.
+     */
+    private static int countChangeSetsThrough(Liquibase liquibase, String migrationFileName)
+            throws LiquibaseException {
+        var changeSets = liquibase.getDatabaseChangeLog().getChangeSets();
+        for (int i = changeSets.size() - 1; i >= 0; i--) {
+            if (changeSets.get(i).getFilePath().endsWith(migrationFileName)) {
+                return i + 1;
+            }
+        }
+        throw new IllegalArgumentException(
+                "No changeset found for migration file '%s' in %s".formatted(migrationFileName,
+                        CLICKHOUSE_CHANGELOG_FILE));
+    }
+
     private static void runDbMigration(String changeLogFile, Map<String, String> parameters,
             DatabaseConnection connection) {
         try {
 
@@ -0,0 +1,187 @@
+package com.comet.opik.db;
+
+import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.util.ArrayList;
+import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/**
+ * A {@code SHOW CREATE}-level snapshot of one ClickHouse table, read from the {@code system} tables so it reflects the
+ * schema the server actually holds rather than the DDL someone believes it applied.
+ *
+ * <p>It carries every aspect a schema change can touch — columns (with their type, DEFAULT/MATERIALIZED kind and
+ * expression, and compression codec), data-skipping indices, the sorting / primary / partition keys, projections, and
+ * the engine — so a guard comparing two tables can assert on all of them instead of on column names alone. Parsing
+ * {@code SHOW CREATE TABLE} text would carry the same information but compare formatting as well as substance; the
+ * {@code system} tables give the same facts already decomposed.
+ *
+ * <p>{@code columns} keeps the server's declaration order ({@code system.columns.position}), which matters because two
+ * tables can hold the same column set in a different order — the trace shard and its shadow do exactly that. Callers
+ * comparing sets should go through {@link #columnNames()} / {@link #storedColumnNames()} rather than comparing the
+ * lists.
+ */
+record TableSchema(
+        String table,
+        String engine,
+        String partitionKey,
+        String sortingKey,
+        String primaryKey,
+        List<Column> columns,
+        List<SkipIndex> skipIndices,
+        List<Projection> projections) {
+
+    /**
+     * @param defaultKind {@code DEFAULT}, {@code MATERIALIZED}, {@code ALIAS}, or empty when the column simply has no
+     *            default. The distinction is load-bearing: only a non-{@code MATERIALIZED}/{@code ALIAS} column can be
+     *            named in an {@code INSERT}, so it is what separates a column the cutover backfill must carry from one
+     *            the destination recomputes for itself.
+     */
+    record Column(String name, String type, String defaultKind, String defaultExpression, String codec) {
+    }
+
+    record SkipIndex(String name, String typeFull, String expression, long granularity) {
+    }
+
+    record Projection(String name, String query) {
+    }
+
+    private static final Set<String> COMPUTED_DEFAULT_KINDS = Set.of("MATERIALIZED", "ALIAS");
+
+    static TableSchema read(Connection connection, String database, String table) throws SQLException {
+        var tableRow = readTableRow(connection, database, table);
+        return new TableSchema(
+                table,
+                tableRow.get("engine_full"),
+                tableRow.get("partition_key"),
+                tableRow.get("sorting_key"),
+                tableRow.get("primary_key"),
+                readColumns(connection, database, table),
+                readSkipIndices(connection, database, table),
+                readProjections(connection, database, table));
+    }
+
+    /** Column names in the server's declaration order. */
+    List<String> columnNames() {
+        return columns.stream().map(Column::name).toList();
+    }
+
+    /**
+     * The columns an {@code INSERT} can name — everything except {@code MATERIALIZED} / {@code ALIAS}, which the server
+     * computes and refuses to accept a value for.
+     */
+    Set<String> storedColumnNames() {
+        return columns.stream()
+                .filter(column -> !COMPUTED_DEFAULT_KINDS.contains(column.defaultKind()))
+                .map(Column::name)
+                .collect(Collectors.toCollection(LinkedHashSet::new));
+    }
+
+    Map<String, Column> columnsByName() {
+        var byName = new LinkedHashMap<String, Column>();
+        columns.forEach(column -> byName.put(column.name(), column));
+        return byName;
+    }
+
+    Map<String, SkipIndex> skipIndicesByName() {
+        var byName = new LinkedHashMap<String, SkipIndex>();
+        skipIndices.forEach(index -> byName.put(index.name(), index));
+        return byName;
+    }
+
+    Set<String> skipIndexNames() {
+        return skipIndicesByName().keySet();
+    }
+
+    Set<String> projectionNames() {
+        var names = new LinkedHashSet<String>();
+        projections.forEach(projection -> names.add(projection.name()));
+        return names;
+    }
+
+    boolean isDistributed() {
+        return engine.startsWith("Distributed");
+    }
+
+    private static Map<String, String> readTableRow(Connection connection, String database, String table)
+            throws SQLException {
+        var sql = """
+                SELECT engine_full, partition_key, sorting_key, primary_key
+                FROM system.tables WHERE database = '%s' AND name = '%s'
+                """.formatted(database, table);
+        try (var statement = connection.createStatement(); var resultSet = statement.executeQuery(sql)) {
+            if (!resultSet.next()) {
+                throw new IllegalStateException("Table '%s.%s' does not exist".formatted(database, table));
+            }
+            return Map.of(
+                    "engine_full", text(resultSet, "engine_full"),
+                    "partition_key", text(resultSet, "partition_key"),
+                    "sorting_key", text(resultSet, "sorting_key"),
+                    "primary_key", text(resultSet, "primary_key"));
+        }
+    }
+
+    private static List<Column> readColumns(Connection connection, String database, String table) throws SQLException {
+        var sql = """
+                SELECT name, type, default_kind, default_expression, compression_codec
+                FROM system.columns WHERE database = '%s' AND table = '%s' ORDER BY position
+                """.formatted(database, table);
+        var columns = new ArrayList<Column>();
+        try (var statement = connection.createStatement(); var resultSet = statement.executeQuery(sql)) {
+            while (resultSet.next()) {
+                columns.add(new Column(
+                        text(resultSet, "name"),
+                        text(resultSet, "type"),
+                        text(resultSet, "default_kind"),
+                        text(resultSet, "default_expression"),
+                        text(resultSet, "compression_codec")));
+            }
+        }
+        return List.copyOf(columns);
+    }
+
+    private static List<SkipIndex> readSkipIndices(Connection connection, String database, String table)
+            throws SQLException {
+        var sql = """
+                SELECT name, type_full, expr, granularity
+                FROM system.data_skipping_indices WHERE database = '%s' AND table = '%s' ORDER BY name
+                """.formatted(database, table);
+        var indices = new ArrayList<SkipIndex>();
+        try (var statement = connection.createStatement(); var resultSet = statement.executeQuery(sql)) {
+            while (resultSet.next()) {
+                indices.add(new SkipIndex(
+                        text(resultSet, "name"),
+                        text(resultSet, "type_full"),
+                        text(resultSet, "expr"),
+                        resultSet.getLong("granularity")));
+            }
+        }
+        return List.copyOf(indices);
+    }
+
+    private static List<Projection> readProjections(Connection connection, String database, String table)
+            throws SQLException {
+        var sql = """
+                SELECT name, query FROM system.projections
+                WHERE database = '%s' AND table = '%s' ORDER BY name
+                """.formatted(database, table);
+        var projections = new ArrayList<Projection>();
+        try (var statement = connection.createStatement(); var resultSet = statement.executeQuery(sql)) {
+            while (resultSet.next()) {
+                projections.add(new Projection(text(resultSet, "name"), text(resultSet, "query")));
+            }
+        }
+        return List.copyOf(projections);
+    }
+
+    /** ClickHouse returns an absent String as {@code ""}; normalise the JDBC {@code null} case to match. */
+    private static String text(ResultSet resultSet, String column) throws SQLException {
+        var value = resultSet.getString(column);
+        return value == null ? "" : value;
+    }
+}