apexbase 1.25.0

"""Tests for CTE (WITH...AS), INSERT...SELECT, and EXPLAIN SQL features."""
import pytest
import tempfile
import os
import shutil
from apexbase import ApexStorage


def get_rows(result):
    """Extract rows from execute() result dict (handles both columns+rows and columns_dict formats)."""
    if 'rows' in result:
        return result['rows']
    cd = result.get('columns_dict')
    if cd:
        cols = list(cd.keys())
        if not cols:
            return []
        n = len(cd[cols[0]])
        return [[cd[c][i] for c in cols] for i in range(n)]
    return []

def get_columns(result):
    """Extract column names from execute() result dict."""
    if 'columns' in result:
        return result['columns']
    cd = result.get('columns_dict')
    if cd:
        return list(cd.keys())
    return []

def get_scalar(result):
    """Extract scalar value (first cell) from result."""
    rows = get_rows(result)
    if rows and rows[0]:
        return rows[0][0]
    return None

def get_col_values(result, col_name):
    """Extract all values from a specific column."""
    cd = result.get('columns_dict')
    if cd and col_name in cd:
        return list(cd[col_name])
    cols = get_columns(result)
    if col_name not in cols:
        return []
    idx = cols.index(col_name)
    return [row[idx] for row in get_rows(result)]


@pytest.fixture
def db_with_data():
    """Create a database with sample data in a unique temp directory."""
    tmpdir = tempfile.mkdtemp(prefix="apextest_cte_")
    storage = ApexStorage(tmpdir)
    storage.execute("CREATE TABLE IF NOT EXISTS users (name TEXT, age INT, city TEXT)")
    storage.execute("TRUNCATE TABLE users")
    storage.execute("INSERT INTO users (name, age, city) VALUES ('Alice', 30, 'NYC')")
    storage.execute("INSERT INTO users (name, age, city) VALUES ('Bob', 25, 'LA')")
    storage.execute("INSERT INTO users (name, age, city) VALUES ('Charlie', 35, 'NYC')")
    storage.execute("INSERT INTO users (name, age, city) VALUES ('Diana', 28, 'LA')")
    storage.execute("INSERT INTO users (name, age, city) VALUES ('Eve', 32, 'NYC')")
    yield storage
    storage.close()
    shutil.rmtree(tmpdir, ignore_errors=True)


# ========== EXPLAIN Tests ==========

class TestExplain:
    def test_explain_simple_select(self, db_with_data):
        result = db_with_data.execute("EXPLAIN SELECT * FROM users")
        plan = get_rows(result)[0][0]
        assert 'Query Plan' in plan
        assert 'Scan' in plan
        assert 'users' in plan

    def test_explain_select_with_where(self, db_with_data):
        result = db_with_data.execute("EXPLAIN SELECT name FROM users WHERE age > 30")
        plan = get_rows(result)[0][0]
        assert 'Filter' in plan
        assert 'WHERE' in plan

    def test_explain_select_with_group_by(self, db_with_data):
        result = db_with_data.execute("EXPLAIN SELECT city, COUNT(*) FROM users GROUP BY city")
        plan = get_rows(result)[0][0]
        assert 'GroupBy' in plan
        assert 'city' in plan

    def test_explain_select_with_order_limit(self, db_with_data):
        result = db_with_data.execute("EXPLAIN SELECT * FROM users ORDER BY age DESC LIMIT 3")
        plan = get_rows(result)[0][0]
        assert 'Sort' in plan
        assert 'DESC' in plan
        assert 'Limit: 3' in plan

    def test_explain_analyze(self, db_with_data):
        result = db_with_data.execute("EXPLAIN ANALYZE SELECT * FROM users WHERE age > 25")
        plan = get_rows(result)[0][0]
        assert 'Actual Time' in plan
        assert 'Actual Rows' in plan
        assert 'Planning Time' in plan
        assert 'Execution Time' in plan

    def test_explain_insert(self, db_with_data):
        result = db_with_data.execute("EXPLAIN INSERT INTO users (name, age, city) VALUES ('Frank', 40, 'SF')")
        plan = get_rows(result)[0][0]
        assert 'Insert' in plan
        assert 'users' in plan

    def test_explain_shows_row_count(self, db_with_data):
        result = db_with_data.execute("EXPLAIN SELECT * FROM users")
        plan = get_rows(result)[0][0]
        assert 'Rows: ~' in plan
        assert 'Columns: 3' in plan

    def test_explain_shows_physical_candidates(self, db_with_data):
        db_with_data.execute(
            "CREATE INDEX IF NOT EXISTS idx_users_age ON users (age) USING BTREE"
        )
        result = db_with_data.execute(
            "EXPLAIN SELECT name FROM users WHERE age > 30"
        )
        plan = get_rows(result)[0][0]
        assert "Chosen Plan:" in plan
        assert "estimated_cost=" in plan
        assert "Candidate:" in plan
        assert "stats=" in plan

        zone_plan = get_rows(
            db_with_data.execute(
                "EXPLAIN SELECT name FROM users WHERE age BETWEEN 25 AND 35"
            )
        )[0][0]
        assert "zone-map-scan(" in zone_plan

    def test_explain_cbo_candidate_invariants(self, db_with_data):
        db_with_data.execute(
            "CREATE INDEX IF NOT EXISTS idx_users_city_age ON users (city, age)"
        )
        result = db_with_data.execute(
            "EXPLAIN SELECT name FROM users WHERE age = 30 AND city = 'NYC'"
        )
        plan = get_rows(result)[0][0]
        assert "composite-index-scan(city,age)" in plan

        db_with_data.execute(
            "CREATE INDEX IF NOT EXISTS idx_users_age ON users (age) USING BTREE"
        )
        result = db_with_data.execute(
            "EXPLAIN SELECT name FROM users WHERE age = 25 OR age = 30"
        )
        plan = get_rows(result)[0][0]
        assert "index-scan(age)" not in plan
        assert "index-range-scan(age)" not in plan
        assert "index-union(deduplicated)" in plan

        db_with_data.execute("ANALYZE users")
        before = get_rows(
            db_with_data.execute("EXPLAIN SELECT name FROM users WHERE age > 30")
        )[0][0]
        assert "stats=available" in before

        db_with_data.execute(
            "INSERT INTO users (name, age, city) VALUES ('Frank', 40, 'SF')"
        )
        after = get_rows(
            db_with_data.execute("EXPLAIN SELECT name FROM users WHERE age > 30")
        )[0][0]
        assert "stats=default" in after


