diff --git a/core/trino-main/src/main/java/io/trino/sql/analyzer/StatementAnalyzer.java b/core/trino-main/src/main/java/io/trino/sql/analyzer/StatementAnalyzer.java index 7945af4ef9b0..9d4cf239aadf 100644 --- a/core/trino-main/src/main/java/io/trino/sql/analyzer/StatementAnalyzer.java +++ b/core/trino-main/src/main/java/io/trino/sql/analyzer/StatementAnalyzer.java @@ -3414,15 +3414,12 @@ protected Scope visitMerge(Merge merge, Optional scope) if (!accessControl.getRowFilters(session.toSecurityContext(), tableName).isEmpty()) { throw semanticException(NOT_SUPPORTED, merge, "Cannot merge into a table with row filters"); } - if (!tableSchema.getTableSchema().getCheckConstraints().isEmpty()) { - // TODO https://github.com/trinodb/trino/issues/15411 Add support for CHECK constraint to MERGE statement - throw semanticException(NOT_SUPPORTED, merge, "Cannot merge into a table with check constraints"); - } - Scope mergeScope = createScope(scope); Scope targetTableScope = analyzer.analyzeForUpdate(relation, Optional.of(mergeScope), UpdateKind.MERGE); Scope sourceTableScope = process(merge.getSource(), mergeScope); Scope joinScope = createAndAssignScope(merge, Optional.of(mergeScope), targetTableScope.getRelationType().joinWith(sourceTableScope.getRelationType())); + analyzeCheckConstraints(table, tableName, targetTableScope, tableSchema.getTableSchema().getCheckConstraints()); + analysis.registerTable(table, redirection.tableHandle(), tableName, session.getIdentity().getUser(), targetTableScope); for (ColumnSchema column : dataColumnSchemas) { if (accessControl.getColumnMask(session.toSecurityContext(), tableName, column.getName(), column.getType()).isPresent()) { 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 index ac422f89bb53..89c760f657ff 100644 --- 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 @@ -683,23 +683,7 @@ public PlanNode plan(Update node) idAllocator.getNextId(), subPlanBuilder.getRoot(), assignments.build())); - - PlanBuilder constraintBuilder = subPlanBuilder.appendProjections(constraints, symbolAllocator, idAllocator); - - List predicates = new ArrayList<>(); - for (Expression constraint : constraints) { - Expression symbol = constraintBuilder.translate(constraint).toSymbolReference(); - - 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), - TRUE_LITERAL, - new Cast(failFunction(plannerContext.getMetadata(), session, CONSTRAINT_VIOLATION, "Check constraint violation: " + constraint), toSqlType(BOOLEAN))); - - predicates.add(predicate); - } - - subPlanBuilder = subPlanBuilder.withNewRoot(new FilterNode(idAllocator.getNextId(), constraintBuilder.getRoot(), and(predicates))); + subPlanBuilder = addCheckConstraints(constraints, subPlanBuilder); } // Build the page, containing: @@ -734,6 +718,26 @@ public PlanNode plan(Update node) return createMergePipeline(table, relationPlan, projectNode, rowIdSymbol, mergeRowSymbol); } + private PlanBuilder addCheckConstraints(List constraints, PlanBuilder subPlanBuilder) + { + PlanBuilder constraintBuilder = subPlanBuilder.appendProjections(constraints, symbolAllocator, idAllocator); + + List predicates = new ArrayList<>(); + for (Expression constraint : constraints) { + Expression symbol = constraintBuilder.translate(constraint).toSymbolReference(); + + 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), + TRUE_LITERAL, + new Cast(failFunction(plannerContext.getMetadata(), session, CONSTRAINT_VIOLATION, "Check constraint violation: " + constraint), toSqlType(BOOLEAN))); + + predicates.add(predicate); + } + + return subPlanBuilder.withNewRoot(new FilterNode(idAllocator.getNextId(), constraintBuilder.getRoot(), and(predicates))); + } + public MergeWriterNode plan(Merge merge) { MergeAnalysis mergeAnalysis = analysis.getMergeAnalysis().orElseThrow(() -> new IllegalArgumentException("analysis.getMergeAnalysis() isn't present")); @@ -773,6 +777,9 @@ public MergeWriterNode plan(Merge merge) PlanBuilder subPlan = newPlanBuilder(joinPlan, analysis, lambdaDeclarationToSymbolMap, session, plannerContext); + FieldReference rowIdReference = analysis.getRowIdField(mergeAnalysis.getTargetTable()); + Symbol rowIdSymbol = planWithPresentColumn.getFieldMappings().get(rowIdReference.getFieldIndex()); + // Build the SearchedCaseExpression that creates the project merge_row Metadata metadata = plannerContext.getMetadata(); List dataColumnSchemas = mergeAnalysis.getDataColumnSchemas(); @@ -790,25 +797,28 @@ public MergeWriterNode plan(Merge merge) } ImmutableList.Builder rowBuilder = ImmutableList.builder(); + Assignments.Builder assignments = Assignments.builder(); List mergeCaseSetColumns = mergeCaseColumnsHandles.get(caseNumber); for (ColumnHandle dataColumnHandle : mergeAnalysis.getDataColumnHandles()) { int index = mergeCaseSetColumns.indexOf(dataColumnHandle); + int fieldNumber = mergeAnalysis.getColumnHandleFieldNumbers().get(dataColumnHandle); + Symbol field = planWithPresentColumn.getFieldMappings().get(fieldNumber); if (index >= 0) { Expression setExpression = mergeCase.getSetExpressions().get(index); subPlan = subqueryPlanner.handleSubqueries(subPlan, setExpression, analysis.getSubqueries(merge)); Expression rewritten = subPlan.rewrite(setExpression); rewritten = coerceIfNecessary(analysis, setExpression, rewritten); if (nonNullableColumnHandles.contains(dataColumnHandle)) { - int fieldIndex = requireNonNull(mergeAnalysis.getColumnHandleFieldNumbers().get(dataColumnHandle), "Could not find fieldIndex for non nullable column"); - ColumnSchema columnSchema = dataColumnSchemas.get(fieldIndex); + ColumnSchema columnSchema = dataColumnSchemas.get(fieldNumber); String columnName = columnSchema.getName(); rewritten = new CoalesceExpression(rewritten, new Cast(failFunction(metadata, session, INVALID_ARGUMENTS, "Assigning NULL to non-null MERGE target table column " + columnName), toSqlType(columnSchema.getType()))); } rowBuilder.add(rewritten); + assignments.put(field, rewritten); } else { - Integer fieldNumber = requireNonNull(mergeAnalysis.getColumnHandleFieldNumbers().get(dataColumnHandle), "Field number for ColumnHandle is null"); - rowBuilder.add(planWithPresentColumn.getFieldMappings().get(fieldNumber).toSymbolReference()); + rowBuilder.add(field.toSymbolReference()); + assignments.putIdentity(field); } } @@ -834,6 +844,19 @@ public MergeWriterNode plan(Merge merge) } whenClauses.add(new WhenClause(condition, new Row(rowBuilder.build()))); + + List constraints = analysis.getCheckConstraints(mergeAnalysis.getTargetTable()); + if (!constraints.isEmpty()) { + assignments.putIdentity(uniqueIdSymbol); + assignments.putIdentity(presentColumn); + assignments.putIdentity(rowIdSymbol); + assignments.putIdentities(source.getFieldMappings()); + subPlan = subPlan.withNewRoot(new ProjectNode( + idAllocator.getNextId(), + subPlan.getRoot(), + assignments.build())); + subPlan = addCheckConstraints(constraints, subPlan.withScope(targetTablePlan.getScope(), targetTablePlan.getFieldMappings())); + } } // Build the "else" clause for the SearchedCaseExpression @@ -848,8 +871,6 @@ public MergeWriterNode plan(Merge merge) SearchedCaseExpression caseExpression = new SearchedCaseExpression(whenClauses.build(), Optional.of(new Row(rowBuilder.build()))); - FieldReference rowIdReference = analysis.getRowIdField(mergeAnalysis.getTargetTable()); - Symbol rowIdSymbol = planWithPresentColumn.getFieldMappings().get(rowIdReference.getFieldIndex()); Symbol mergeRowSymbol = symbolAllocator.newSymbol("merge_row", mergeAnalysis.getMergeRowType()); Symbol caseNumberSymbol = symbolAllocator.newSymbol("case_number", INTEGER); 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 index cd22fa1d4126..0fae78fa340e 100644 --- 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 @@ -187,34 +187,34 @@ public void testInsert() public void testMergeInsert() { // Within allowed check constraint - assertThatThrownBy(() -> assertions.query(""" + assertThat(assertions.query(""" MERGE INTO mock.tiny.nation USING (VALUES 42) t(dummy) ON false WHEN NOT MATCHED THEN INSERT VALUES (101, 'POLAND', 0, 'No comment') """)) - .hasMessage("line 1:1: Cannot merge into a table with check constraints"); + .matches("SELECT BIGINT '1'"); - // Outside allowed check constraint - assertThatThrownBy(() -> assertions.query(""" + assertThat(assertions.query(""" MERGE INTO mock.tiny.nation USING (VALUES 42) t(dummy) ON false - WHEN NOT MATCHED THEN INSERT VALUES (26, 'POLAND', 0, 'No comment') + WHEN NOT MATCHED THEN INSERT (nationkey) VALUES (NULL) """)) - .hasMessage("line 1:1: Cannot merge into a table with check constraints"); - assertThatThrownBy(() -> assertions.query(""" - MERGE INTO mock.tiny.nation USING (VALUES (26, 'POLAND', 0, 'No comment'), (27, 'HOLLAND', 0, 'A comment')) t(a,b,c,d) ON nationkey = a - WHEN NOT MATCHED THEN INSERT VALUES (a,b,c,d) + .matches("SELECT BIGINT '1'"); + assertThat(assertions.query(""" + MERGE INTO mock.tiny.nation USING (VALUES 42) t(dummy) ON false + WHEN NOT MATCHED THEN INSERT (nationkey) VALUES (0) """)) - .hasMessage("line 1:1: Cannot merge into a table with check constraints"); + .matches("SELECT BIGINT '1'"); + // Outside allowed check constraint assertThatThrownBy(() -> assertions.query(""" MERGE INTO mock.tiny.nation USING (VALUES 42) t(dummy) ON false - WHEN NOT MATCHED THEN INSERT (nationkey) VALUES (NULL) + WHEN NOT MATCHED THEN INSERT VALUES (101, 'POLAND', 10, 'No comment') """)) - .hasMessage("line 1:1: Cannot merge into a table with check constraints"); + .hasMessage("Check constraint violation: (regionkey < 10)"); assertThatThrownBy(() -> assertions.query(""" - MERGE INTO mock.tiny.nation USING (VALUES 42) t(dummy) ON false - WHEN NOT MATCHED THEN INSERT (nationkey) VALUES (0) + MERGE INTO mock.tiny.nation USING (VALUES (26, 'POLAND', 10, 'No comment'), (27, 'HOLLAND', 10, 'A comment')) t(a,b,c,d) ON nationkey = a + WHEN NOT MATCHED THEN INSERT VALUES (a,b,c,d) """)) - .hasMessage("line 1:1: Cannot merge into a table with check constraints"); + .hasMessage("Check constraint violation: (regionkey < 10)"); } @Test @@ -319,28 +319,25 @@ public void testDelete() public void testMergeDelete() { // Within allowed check constraint - assertThatThrownBy(() -> assertions.query(""" + assertThat(assertions.query(""" MERGE INTO mock.tiny.nation USING (VALUES 1,2) t(x) ON nationkey = x WHEN MATCHED THEN DELETE """)) - .hasMessage("line 1:1: Cannot merge into a table with check constraints"); + .matches("SELECT BIGINT '2'"); - // Outside allowed check constraint - assertThatThrownBy(() -> assertions.query(""" - MERGE INTO mock.tiny.nation USING (VALUES 1,2,3,4,5) t(x) ON regionkey = x + // Source values outside allowed check constraint should not cause failure + assertThat(assertions.query(""" + MERGE INTO mock.tiny.nation USING (VALUES 1,2,3,4,11) t(x) ON regionkey = x WHEN MATCHED THEN DELETE """)) - .hasMessage("line 1:1: Cannot merge into a table with check constraints"); - assertThatThrownBy(() -> assertions.query(""" + .matches("SELECT BIGINT '20'"); + + // No check constraining column in query + assertThat(assertions.query(""" MERGE INTO mock.tiny.nation USING (VALUES 1,11) t(x) ON nationkey = x WHEN MATCHED THEN DELETE """)) - .hasMessage("line 1:1: Cannot merge into a table with check constraints"); - assertThatThrownBy(() -> assertions.query(""" - MERGE INTO mock.tiny.nation USING (VALUES 11,12,13,14,15) t(x) ON nationkey = x - WHEN MATCHED THEN DELETE - """)) - .hasMessage("line 1:1: Cannot merge into a table with check constraints"); + .matches("SELECT BIGINT '2'"); } /** @@ -475,56 +472,247 @@ public void testUpdateNotDeterministic() public void testMergeUpdate() { // Within allowed check constraint - assertThatThrownBy(() -> assertions.query(""" + assertThat(assertions.query(""" MERGE INTO mock.tiny.nation USING (VALUES 5) t(x) ON nationkey = x WHEN MATCHED THEN UPDATE SET regionkey = regionkey * 2 """)) - .hasMessage("line 1:1: Cannot merge into a table with check constraints"); + .matches("SELECT BIGINT '1'"); - // Outside allowed check constraint + // Merge column within allowed check constraint, but updated rows are outside the check constraint assertThatThrownBy(() -> assertions.query(""" MERGE INTO mock.tiny.nation USING (VALUES 1,2,3,4,5,6) t(x) ON regionkey = x - WHEN MATCHED THEN UPDATE SET regionkey = regionkey * 2 + WHEN MATCHED THEN UPDATE SET regionkey = regionkey * 5 """)) - .hasMessage("line 1:1: Cannot merge into a table with check constraints"); + .hasMessage("Check constraint violation: (regionkey < 10)"); assertThatThrownBy(() -> assertions.query(""" MERGE INTO mock.tiny.nation USING (VALUES 1, 11) t(x) ON nationkey = x - WHEN MATCHED THEN UPDATE SET regionkey = regionkey * 2 + WHEN MATCHED THEN UPDATE SET regionkey = regionkey * 5 """)) - .hasMessage("line 1:1: Cannot merge into a table with check constraints"); + .hasMessage("Check constraint violation: (regionkey < 10)"); + assertThatThrownBy(() -> assertions.query(""" - MERGE INTO mock.tiny.nation USING (VALUES 11) t(x) ON nationkey = x + MERGE INTO mock.tiny.nation t USING mock.tiny.nation s ON t.nationkey = s.nationkey + WHEN MATCHED THEN UPDATE SET regionkey = 10 + """)) + .hasMessage("Check constraint violation: (regionkey < 10)"); + + // Merge column outside allowed check constraint and updated rows within allowed check constraint + assertThat(assertions.query(""" + MERGE INTO mock.tiny.nation USING (VALUES 1, 11) t(x) ON regionkey = x WHEN MATCHED THEN UPDATE SET regionkey = regionkey * 2 """)) - .hasMessage("line 1:1: Cannot merge into a table with check constraints"); + .matches("SELECT BIGINT '5'"); - // Within allowed check constraint, but updated rows are outside the check constraint + // Merge column outside allowed check constraint and updated rows are outside the check constraint assertThatThrownBy(() -> assertions.query(""" - MERGE INTO mock.tiny.nation USING (VALUES 1,2,3) t(x) ON nationkey = x - WHEN MATCHED THEN UPDATE SET nationkey = 10 + MERGE INTO mock.tiny.nation USING (VALUES 11) t(x) ON nationkey = x + WHEN MATCHED THEN UPDATE SET regionkey = regionkey * 5 """)) - .hasMessage("line 1:1: Cannot merge into a table with check constraints"); - assertThatThrownBy(() -> assertions.query(""" + .hasMessage("Check constraint violation: (regionkey < 10)"); + + // No check constraining column in query + assertThat(assertions.query(""" MERGE INTO mock.tiny.nation USING (VALUES 1,2,3) t(x) ON nationkey = x WHEN MATCHED THEN UPDATE SET nationkey = NULL """)) - .hasMessage("line 1:1: Cannot merge into a table with check constraints"); - - // Outside allowed check constraint, but updated rows are outside the check constraint - assertThatThrownBy(() -> assertions.query(""" + .matches("SELECT BIGINT '3'"); + assertThat(assertions.query(""" MERGE INTO mock.tiny.nation USING (VALUES 10) t(x) ON nationkey = x WHEN MATCHED THEN UPDATE SET nationkey = 13 """)) - .hasMessage("line 1:1: Cannot merge into a table with check constraints"); - assertThatThrownBy(() -> assertions.query(""" + .matches("SELECT BIGINT '1'"); + assertThat(assertions.query(""" MERGE INTO mock.tiny.nation USING (VALUES 10) t(x) ON nationkey = x WHEN MATCHED THEN UPDATE SET nationkey = NULL """)) - .hasMessage("line 1:1: Cannot merge into a table with check constraints"); - assertThatThrownBy(() -> assertions.query(""" + .matches("SELECT BIGINT '1'"); + assertThat(assertions.query(""" MERGE INTO mock.tiny.nation USING (VALUES 10) t(x) ON nationkey IS NULL WHEN MATCHED THEN UPDATE SET nationkey = 13 """)) - .hasMessage("line 1:1: Cannot merge into a table with check constraints"); + .matches("SELECT BIGINT '0'"); + } + + @Test + public void testComplexMerge() + { + // Within allowed check constraint + assertThat(assertions.query(""" + MERGE INTO mock.tiny.nation USING (VALUES 1,2,3,4,5,6) t(x) ON regionkey = x + WHEN MATCHED AND t.x = 1 THEN DELETE + WHEN MATCHED THEN UPDATE SET regionkey = 9 + WHEN NOT MATCHED THEN INSERT VALUES (101, 'POLAND', 0, 'No comment') + """)) + .matches("SELECT BIGINT '22'"); + + // Outside allowed check constraint + assertThatThrownBy(() -> assertions.query(""" + MERGE INTO mock.tiny.nation USING (VALUES 1,2,3,4,5,6) t(x) ON regionkey = x + WHEN MATCHED AND t.x = 1 THEN DELETE + WHEN MATCHED THEN UPDATE SET regionkey = 10 + WHEN NOT MATCHED THEN INSERT VALUES (101, 'POLAND', 9, 'No comment') + """)) + .hasMessage("Check constraint violation: (regionkey < 10)"); + + assertThatThrownBy(() -> assertions.query(""" + MERGE INTO mock.tiny.nation USING (VALUES 1,2,3,4,5,6) t(x) ON regionkey = x + WHEN MATCHED AND t.x = 1 THEN DELETE + WHEN MATCHED THEN UPDATE SET regionkey = 9 + WHEN NOT MATCHED THEN INSERT VALUES (101, 'POLAND', 10, 'No comment') + """)) + .hasMessage("Check constraint violation: (regionkey < 10)"); + } + + @Test + public void testMergeCheckMultipleColumns() + { + // Within allowed check constraint + assertThat(assertions.query(""" + MERGE INTO mock.tiny.nation_multiple_column_constraint USING (VALUES 1,2,3,4,5,6) t(x) ON regionkey = x + WHEN MATCHED AND t.x = 1 THEN DELETE + WHEN MATCHED THEN UPDATE SET regionkey = 49 + WHEN NOT MATCHED THEN INSERT VALUES (99, 'POLAND', 49, 'No comment') + """)) + .matches("SELECT BIGINT '22'"); + + // Outside allowed check constraint (regionkey in UPDATE) + assertThatThrownBy(() -> assertions.query(""" + MERGE INTO mock.tiny.nation_multiple_column_constraint USING (VALUES 1,2,3,4,5,6) t(x) ON regionkey = x + WHEN MATCHED AND t.x = 1 THEN DELETE + WHEN MATCHED THEN UPDATE SET regionkey = 50 + WHEN NOT MATCHED THEN INSERT VALUES (99, 'POLAND', 49, 'No comment') + """)) + .hasMessage("Check constraint violation: ((nationkey < 100) AND (regionkey < 50))"); + + // Outside allowed check constraint (regionkey in INSERT) + assertThatThrownBy(() -> assertions.query(""" + MERGE INTO mock.tiny.nation_multiple_column_constraint USING (VALUES 1,2,3,4,5,6) t(x) ON regionkey = x + WHEN MATCHED AND t.x = 1 THEN DELETE + WHEN MATCHED THEN UPDATE SET regionkey = 49 + WHEN NOT MATCHED THEN INSERT VALUES (99, 'POLAND', 50, 'No comment') + """)) + .hasMessage("Check constraint violation: ((nationkey < 100) AND (regionkey < 50))"); + + // Outside allowed check constraint (nationkey in UPDATE) + assertThatThrownBy(() -> assertions.query(""" + MERGE INTO mock.tiny.nation_multiple_column_constraint USING (VALUES 1,2,3,4,5,6) t(x) ON regionkey = x + WHEN MATCHED AND t.x = 1 THEN DELETE + WHEN MATCHED THEN UPDATE SET nationkey = 100 + WHEN NOT MATCHED THEN INSERT VALUES (99, 'POLAND', 49, 'No comment') + """)) + .hasMessage("Check constraint violation: ((nationkey < 100) AND (regionkey < 50))"); + + // Outside allowed check constraint (nationkey in INSERT) + assertThatThrownBy(() -> assertions.query(""" + MERGE INTO mock.tiny.nation_multiple_column_constraint USING (VALUES 1,2,3,4,5,6) t(x) ON regionkey = x + WHEN MATCHED AND t.x = 1 THEN DELETE + WHEN MATCHED THEN UPDATE SET nationkey = 99 + WHEN NOT MATCHED THEN INSERT VALUES (100, 'POLAND', 50, 'No comment') + """)) + .hasMessage("Check constraint violation: ((nationkey < 100) AND (regionkey < 50))"); + } + + @Test + public void testMergeSubquery() + { + // TODO https://github.com/trinodb/trino/issues/18230 Support subqueries for MERGE statement in check constraint + assertThatThrownBy(() -> assertions.query(""" + MERGE INTO mock.tiny.nation_subquery USING (VALUES 1,2,3,4,5,6) t(x) ON regionkey = x + WHEN MATCHED AND t.x = 1 THEN DELETE + WHEN MATCHED THEN UPDATE SET regionkey = 9 + WHEN NOT MATCHED THEN INSERT VALUES (101, 'POLAND', 0, 'No comment') + """)) + .hasMessageContaining("Unexpected subquery expression in logical plan"); + } + + @Test + public void testMergeUnsupportedCurrentDate() + { + assertThatThrownBy(() -> assertions.query(""" + MERGE INTO mock.tiny.nation_current_date USING (VALUES 1,2,3,4,5,6) t(x) ON regionkey = x + WHEN MATCHED AND t.x = 1 THEN DELETE + WHEN MATCHED THEN UPDATE SET regionkey = 9 + WHEN NOT MATCHED THEN INSERT VALUES (101, 'POLAND', 0, 'No comment') + """)) + .hasMessageContaining("Check constraint expression should not contain temporal expression"); + } + + @Test + public void testMergeUnsupportedCurrentTime() + { + assertThatThrownBy(() -> assertions.query(""" + MERGE INTO mock.tiny.nation_current_time USING (VALUES 1,2,3,4,5,6) t(x) ON regionkey = x + WHEN MATCHED AND t.x = 1 THEN DELETE + WHEN MATCHED THEN UPDATE SET regionkey = 9 + WHEN NOT MATCHED THEN INSERT VALUES (101, 'POLAND', 0, 'No comment') + """)) + .hasMessageContaining("Check constraint expression should not contain temporal expression"); + } + + @Test + public void testMergeUnsupportedCurrentTimestamp() + { + assertThatThrownBy(() -> assertions.query(""" + MERGE INTO mock.tiny.nation_current_timestamp USING (VALUES 1,2,3,4,5,6) t(x) ON regionkey = x + WHEN MATCHED AND t.x = 1 THEN DELETE + WHEN MATCHED THEN UPDATE SET regionkey = 9 + WHEN NOT MATCHED THEN INSERT VALUES (101, 'POLAND', 0, 'No comment') + """)) + .hasMessageContaining("Check constraint expression should not contain temporal expression"); + } + + @Test + public void testMergeUnsupportedLocaltime() + { + assertThatThrownBy(() -> assertions.query(""" + MERGE INTO mock.tiny.nation_localtime USING (VALUES 1,2,3,4,5,6) t(x) ON regionkey = x + WHEN MATCHED AND t.x = 1 THEN DELETE + WHEN MATCHED THEN UPDATE SET regionkey = 9 + WHEN NOT MATCHED THEN INSERT VALUES (101, 'POLAND', 0, 'No comment') + """)) + .hasMessageContaining("Check constraint expression should not contain temporal expression"); + } + + @Test + public void testMergeUnsupportedLocaltimestamp() + { + assertThatThrownBy(() -> assertions.query(""" + MERGE INTO mock.tiny.nation_localtimestamp USING (VALUES 1,2,3,4,5,6) t(x) ON regionkey = x + WHEN MATCHED AND t.x = 1 THEN DELETE + WHEN MATCHED THEN UPDATE SET regionkey = 9 + WHEN NOT MATCHED THEN INSERT VALUES (101, 'POLAND', 0, 'No comment') + """)) + .hasMessageContaining("Check constraint expression should not contain temporal expression"); + } + + @Test + public void testMergeUnsupportedConstraint() + { + assertThatThrownBy(() -> assertions.query(""" + MERGE INTO mock.tiny.nation_invalid_function USING (VALUES 1,2,3,4,5,6) t(x) ON regionkey = x + WHEN MATCHED AND t.x = 1 THEN DELETE + WHEN MATCHED THEN UPDATE SET regionkey = 9 + WHEN NOT MATCHED THEN INSERT VALUES (101, 'POLAND', 0, 'No comment') + """)) + .hasMessageContaining("Function 'invalid_function' not registered"); + assertThatThrownBy(() -> assertions.query(""" + MERGE INTO mock.tiny.nation_not_boolean_expression USING (VALUES 1,2,3,4,5,6) t(x) ON regionkey = x + WHEN MATCHED AND t.x = 1 THEN DELETE + WHEN MATCHED THEN UPDATE SET regionkey = 9 + WHEN NOT MATCHED THEN INSERT VALUES (101, 'POLAND', 0, 'No comment') + """)) + .hasMessageContaining("to be of type BOOLEAN, but was integer"); + } + + @Test + public void testMergeNotDeterministic() + { + assertThatThrownBy(() -> assertions.query(""" + MERGE INTO mock.tiny.nation_not_deterministic USING (VALUES 1,2,3,4,5,6) t(x) ON regionkey = x + WHEN MATCHED AND t.x = 1 THEN DELETE + WHEN MATCHED THEN UPDATE SET regionkey = 9 + WHEN NOT MATCHED THEN INSERT VALUES (101, 'POLAND', 0, 'No comment') + """)) + .hasMessageContaining("Check constraint expression should be deterministic"); } } diff --git a/testing/trino-product-tests/src/main/java/io/trino/tests/product/deltalake/TestDeltaLakeCheckConstraintCompatibility.java b/testing/trino-product-tests/src/main/java/io/trino/tests/product/deltalake/TestDeltaLakeCheckConstraintCompatibility.java index edd6fab5ce84..58a09e4016f2 100644 --- a/testing/trino-product-tests/src/main/java/io/trino/tests/product/deltalake/TestDeltaLakeCheckConstraintCompatibility.java +++ b/testing/trino-product-tests/src/main/java/io/trino/tests/product/deltalake/TestDeltaLakeCheckConstraintCompatibility.java @@ -137,6 +137,69 @@ public void testCheckConstraintUpdateCompatibility() } } + // TODO: Add DELTA_LAKE_DATABRICKS and DELTA_LAKE_EXCLUDE_73 groups once flakiness on Databricks is resolved + @Test(groups = {DELTA_LAKE_OSS, PROFILE_SPECIFIC_TESTS}) + public void testCheckConstraintMergeCompatibility() + { + String tableName1 = "test_check_constraint_merge_" + randomNameSuffix(); + String tableName2 = "test_check_constraint_merge_" + randomNameSuffix(); + + onDelta().executeQuery("CREATE TABLE default." + tableName1 + + "(a INT) " + + "USING DELTA " + + "LOCATION 's3://" + bucketName + "/databricks-compatibility-test-" + tableName1 + "'"); + + onDelta().executeQuery("CREATE TABLE default." + tableName2 + + "(a INT) " + + "USING DELTA " + + "LOCATION 's3://" + bucketName + "/databricks-compatibility-test-" + tableName2 + "'"); + + try { + onDelta().executeQuery("ALTER TABLE default." + tableName1 + " ADD CONSTRAINT a_constraint CHECK (a < 3)"); + + onTrino().executeQuery("INSERT INTO delta.default." + tableName1 + " VALUES (1)"); + onTrino().executeQuery("INSERT INTO delta.default." + tableName2 + " VALUES (1), (2)"); + + assertThat(onTrino().executeQuery("SELECT * FROM delta.default." + tableName1)) + .containsOnly(row(1)); + + onTrino().executeQuery("MERGE INTO delta.default." + tableName1 + " t USING delta.default." + tableName2 + " s " + + "ON (t.a = s.a) " + + "WHEN MATCHED " + + "THEN UPDATE SET a = 2 " + + "WHEN NOT MATCHED " + + "THEN INSERT (a) VALUES (2)"); + + assertThat(onTrino().executeQuery("SELECT * FROM delta.default." + tableName1)) + .containsOnly( + row(2), + row(2)); + + assertThatThrownBy(() -> onDelta().executeQuery("MERGE INTO default." + tableName1 + " t USING default." + tableName2 + " s " + + "ON (t.a = s.a) " + + "WHEN MATCHED " + + "THEN UPDATE SET a = 3 " + + "WHEN NOT MATCHED " + + "THEN INSERT (a) VALUES (3)")) + .hasMessageMatching("(?s).* CHECK constraint .* violated by row with values.*"); + assertThatThrownBy(() -> onTrino().executeQuery("MERGE INTO delta.default." + tableName1 + " t USING delta.default." + tableName2 + " s " + + "ON (t.a = s.a) " + + "WHEN MATCHED " + + "THEN UPDATE SET a = 3 " + + "WHEN NOT MATCHED " + + "THEN INSERT (a) VALUES (3)")) + .hasMessageContaining("Check constraint violation"); + assertThat(onTrino().executeQuery("SELECT * FROM delta.default." + tableName1)) + .containsOnly( + row(2), + row(2)); + } + finally { + dropDeltaTableWithRetry("default." + tableName1); + dropDeltaTableWithRetry("default." + tableName2); + } + } + @Test(groups = {DELTA_LAKE_DATABRICKS, DELTA_LAKE_OSS, DELTA_LAKE_EXCLUDE_73, PROFILE_SPECIFIC_TESTS}) @Flaky(issue = DATABRICKS_COMMUNICATION_FAILURE_ISSUE, match = DATABRICKS_COMMUNICATION_FAILURE_MATCH) public void testCheckConstraintUnknownCondition() @@ -192,26 +255,6 @@ public void testCheckConstraintAcrossColumns() } } - @Test(groups = {DELTA_LAKE_OSS, DELTA_LAKE_DATABRICKS, DELTA_LAKE_EXCLUDE_73, PROFILE_SPECIFIC_TESTS}) - @Flaky(issue = DATABRICKS_COMMUNICATION_FAILURE_ISSUE, match = DATABRICKS_COMMUNICATION_FAILURE_MATCH) - public void testWritesToTableWithCheckConstraintFails() - { - String tableName = "test_check_constraint_unsupported_operatins_" + randomNameSuffix(); - - onDelta().executeQuery("CREATE TABLE default." + tableName + " (a INT, b INT) " + - "USING DELTA " + - "LOCATION 's3://" + bucketName + "/databricks-compatibility-test-" + tableName + "'"); - onDelta().executeQuery("ALTER TABLE default." + tableName + " ADD CONSTRAINT aIsPositive CHECK (a > 0)"); - try { - assertQueryFailure(() -> onTrino().executeQuery("MERGE INTO delta.default." + tableName + " t USING delta.default." + tableName + " s " + - "ON (t.a = s.a) WHEN MATCHED THEN UPDATE SET b = 42")) - .hasMessageContaining("Cannot merge into a table with check constraints"); - } - finally { - dropDeltaTableWithRetry("default." + tableName); - } - } - @Test(groups = {DELTA_LAKE_OSS, DELTA_LAKE_DATABRICKS, DELTA_LAKE_EXCLUDE_73, PROFILE_SPECIFIC_TESTS}) @Flaky(issue = DATABRICKS_COMMUNICATION_FAILURE_ISSUE, match = DATABRICKS_COMMUNICATION_FAILURE_MATCH) public void testMetadataOperationsRetainCheckConstraint() @@ -252,7 +295,7 @@ public void testUnsupportedCheckConstraintExpression() onTrino().executeQuery("DELETE FROM delta.default." + tableName + " WHERE a = 2"); assertQueryFailure(() -> onTrino().executeQuery("MERGE INTO delta.default." + tableName + " t USING delta.default." + tableName + " s " + "ON (t.a = s.a) WHEN MATCHED THEN UPDATE SET b = -1")) - .hasMessageContaining("Cannot merge into a table with check constraints"); + .hasMessageContaining("Failed to convert Delta check constraints to Trino expression"); // Verify these operations succeed even if check constraints have unsupported expressions onTrino().executeQuery("ALTER TABLE delta.default." + tableName + " ADD COLUMN c INT");