connect/tests/unit_tests/test_driver/test_cursor.py

from types import SimpleNamespace
from unittest.mock import Mock

import pytest

from clickhouse_connect.dbapi.cursor import Cursor
from clickhouse_connect.driver.exceptions import ProgrammingError


def create_mock_client(result_data):
    """Helper to create a mock client with query result"""
    client = Mock()
    query_result = Mock()
    query_result.result_set = result_data
    query_result.column_names = ["col1", "col2", "col3"]
    query_result.column_types = [Mock(name="String")] * 3
    query_result.summary = {"rows": len(result_data)}
    client.query.return_value = query_result
    return client


def create_mock_query_result(result_data, column_names=None, column_types=None):
    """Create a mock query result with optional metadata."""
    query_result = Mock()
    query_result.result_set = result_data
    query_result.column_names = column_names or []
    query_result.column_types = column_types or []
    query_result.summary = {"rows": len(result_data)}
    return query_result


def test_fetchall_respects_cursor_position():
    """Test that fetchall() returns only unread rows and respects cursor position"""
    test_data = [
        ("row1_col1", "row1_col2", "row1_col3"),
        ("row2_col1", "row2_col2", "row2_col3"),
        ("row3_col1", "row3_col2", "row3_col3"),
        ("row4_col1", "row4_col2", "row4_col3"),
        ("row5_col1", "row5_col2", "row5_col3"),
    ]

    client = create_mock_client(test_data)
    cursor = Cursor(client)

    # Execute a query to populate cursor data
    cursor.execute("SELECT * FROM test_table")

    # Fetch first two rows
    row1 = cursor.fetchone()
    row2 = cursor.fetchone()

    assert row1 == test_data[0]
    assert row2 == test_data[1]
    assert cursor._ix == 2  # Cursor should be at position 2

    # fetchall() should return remaining rows, not all rows
    remaining_rows = cursor.fetchall()

    # Should only get rows 3, 4, and 5 (indices 2, 3, 4)
    expected_remaining = test_data[2:]
    assert remaining_rows == expected_remaining
    assert len(remaining_rows) == 3

    # Cursor should now be at the end
    assert cursor._ix == cursor._rowcount

    # Another fetchall() should return empty list since all rows consumed
    empty_result = cursor.fetchall()
    assert empty_result == []


def test_fetchmany_respects_size_parameter():
    """Test that fetchmany() correctly handles the size parameter"""
    test_data = [
        ("row1",),
        ("row2",),
        ("row3",),
        ("row4",),
        ("row5",),
        ("row6",),
        ("row7",),
        ("row8",),
        ("row9",),
        ("row10",),
    ]

    client = create_mock_client(test_data)
    cursor = Cursor(client)
    cursor.execute("SELECT * FROM test_table")

    # Test fetchmany with explicit size
    batch1 = cursor.fetchmany(size=3)
    assert len(batch1) == 3
    assert batch1 == test_data[0:3]
    assert cursor._ix == 3

    # Test fetchmany with size larger than remaining rows
    batch2 = cursor.fetchmany(size=10)
    assert len(batch2) == 7  # Only 7 rows remaining
    assert batch2 == test_data[3:10]
    assert cursor._ix == 10

    # Test fetchmany when no rows remain
    batch3 = cursor.fetchmany(size=5)
    assert batch3 == []
    assert cursor._ix == 10


def test_fetchmany_negative_values():
    """Test fetchmany with various negative values"""
    test_data = [("row1",), ("row2",), ("row3",), ("row4",), ("row5",)]

    client = create_mock_client(test_data)
    cursor = Cursor(client)
    cursor.execute("SELECT * FROM test_table")

    # Advance cursor partway
    cursor.fetchone()  # Now at index 1

    # Any negative value should fetch all remaining
    remaining = cursor.fetchmany(-999)
    assert len(remaining) == 4
    assert remaining == test_data[1:]


def test_fetchmany_w_no_size_parameter_fetches_all_remaining():
    """Test default behavior or fetchmany"""
    test_data = [("A", 1), ("B", 2), ("C", 3), ("D", 4), ("E", 5), ("F", 6)]

    client = create_mock_client(test_data)
    cursor = Cursor(client)
    cursor.execute("SELECT * FROM test_table")

    # Fetch many (no size parameter)
    batch = cursor.fetchmany()
    assert batch == test_data

    # Reset cursor
    cursor.execute("SELECT * FROM test_table")

    # Fetch one
    row1 = cursor.fetchone()
    assert row1 == test_data[0]

    # Fetch remaining (fetchmany with no size parameter)
    batch = cursor.fetchmany()
    assert batch == test_data[1:]


