Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -14,6 +17,7 @@
import ru.yandex.clickhouse.ClickHouseConnectionImpl;

import java.sql.SQLException;
import java.util.List;
import java.util.Map;

@UtilityClass
Expand Down Expand Up @@ -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 {
Expand Down
187 changes: 187 additions & 0 deletions apps/opik-backend/src/test/java/com/comet/opik/db/TableSchema.java
Original file line number Diff line number Diff line change
@@ -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) {
Comment thread
thiagohora marked this conversation as resolved.

/**
* @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;
}
}
Loading
Loading