diff --git a/ydb-trino-adapter/ROADMAP.md b/ydb-trino-adapter/ROADMAP.md index 546927b0..c2021dce 100644 --- a/ydb-trino-adapter/ROADMAP.md +++ b/ydb-trino-adapter/ROADMAP.md @@ -88,8 +88,9 @@ failure. changes as atomic delete+insert row changes, or reject that statement with a documented `NOT_SUPPORTED` error. See the [YQL UPDATE contract](https://ydb.tech/docs/en/yql/reference/syntax/update). -4. Add focused tests for composite primary keys, a non-unique first visible - column, physical-key updates, rollback/close, and fresh-state retries. +4. A failed `Connection.commit()` has an ambiguous outcome, so the connector + never replays that MERGE attempt. Add further focused tests for composite + primary keys, physical-key updates, rollback/close, and fresh-state retries. **Exit criterion:** retain the current green inherited `testMerge*` suite and add a bounded-memory benchmark that demonstrates acceptable production-scale @@ -114,7 +115,7 @@ runtime for the set-based implementation. - Do not replay buffered INSERT pages after `JdbcPageSink` may already have committed an internal batch. - Add unit tests for status classification, interrupted backoff, rollback - failure suppression, connection cleanup, and a failure after commit. + failure suppression, and connection cleanup. - Define memory/backpressure limits for buffered merge pages; memory usage must not remain unreported. diff --git a/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbMergeSink.java b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbMergeSink.java index d1f29ddd..07d33a86 100644 --- a/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbMergeSink.java +++ b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbMergeSink.java @@ -96,10 +96,12 @@ public CompletableFuture> finish() { for (int attempt = 0; attempt < maxAttempts; attempt++) { Connection connection = null; + boolean commitAttempted = false; boolean committed = false; try { connection = openConnection(); executeMergeInTransaction(connection); + commitAttempted = true; connection.commit(); committed = true; connection.close(); @@ -131,6 +133,13 @@ public CompletableFuture> finish() { throw new TrinoException(JDBC_ERROR, "YDB MERGE committed, but closing its connection failed", e); } + if (commitAttempted) { + throw new TrinoException( + JDBC_ERROR, + "YDB MERGE commit outcome is unknown; operation was not retried to avoid duplicate writes", + e); + } + if (!YdbRetryUtils.isRetryable(e) && !isRejectedDriverSessionAcquisition(e)) { throw new TrinoException(JDBC_ERROR, e); } diff --git a/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestYdbMergeSink.java b/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestYdbMergeSink.java new file mode 100644 index 00000000..fd448a1e --- /dev/null +++ b/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestYdbMergeSink.java @@ -0,0 +1,114 @@ +package tech.ydb.trino; + +import io.trino.plugin.jdbc.JdbcClient; +import io.trino.plugin.jdbc.JdbcMergeTableHandle; +import io.trino.plugin.jdbc.JdbcOutputTableHandle; +import io.trino.plugin.jdbc.JdbcTableHandle; +import io.trino.plugin.jdbc.RemoteTableName; +import io.trino.plugin.jdbc.WriteMapping; +import io.trino.plugin.jdbc.logging.RemoteQueryModifier; +import io.trino.spi.Page; +import io.trino.spi.TrinoException; +import io.trino.spi.block.Block; +import io.trino.spi.block.BlockBuilder; +import io.trino.spi.connector.SchemaTableName; +import io.trino.spi.type.Type; +import org.junit.jupiter.api.Test; +import tech.ydb.core.Status; +import tech.ydb.core.StatusCode; +import tech.ydb.jdbc.exception.YdbStatusable; + +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.SQLException; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicInteger; + +import static io.trino.spi.connector.ConnectorMergeSink.INSERT_OPERATION_NUMBER; +import static io.trino.spi.type.BigintType.BIGINT; +import static io.trino.spi.type.IntegerType.INTEGER; +import static io.trino.spi.type.TinyintType.TINYINT; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +class TestYdbMergeSink { + @Test + void testCommitFailureIsNotRetried() { + AtomicInteger connections = new AtomicInteger(); + AtomicInteger executions = new AtomicInteger(); + AtomicInteger commits = new AtomicInteger(); + SQLException commitFailure = new RetryableSQLException(); + PreparedStatement statement = proxy(PreparedStatement.class, (_, method, _) -> { + if (method.getName().equals("executeBatch")) { + executions.incrementAndGet(); + return new int[] {1}; + } + return null; + }); + Connection connection = proxy(Connection.class, (_, method, _) -> switch (method.getName()) { + case "prepareStatement" -> statement; + case "commit" -> { + commits.incrementAndGet(); + throw commitFailure; + } + default -> null; + }); + JdbcClient jdbcClient = proxy(JdbcClient.class, (_, method, _) -> switch (method.getName()) { + case "getConnection" -> { + connections.incrementAndGet(); + yield connection; + } + case "toWriteMapping" -> WriteMapping.longMapping("Int64", PreparedStatement::setLong); + case "buildInsertSql" -> "INSERT INTO `target` (`value`) VALUES (?)"; + default -> null; + }); + YdbMergeSink sink = new YdbMergeSink( + null, null, mergeHandle(), jdbcClient, () -> 1, RemoteQueryModifier.NONE, null); + sink.storeMergedRows(new Page( + block(BIGINT, 42), block(TINYINT, INSERT_OPERATION_NUMBER), block(INTEGER, 0), block(BIGINT, 42))); + + assertThatThrownBy(sink::finish) + .isInstanceOfSatisfying(TrinoException.class, exception -> { + assertThat(exception).hasMessageContaining("commit outcome is unknown"); + assertThat(exception.getCause()).isSameAs(commitFailure); + }); + assertThat(connections).hasValue(1); + assertThat(executions).hasValue(1); + assertThat(commits).hasValue(1); + } + + private static JdbcMergeTableHandle mergeHandle() { + RemoteTableName tableName = new RemoteTableName(Optional.empty(), Optional.empty(), "target"); + JdbcOutputTableHandle outputHandle = new JdbcOutputTableHandle( + tableName, List.of("value"), List.of(BIGINT), Optional.empty(), Optional.empty(), Optional.empty()); + return new JdbcMergeTableHandle( + new JdbcTableHandle(new SchemaTableName("default", "target"), tableName, Optional.empty()), + outputHandle, + Map.of(), + Optional.empty(), + List.of(), + List.of(), + Map.of()); + } + + private static Block block(Type type, long value) { + BlockBuilder builder = type.createBlockBuilder(null, 1); + type.writeLong(builder, value); + return builder.build(); + } + + private static T proxy(Class type, InvocationHandler handler) { + return type.cast(Proxy.newProxyInstance(type.getClassLoader(), new Class[] {type}, handler)); + } + + private static final class RetryableSQLException extends SQLException implements YdbStatusable { + @Override + public Status getStatus() { + return Status.of(StatusCode.ABORTED); + } + } +}