def test_mixed_fetch_operations():
    """Test mixing different fetch operations"""
    test_data = [("A", 1), ("B", 2), ("C", 3), ("D", 4), ("E", 5), ("F", 6)]

    client = create_mock_client(test_data)
    cursor = Cursor(client)
    cursor.execute("SELECT * FROM test_table")

    # Fetch one
    row1 = cursor.fetchone()
    assert row1 == test_data[0]

    # Fetch many
    batch = cursor.fetchmany(2)
    assert batch == test_data[1:3]

    # Fetch all remaining
    remaining = cursor.fetchall()
    assert remaining == test_data[3:6]

    # All subsequent fetches should return empty/None
    assert cursor.fetchone() is None
    assert cursor.fetchone() is None  # Should continue returning None
    assert cursor.fetchmany(10) == []
    assert cursor.fetchall() == []


def test_cursor_reset_on_new_execute():
    """Test that cursor position resets on new execute"""
    test_data = [("row1",), ("row2",), ("row3",)]

    client = create_mock_client(test_data)
    cursor = Cursor(client)

    # First query
    cursor.execute("SELECT * FROM test_table")
    cursor.fetchmany(2)
    assert cursor._ix == 2

    # New query should reset cursor
    cursor.execute("SELECT * FROM test_table")
    assert cursor._ix == 0

    # Should be able to fetch all rows again
    all_rows = cursor.fetchall()
    assert len(all_rows) == 3
    assert all_rows == test_data


def test_check_valid():
    """Test that operations fail when cursor is not valid"""
    client = Mock()
    cursor = Cursor(client)

    # Cursor should be invalid before execute
    with pytest.raises(ProgrammingError):
        cursor.fetchone()

    with pytest.raises(ProgrammingError):
        cursor.fetchall()

    with pytest.raises(ProgrammingError):
        cursor.fetchmany()


def test_empty_result_set():
    """Test cursor behavior with empty result set"""
    client = create_mock_client([])
    cursor = Cursor(client)
    cursor.execute("SELECT * FROM empty_table")

    assert cursor.rowcount == 0
    assert cursor.fetchone() is None
    assert cursor.fetchall() == []
    assert cursor.fetchmany(5) == []


def test_executemany_bulk_insert_with_tuple_rows():
    """Sequence rows (PEP 249 seq_of_parameters) must use the bulk insert path
    with column names taken from the INSERT statement.
    """
    client = Mock()
    cursor = Cursor(client)

    rows = [(13, "user_1"), (79, "user_2")]
    cursor.executemany("INSERT INTO test_table (id, name) VALUES (%s, %s)", rows)

    client.insert.assert_called_once_with("test_table", rows, ["id", "name"])
    client.query.assert_not_called()


def test_executemany_bulk_insert_tuple_rows_without_column_list():
    """Without an explicit column list, sequence rows insert into all columns."""
    client = Mock()
    cursor = Cursor(client)

    rows = [(13, "user_1"), (79, "user_2")]
    cursor.executemany("INSERT INTO test_table VALUES (%s, %s)", rows)

    client.insert.assert_called_once_with("test_table", rows, "*")
    client.query.assert_not_called()


def test_executemany_bulk_insert_with_dict_rows():
    """Mapping rows keep the existing bulk insert behavior."""
    client = Mock()
    cursor = Cursor(client)

    rows = [{"id": 13, "name": "user_1"}, {"id": 79, "name": "user_2"}]
    cursor.executemany("INSERT INTO test_table (id, name) VALUES (%(id)s, %(name)s)", rows)

    client.insert.assert_called_once_with("test_table", [[13, "user_1"], [79, "user_2"]], ["id", "name"])
    client.query.assert_not_called()


def test_execute_unescapes_double_percents_without_parameters():
    """Test that cursor.execute unescapes %% to % when no parameters are given.

    This is required by the PEP 249 pyformat paramstyle contract: callers
    (e.g. SQLAlchemy) escape literal percent signs as %% in the operation
    string.  When there are no parameters, the cursor must unescape them.
    See https://github.com/ClickHouse/clickhouse-connect/issues/297
    """
    client = create_mock_client([])
    cursor = Cursor(client)

    # Simulate what SQLAlchemy sends for:
    #   text("SELECT formatDateTime(toDate('2010-01-04'), '%g')")
    # with _double_percents=True (pyformat paramstyle)
    cursor.execute("SELECT formatDateTime(toDate('2010-01-04'), '%%g')")

    # The query passed to client.query should have %% unescaped to %
    actual_query = client.query.call_args[0][0]
    assert actual_query == "SELECT formatDateTime(toDate('2010-01-04'), '%g')"
    assert "%%" not in actual_query


