Add support create table with auto increment column in MySQL

Author: hqbhoho

docs/src/main/sphinx/connector/mysql.md

@@ -376,6 +376,7 @@ following features:
 - [](/sql/drop-table)
 - [](/sql/create-schema)
 - [](/sql/drop-schema)
+- [](mysql-schema-and-table-management)
 - [](mysql-procedures)
 - [](mysql-table-functions)
 
@@ -395,6 +396,41 @@ following features:
 ```{include} non-transactional-merge.fragment
 ```
 
+(mysql-schema-and-table-management)=
+### Schema and table management
+
+#### Column properties
+
+Column property usage example:
+
+```
+CREATE TABLE person (
+  id INT NOT NULL WITH (auto_increment = true),
+  name VARCHAR,
+  age INT,
+  birthday DATE
+)
+WITH (
+  primary_key = ARRAY['id']
+);
+```
+
+The following are supported MySQL column properties:
+
+:::{list-table}
+:widths: 30, 10, 60
+:header-rows: 1
+
+* - Property name
+  - Required
+  - Description
+* - `auto_increment`
+  - No
+  - Auto generate a unique identity for new rows. There can be only one auto increment column
+    and must be defined as the first key. Only applies to integer types (`TINYINT`,
+    `SMALLINT`, `INTEGER`, `BIGINT`) column.
+:::
+
 (mysql-procedures)=
 ### Procedures
 

plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/BaseJdbcClient.java

