Fix failure when check constraint is null

The planner was incorrectly trying to derive the necessary coercions based on the rewritten expression instead of the original one.
trinodb · Mar 1, 2024 · a3d1fe8 · a3d1fe8
1 parent 5bce777
commit a3d1fe8
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 1 deletion.
diff --git a/core/trino-main/src/main/java/io/trino/sql/planner/QueryPlanner.java b/core/trino-main/src/main/java/io/trino/sql/planner/QueryPlanner.java
@@ -727,7 +727,7 @@ private PlanBuilder addCheckConstraints(List<Expression> constraints, PlanBuilde
 
             Expression predicate = new IfExpression(
                     // When predicate evaluates to UNKNOWN (e.g. NULL > 100), it should not violate the check constraint.
-                    new CoalesceExpression(coerceIfNecessary(analysis, symbol, symbol), TRUE_LITERAL),
+                    new CoalesceExpression(coerceIfNecessary(analysis, constraint, symbol), TRUE_LITERAL),
                     TRUE_LITERAL,
                     new Cast(failFunction(plannerContext.getMetadata(), CONSTRAINT_VIOLATION, "Check constraint violation: " + constraint), toSqlType(BOOLEAN)));
 

diff --git a/core/trino-main/src/test/java/io/trino/sql/query/TestCheckConstraint.java b/core/trino-main/src/test/java/io/trino/sql/query/TestCheckConstraint.java
@@ -65,6 +65,9 @@ public TestCheckConstraint()
                     if (schemaTableName.equals(new SchemaTableName("tiny", "nation"))) {
                         return TPCH_NATION_SCHEMA;
                     }
+                    if (schemaTableName.equals(new SchemaTableName("tiny", "nation_with_null_check"))) {
+                        return TPCH_NATION_SCHEMA;
+                    }
                     if (schemaTableName.equals(new SchemaTableName("tiny", "nation_multiple_column_constraint"))) {
                         return TPCH_NATION_SCHEMA;
                     }
@@ -101,6 +104,9 @@ public TestCheckConstraint()
                     if (schemaTableName.equals(new SchemaTableName("tiny", "nation"))) {
                         return ImmutableList.of("regionkey < 10");
                     }
+                    if (schemaTableName.equals(new SchemaTableName("tiny", "nation_with_null_check"))) {
+                        return ImmutableList.of("null");
+                    }
                     if (schemaTableName.equals(new SchemaTableName("tiny", "nation_multiple_column_constraint"))) {
                         return ImmutableList.of("nationkey < 100 AND regionkey < 50");
                     }
@@ -137,6 +143,9 @@ public TestCheckConstraint()
                     if (schemaTableName.equals(new SchemaTableName("tiny", "nation"))) {
                         return TPCH_NATION_DATA;
                     }
+                    if (schemaTableName.equals(new SchemaTableName("tiny", "nation_with_null_check"))) {
+                        return TPCH_NATION_DATA;
+                    }
                     if (schemaTableName.equals(new SchemaTableName("tiny", "nation_multiple_column_constraint"))) {
                         return TPCH_NATION_DATA;
                     }
@@ -226,6 +235,34 @@ public void testInsertAllowUnknown()
                 .matches("SELECT BIGINT '1'");
     }
 
+    @Test
+    void testNullConstraint()
+    {
+        assertThat(assertions.query("""
+                MERGE INTO mock.tiny.nation_with_null_check USING (VALUES 1,2) t(x) ON nationkey = x
+                WHEN MATCHED THEN DELETE
+                """))
+                .matches("SELECT BIGINT '2'");
+
+        assertThat(assertions.query("""
+                MERGE INTO mock.tiny.nation_with_null_check USING (VALUES 5) t(x) ON nationkey = x
+                WHEN MATCHED THEN UPDATE SET regionkey = regionkey * 2
+                """))
+                .matches("SELECT BIGINT '1'");
+
+        assertThat(assertions.query("""
+                MERGE INTO mock.tiny.nation_with_null_check USING (VALUES 42) t(dummy) ON false
+                WHEN NOT MATCHED THEN INSERT VALUES (101, 'POLAND', 0, 'No comment')
+                """))
+                .matches("SELECT BIGINT '1'");
+
+        assertThat(assertions.query("INSERT INTO mock.tiny.nation_with_null_check VALUES (99, 'POLAND', 49, 'No comment')"))
+                .matches("SELECT BIGINT '1'");
+
+        assertThat(assertions.query("DELETE FROM mock.tiny.nation_with_null_check WHERE nationkey < 3"))
+                .matches("SELECT BIGINT '3'");
+    }
+
     @Test
     public void testInsertCheckMultipleColumns()
     {