fix: Fill in uuid and auth columns before building mutation message keys

agent/skills/implement-sqrl/SKILL.md

@@ -141,6 +141,7 @@ The internal table does not need `WATERMARK` declaration. SQRL generates the wat
 Declare a `WATERMARK` explicitely on an internal table only when the user asks for one.
 
 With `engine(kafka)`, `PARTITIONED BY (col, ...)` on an internal table without a `PRIMARY KEY` keys each record by those columns, both the records the server writes for mutations and the records the pipeline exports into the table, so records with the same key stay on one Kafka partition in order.
+The key columns can be `uuid` or `auth` metadata columns, but not the `timestamp` metadata column.
 The Flink source ignores that key and reads every column from the message value, so adding `PARTITIONED BY` to an existing internal table is safe: records written before the table had a key are still read, also when the topic is read again from the earliest offset.
 This relies on the default `kafka-safe` connector. Do not switch internal Kafka tables to the plain `kafka` connector, which drops records that have no key.
 Use `PARTITIONED BY` when records only need a Kafka key for partitioning and ordering. A `PRIMARY KEY` also keys the records, but it makes the table an upsert table (`VERSIONED_STATE`), and its upsert source drops records that were written without a key.

documentation/docs/sqrl-language.md

@@ -113,7 +113,7 @@ Internal tables support `kafka` and `iceberg` engines.
 | Watermark spec             | **generated**                                                   | **required**                 |
 | Primary key                | *Unenforced* upsert semantics                                   | Same as Flink                |
 
-With the `kafka` engine, `PARTITIONED BY (col, ...)` on an internal table without a primary key sets the Kafka message key of each record, both for the records that the server writes for a mutation and for the records that the pipeline exports into the table. Records with the same key go to the same Kafka partition, so their order is kept. The clause only declares the key; the table is not partitioned in Flink.
+With the `kafka` engine, `PARTITIONED BY (col, ...)` on an internal table without a primary key sets the Kafka message key of each record, both for the records that the server writes for a mutation and for the records that the pipeline exports into the table. Records with the same key go to the same Kafka partition, so their order is kept. The clause only declares the key; the table is not partitioned in Flink. The key can contain `METADATA FROM 'uuid'` and `METADATA FROM 'auth.<claim>'` columns, which the server fills in before it writes the record, but not the `METADATA FROM 'timestamp'` column, which is only known after the record is written.
 The Flink source of the table reads every column from the message value and ignores the key. Adding `PARTITIONED BY` to an existing table therefore keeps the records written before the table had a key, also when the pipeline reads the topic again from the earliest offset. This relies on the default `kafka-safe` connector; the plain `kafka` connector drops records that have no key.
 A table with a primary key is keyed by its primary key instead, and its upsert source drops records that were written without a key.
 

sqrl-planner/src/main/java/com/datasqrl/planner/Sqrl2FlinkSQLTranslator.java

@@ -121,6 +121,7 @@
 import org.apache.flink.table.api.bridge.java.internal.StreamTableEnvironmentImpl;
 import org.apache.flink.table.api.internal.TableResultInternal;
 import org.apache.flink.table.catalog.CatalogManager;
+import org.apache.flink.table.catalog.Column;
 import org.apache.flink.table.catalog.Column.ComputedColumn;
 import org.apache.flink.table.catalog.Column.MetadataColumn;
 import org.apache.flink.table.catalog.Column.PhysicalColumn;
@@ -953,6 +954,7 @@ Note, that this requires we replace the table with the actual table (and the cor
             .orElse(PrimaryKeyMap.UNDEFINED);
     // Finish building mutation query by building input and output types from table schema
     if (mutationBld != null) {
+      validateMutationPartitionKey(flinkSchema, mutationBld.build().getPartitionKey());
       var fields = convertSchema2RelDataType(flinkSchema);
       var inputType = CalciteUtil.getRelTypeBuilder(typeFactory);
       var outputType = CalciteUtil.getRelTypeBuilder(typeFactory);
@@ -1175,4 +1177,28 @@ private static void validateMutationHints(
               + " Supported engines: \"kafka\", \"iceberg\".");
     }
   }
