Disallow writes into hive tables that have any enforced constraints

prestodb · Feb 26, 2024 · 492aeeb · 492aeeb
1 parent b0f28f5
commit 492aeeb
Show file tree

Hide file tree

Showing 3 changed files with 75 additions and 2 deletions.
diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/HiveMetadata.java b/presto-hive/src/main/java/com/facebook/presto/hive/HiveMetadata.java
@@ -1868,7 +1868,7 @@ private HiveInsertTableHandle beginInsertInternal(ConnectorSession session, Conn
         Table table = metastore.getTable(metastoreContext, tableName.getSchemaName(), tableName.getTableName())
                 .orElseThrow(() -> new TableNotFoundException(tableName));
 
-        checkTableIsWritable(table, writesToNonManagedTablesEnabled);
+        checkTableIsWritable(table, writesToNonManagedTablesEnabled, metastore.getTableConstraints(metastoreContext, tableName.getSchemaName(), tableName.getTableName()));
 
         for (Column column : table.getDataColumns()) {
             if (!isWritableType(column.getType())) {

diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/HiveWriteUtils.java b/presto-hive/src/main/java/com/facebook/presto/hive/HiveWriteUtils.java
@@ -42,6 +42,7 @@
 import com.facebook.presto.spi.PrestoException;
 import com.facebook.presto.spi.SchemaNotFoundException;
 import com.facebook.presto.spi.SchemaTableName;
+import com.facebook.presto.spi.constraints.TableConstraint;
 import com.facebook.presto.spi.security.ConnectorIdentity;
 import com.google.common.collect.ImmutableList;
 import com.google.common.primitives.Shorts;
@@ -300,7 +301,7 @@ else if (isRowType(type)) {
         throw new IllegalArgumentException("unsupported type: " + type);
     }
 
-    public static void checkTableIsWritable(Table table, boolean writesToNonManagedTablesEnabled)
+    public static void checkTableIsWritable(Table table, boolean writesToNonManagedTablesEnabled, List<TableConstraint<String>> constraints)
     {
         PrestoTableType tableType = table.getTableType();
         if (!writesToNonManagedTablesEnabled
@@ -310,6 +311,10 @@ public static void checkTableIsWritable(Table table, boolean writesToNonManagedT
             throw new PrestoException(NOT_SUPPORTED, "Cannot write to non-managed Hive table");
         }
 
+        if (constraints.stream().anyMatch(TableConstraint::isEnforced)) {
+            throw new PrestoException(NOT_SUPPORTED, format("Cannot write to table %s since it has table constraints that are enforced", table.getSchemaTableName().toString()));
+        }
+
         checkWritable(
                 table.getSchemaTableName(),
                 Optional.empty(),

diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveIntegrationSmokeTest.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveIntegrationSmokeTest.java
@@ -6429,6 +6429,74 @@ public void testCreateTableWithConstraints()
         assertQueryFails(createTableWithTwoConstraintsSql, "Multiple primary key constraints are not allowed");
     }
 
+    @Test
+    public void testDMLWithEnforcedConstraints()
+    {
+        String tableName = "test_table_with_constraints";
+
+        String createTableWithOneConstraintFormat = "CREATE TABLE %s.%s.%s (\n" +
+                "   %s bigint,\n" +
+                "   %s double,\n" +
+                "   %s varchar,\n" +
+                "   %s bigint,\n" +
+                "   %s\n" +
+                ")\n" +
+                "WITH (\n" +
+                "   format = 'ORC'\n" +
+                ")";
+
+        String createTableWithOneConstraintSql = format(
+                createTableWithOneConstraintFormat,
+                getSession().getCatalog().get(),
+                getSession().getSchema().get(),
+                tableName,
+                "c1",
+                "c2",
+                "c3",
+                "c4",
+                "CONSTRAINT pk PRIMARY KEY (c4)");
+
+        String expectedcreateTableWithOneConstraint = format(
+                createTableWithOneConstraintFormat,
+                getSession().getCatalog().get(),
+                getSession().getSchema().get(),
+                tableName,
+                "\"c1\"",
+                "\"c2\"",
+                "\"c3\"",
+                "\"c4\"",
+                "CONSTRAINT pk PRIMARY KEY (c4)");
+
+        assertUpdate(getSession(), createTableWithOneConstraintSql);
+        MaterializedResult actualResult = computeActual("SHOW CREATE TABLE " + tableName);
+        assertEquals(getOnlyElement(actualResult.getOnlyColumnAsSet()), expectedcreateTableWithOneConstraint);
+
+        String insertStmt = format("INSERT INTO %s VALUES (1, 2.3, 'abc', 4)", tableName);
+        assertQueryFails(insertStmt, format("Cannot write to table %s.%s since it has table constraints that are enforced", getSession().getSchema().get(), tableName));
+
+        String dropConstraintStmt = format("ALTER TABLE %s.%s.%s DROP CONSTRAINT pk", getSession().getCatalog().get(), getSession().getSchema().get(), tableName);
+        assertUpdate(getSession(), dropConstraintStmt);
+
+        String addPrimaryKeyStmt = "ALTER TABLE " + tableName + " ADD CONSTRAINT pk PRIMARY KEY (c2, c3) NOT ENFORCED";
+        assertUpdate(getSession(), addPrimaryKeyStmt);
+        assertUpdate(getSession(), insertStmt, 1);
+
+        String addUniqueEnforced = "ALTER TABLE " + tableName + " ADD CONSTRAINT uq1 UNIQUE (c1) ENFORCED";
+        assertUpdate(getSession(), addUniqueEnforced);
+        assertQueryFails(insertStmt, format("Cannot write to table %s.%s since it has table constraints that are enforced", getSession().getSchema().get(), tableName));
+
+        String addSecondUniqueNotEnforced = "ALTER TABLE " + tableName + " ADD CONSTRAINT uq2 UNIQUE (c4) NOT ENFORCED";
+        assertUpdate(getSession(), addSecondUniqueNotEnforced);
+        assertQueryFails(insertStmt, format("Cannot write to table %s.%s since it has table constraints that are enforced", getSession().getSchema().get(), tableName));
+
+        dropConstraintStmt = format("ALTER TABLE %s.%s.%s DROP CONSTRAINT uq1", getSession().getCatalog().get(), getSession().getSchema().get(), tableName);
+        assertUpdate(getSession(), dropConstraintStmt);
+        assertUpdate(getSession(), insertStmt, 1);
+
+        String dropTableStmt = format("DROP TABLE %s.%s.%s", getSession().getCatalog().get(), getSession().getSchema().get(), tableName);
+        assertUpdate(getSession(), dropTableStmt);
+    }
+
     protected String retentionDays(int days)
     {
         return "";