From e3ec4f7ea25c6b331a4fec38f65b9309ac2bb139 Mon Sep 17 00:00:00 2001
From: Pekka Enberg <penberg@iki.fi>
Date: Mon, 8 Sep 2025 11:11:57 +0300
Subject: [PATCH] sql_generation: Add support for predefined columns

---
 sql_generation/generation/table.rs | 31 ++++++++++++++++++++++++------
 1 file changed, 25 insertions(+), 6 deletions(-)

diff --git a/sql_generation/generation/table.rs b/sql_generation/generation/table.rs
index 21c89a179..e8b722180 100644
--- a/sql_generation/generation/table.rs
+++ b/sql_generation/generation/table.rs
@@ -16,21 +16,33 @@ impl Arbitrary for Name {
     }
 }
 
-impl Arbitrary for Table {
-    fn arbitrary<R: Rng + ?Sized, C: GenerationContext>(rng: &mut R, context: &C) -> Self {
+impl Table {
+    /// Generate a table with some predefined columns
+    pub fn arbitrary_with_columns<R: Rng + ?Sized, C: GenerationContext>(
+        rng: &mut R,
+        context: &C,
+        name: String,
+        predefined_columns: Vec<Column>,
+    ) -> Self {
         let opts = context.opts().table.clone();
-        let name = Name::arbitrary(rng, context).0;
         let large_table =
             opts.large_table.enable && rng.random_bool(opts.large_table.large_table_prob);
-        let column_size = if large_table {
+        let target_column_size = if large_table {
             rng.random_range(opts.large_table.column_range)
         } else {
             rng.random_range(opts.column_range)
         } as usize;
-        let mut column_set = IndexSet::with_capacity(column_size);
+
+        // Start with predefined columns
+        let mut column_set = IndexSet::with_capacity(target_column_size);
+        for col in predefined_columns {
+            column_set.insert(col);
+        }
+
+        // Generate additional columns to reach target size
         for col in std::iter::repeat_with(|| Column::arbitrary(rng, context)) {
             column_set.insert(col);
-            if column_set.len() == column_size {
+            if column_set.len() >= target_column_size {
                 break;
             }
         }
@@ -44,6 +56,13 @@ impl Arbitrary for Table {
     }
 }
 
+impl Arbitrary for Table {
+    fn arbitrary<R: Rng + ?Sized, C: GenerationContext>(rng: &mut R, context: &C) -> Self {
+        let name = Name::arbitrary(rng, context).0;
+        Table::arbitrary_with_columns(rng, context, name, vec![])
+    }
+}
+
 impl Arbitrary for Column {
     fn arbitrary<R: Rng + ?Sized, C: GenerationContext>(rng: &mut R, context: &C) -> Self {
         let name = Name::arbitrary(rng, context).0;