+
+  /**
+   * Partition columns must be physical, which includes the 'uuid' and 'auth' metadata columns the
+   * server fills in. Flink does not check this for engines that remove the partition from the
+   * table, e.g. Kafka uses it as the message key, which the server builds before the record is
+   * written.
+   */
+  private static void validateMutationPartitionKey(
+      ResolvedSchema schema, List<String> partitionKey) {
+
+    for (var colName : partitionKey) {
+      if (!schema.getColumn(colName).map(Column::isPhysical).orElse(false)) {
+        var availableCols =
+            schema.getColumns().stream()
+                .filter(Column::isPhysical)
+                .map(Column::getName)
+                .collect(Collectors.joining(", "));
+
+        throw new NoLocationStatementParserException(
+            "Invalid PARTITIONED BY column [%s]: partition columns must be regular columns or 'uuid' or 'auth' metadata columns. Available columns: [%s]"
+                .formatted(colName, availableCols));
+      }
+    }
+  }
 }

sqrl-server/sqrl-server-vertx-base/src/main/java/com/datasqrl/server/MutationConfigurationImpl.java

@@ -23,6 +23,7 @@
 import com.datasqrl.server.io.SinkProducer;
 import com.datasqrl.server.kafka.KafkaSinkProducer;
 import com.datasqrl.server.kafka.SaslRetryingSinkProducer;
+import com.google.common.annotations.VisibleForTesting;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
 import io.vertx.core.CompositeFuture;
@@ -122,7 +123,8 @@ interface ComputeInputColumns {
     Object compute(DataFetchingEnvironment env);
   }
 
