Fix failure when setting NULL in UPDATE statement in JDBC-based connectors

plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/DefaultQueryBuilder.java

@@ -288,25 +288,31 @@ public PreparedStatement prepareStatement(
             WriteFunction writeFunction = parameter.getJdbcType()
                     .map(jdbcType -> getWriteFunction(client, session, connection, jdbcType, parameter.getType()))
                     .orElseGet(() -> getWriteFunction(client, session, parameter.getType()));
-            Class<?> javaType = writeFunction.getJavaType();
-            Object value = parameter.getValue()
-                    // The value must be present, since DefaultQueryBuilder never creates null parameters. Values coming from Domain's ValueSet are non-null, and
-                    // nullable domains are handled explicitly, with SQL syntax.
-                    .orElseThrow(() -> new VerifyException("Value is missing"));
-            if (javaType == boolean.class) {
-                ((BooleanWriteFunction) writeFunction).set(statement, parameterIndex, (boolean) value);
-            }
-            else if (javaType == long.class) {
-                ((LongWriteFunction) writeFunction).set(statement, parameterIndex, (long) value);
-            }
-            else if (javaType == double.class) {
-                ((DoubleWriteFunction) writeFunction).set(statement, parameterIndex, (double) value);
-            }
-            else if (javaType == Slice.class) {
-                ((SliceWriteFunction) writeFunction).set(statement, parameterIndex, (Slice) value);
+
+            if (parameter.getValue().isEmpty()) {
+                writeFunction.setNull(statement, parameterIndex);
             }
             else {
-                ((ObjectWriteFunction) writeFunction).set(statement, parameterIndex, value);
+                Class<?> javaType = writeFunction.getJavaType();
+                Object value = parameter.getValue()
+                        // The value must be present, since DefaultQueryBuilder never creates null parameters. Values coming from Domain's ValueSet are non-null, and
+                        // nullable domains are handled explicitly, with SQL syntax.
+                        .orElseThrow(() -> new VerifyException("Value is missing"));
+                if (javaType == boolean.class) {
+                    ((BooleanWriteFunction) writeFunction).set(statement, parameterIndex, (boolean) value);
+                }
+                else if (javaType == long.class) {
+                    ((LongWriteFunction) writeFunction).set(statement, parameterIndex, (long) value);
+                }
+                else if (javaType == double.class) {
+                    ((DoubleWriteFunction) writeFunction).set(statement, parameterIndex, (double) value);
+                }
+                else if (javaType == Slice.class) {
+                    ((SliceWriteFunction) writeFunction).set(statement, parameterIndex, (Slice) value);
+                }
+                else {
+                    ((ObjectWriteFunction) writeFunction).set(statement, parameterIndex, value);
+                }
             }
         }
 

plugin/trino-hive/src/test/java/io/trino/plugin/hive/BaseHiveConnectorTest.java

@@ -351,6 +351,14 @@ public void testUpdateMultipleCondition()
                 .hasMessage(MODIFYING_NON_TRANSACTIONAL_TABLE_MESSAGE);
     }
 
+    @Test
+    @Override
+    public void testUpdateWithNullValues()
+    {
+        assertThatThrownBy(super::testUpdateWithNullValues)
+                .hasMessage(MODIFYING_NON_TRANSACTIONAL_TABLE_MESSAGE);
+    }
+
     @Test
     @Override
     public void testRowLevelUpdate()

plugin/trino-kudu/src/test/java/io/trino/plugin/kudu/TestKuduConnectorTest.java

@@ -978,6 +978,30 @@ public void testUpdate()
     @Override
     public void testUpdateMultipleCondition() {}
 
+    /**
+     * The test is overridden because Kudu requires nullable columns to be explicitly specified in the creation statement using `WITH (nullable=true)`.
+     * Additionally, the first column will be the primary key in the table, override to use `regionkey` to test.
+     */
+    @Test
+    @Override
+    public void testUpdateWithNullValues()
+    {
+        withTableName("test_update_nulls", tableName -> {
+            assertUpdate(createKuduTableForWrites("CREATE TABLE %s (nationkey bigint, name varchar(25), regionkey bigint WITH (nullable=true), comment varchar(152))".formatted(tableName)));
+            assertUpdate("INSERT INTO " + tableName + " SELECT * FROM nation", 25);
+
+            assertQuery("SELECT count(*) FROM " + tableName + " WHERE regionkey IS NULL", "VALUES 0");
+            assertUpdate("UPDATE " + tableName + " SET regionkey = NULL WHERE nationkey > 20", 4);
+            assertQuery("SELECT count(*) FROM " + tableName + " WHERE regionkey IS NULL", "VALUES 4");
+
+            // Kudu connector does not have ConnectorCapabilities with `SUPPORTS_NOT_NULL_CONSTRAINT`, but column definition not null by default
+            // Here verify set null to the not null column will fail in Kudu
+            assertThatThrownBy(() -> getQueryRunner().execute("UPDATE " + tableName + " SET nationkey = NULL WHERE nationkey > 20"))
+                    .isInstanceOf(RuntimeException.class)
+                    .hasMessage("nationkey cannot be set to null");
+        });
+    }
+
     /**
      * This test fails intermittently because Kudu doesn't have strong enough
      * semantics to support writing from multiple threads.

testing/trino-testing/src/main/java/io/trino/testing/BaseConnectorTest.java

@@ -4999,6 +4999,40 @@ public void testUpdateMultipleCondition()
         }
     }
 
+    @Test
+    public void testUpdateWithNullValues()
+    {
+        skipTestUnless(hasBehavior(SUPPORTS_UPDATE));
+
+        try (TestTable table = new TestTable(getQueryRunner()::execute, "test_update_nulls", "AS SELECT * FROM nation")) {
+            String tableName = table.getName();
+
+            assertQuery("SELECT count(*) FROM " + tableName + " WHERE nationkey IS NULL", "VALUES 0");
+            assertUpdate("UPDATE " + tableName + " SET nationkey = NULL WHERE regionkey = 2", 5);
+            assertQuery("SELECT count(*) FROM " + tableName + " WHERE nationkey IS NULL", "VALUES 5");
+        }
+
+        if (!hasBehavior(SUPPORTS_NOT_NULL_CONSTRAINT)) {
+            return;
+        }
+
+        try (TestTable table = createTestTableForWrites("test_update_nulls", "(nullable_col INTEGER, not_null_col INTEGER NOT NULL, pk INT)", "pk")) {
+            String tableName = table.getName();
+            assertUpdate("INSERT INTO " + tableName + " VALUES (1, 1, 1), (2, 2, 2)", 2);
+            assertQuery("SELECT * FROM " + tableName, "VALUES (1, 1, 1), (2, 2, 2)");
+
+            assertUpdate("UPDATE " + tableName + " SET nullable_col = null WHERE not_null_col = 1", 1);
+            assertQuery("SELECT * FROM " + tableName, "VALUES (null, 1, 1), (2, 2, 2)");
+
+            if (hasBehavior(SUPPORTS_ROW_LEVEL_UPDATE)) {
+                // Covered by testUpdateNotNullColumn
+                return;
+            }
+
+            assertQueryFails("UPDATE " + tableName + " SET not_null_col = TRY(1 / 0) WHERE not_null_col = 2", MODIFYING_ROWS_MESSAGE);
+        }
+    }
+
     @Test
     public void testRowLevelUpdate()
     {