Support explicit primary keys for YDB table creation

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` без

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,44 @@ private long executeReturningDml(
 
     @Override
     protected List<String> createTableSqls(RemoteTableName remoteTableName, List<String> 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<String> 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.contains(null)) {
+            throw new TrinoException(INVALID_TABLE_PROPERTY, "Table property 'primary_key' must not contain null columns");
+        }
+        if (primaryKeys.stream().distinct().count() != primaryKeys.size()) {
+            throw new TrinoException(INVALID_TABLE_PROPERTY, "Table property 'primary_key' contains duplicate columns");
+        }
+
+        Set<String> 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<String, Object> getTableProperties(ConnectorSession session, JdbcTableHandle tableHandle) {
+        List<String> primaryKeys = getPrimaryKeys(session, tableHandle.getRequiredNamedRelation().getRemoteTableName()).stream()
+                .map(JdbcColumnHandle::getColumnName)
+                .toList();
+        return Map.of(PRIMARY_KEY_PROPERTY, primaryKeys);
     }
 }

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);
     }
 

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<PropertyMetadata<?>> 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<PropertyMetadata<?>> getTableProperties() {
+        return tableProperties;
+    }
+
+    public static List<String> getPrimaryKey(Map<String, Object> tableProperties) {
+        requireNonNull(tableProperties, "tableProperties is null");
+        return ((List<?>) tableProperties.get(PRIMARY_KEY_PROPERTY)).stream()
+                .map(String.class::cast)
+                .toList();
+    }
+}

ydb-trino-adapter/src/test/java/tech/ydb/trino/TestYdbCreateTable.java

@@ -0,0 +1,76 @@
+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 (CAST(1 AS BIGINT), 'a'), (CAST(2 AS BIGINT), 'b')");
+
+            var tablesBefore = computeActual("SHOW TABLES").getOnlyColumnAsSet();
+            assertThat(query("CREATE TABLE " + duplicateTableName + " (id, payload) " +
+                    "WITH (primary_key = ARRAY['id']) AS VALUES (BIGINT '1', 'a'), (BIGINT '1', 'b')"))
+                    .failure();
+            assertThat(computeActual("SHOW TABLES").getOnlyColumnAsSet()).isEqualTo(tablesBefore);
+            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");
+        assertQueryFails("CREATE TABLE " + tableName + " (id bigint) WITH (primary_key = ARRAY[CAST(NULL AS VARCHAR)])",
+                ".*Table property 'primary_key' must not contain null columns");
+        assertThat(getQueryRunner().tableExists(getSession(), tableName)).isFalse();
+    }
+}

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<JdbcColumnHandle> getColumnsForPrimaryKeyLookup(
         return super.getColumns(session, schemaTableName, remoteTableName);
     }
 
+    @Override
+    public Map<String, Object> getTableProperties(ConnectorSession session, JdbcTableHandle tableHandle) {
+        Map<String, Object> properties = super.getTableProperties(session, tableHandle);
+        if (YdbTableProperties.getPrimaryKey(properties).equals(List.of(YDB_HIDDEN_PK_COLUMN))) {
+            return Map.of();
+        }
+        return properties;
+    }
+
     @Override
     protected List<String> createTableSqls(RemoteTableName remoteTableName, List<String> columns, ConnectorTableMetadata tableMetadata) {
         if (tableMetadata.getComment().isPresent()) {

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);
     }
 

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<Builder> {
         private final Map<String, String> connectorProperties = new HashMap<>();
         private List<TpchTable<?>> initialTables = ImmutableList.of();
+        private boolean useTestingClient = true;
 
         private Builder() {
             super(testSessionBuilder()
@@ -65,14 +66,19 @@ public Builder setInitialTables(List<TpchTable<?>> tables) {
             return this;
         }
 
+        public Builder useProductionClient() {
+            this.useTestingClient = false;
+            return this;
+        }
+
         @Override
         public DistributedQueryRunner build() throws Exception {
             DistributedQueryRunner queryRunner = super.build();
             try {
                 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) {