-  private List<SinkProducer.Record> getRecords(
+  @VisibleForTesting
+  static List<SinkProducer.Record> getRecords(
       DataFetchingEnvironment env,
       List<String> keyColumns,
       Map<String, ComputeInputColumns> computedColumns) {
@@ -140,20 +142,6 @@ private List<SinkProducer.Record> getRecords(
       entries = List.of((Map) argument);
     }
 
-    // Construct keys if necessary (null if no key cols)
-    var keys = new Map[entries.size()];
-    if (!keyColumns.isEmpty()) {
-      for (int i = 0; i < entries.size(); i++) {
-        var filtered = new LinkedHashMap();
-        for (var key : keyColumns) {
-          if (entries.get(i).containsKey(key)) {
-            filtered.put(key, entries.get(i).get(key));
-          }
-        }
-        keys[i] = filtered;
-      }
-    }
-
     if (!computedColumns.isEmpty()) {
       // Add UUID for event for the computed uuid columns
       entries.forEach(
@@ -166,14 +154,28 @@ private List<SinkProducer.Record> getRecords(
           });
     }
 
+    // Keys are built after the computed columns are added, since those can be key columns too
     var records = new ArrayList<SinkProducer.Record>();
-    for (int i = 0; i < entries.size(); i++) {
-      records.add(new SinkProducer.Record(keys[i], entries.get(i)));
+    for (var entry : entries) {
+      records.add(new SinkProducer.Record(createKey(entry, keyColumns), entry));
     }
 
     return records;
   }
 
+  /**
+   * Keeps every key column, also when the input omits it, so that the key matches the one Flink
+   * writes for the same record (null if no key cols)
+   */
+  private static Map createKey(Map entry, List<String> keyColumns) {
+    if (keyColumns.isEmpty()) {
+      return null;
+    }
+    var key = new LinkedHashMap();
+    keyColumns.forEach(col -> key.put(col, entry.get(col)));
+    return key;
+  }
+
   private void completeWithResults(
       CompletableFuture<Object> cf, List<Object> results, boolean returnList) {
     if (returnList) {

sqrl-server/sqrl-server-vertx-base/src/test/java/com/datasqrl/server/MutationConfigurationImplTest.java

@@ -0,0 +1,78 @@
+/*
+ * Copyright © 2021 DataSQRL (contact@datasqrl.com)
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.datasqrl.server;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.assertj.core.api.Assertions.assertThat;
+
+import com.datasqrl.server.MutationConfigurationImpl.ComputeInputColumns;
+import com.datasqrl.server.io.SinkProducer;
+import com.datasqrl.server.kafka.JsonSerializer;
+import graphql.schema.DataFetchingEnvironment;
+import graphql.schema.DataFetchingEnvironmentImpl;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.junit.jupiter.api.Test;
+
+class MutationConfigurationImplTest {
+
+  private static final JsonSerializer KEY_SERIALIZER = new JsonSerializer();
+
+  @Test
+  void givenComputedKeyColumn_whenGettingRecords_thenKeyContainsComputedValue() {
+    var env = mutationOf(new HashMap<>(Map.of("sensorid", 1)));
+    Map<String, ComputeInputColumns> computedColumns = Map.of("tenantId", e -> "tenant-1");
+
+    var records =
+        MutationConfigurationImpl.getRecords(env, List.of("sensorid", "tenantId"), computedColumns);
+
+    assertThat(serializedKeys(records))
+        .containsExactly("{\"sensorid\":1,\"tenantId\":\"tenant-1\"}");
+  }
+
+  @Test
+  void givenOmittedAndNullKeyColumn_whenGettingRecords_thenBothKeysContainNull() {
+    var explicitNull = new HashMap<String, Object>();
+    explicitNull.put("zone", null);
+    var env = mutationOf(List.of(new HashMap<>(), explicitNull));
+
+    var records = MutationConfigurationImpl.getRecords(env, List.of("zone"), Map.of());
+
+    assertThat(serializedKeys(records)).containsExactly("{\"zone\":null}", "{\"zone\":null}");
+  }
+
+  @Test
+  void givenNoKeyColumns_whenGettingRecords_thenRecordHasNoKey() {
+    var env = mutationOf(new HashMap<>(Map.of("sensorid", 1)));
+
+    var records = MutationConfigurationImpl.getRecords(env, List.of(), Map.of());
+
+    assertThat(records).singleElement().extracting(SinkProducer.Record::key).isNull();
+  }
+
+  private static DataFetchingEnvironment mutationOf(Object input) {
+    return DataFetchingEnvironmentImpl.newDataFetchingEnvironment()
+        .arguments(Map.of("event", input))
+        .build();
+  }
+
+  private static List<String> serializedKeys(List<SinkProducer.Record> records) {
+    return records.stream()
+        .map(rec -> new String(KEY_SERIALIZER.serialize("topic", rec.key()), UTF_8))
+        .toList();
+  }
+}

sqrl-testing/sqrl-testing-integration/src/test/resources/dagplanner/partitionedKafkaMutationTimestampKey-fail.sqrl

@@ -0,0 +1,6 @@
+/*+engine(kafka) */
+CREATE TABLE Events (
+    tenantId STRING NOT NULL METADATA FROM 'auth.tenant_id',
+    payload STRING NOT NULL,
+    event_time TIMESTAMP_LTZ(3) NOT NULL METADATA FROM 'timestamp'
+) PARTITIONED BY (tenantId, event_time);

sqrl-testing/sqrl-testing-integration/src/test/resources/snapshots/com/datasqrl/DAGPlannerTest/partitionedKafkaMutationTimestampKey-fail.txt

@@ -0,0 +1,6 @@
+[FATAL] Invalid PARTITIONED BY column [event_time]: partition columns must be regular columns or 'uuid' or 'auth' metadata columns. Available columns: [tenantId, payload]
+in script:partitionedKafkaMutationTimestampKey-fail.sqrl [2:1]:
+/*+engine(kafka) */
+CREATE TABLE Events (
+^
+