Fix failure when inserting into Kudu table having row uuid column

trinodb · Jun 21, 2022 · 1b60552 · 1b60552
1 parent dec1ed6
commit 1b60552
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 3 deletions.
diff --git a/plugin/trino-kudu/src/main/java/io/trino/plugin/kudu/KuduInsertTableHandle.java b/plugin/trino-kudu/src/main/java/io/trino/plugin/kudu/KuduInsertTableHandle.java
@@ -30,23 +30,27 @@ public class KuduInsertTableHandle
 {
     private final SchemaTableName schemaTableName;
     private final List<Type> columnTypes;
+    private final boolean generateUUID;
     private transient KuduTable table;
 
     @JsonCreator
     public KuduInsertTableHandle(
             @JsonProperty("schemaTableName") SchemaTableName schemaTableName,
-            @JsonProperty("columnTypes") List<Type> columnTypes)
+            @JsonProperty("columnTypes") List<Type> columnTypes,
+            @JsonProperty("generateUUID") boolean generateUUID)
     {
-        this(schemaTableName, columnTypes, null);
+        this(schemaTableName, columnTypes, generateUUID, null);
     }
 
     public KuduInsertTableHandle(
             SchemaTableName schemaTableName,
             List<Type> columnTypes,
+            boolean generateUUID,
             KuduTable table)
     {
         this.schemaTableName = requireNonNull(schemaTableName, "schemaTableName is null");
         this.columnTypes = ImmutableList.copyOf(requireNonNull(columnTypes, "columnTypes is null"));
+        this.generateUUID = generateUUID;
         this.table = table;
     }
 
@@ -70,9 +74,10 @@ public List<Type> getOriginalColumnTypes()
     }
 
     @Override
+    @JsonProperty
     public boolean isGenerateUUID()
     {
-        return false;
+        return generateUUID;
     }
 
     public KuduTable getTable(KuduClientSession session)

diff --git a/plugin/trino-kudu/src/main/java/io/trino/plugin/kudu/KuduMetadata.java b/plugin/trino-kudu/src/main/java/io/trino/plugin/kudu/KuduMetadata.java
@@ -313,6 +313,8 @@ public ConnectorInsertTableHandle beginInsert(ConnectorSession session, Connecto
         return new KuduInsertTableHandle(
                 tableHandle.getSchemaTableName(),
                 columnTypes,
+                columns.stream()
+                        .anyMatch(column -> column.getName().equals(ROW_ID)),
                 table);
     }
 

diff --git a/plugin/trino-kudu/src/test/java/io/trino/plugin/kudu/AbstractKuduConnectorTest.java b/plugin/trino-kudu/src/test/java/io/trino/plugin/kudu/AbstractKuduConnectorTest.java
@@ -333,6 +333,17 @@ public void testInsert()
         throw new SkipException("TODO");
     }
 
+    @Test
+    public void testInsertIntoTableHavingRowUuid()
+    {
+        try (TestTable table = new TestTable(getQueryRunner()::execute, "test_insert_", " AS SELECT * FROM region WITH NO DATA")) {
+            assertUpdate("INSERT INTO " + table.getName() + " SELECT * FROM region", 5);
+
+            assertThat(query("SELECT * FROM " + table.getName()))
+                    .matches("SELECT * FROM region");
+        }
+    }
+
     @Test
     @Override
     public void testInsertUnicode()