@@ -336,13 +336,16 @@ public List<JdbcColumnHandle> getColumns(ConnectorSession session, SchemaTableNa
                 boolean nullable = (resultSet.getInt("NULLABLE") != columnNoNulls);
                 // Note: some databases (e.g. SQL Server) do not return column remarks/comment here.
                 Optional<String> comment = Optional.ofNullable(emptyToNull(resultSet.getString("REMARKS")));
+                // Note: the default implementation returns an empty map, and each connector needs to provide its own implementation.
+                Map<String, Object> columnProperties = getColumnProperties(resultSet);
                 // skip unsupported column types
                 columnMapping.ifPresent(mapping -> columns.add(JdbcColumnHandle.builder()
                         .setColumnName(columnName)
                         .setJdbcTypeHandle(typeHandle)
                         .setColumnType(mapping.getType())
                         .setNullable(nullable)
                         .setComment(comment)
+                        .setColumnProperties(columnProperties)
                         .build()));
                 if (columnMapping.isEmpty()) {
                     UnsupportedTypeHandling unsupportedTypeHandling = getUnsupportedTypeHandling(session);
@@ -366,6 +369,12 @@ public List<JdbcColumnHandle> getColumns(ConnectorSession session, SchemaTableNa
         }
     }
 
+    public Map<String, Object> getColumnProperties(ResultSet resultSet)
+            throws SQLException
+    {
+        return emptyMap();
+    }
+
     @Override
     public Iterator<RelationColumnsMetadata> getAllTableColumns(ConnectorSession session, Optional<String> schema)
     {

plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/JdbcColumnHandle.java

@@ -16,18 +16,24 @@
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.base.Joiner;
+import com.google.common.collect.ImmutableMap;
 import io.airlift.slice.SizeOf;
+import io.trino.spi.TrinoException;
 import io.trino.spi.connector.ColumnHandle;
 import io.trino.spi.connector.ColumnMetadata;
 import io.trino.spi.connector.ColumnSchema;
 import io.trino.spi.type.Type;
 
+import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 
 import static io.airlift.slice.SizeOf.estimatedSizeOf;
 import static io.airlift.slice.SizeOf.instanceSize;
 import static io.airlift.slice.SizeOf.sizeOf;
+import static io.trino.spi.StandardErrorCode.INVALID_COLUMN_PROPERTY;
+import static java.util.Collections.emptyMap;
 import static java.util.Objects.requireNonNull;
 
 public final class JdbcColumnHandle
@@ -40,11 +46,12 @@ public final class JdbcColumnHandle
     private final Type columnType;
     private final boolean nullable;
     private final Optional<String> comment;
+    private final Map<String, Object> columnProperties;
 
     // All and only required fields
     public JdbcColumnHandle(String columnName, JdbcTypeHandle jdbcTypeHandle, Type columnType)
     {
-        this(columnName, jdbcTypeHandle, columnType, true, Optional.empty());
+        this(columnName, jdbcTypeHandle, columnType, true, Optional.empty(), emptyMap());
     }
 
     /**
@@ -57,13 +64,15 @@ public JdbcColumnHandle(
             @JsonProperty("jdbcTypeHandle") JdbcTypeHandle jdbcTypeHandle,
             @JsonProperty("columnType") Type columnType,
             @JsonProperty("nullable") boolean nullable,
-            @JsonProperty("comment") Optional<String> comment)
+            @JsonProperty("comment") Optional<String> comment,
+            @JsonProperty("columnProperties") Map<String, Object> columnProperties)
     {
         this.columnName = requireNonNull(columnName, "columnName is null");
         this.jdbcTypeHandle = requireNonNull(jdbcTypeHandle, "jdbcTypeHandle is null");
         this.columnType = requireNonNull(columnType, "columnType is null");
         this.nullable = nullable;
         this.comment = requireNonNull(comment, "comment is null");
+        this.columnProperties = ImmutableMap.copyOf(columnProperties);
     }
 
     @JsonProperty
@@ -96,13 +105,20 @@ public Optional<String> getComment()
         return comment;
     }
 
+    @JsonProperty
+    public Map<String, Object> getColumnProperties()
+    {
+        return columnProperties;
+    }
+
     public ColumnMetadata getColumnMetadata()
     {
         return ColumnMetadata.builder()
                 .setName(columnName)
                 .setType(columnType)
                 .setNullable(nullable)
                 .setComment(comment)
+                .setProperties(columnProperties)
                 .build();
     }
 
@@ -149,7 +165,25 @@ public long getRetainedSizeInBytes()
                 + sizeOf(nullable)
                 + estimatedSizeOf(columnName)
                 + sizeOf(comment, SizeOf::estimatedSizeOf)
-                + jdbcTypeHandle.getRetainedSizeInBytes();
+                + jdbcTypeHandle.getRetainedSizeInBytes()
+                + estimatedSizeOf(columnProperties, SizeOf::estimatedSizeOf, JdbcColumnHandle::estimatedObjectSizeOf);
+    }
+
+    /**
+     * Returns the estimated size of a property value.
+     * <p>
+     * Supported types are the same as those accepted by
+     * {@link io.trino.metadata.PropertyUtil#toExpression}.
+     */
+    private static long estimatedObjectSizeOf(Object value)
+    {
+        return switch (value) {
+            case String stringValue -> estimatedSizeOf(stringValue);
+            case Boolean _, Integer _, Long _, Double _ -> instanceSize(value.getClass());
+            case List<?> list -> estimatedSizeOf(list, JdbcColumnHandle::estimatedObjectSizeOf);
+            case null -> throw new TrinoException(INVALID_COLUMN_PROPERTY, "Property value is null");
+            default -> throw new TrinoException(INVALID_COLUMN_PROPERTY, "Unsupported property value type: " + value.getClass().getName());
+        };
     }
 
     public static Builder builder()
@@ -169,6 +203,7 @@ public static final class Builder
         private Type columnType;
         private boolean nullable = true;
         private Optional<String> comment = Optional.empty();
+        private Map<String, Object> columnProperties = emptyMap();
 
         public Builder() {}
 
@@ -179,6 +214,7 @@ private Builder(JdbcColumnHandle handle)
             this.columnType = handle.getColumnType();
             this.nullable = handle.isNullable();
             this.comment = handle.getComment();
+            this.columnProperties = handle.getColumnProperties();
         }
 
         public Builder setColumnName(String columnName)
@@ -211,14 +247,21 @@ public Builder setComment(Optional<String> comment)
             return this;
         }
 
