From 4225231fba1face72b3c1a3f14b57837f82c8917 Mon Sep 17 00:00:00 2001 From: KirillKurdyukov Date: Tue, 15 Sep 2026 11:13:51 +0300 Subject: [PATCH 1/3] Add explicit primary keys for YDB table creation --- ydb-trino-adapter/README.md | 19 +++++ .../main/java/tech/ydb/trino/YdbClient.java | 39 +++++++++- .../java/tech/ydb/trino/YdbClientModule.java | 2 + .../tech/ydb/trino/YdbTableProperties.java | 38 ++++++++++ .../tech/ydb/trino/TestYdbCreateTable.java | 72 +++++++++++++++++++ .../tech/ydb/trino/TestingYdbJdbcClient.java | 11 +++ .../tech/ydb/trino/TestingYdbJdbcModule.java | 3 + .../java/tech/ydb/trino/YdbQueryRunner.java | 8 ++- 8 files changed, 190 insertions(+), 2 deletions(-) create mode 100644 ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbTableProperties.java create mode 100644 ydb-trino-adapter/src/test/java/tech/ydb/trino/TestYdbCreateTable.java diff --git a/ydb-trino-adapter/README.md b/ydb-trino-adapter/README.md index 688b22d7..e4aad3c5 100644 --- a/ydb-trino-adapter/README.md +++ b/ydb-trino-adapter/README.md @@ -45,6 +45,25 @@ SELECT * FROM local.default.orders; Прежние обращения `catalog.ydb.table` нужно заменить на `catalog.default.table`. +## Создание таблиц + +YDB требует первичный ключ, поэтому `CREATE TABLE` и `CREATE TABLE AS` должны +задавать упорядоченное свойство `primary_key`. Адаптер не добавляет скрытый ключ: + +```sql +CREATE TABLE events (tenant bigint, event_id bigint, payload varchar) +WITH (primary_key = ARRAY['tenant', 'event_id']); + +CREATE TABLE events_copy +WITH (primary_key = ARRAY['tenant', 'event_id']) +AS SELECT tenant, event_id, payload FROM events; +``` + +Имена ключей в CTAS относятся к выходным столбцам запроса. Отсутствующий, +пустой, повторяющийся или неизвестный `primary_key` отклоняется. Текущая +integration-проверка CTAS использует `insert.non-transactional-insert.enabled=true`; +transactional staging этой проверкой не подтверждается. + ## Текст и байты YDB `Text` отображается в Trino как `varchar`, а `Bytes` — как `varbinary` без diff --git a/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbClient.java b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbClient.java index 546f003f..628325bb 100644 --- a/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbClient.java +++ b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbClient.java @@ -100,6 +100,7 @@ import static io.trino.plugin.jdbc.StandardColumnMappings.varcharColumnMapping; import static io.trino.plugin.jdbc.StandardColumnMappings.varcharReadFunction; import static io.trino.plugin.jdbc.StandardColumnMappings.varcharWriteFunction; +import static io.trino.spi.StandardErrorCode.INVALID_TABLE_PROPERTY; import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.BooleanType.BOOLEAN; @@ -124,6 +125,7 @@ import static tech.ydb.trino.YdbColumnMappings.dateWriteFunction; import static tech.ydb.trino.YdbColumnMappings.timestampColumnMapping; import static tech.ydb.trino.YdbColumnMappings.timestampWriteFunction; +import static tech.ydb.trino.YdbTableProperties.PRIMARY_KEY_PROPERTY; public class YdbClient extends BaseJdbcClient { static final String DEFAULT_SCHEMA = "default"; @@ -730,6 +732,41 @@ private long executeReturningDml( @Override protected List createTableSqls(RemoteTableName remoteTableName, List columns, ConnectorTableMetadata tableMetadata) { - return super.createTableSqls(remoteTableName, columns, tableMetadata); + if (tableMetadata.getComment().isPresent()) { + throw new TrinoException(NOT_SUPPORTED, "This connector does not support creating tables with table comment"); + } + + List primaryKeys = YdbTableProperties.getPrimaryKey(tableMetadata.getProperties()); + if (primaryKeys.isEmpty()) { + throw new TrinoException(INVALID_TABLE_PROPERTY, "Table property 'primary_key' must contain at least one column"); + } + if (primaryKeys.stream().distinct().count() != primaryKeys.size()) { + throw new TrinoException(INVALID_TABLE_PROPERTY, "Table property 'primary_key' contains duplicate columns"); + } + + Set columnNames = tableMetadata.getColumns().stream() + .map(ColumnMetadata::getName) + .collect(Collectors.toSet()); + primaryKeys.stream() + .filter(primaryKey -> !columnNames.contains(primaryKey)) + .findFirst() + .ifPresent(primaryKey -> { + throw new TrinoException( + INVALID_TABLE_PROPERTY, + "Column '%s' specified in table property '%s' does not exist".formatted(primaryKey, PRIMARY_KEY_PROPERTY)); + }); + + return List.of("CREATE TABLE %s (%s, PRIMARY KEY (%s))".formatted( + quoted(remoteTableName), + String.join(", ", columns), + primaryKeys.stream().map(this::quoted).collect(joining(", ")))); + } + + @Override + public Map getTableProperties(ConnectorSession session, JdbcTableHandle tableHandle) { + List primaryKeys = getPrimaryKeys(session, tableHandle.getRequiredNamedRelation().getRemoteTableName()).stream() + .map(JdbcColumnHandle::getColumnName) + .toList(); + return Map.of(PRIMARY_KEY_PROPERTY, primaryKeys); } } diff --git a/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbClientModule.java b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbClientModule.java index 53b2da8c..46d8f723 100644 --- a/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbClientModule.java +++ b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbClientModule.java @@ -18,6 +18,7 @@ import tech.ydb.jdbc.YdbDriver; import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder; +import static io.trino.plugin.jdbc.JdbcModule.bindTablePropertiesProvider; public class YdbClientModule implements Module { @@ -28,6 +29,7 @@ public void configure(Binder binder) { .to(YdbMetadataFactory.class) .in(Scopes.SINGLETON); + bindTablePropertiesProvider(binder, YdbTableProperties.class); binder.bind(YdbConnector.class).in(Scopes.SINGLETON); } diff --git a/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbTableProperties.java b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbTableProperties.java new file mode 100644 index 00000000..621d3b43 --- /dev/null +++ b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbTableProperties.java @@ -0,0 +1,38 @@ +package tech.ydb.trino; + +import io.trino.plugin.jdbc.TablePropertiesProvider; +import io.trino.spi.session.PropertyMetadata; +import io.trino.spi.type.ArrayType; + +import java.util.List; +import java.util.Map; + +import static io.trino.spi.type.VarcharType.VARCHAR; +import static java.util.Objects.requireNonNull; + +public class YdbTableProperties implements TablePropertiesProvider { + public static final String PRIMARY_KEY_PROPERTY = "primary_key"; + + private final List> tableProperties = List.of( + new PropertyMetadata<>( + PRIMARY_KEY_PROPERTY, + "Columns forming the table primary key, in order", + new ArrayType(VARCHAR), + List.class, + List.of(), + false, + value -> (List) value, + value -> value)); + + @Override + public List> getTableProperties() { + return tableProperties; + } + + public static List getPrimaryKey(Map tableProperties) { + requireNonNull(tableProperties, "tableProperties is null"); + return ((List) tableProperties.get(PRIMARY_KEY_PROPERTY)).stream() + .map(String.class::cast) + .toList(); + } +} diff --git a/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestYdbCreateTable.java b/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestYdbCreateTable.java new file mode 100644 index 00000000..080042ef --- /dev/null +++ b/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestYdbCreateTable.java @@ -0,0 +1,72 @@ +package tech.ydb.trino; + +import io.trino.testing.AbstractTestQueryFramework; +import io.trino.testing.QueryRunner; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; +import tech.ydb.test.junit5.YdbHelperExtension; + +import static org.assertj.core.api.Assertions.assertThat; + +public class TestYdbCreateTable extends AbstractTestQueryFramework { + @RegisterExtension + static final YdbHelperExtension ydb = new YdbHelperExtension(); + + @Override + protected QueryRunner createQueryRunner() throws Exception { + return YdbQueryRunner.builder(ydb) + .useProductionClient() + .build(); + } + + @Test + public void testCreateTableWithPrimaryKey() { + String tableName = "production_create_with_pk"; + try { + assertUpdate("CREATE TABLE " + tableName + " (tenant bigint, event_id bigint, payload varchar) " + + "WITH (primary_key = ARRAY['event_id', 'tenant'])"); + assertUpdate("INSERT INTO " + tableName + " VALUES (1, 10, 'a'), (2, NULL, 'b')", 2); + assertQuery("SELECT * FROM " + tableName, "VALUES (1, 10, 'a'), (2, NULL, 'b')"); + assertThat((String) computeScalar("SHOW CREATE TABLE " + tableName)) + .contains("primary_key = ARRAY['event_id', 'tenant']"); + } + finally { + assertUpdate("DROP TABLE IF EXISTS " + tableName); + } + } + + @Test + public void testCreateTableAsSelectWithPrimaryKey() { + String tableName = "production_ctas_with_pk"; + String duplicateTableName = "production_ctas_duplicate_pk"; + try { + assertUpdate("CREATE TABLE " + tableName + " (id, payload) WITH (primary_key = ARRAY['id']) " + + "AS VALUES (BIGINT '1', 'a'), (BIGINT '2', 'b')", 2); + assertQuery("SELECT * FROM " + tableName, "VALUES (BIGINT '1', 'a'), (BIGINT '2', 'b')"); + + assertThat(query("CREATE TABLE " + duplicateTableName + " (id, payload) " + + "WITH (primary_key = ARRAY['id']) AS VALUES (BIGINT '1', 'a'), (BIGINT '1', 'b')")) + .failure(); + assertThat(getQueryRunner().tableExists(getSession(), duplicateTableName)).isFalse(); + } + finally { + assertUpdate("DROP TABLE IF EXISTS " + tableName); + assertUpdate("DROP TABLE IF EXISTS " + duplicateTableName); + } + } + + @Test + public void testInvalidPrimaryKeyProperties() { + String tableName = "production_create_invalid_pk"; + assertQueryFails("CREATE TABLE " + tableName + " (id bigint)", + ".*Table property 'primary_key' must contain at least one column"); + assertQueryFails("CREATE TABLE " + tableName + " (id bigint) " + + "WITH (primary_key = CAST(ARRAY[] AS ARRAY(VARCHAR)))", + ".*Table property 'primary_key' must contain at least one column"); + assertQueryFails("CREATE TABLE " + tableName + " (id bigint) WITH (primary_key = ARRAY['missing'])", + ".*Column 'missing' specified in table property 'primary_key' does not exist"); + assertQueryFails("CREATE TABLE " + tableName + " (id bigint) WITH (primary_key = ARRAY['id', 'id'])", + ".*Table property 'primary_key' contains duplicate columns"); + assertThat(getQueryRunner().tableExists(getSession(), tableName)).isFalse(); + } +} diff --git a/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestingYdbJdbcClient.java b/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestingYdbJdbcClient.java index 40a7e0ae..82e4f1ff 100644 --- a/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestingYdbJdbcClient.java +++ b/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestingYdbJdbcClient.java @@ -4,6 +4,7 @@ import io.trino.plugin.jdbc.BaseJdbcConfig; import io.trino.plugin.jdbc.ConnectionFactory; import io.trino.plugin.jdbc.JdbcColumnHandle; +import io.trino.plugin.jdbc.JdbcTableHandle; import io.trino.plugin.jdbc.QueryBuilder; import io.trino.plugin.jdbc.RemoteTableName; import io.trino.plugin.jdbc.logging.RemoteQueryModifier; @@ -13,6 +14,7 @@ import io.trino.spi.connector.SchemaTableName; import java.util.List; +import java.util.Map; import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; @@ -47,6 +49,15 @@ protected List getColumnsForPrimaryKeyLookup( return super.getColumns(session, schemaTableName, remoteTableName); } + @Override + public Map getTableProperties(ConnectorSession session, JdbcTableHandle tableHandle) { + Map properties = super.getTableProperties(session, tableHandle); + if (YdbTableProperties.getPrimaryKey(properties).equals(List.of(YDB_HIDDEN_PK_COLUMN))) { + return Map.of(); + } + return properties; + } + @Override protected List createTableSqls(RemoteTableName remoteTableName, List columns, ConnectorTableMetadata tableMetadata) { if (tableMetadata.getComment().isPresent()) { diff --git a/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestingYdbJdbcModule.java b/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestingYdbJdbcModule.java index dc71a297..9b380331 100644 --- a/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestingYdbJdbcModule.java +++ b/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestingYdbJdbcModule.java @@ -14,7 +14,9 @@ import io.trino.plugin.jdbc.QueryBuilder; import io.trino.plugin.jdbc.credential.CredentialProvider; import io.trino.plugin.jdbc.logging.RemoteQueryModifier; + import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder; +import static io.trino.plugin.jdbc.JdbcModule.bindTablePropertiesProvider; public class TestingYdbJdbcModule implements Module { @@ -25,6 +27,7 @@ public void configure(Binder binder) { .to(YdbMetadataFactory.class) .in(Scopes.SINGLETON); + bindTablePropertiesProvider(binder, YdbTableProperties.class); binder.bind(YdbConnector.class).in(Scopes.SINGLETON); } diff --git a/ydb-trino-adapter/src/test/java/tech/ydb/trino/YdbQueryRunner.java b/ydb-trino-adapter/src/test/java/tech/ydb/trino/YdbQueryRunner.java index dc1e4a1d..4046661b 100644 --- a/ydb-trino-adapter/src/test/java/tech/ydb/trino/YdbQueryRunner.java +++ b/ydb-trino-adapter/src/test/java/tech/ydb/trino/YdbQueryRunner.java @@ -47,6 +47,7 @@ static String buildJdbcUrl(YdbHelperExtension ydb) { public static class Builder extends DistributedQueryRunner.Builder { private final Map connectorProperties = new HashMap<>(); private List> initialTables = ImmutableList.of(); + private boolean useTestingClient = true; private Builder() { super(testSessionBuilder() @@ -65,6 +66,11 @@ public Builder setInitialTables(List> tables) { return this; } + public Builder useProductionClient() { + this.useTestingClient = false; + return this; + } + @Override public DistributedQueryRunner build() throws Exception { DistributedQueryRunner queryRunner = super.build(); @@ -72,7 +78,7 @@ public DistributedQueryRunner build() throws Exception { queryRunner.installPlugin(new io.trino.plugin.tpch.TpchPlugin()); queryRunner.createCatalog("tpch", "tpch"); - queryRunner.installPlugin(new YdbPlugin(new TestingYdbJdbcModule())); + queryRunner.installPlugin(new YdbPlugin(useTestingClient ? new TestingYdbJdbcModule() : new YdbClientModule())); queryRunner.createCatalog("local", "ydb", ImmutableMap.copyOf(connectorProperties)); for (TpchTable table : initialTables) { From d0c3f6ff44012a6ec0cea0c2dc1c26aa9dffc5d5 Mon Sep 17 00:00:00 2001 From: KirillKurdyukov Date: Tue, 15 Sep 2026 11:25:01 +0300 Subject: [PATCH 2/3] Harden primary key validation and CTAS cleanup test --- ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbClient.java | 3 +++ .../src/test/java/tech/ydb/trino/TestYdbCreateTable.java | 4 ++++ 2 files changed, 7 insertions(+) diff --git a/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbClient.java b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbClient.java index 628325bb..87d7f8ae 100644 --- a/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbClient.java +++ b/ydb-trino-adapter/src/main/java/tech/ydb/trino/YdbClient.java @@ -740,6 +740,9 @@ protected List createTableSqls(RemoteTableName remoteTableName, List Date: Sun, 20 Sep 2026 22:27:18 +0300 Subject: [PATCH 3/3] Fix YDB create table test expectations --- .../src/test/java/tech/ydb/trino/TestYdbCreateTable.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestYdbCreateTable.java b/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestYdbCreateTable.java index 1a3c4807..6b14a3da 100644 --- a/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestYdbCreateTable.java +++ b/ydb-trino-adapter/src/test/java/tech/ydb/trino/TestYdbCreateTable.java @@ -28,7 +28,7 @@ public void testCreateTableWithPrimaryKey() { assertUpdate("INSERT INTO " + tableName + " VALUES (1, 10, 'a'), (2, NULL, 'b')", 2); assertQuery("SELECT * FROM " + tableName, "VALUES (1, 10, 'a'), (2, NULL, 'b')"); assertThat((String) computeScalar("SHOW CREATE TABLE " + tableName)) - .contains("primary_key = ARRAY['event_id', 'tenant']"); + .contains("primary_key = ARRAY['event_id','tenant']"); } finally { assertUpdate("DROP TABLE IF EXISTS " + tableName); @@ -42,7 +42,7 @@ public void testCreateTableAsSelectWithPrimaryKey() { try { assertUpdate("CREATE TABLE " + tableName + " (id, payload) WITH (primary_key = ARRAY['id']) " + "AS VALUES (BIGINT '1', 'a'), (BIGINT '2', 'b')", 2); - assertQuery("SELECT * FROM " + tableName, "VALUES (BIGINT '1', 'a'), (BIGINT '2', 'b')"); + assertQuery("SELECT * FROM " + tableName, "VALUES (CAST(1 AS BIGINT), 'a'), (CAST(2 AS BIGINT), 'b')"); var tablesBefore = computeActual("SHOW TABLES").getOnlyColumnAsSet(); assertThat(query("CREATE TABLE " + duplicateTableName + " (id, payload) " +