From 02ccdd71940474c6a33252bd2b6a93b734b9e178 Mon Sep 17 00:00:00 2001 From: Pekka Enberg Date: Mon, 21 Jul 2025 13:01:21 +0300 Subject: [PATCH] antithesis: Add `CREATE INDEX` parallel driver --- .../parallel_driver_create_index.py | 131 ++++++++++++++++++ 1 file changed, 131 insertions(+) create mode 100755 antithesis-tests/stress-composer/parallel_driver_create_index.py diff --git a/antithesis-tests/stress-composer/parallel_driver_create_index.py b/antithesis-tests/stress-composer/parallel_driver_create_index.py new file mode 100755 index 000000000..26546a7a2 --- /dev/null +++ b/antithesis-tests/stress-composer/parallel_driver_create_index.py @@ -0,0 +1,131 @@ +#!/usr/bin/env -S python3 -u + +import json +import string + +import turso +from antithesis.random import get_random, random_choice +from utils import generate_random_value + +# Get initial state +try: + con_init = turso.connect("init_state.db") +except Exception as e: + print(f"Error connecting to database: {e}") + exit(0) + +cur_init = con_init.cursor() + +# Get table count +tbl_len = cur_init.execute("SELECT count FROM tables").fetchone()[0] +if tbl_len == 0: + print("No tables available for index creation") + exit(0) + +# Select a random table +selected_tbl = get_random() % tbl_len +tbl_schema = json.loads(cur_init.execute(f"SELECT schema FROM schemas WHERE tbl = {selected_tbl}").fetchone()[0]) + +# Connect to the main database +try: + con = turso.connect("stress_composer.db", experimental_indexes=True) +except Exception as e: + print(f"Failed to open stress_composer.db. Exiting... {e}") + exit(0) + +cur = con.cursor() + +# Check existing indexes on this table +existing_indexes = cur.execute(f""" + SELECT name FROM sqlite_master + WHERE type = 'index' AND tbl_name = 'tbl_{selected_tbl}' + AND sql IS NOT NULL +""").fetchall() +existing_index_names = {idx[0] for idx in existing_indexes} + +print(f"Selected table: tbl_{selected_tbl} with {tbl_schema['colCount']} columns") +print(f"Existing indexes: {len(existing_indexes)}") + +# Decide whether to create a single-column or composite index +create_composite = tbl_schema["colCount"] > 2 and get_random() % 3 == 0 # 33% chance for composite + +if create_composite: + # Create composite index + num_cols = 2 + (get_random() % min(2, tbl_schema["colCount"] - 1)) # 2-3 columns + selected_cols = [] + available_cols = list(range(tbl_schema["colCount"])) + + for _ in range(min(num_cols, len(available_cols))): + if available_cols: + idx = get_random() % len(available_cols) + selected_cols.append(available_cols.pop(idx)) + + col_names = [f"col_{col}" for col in sorted(selected_cols)] + cols_suffix = "_".join(str(c) for c in sorted(selected_cols)) + random_suffix = "".join(random_choice(string.ascii_lowercase) for _ in range(4)) + index_name = f"idx_tbl{selected_tbl}_{cols_suffix}_{random_suffix}" + + # Check if this combination already has an index + if index_name not in existing_index_names: + try: + create_stmt = f"CREATE INDEX {index_name} ON tbl_{selected_tbl} ({', '.join(col_names)})" + print(f"Creating composite index: {create_stmt}") + cur.execute(create_stmt) + + # Store index information in init_state.db + cur_init.execute(f""" + INSERT INTO indexes (idx_name, tbl_name, idx_type, cols) + VALUES ('{index_name}', 'tbl_{selected_tbl}', 'composite', '{", ".join(col_names)}') + """) + con_init.commit() + print(f"Successfully created composite index: {index_name}") + except turso.OperationalError as e: + print(f"Failed to create composite index: {e}") + con.rollback() + else: + print(f"Index {index_name} already exists, skipping") +else: + # Create single-column index + # Select a random column (avoiding primary key which is usually col_0) + available_cols = list(range(1, tbl_schema["colCount"])) + if not available_cols: + available_cols = [0] # Fall back to first column if only one exists + + selected_col = available_cols[get_random() % len(available_cols)] + col_name = f"col_{selected_col}" + + # Determine index type based on column data type + col_type = tbl_schema[col_name]["data_type"] + index_suffix = "".join(random_choice(string.ascii_lowercase) for _ in range(4)) + + if col_type == "TEXT" and tbl_schema[col_name].get("unique", False): + # Create unique index for unique text columns + index_name = f"idx_tbl{selected_tbl}_col{selected_col}_unique_{index_suffix}" + create_stmt = f"CREATE UNIQUE INDEX {index_name} ON tbl_{selected_tbl} ({col_name})" + else: + # Create regular index + index_name = f"idx_tbl{selected_tbl}_col{selected_col}_{index_suffix}" + create_stmt = f"CREATE INDEX {index_name} ON tbl_{selected_tbl} ({col_name})" + + if index_name not in existing_index_names: + try: + print(f"Creating single-column index: {create_stmt}") + cur.execute(create_stmt) + + # Store index information in init_state.db + idx_type = "unique" if "UNIQUE" in create_stmt else "single" + cur_init.execute(f""" + INSERT INTO indexes (idx_name, tbl_name, idx_type, cols) + VALUES ('{index_name}', 'tbl_{selected_tbl}', '{idx_type}', '{col_name}') + """) + con_init.commit() + print(f"Successfully created {idx_type} index: {index_name}") + except turso.OperationalError as e: + print(f"Failed to create index: {e}") + con.rollback() + else: + print(f"Index {index_name} already exists, skipping") + +con.commit() +con.close() +con_init.close()