source: Fix key-value load gen w/o primary export

This commit fixes the key-value load generator source when the primary export is removed via the force_source_table_syntax flag. Fixes #MaterializeInc/database-issues/issues/8904
MaterializeInc · Jan 22, 2025 · c4d2649 · c4d2649
1 parent 411e6c9
commit c4d2649
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 42 deletions.
diff --git a/src/storage/src/source/generator.rs b/src/storage/src/source/generator.rs
@@ -148,6 +148,10 @@ impl GeneratorKind {
     ) {
         // figure out which output types from the generator belong to which output indexes
         let mut output_map = BTreeMap::new();
+        // Make sure that there's an entry for the default output, even if there are no exports
+        // that need data output. Certain implementations rely on it (at the time of this comment
+        // that includes the key-value load gen source).
+        output_map.insert(LoadGeneratorOutput::Default, Vec::new());
         for (idx, (_, export)) in config.source_exports.iter().enumerate() {
             let output_type = match &export.details {
                 SourceExportDetails::LoadGenerator(details) => details.output,

diff --git a/test/testdrive/force-source-tables.td b/test/testdrive/force-source-tables.td
@@ -464,49 +464,47 @@ avro_table_upsert ""
 # Key-value load generator source using source-fed tables
 #
 
-# TODO(database-issues#8904): Re-enable when key-value load gen sources work with force_source_table_syntax.
+$ postgres-execute connection=postgres://mz_system:materialize@${testdrive.materialize-internal-sql-addr}
+ALTER SYSTEM SET enable_create_table_from_source = true
+ALTER SYSTEM SET enable_load_generator_key_value = true
+ALTER SYSTEM SET force_source_table_syntax = true
 
-# $ postgres-execute connection=postgres://mz_system:materialize@${testdrive.materialize-internal-sql-addr}
-# ALTER SYSTEM SET enable_create_table_from_source = true
-# ALTER SYSTEM SET enable_load_generator_key_value = true
-# ALTER SYSTEM SET force_source_table_syntax = true
-#
-# > CREATE SOURCE keyvalue
-#   IN CLUSTER ${arg.single-replica-cluster}
-#   FROM LOAD GENERATOR KEY VALUE (
-#     KEYS 16,
-#     PARTITIONS 4,
-#     SNAPSHOT ROUNDS 3,
-#     SEED 123,
-#     VALUE SIZE 10,
-#     BATCH SIZE 2,
-#     TICK INTERVAL '1s'
-#   );
-#
-# > CREATE TABLE kv_1 FROM SOURCE keyvalue INCLUDE KEY ENVELOPE UPSERT;
-#
-# > CREATE TABLE kv_2 FROM SOURCE keyvalue INCLUDE KEY ENVELOPE NONE;
-#
-# > SELECT partition, count(*) FROM kv_1 GROUP BY partition
-# 0   4
-# 1   4
-# 2   4
-# 3   4
-#
-# > SELECT status FROM mz_internal.mz_source_statuses WHERE name = 'kv_1';
-# running
-#
-# > SELECT partition, count(*) > 10 FROM kv_2 GROUP BY partition
-# 0   true
-# 1   true
-# 2   true
-# 3   true
-#
-# > SHOW TABLES ON keyvalue;
-# kv_1 ""
-# kv_2 ""
-#
-# > DROP SOURCE keyvalue CASCADE;
+> CREATE SOURCE keyvalue
+  IN CLUSTER ${arg.single-replica-cluster}
+  FROM LOAD GENERATOR KEY VALUE (
+    KEYS 16,
+    PARTITIONS 4,
+    SNAPSHOT ROUNDS 3,
+    SEED 123,
+    VALUE SIZE 10,
+    BATCH SIZE 2,
+    TICK INTERVAL '1s'
+  );
+
+> CREATE TABLE kv_1 FROM SOURCE keyvalue INCLUDE KEY ENVELOPE UPSERT;
+
+> CREATE TABLE kv_2 FROM SOURCE keyvalue INCLUDE KEY ENVELOPE NONE;
+
+> SELECT partition, count(*) FROM kv_1 GROUP BY partition
+0   4
+1   4
+2   4
+3   4
+
+> SELECT status FROM mz_internal.mz_source_statuses WHERE name = 'kv_1';
+running
+
+> SELECT partition, count(*) > 10 FROM kv_2 GROUP BY partition
+0   true
+1   true
+2   true
+3   true
+
+> SHOW TABLES ON keyvalue;
+kv_1 ""
+kv_2 ""
+
+> DROP SOURCE keyvalue CASCADE;
 
 #
 # Force usage of the new syntax and check that old statements are disallowed