From 0762c8f780aafcbf4b12e59bb9133c9b92c318ff Mon Sep 17 00:00:00 2001
From: Jussi Saurio <jussi.saurio@gmail.com>
Date: Mon, 7 Jul 2025 11:09:38 +0300
Subject: [PATCH] stress: add a way to run stress with indexes enabled

---
 stress/Cargo.toml |  1 +
 stress/main.rs    | 52 ++++++++++++++++++++++++++++++++++++-----------
 2 files changed, 41 insertions(+), 12 deletions(-)

diff --git a/stress/Cargo.toml b/stress/Cargo.toml
index 65930f86b..47fa7331e 100644
--- a/stress/Cargo.toml
+++ b/stress/Cargo.toml
@@ -17,6 +17,7 @@ path = "main.rs"
 [features]
 default = []
 antithesis = ["turso/antithesis"]
+experimental_indexes = ["turso/experimental_indexes"]
 
 [dependencies]
 anarchist-readable-name-generator-lib = "0.1.0"
diff --git a/stress/main.rs b/stress/main.rs
index daab218ee..8531997cc 100644
--- a/stress/main.rs
+++ b/stress/main.rs
@@ -33,7 +33,7 @@ pub struct Column {
 }
 
 /// Represents SQLite data types
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq)]
 pub enum DataType {
     Integer,
     Real,
@@ -47,6 +47,7 @@ pub enum DataType {
 pub enum Constraint {
     PrimaryKey,
     NotNull,
+    #[cfg(feature = "experimental_indexes")]
     Unique,
 }
 
@@ -79,17 +80,20 @@ fn generate_random_data_type() -> DataType {
 }
 
 fn generate_random_constraint() -> Constraint {
+    #[cfg(feature = "experimental_indexes")]
     match get_random() % 2 {
         0 => Constraint::NotNull,
         _ => Constraint::Unique,
     }
+    #[cfg(not(feature = "experimental_indexes"))]
+    Constraint::NotNull
 }
 
 fn generate_random_column() -> Column {
     let name = generate_random_identifier();
     let data_type = generate_random_data_type();
 
-    let constraint_count = (get_random() % 3) as usize;
+    let constraint_count = (get_random() % 2) as usize;
     let mut constraints = Vec::with_capacity(constraint_count);
 
     for _ in 0..constraint_count {
@@ -122,11 +126,37 @@ fn generate_random_table() -> Table {
         columns.push(column);
     }
 
-    // Then, randomly select one column to be the primary key
-    let pk_index = (get_random() % column_count as u64) as usize;
-    columns[pk_index].constraints.push(Constraint::PrimaryKey);
+    #[cfg(feature = "experimental_indexes")]
+    {
+        // Then, randomly select one column to be the primary key
+        let pk_index = (get_random() % column_count as u64) as usize;
+        columns[pk_index].constraints.push(Constraint::PrimaryKey);
+        Table { name, columns }
+    }
+    #[cfg(not(feature = "experimental_indexes"))]
+    {
+        // Pick a random column that is exactly INTEGER type to be the primary key (INTEGER PRIMARY KEY does not require indexes,
+        // as it becomes an alias for the ROWID).
+        let pk_candidates = columns
+            .iter()
+            .enumerate()
+            .filter(|(_, col)| col.data_type == DataType::Integer)
+            .map(|(i, _)| i)
+            .collect::<Vec<_>>();
+        if pk_candidates.is_empty() {
+            // if there are no INTEGER columns, make a random column INTEGER and set it as PRIMARY KEY
+            let col_id = (get_random() % column_count as u64) as usize;
+            columns[col_id].data_type = DataType::Integer;
+            columns[col_id].constraints.push(Constraint::PrimaryKey);
+            return Table { name, columns };
+        }
+        let pk_index = pk_candidates
+            .get((get_random() % pk_candidates.len() as u64) as usize)
+            .unwrap();
+        columns[*pk_index].constraints.push(Constraint::PrimaryKey);
 
-    Table { name, columns }
+        Table { name, columns }
+    }
 }
 
 pub fn gen_bool(probability_true: f64) -> bool {
@@ -165,12 +195,9 @@ impl ArbitrarySchema {
                     .map(|col| {
                         let mut col_def =
                             format!("  {} {}", col.name, data_type_to_sql(&col.data_type));
-                        if false {
-                            /* FIXME */
-                            for constraint in &col.constraints {
-                                col_def.push(' ');
-                                col_def.push_str(&constraint_to_sql(constraint));
-                            }
+                        for constraint in &col.constraints {
+                            col_def.push(' ');
+                            col_def.push_str(&constraint_to_sql(constraint));
                         }
                         col_def
                     })
@@ -197,6 +224,7 @@ fn constraint_to_sql(constraint: &Constraint) -> String {
     match constraint {
         Constraint::PrimaryKey => "PRIMARY KEY".to_string(),
         Constraint::NotNull => "NOT NULL".to_string(),
+        #[cfg(feature = "experimental_indexes")]
         Constraint::Unique => "UNIQUE".to_string(),
     }
 }