turso/antithesis-tests/stress-composer/parallel_driver_create_index.py

#!/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()