# ========== INSERT...SELECT Tests ==========

class TestInsertSelect:
    def test_insert_select_same_table(self, db_with_data):
        db_with_data.execute("CREATE TABLE IF NOT EXISTS users_backup (name TEXT, age INT, city TEXT)")
        db_with_data.execute("TRUNCATE TABLE users_backup")
        db_with_data.execute("INSERT INTO users_backup (name, age, city) SELECT name, age, city FROM users")
        result = db_with_data.execute("SELECT COUNT(*) FROM users_backup")
        assert get_scalar(result) == 5

    def test_insert_select_with_filter(self, db_with_data):
        db_with_data.execute("CREATE TABLE IF NOT EXISTS nyc_users (name TEXT, age INT, city TEXT)")
        db_with_data.execute("TRUNCATE TABLE nyc_users")
        db_with_data.execute("INSERT INTO nyc_users (name, age, city) SELECT name, age, city FROM users WHERE city = 'NYC'")
        result = db_with_data.execute("SELECT COUNT(*) FROM nyc_users")
        assert get_scalar(result) == 3

    def test_insert_select_with_order_limit(self, db_with_data):
        db_with_data.execute("CREATE TABLE IF NOT EXISTS top_users (name TEXT, age INT)")
        db_with_data.execute("TRUNCATE TABLE top_users")
        db_with_data.execute("INSERT INTO top_users (name, age) SELECT name, age FROM users ORDER BY age DESC LIMIT 2")
        result = db_with_data.execute("SELECT * FROM top_users ORDER BY age DESC")
        rows = get_rows(result)
        assert len(rows) == 2
        ages = get_col_values(result, 'age')
        assert ages[0] == 35  # Charlie
        assert ages[1] == 32  # Eve

    def test_insert_select_aggregation(self, db_with_data):
        db_with_data.execute("CREATE TABLE IF NOT EXISTS city_stats (city TEXT, cnt INT)")
        db_with_data.execute("TRUNCATE TABLE city_stats")
        db_with_data.execute("INSERT INTO city_stats (city, cnt) SELECT city, COUNT(*) FROM users GROUP BY city")
        result = db_with_data.execute("SELECT * FROM city_stats ORDER BY city")
        rows = get_rows(result)
        assert len(rows) == 2
        cities = get_col_values(result, 'city')
        cnts = get_col_values(result, 'cnt')
        la_idx = cities.index('LA')
        nyc_idx = cities.index('NYC')
        assert cnts[la_idx] == 2
        assert cnts[nyc_idx] == 3

    def test_insert_select_empty_result(self, db_with_data):
        db_with_data.execute("CREATE TABLE IF NOT EXISTS empty_table (name TEXT, age INT)")
        db_with_data.execute("TRUNCATE TABLE empty_table")
        db_with_data.execute("INSERT INTO empty_table (name, age) SELECT name, age FROM users WHERE age > 100")
        result = db_with_data.execute("SELECT COUNT(*) FROM empty_table")
        assert get_scalar(result) == 0