def test_execute_preserves_percent_with_parameters():
    """Test that cursor.execute does NOT manually unescape %% when parameters
    are provided, since finalize_query handles it via Python's % operator.
    """
    client = create_mock_client([])
    cursor = Cursor(client)

    # Simulate what SQLAlchemy sends for:
    #   text("SELECT formatDateTime(toDate(:d), '%g')")
    # with _double_percents=True and bound parameter d
    cursor.execute("SELECT formatDateTime(toDate(%(d)s), '%%g')", {"d": "2010-01-04"})

    # Parameters are passed through to client.query; finalize_query handles
    # the %% -> % unescaping via the % operator during parameter substitution.
    actual_query = client.query.call_args[0][0]
    actual_params = client.query.call_args[0][1]
    assert actual_query == "SELECT formatDateTime(toDate(%(d)s), '%%g')"
    assert actual_params == {"d": "2010-01-04"}


def test_execute_unescapes_multiple_percents():
    """Test unescaping multiple %% occurrences in a single query."""
    client = create_mock_client([])
    cursor = Cursor(client)

    cursor.execute("SELECT formatDateTime(now(), '%%Y-%%m-%%d %%H:%%M:%%S')")

    actual_query = client.query.call_args[0][0]
    assert actual_query == "SELECT formatDateTime(now(), '%Y-%m-%d %H:%M:%S')"


def test_execute_empty_result_fetches_metadata_with_parameters():
    """Empty SELECT results should still populate description metadata."""
    client = Mock()
    client.query.side_effect = [
        create_mock_query_result([]),
        create_mock_query_result(
            [],
            column_names=["value_1"],
            column_types=[SimpleNamespace(name="UInt64")],
        ),
    ]
    cursor = Cursor(client)

    cursor.execute(
        "SELECT value_1 FROM test_table WHERE value_1 = %(value_1)s LIMIT %(param_1)s",
        {"value_1": 13, "param_1": 1},
    )

    assert cursor.description == [("value_1", "UInt64", None, None, None, None, True)]
    assert client.query.call_args_list[1].args == (
        "SELECT * FROM (SELECT value_1 FROM test_table WHERE value_1 = %(value_1)s LIMIT %(param_1)s) LIMIT 0",
        {"value_1": 13, "param_1": 1},
    )


def test_execute_empty_with_query_fetches_metadata():
    """CTE queries should use the same metadata fallback."""
    client = Mock()
    client.query.side_effect = [
        create_mock_query_result([]),
        create_mock_query_result(
            [],
            column_names=["value_1"],
            column_types=[SimpleNamespace(name="UInt64")],
        ),
    ]
    cursor = Cursor(client)

    cursor.execute("WITH value_1 AS 13 SELECT value_1 WHERE value_1 = 79")

    assert cursor.description == [("value_1", "UInt64", None, None, None, None, True)]
    assert client.query.call_args_list[1].args == (
        "SELECT * FROM (WITH value_1 AS 13 SELECT value_1 WHERE value_1 = 79) LIMIT 0",
        None,
    )


def _mock_insert_client(written_rows: int, summary_extra: dict | None = None):
    """Return a mock client whose insert() returns a QuerySummary-like object."""
    summary_dict = {"written_rows": str(written_rows), **(summary_extra or {})}
    insert_summary = Mock()
    insert_summary.written_rows = written_rows
    insert_summary.summary = summary_dict
    client = Mock()
    client.insert.return_value = insert_summary
    return client, summary_dict


def test_executemany_bulk_insert_rowcount_equals_written_rows():
    """rowcount after a bulk executemany reflects the actual number of inserted rows."""
    client, _ = _mock_insert_client(written_rows=2)
    cursor = Cursor(client)

    rows = [(13, "user_1"), (79, "user_2")]
    cursor.executemany("INSERT INTO test_table (id, name) VALUES (%s, %s)", rows)

    assert cursor.rowcount == 2


def test_executemany_bulk_insert_appends_summary():
    """summary is populated from the insert response after a bulk executemany."""
    client, summary_dict = _mock_insert_client(written_rows=2, summary_extra={"written_bytes": "64"})
    cursor = Cursor(client)

    rows = [(13, "user_1"), (79, "user_2")]
    cursor.executemany("INSERT INTO test_table (id, name) VALUES (%s, %s)", rows)

    assert cursor.summary == [summary_dict]


def test_executemany_generator_falls_through_to_row_by_row():
    """A generator passed to executemany falls through to the row-by-row path without raising TypeError."""
    client = Mock()
    client.query.return_value = create_mock_query_result([])

    cursor = Cursor(client)

    def row_generator():
        yield {"id": 1, "name": "user_1"}
        yield {"id": 2, "name": "user_2"}

    cursor.executemany("INSERT INTO test_table (id, name) VALUES (%(id)s, %(name)s)", row_generator())

    client.insert.assert_not_called()
    assert client.query.call_count == 2