Restrict insert overwrite only to auto commit context

trinodb · Oct 22, 2021 · efcef56 · efcef56
1 parent ce7b034
commit efcef56
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 6 deletions.
diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMetadata.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMetadata.java
@@ -1681,9 +1681,15 @@ public HiveInsertTableHandle beginInsert(ConnectorSession session, ConnectorTabl
 
         WriteInfo writeInfo = locationService.getQueryWriteInfo(locationHandle);
         if (getInsertExistingPartitionsBehavior(session) == InsertExistingPartitionsBehavior.OVERWRITE
-                && isTransactionalTable(table.getParameters())
                 && writeInfo.getWriteMode() == DIRECT_TO_TARGET_EXISTING_DIRECTORY) {
-            throw new TrinoException(NOT_SUPPORTED, "Overwriting existing partition in transactional tables doesn't support DIRECT_TO_TARGET_EXISTING_DIRECTORY write mode");
+            if (isTransactionalTable(table.getParameters())) {
+                throw new TrinoException(NOT_SUPPORTED, "Overwriting existing partition in transactional tables doesn't support DIRECT_TO_TARGET_EXISTING_DIRECTORY write mode");
+            }
+            // This check is required to prevent using partition overwrite operation during user managed transactions
+            // Partition overwrite operation is nonatomic thus can't and shouldn't be used in non autocommit context.
+            if (!session.isAutoCommitContext()) {
+                throw new TrinoException(NOT_SUPPORTED, "Overwriting existing partition in non auto commit context doesn't support DIRECT_TO_TARGET_EXISTING_DIRECTORY write mode");
+            }
         }
         metastore.declareIntentionToWrite(session, writeInfo.getWriteMode(), writeInfo.getWritePath(), tableName);
         return result;

diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/BaseTestHiveInsertOverwrite.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/BaseTestHiveInsertOverwrite.java
@@ -14,6 +14,7 @@
 package io.trino.plugin.hive;
 
 import com.google.common.collect.ImmutableMap;
+import io.trino.Session;
 import io.trino.plugin.hive.containers.HiveMinioDataLake;
 import io.trino.plugin.hive.s3.S3HiveQueryRunner;
 import io.trino.testing.AbstractTestQueryFramework;
@@ -29,6 +30,7 @@
 import static java.lang.String.format;
 import static java.util.Objects.requireNonNull;
 import static java.util.stream.Collectors.joining;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 public abstract class BaseTestHiveInsertOverwrite
         extends AbstractTestQueryFramework
@@ -75,6 +77,20 @@ public void setUp()
                 bucketName));
     }
 
+    @Test
+    public void testInsertOverwriteInTransaction()
+    {
+        String testTable = getTestTableName();
+        computeActual(getCreateTableStatement(testTable, "partitioned_by=ARRAY['regionkey']"));
+        assertThatThrownBy(
+                () -> newTransaction()
+                        .execute(getSession(), session -> {
+                            getQueryRunner().execute(session, createInsertStatement(testTable));
+                        }))
+                .hasMessage("Overwriting existing partition in non auto commit context doesn't support DIRECT_TO_TARGET_EXISTING_DIRECTORY write mode");
+        computeActual(format("DROP TABLE %s", testTable));
+    }
+
     @Test
     public void testInsertOverwriteNonPartitionedTable()
     {
@@ -149,15 +165,26 @@ public void testInsertOverwritePartitionedAndBucketedExternalTable()
     }
 
     protected void assertInsertFailure(String testTable, String expectedMessageRegExp)
+    {
+        assertInsertFailure(getSession(), testTable, expectedMessageRegExp);
+    }
+
+    protected void assertInsertFailure(Session session, String testTable, String expectedMessageRegExp)
     {
         assertQueryFails(
-                format("INSERT INTO %s " +
-                                "SELECT name, comment, nationkey, regionkey " +
-                                "FROM tpch.tiny.nation",
-                        testTable),
+                session,
+                createInsertStatement(testTable),
                 expectedMessageRegExp);
     }
 
+    private String createInsertStatement(String testTable)
+    {
+        return format("INSERT INTO %s " +
+                        "SELECT name, comment, nationkey, regionkey " +
+                        "FROM tpch.tiny.nation",
+                testTable);
+    }
+
     protected void assertOverwritePartition(String testTable)
     {
         computeActual(format(