+        public Builder setColumnProperties(Map<String, Object> columnProperties)
+        {
+            this.columnProperties = columnProperties;
+            return this;
+        }
+
         public JdbcColumnHandle build()
         {
             return new JdbcColumnHandle(
                     columnName,
                     jdbcTypeHandle,
                     columnType,
                     nullable,
-                    comment);
+                    comment,
+                    columnProperties);
         }
     }
 }

plugin/trino-mysql/src/main/java/io/trino/plugin/mysql/MySqlClient.java

@@ -83,14 +83,18 @@
 import io.trino.spi.statistics.ColumnStatistics;
 import io.trino.spi.statistics.Estimate;
 import io.trino.spi.statistics.TableStatistics;
+import io.trino.spi.type.BigintType;
 import io.trino.spi.type.CharType;
 import io.trino.spi.type.DecimalType;
 import io.trino.spi.type.Decimals;
+import io.trino.spi.type.IntegerType;
 import io.trino.spi.type.LongTimestampWithTimeZone;
+import io.trino.spi.type.SmallintType;
 import io.trino.spi.type.StandardTypes;
 import io.trino.spi.type.TimeType;
 import io.trino.spi.type.TimestampType;
 import io.trino.spi.type.TimestampWithTimeZoneType;
+import io.trino.spi.type.TinyintType;
 import io.trino.spi.type.Type;
 import io.trino.spi.type.TypeManager;
 import io.trino.spi.type.TypeSignature;
@@ -182,6 +186,7 @@
 import static io.trino.plugin.jdbc.StandardColumnMappings.varcharWriteFunction;
 import static io.trino.plugin.jdbc.TypeHandlingJdbcSessionProperties.getUnsupportedTypeHandling;
 import static io.trino.plugin.jdbc.UnsupportedTypeHandling.CONVERT_TO_VARCHAR;
+import static io.trino.plugin.mysql.MySqlColumnProperties.AUTO_INCREMENT;
 import static io.trino.plugin.mysql.MySqlTableProperties.PRIMARY_KEY_PROPERTY;
 import static io.trino.spi.StandardErrorCode.ALREADY_EXISTS;
 import static io.trino.spi.StandardErrorCode.INVALID_TABLE_PROPERTY;
@@ -487,15 +492,30 @@ protected List<String> createTableSqls(RemoteTableName remoteTableName, List<Str
         columnDefinitions.addAll(columns);
 
         List<String> primaryKeys = MySqlTableProperties.getPrimaryKey(tableMetadata.getProperties());
+        verifyColumnMetadata(primaryKeys, tableMetadata.getColumns());
         if (!primaryKeys.isEmpty()) {
-            verifyPrimaryKey(primaryKeys, tableMetadata.getColumns());
             columnDefinitions.add("PRIMARY KEY (" + primaryKeys.stream().map(this::quoted).collect(joining(", ")) + ")");
         }
         return ImmutableList.of(format("CREATE TABLE %s (%s) COMMENT %s", quoted(remoteTableName), join(", ", columnDefinitions.build()), mysqlVarcharLiteral(tableMetadata.getComment().orElse(NO_COMMENT))));
     }
 
