From a149e5520939ee70e2c72b8a1130244be0ec3716 Mon Sep 17 00:00:00 2001
From: pedrocarlo <pedro.muniz.carlo@gmail.com>
Date: Thu, 9 Oct 2025 19:29:14 -0300
Subject: [PATCH] fix Drop Column to only be generated if no columns conflict
 in Indexes

---
 sql_generation/generation/query.rs | 105 ++++++++++++++++++++++-------
 1 file changed, 80 insertions(+), 25 deletions(-)

diff --git a/sql_generation/generation/query.rs b/sql_generation/generation/query.rs
index ef035e0be..94aafaae9 100644
--- a/sql_generation/generation/query.rs
+++ b/sql_generation/generation/query.rs
@@ -1,6 +1,6 @@
 use crate::generation::{
-    gen_random_text, pick_n_unique, pick_unique, Arbitrary, ArbitraryFrom, ArbitrarySized,
-    GenerationContext,
+    gen_random_text, pick_index, pick_n_unique, pick_unique, Arbitrary, ArbitraryFrom,
+    ArbitrarySized, GenerationContext,
 };
 use crate::model::query::alter_table::{AlterTable, AlterTableType, AlterTableTypeDiscriminants};
 use crate::model::query::predicate::Predicate;
@@ -392,26 +392,52 @@ impl Arbitrary for Update {
     }
 }
 
-impl Arbitrary for AlterTable {
-    fn arbitrary<R: Rng + ?Sized, C: GenerationContext>(rng: &mut R, context: &C) -> Self {
-        let table = pick(context.tables(), rng);
-        let choices: &'static [AlterTableTypeDiscriminants] = if table.columns.len() > 1 {
-            &[
-                AlterTableTypeDiscriminants::RenameTo,
-                AlterTableTypeDiscriminants::AddColumn,
-                // AlterTableTypeDiscriminants::AlterColumn,
-                AlterTableTypeDiscriminants::RenameColumn,
-                AlterTableTypeDiscriminants::DropColumn,
-            ]
-        } else {
-            &[
-                AlterTableTypeDiscriminants::RenameTo,
-                AlterTableTypeDiscriminants::AddColumn,
-                // AlterTableTypeDiscriminants::AlterColumn,
-                AlterTableTypeDiscriminants::RenameColumn,
-            ]
-        };
-        let alter_table_type = match choices.choose(rng).unwrap() {
+const ALTER_TABLE_ALL: &[AlterTableTypeDiscriminants] = &[
+    AlterTableTypeDiscriminants::RenameTo,
+    AlterTableTypeDiscriminants::AddColumn,
+    // AlterTableTypeDiscriminants::AlterColumn,
+    AlterTableTypeDiscriminants::RenameColumn,
+    AlterTableTypeDiscriminants::DropColumn,
+];
+const ALTER_TABLE_NO_DROP: &[AlterTableTypeDiscriminants] = &[
+    AlterTableTypeDiscriminants::RenameTo,
+    AlterTableTypeDiscriminants::AddColumn,
+    // AlterTableTypeDiscriminants::AlterColumn,
+    AlterTableTypeDiscriminants::RenameColumn,
+];
+
+// TODO: Unfortunately this diff strategy allocates a couple of IndexSet's
+// in the future maybe change this to be more efficient. This is currently acceptable because this function
+// is only called for `DropColumn`
+fn get_column_diff(table: &Table) -> IndexSet<&str> {
+    // Columns that are referenced in INDEXES cannot be dropped
+    let column_cannot_drop = table
+        .indexes
+        .iter()
+        .flat_map(|index| index.columns.iter().map(|(col_name, _)| col_name.as_str()))
+        .collect::<IndexSet<_>>();
+    if column_cannot_drop.len() == table.columns.len() {
+        // Optimization: all columns are present in indexes so we do not need to but the table column set
+        return IndexSet::new();
+    }
+
+    let column_set: IndexSet<_, std::hash::RandomState> =
+        IndexSet::from_iter(table.columns.iter().map(|col| col.name.as_str()));
+
+    let diff = column_set
+        .difference(&column_cannot_drop)
+        .copied()
+        .collect::<IndexSet<_, std::hash::RandomState>>();
+    diff
+}
+
+impl ArbitraryFrom<(&Table, &[AlterTableTypeDiscriminants])> for AlterTableType {
+    fn arbitrary_from<R: Rng + ?Sized, C: GenerationContext>(
+        rng: &mut R,
+        context: &C,
+        (table, choices): (&Table, &[AlterTableTypeDiscriminants]),
+    ) -> Self {
+        match choices.choose(rng).unwrap() {
             AlterTableTypeDiscriminants::RenameTo => AlterTableType::RenameTo {
                 new_name: Name::arbitrary(rng, context).0,
             },
@@ -425,10 +451,39 @@ impl Arbitrary for AlterTable {
                 old: pick(&table.columns, rng).name.clone(),
                 new: Name::arbitrary(rng, context).0,
             },
-            AlterTableTypeDiscriminants::DropColumn => AlterTableType::DropColumn {
-                column_name: pick(&table.columns, rng).name.clone(),
-            },
+            AlterTableTypeDiscriminants::DropColumn => {
+                let col_diff = get_column_diff(table);
+
+                if col_diff.is_empty() {
+                    // Generate a DropColumn if we can drop a column
+                    return AlterTableType::arbitrary_from(
+                        rng,
+                        context,
+                        (table, ALTER_TABLE_NO_DROP),
+                    );
+                }
+
+                let col_idx = pick_index(col_diff.len(), rng);
+                let col_name = col_diff.get_index(col_idx).unwrap();
+
+                AlterTableType::DropColumn {
+                    column_name: col_name.to_string(),
+                }
+            }
+        }
+    }
+}
+
+impl Arbitrary for AlterTable {
+    fn arbitrary<R: Rng + ?Sized, C: GenerationContext>(rng: &mut R, context: &C) -> Self {
+        let table = pick(context.tables(), rng);
+        let choices: &'static [AlterTableTypeDiscriminants] = if table.columns.len() > 1 {
+            ALTER_TABLE_ALL
+        } else {
+            ALTER_TABLE_NO_DROP
         };
+
+        let alter_table_type = AlterTableType::arbitrary_from(rng, context, (table, choices));
         Self {
             table_name: table.name.clone(),
             alter_table_type,