# ========== CTE (WITH...AS) Tests ==========

class TestCTE:
    def test_simple_cte(self, db_with_data):
        result = db_with_data.execute("""
            WITH seniors AS (SELECT name, age FROM users WHERE age >= 30)
            SELECT * FROM seniors ORDER BY age
        """)
        rows = get_rows(result)
        assert len(rows) == 3
        names = get_col_values(result, 'name')
        assert names == ['Alice', 'Eve', 'Charlie']

    def test_cte_with_aggregation(self, db_with_data):
        result = db_with_data.execute("""
            WITH city_counts AS (SELECT city, COUNT(*) AS cnt FROM users GROUP BY city)
            SELECT * FROM city_counts ORDER BY cnt DESC
        """)
        rows = get_rows(result)
        assert len(rows) == 2
        cnts = get_col_values(result, 'cnt')
        assert cnts[0] == 3  # NYC
        assert cnts[1] == 2  # LA

    def test_cte_with_filter_on_cte(self, db_with_data):
        result = db_with_data.execute("""
            WITH young AS (SELECT name, age, city FROM users WHERE age < 30)
            SELECT name FROM young WHERE city = 'LA'
        """)
        rows = get_rows(result)
        assert len(rows) == 2
        names = set(get_col_values(result, 'name'))
        assert 'Bob' in names
        assert 'Diana' in names

    def test_cte_preserves_original_data(self, db_with_data):
        """CTE should not modify the original table."""
        db_with_data.execute("""
            WITH temp AS (SELECT name FROM users WHERE age > 30)
            SELECT * FROM temp
        """)
        result = db_with_data.execute("SELECT COUNT(*) FROM users")
        assert get_scalar(result) == 5  # Original data unchanged

    def test_shared_cte_reaches_non_adjacent_consumers(self, db_with_data):
        """A shared CTE remains visible through all later chained CTEs."""
        result = db_with_data.execute("""
            WITH base AS (
                SELECT name, age, city FROM users
            ),
            seniors AS (
                SELECT COUNT(*) AS senior_cnt
                FROM base WHERE age >= 30
            ),
            young AS (
                SELECT COUNT(*) AS young_cnt
                FROM base WHERE age < 30
            )
            SELECT s.senior_cnt, y.young_cnt
            FROM seniors s
            CROSS JOIN young y
        """)

        assert get_col_values(result, 'senior_cnt') == [3]
        assert get_col_values(result, 'young_cnt') == [2]

    def test_cte_pruning_preserves_aggregate_source_columns(self, db_with_data):
        db_with_data.execute(
            "CREATE TABLE IF NOT EXISTS cte_orders "
            "(order_id INT, user_id INT, status TEXT, order_day INT)"
        )
        db_with_data.execute("TRUNCATE TABLE cte_orders")
        db_with_data.execute(
            "INSERT INTO cte_orders VALUES "
            "(1, 10, 'paid', 1), (2, 10, 'paid', 2), (3, 20, 'paid', 3)"
        )
        db_with_data.execute(
            "CREATE TABLE IF NOT EXISTS cte_payments "
            "(order_id INT, paid_amount DOUBLE)"
        )
        db_with_data.execute("TRUNCATE TABLE cte_payments")
        db_with_data.execute(
            "INSERT INTO cte_payments VALUES (1, 100.0), (2, 50.0), (3, 80.0)"
        )
        db_with_data.execute(
            "CREATE TABLE IF NOT EXISTS cte_items (order_id INT, qty INT)"
        )
        db_with_data.execute("TRUNCATE TABLE cte_items")
        db_with_data.execute(
            "INSERT INTO cte_items VALUES (1, 2), (1, 3), (3, 4)"
        )

        result = db_with_data.execute("""
            WITH paid_orders AS (
                SELECT o.order_id, o.user_id, o.order_day, p.paid_amount
                FROM cte_orders o
                JOIN cte_payments p ON o.order_id = p.order_id
                WHERE o.status = 'paid'
            ),
            item_agg AS (
                SELECT order_id, SUM(qty) AS item_count
                FROM cte_items
                GROUP BY order_id
            )
            SELECT
                po.user_id,
                COUNT(*) AS order_count,
                SUM(po.paid_amount) AS gmv,
                AVG(po.paid_amount) AS avg_paid,
                MAX(po.order_day) AS last_order_day,
                SUM(COALESCE(ia.item_count, 0)) AS total_items
            FROM paid_orders po
            LEFT JOIN item_agg ia ON po.order_id = ia.order_id
            GROUP BY po.user_id
            ORDER BY po.user_id
        """)

        assert get_rows(result) == [
            [10, 2, 150.0, 75.0, 2, 5.0],
            [20, 1, 80.0, 80.0, 3, 4.0],
        ]