-    private static void verifyPrimaryKey(List<String> primaryKeys, List<ColumnMetadata> columns)
+    private static void verifyColumnMetadata(List<String> primaryKeys, List<ColumnMetadata> columns)
     {
+        Set<String> autoIncrementColumnNames = columns.stream()
+                .filter(column -> (boolean) column.getProperties().getOrDefault(AUTO_INCREMENT, false))
+                .map(ColumnMetadata::getName)
+                .collect(toImmutableSet());
+
+        if (autoIncrementColumnNames.size() > 1) {
+            throw new TrinoException(NOT_SUPPORTED, "There can be only one auto increment column in MySQL");
+        }
+
+        if (autoIncrementColumnNames.size() == 1) {
+            if (primaryKeys.isEmpty() || !autoIncrementColumnNames.contains(primaryKeys.getFirst())) {
+                throw new TrinoException(NOT_SUPPORTED, "Auto increment column must be defined as the first key in MySQL");
+            }
+        }
+
         Set<String> columnNames = columns.stream()
                 .map(column -> {
                     String columnName = column.getName();
@@ -521,10 +541,23 @@ protected String getColumnDefinitionSql(ConnectorSession session, ColumnMetadata
             throw new TrinoException(NOT_SUPPORTED, "This connector does not support creating tables with column comment");
         }
 
-        return "%s %s %s".formatted(
+        return "%s %s %s %s".formatted(
                 quoted(columnName),
                 toWriteMapping(session, column.getType()).getDataType(),
-                column.isNullable() ? "NULL" : "NOT NULL");
+                column.isNullable() ? "NULL" : "NOT NULL",
+                isAutoIncrement(column) ? "AUTO_INCREMENT" : "");
+    }
+
+    private static boolean isAutoIncrement(ColumnMetadata column)
+    {
+        boolean isAutoIncrement = (boolean) column.getProperties().getOrDefault(AUTO_INCREMENT, false);
+        if (isAutoIncrement) {
+            Type type = column.getType();
+            if (!(type instanceof TinyintType || type instanceof SmallintType || type instanceof IntegerType || type instanceof BigintType)) {
+                throw new TrinoException(NOT_SUPPORTED, "Unsupported column type for AUTO_INCREMENT: " + type);
+            }
+        }
+        return isAutoIncrement;
     }
 
     private static String mysqlVarcharLiteral(String value)
@@ -644,6 +677,18 @@ public Optional<ColumnMapping> toColumnMapping(ConnectorSession session, Connect
         return Optional.empty();
     }
 
+    @Override
+    public Map<String, Object> getColumnProperties(ResultSet resultSet)
+            throws SQLException
+    {
+        ImmutableMap.Builder<String, Object> columnPropertiesBuilder = ImmutableMap.builder();
+        boolean autoIncrement = "YES".equals(resultSet.getString("IS_AUTOINCREMENT"));
+        if (autoIncrement) {
+            columnPropertiesBuilder.put(AUTO_INCREMENT, autoIncrement);
+        }
+        return columnPropertiesBuilder.buildOrThrow();
+    }
+
     private static ColumnMapping mySqlDefaultVarcharColumnMapping(int columnSize, Optional<CaseSensitivity> caseSensitivity)
     {
         if (columnSize > VarcharType.MAX_LENGTH) {
@@ -963,6 +1008,10 @@ private void addColumn(ConnectorSession session, RemoteTableName table, ColumnMe
             throw new TrinoException(NOT_SUPPORTED, "This connector does not support adding columns with comments");
         }
 
+        if (!column.getProperties().isEmpty()) {
+            throw new TrinoException(NOT_SUPPORTED, "This connector does not support adding columns with column properties");
+        }
+
         try (Connection connection = connectionFactory.openConnection(session)) {
             verify(connection.getAutoCommit());
             String columnName = column.getName();

plugin/trino-mysql/src/main/java/io/trino/plugin/mysql/MySqlClientModule.java

@@ -40,6 +40,7 @@
 import static com.google.inject.multibindings.Multibinder.newSetBinder;
 import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder;
 import static io.airlift.configuration.ConfigBinder.configBinder;
+import static io.trino.plugin.jdbc.JdbcModule.bindColumnPropertiesProvider;
 import static io.trino.plugin.jdbc.JdbcModule.bindTablePropertiesProvider;
 
 public class MySqlClientModule
@@ -55,6 +56,7 @@ protected void setup(Binder binder)
         configBinder(binder).bindConfig(MySqlConfig.class);
         configBinder(binder).bindConfig(JdbcStatisticsConfig.class);
         bindTablePropertiesProvider(binder, MySqlTableProperties.class);
+        bindColumnPropertiesProvider(binder, MySqlColumnProperties.class);
         install(new DecimalModule());
         install(new JdbcJoinPushdownSupportModule());
         newSetBinder(binder, ConnectorTableFunction.class).addBinding().toProvider(Query.class).in(Scopes.SINGLETON);