diff --git a/docs/changelog.rst b/docs/changelog.rst index f9e152baf..6fb248473 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -21,6 +21,11 @@ v0.54.0 - SQL processing correctness and cleanup statement config plus normalized driver-feature dictionary across backends. * MySQL-family adapter config, driver, and pool modules now resolve runtime vendor symbols through adapter-local typing modules. +* Oracle LOB fetches now default to direct string/byte materialization where + python-oracledb supports it. Pass ``fetch_lobs=True`` when application code + needs native Oracle LOB locators. Unconstrained LOB contents are no longer + parsed with content heuristics; native ``JSON``, ``IS JSON`` CLOB/BLOB, and + OSON-capable values still decode through Oracle JSON metadata. * Driver statement-object caches are now bounded by the configured statement cache size, and cached named-parameter rebinding reuses driver-owned processing state. diff --git a/docs/reference/adapters/oracledb.rst b/docs/reference/adapters/oracledb.rst index 5feaffa17..a426e43d3 100644 --- a/docs/reference/adapters/oracledb.rst +++ b/docs/reference/adapters/oracledb.rst @@ -6,6 +6,32 @@ Sync and async Oracle adapter using `python-oracledb None: + self._driver = driver + self._sql = sql + self._parameters = parameters + self._cursor_manager: AdbcCursor | None = None + self._reader: Any = None + + def start(self) -> None: + cursor_manager = self._driver.with_cursor(self._driver.connection) + reader: Any = None + try: + cursor = cursor_manager.__enter__() + handler = self._driver.handle_database_exceptions() + with handler: + execute_parameters = normalize_postgres_empty_parameters(self._driver._dialect_name, self._parameters) + cursor.execute(self._sql, parameters=execute_parameters) + reader = _fetch_record_batch(cursor) + self._driver._check_pending_exception(handler) + except BaseException: + with contextlib.suppress(Exception): + cursor_manager.__exit__(None, None, None) + raise + + if reader is None: + msg = "ADBC did not return a record batch reader." + raise SQLSpecError(msg) + self._cursor_manager = cursor_manager + self._reader = iter(reader) + + def fetch_chunk(self) -> "list[dict[str, Any]]": + reader = self._reader + if reader is None: + return [] + while True: + try: + batch = next(reader) + except StopIteration: + return [] + rows = cast("list[dict[str, Any]]", batch.to_pylist()) + if rows: + return rows + + def close(self) -> None: + self._reader = None + cursor_manager = self._cursor_manager + self._cursor_manager = None + if cursor_manager is not None: + with contextlib.suppress(Exception): + cursor_manager.__exit__(None, None, None) + + +def _fetch_record_batch(cursor: Any) -> Any: + fetch_record_batch = getattr(cursor, "fetch_record_batch", None) + if fetch_record_batch is None: + msg = "ADBC cursor does not expose fetch_record_batch() for native row streaming." + raise SQLSpecError(msg) + return fetch_record_batch() + + @final class AdbcDriver(SyncDriverAdapterBase): """ADBC driver for Arrow Database Connectivity. @@ -151,19 +216,16 @@ def dispatch_execute(self, cursor: "AdbcRawCursor", statement: SQL) -> "Executio is_select_like = statement.returns_rows() or self._should_force_select(statement, cursor) if is_select_like: - fetched_data = cursor.fetchall() - column_names = self._resolve_column_names(cursor.description) - data, column_names = collect_rows( - cast("list[Any] | None", fetched_data), cursor.description, column_names=column_names - ) - row_format = "dict" if data and isinstance(data[0], dict) else "tuple" + arrow_table = cursor.fetch_arrow_table() + data = arrow_table.to_pylist() + column_names = list(arrow_table.column_names) return self.create_execution_result( cursor, selected_data=data, column_names=column_names, data_row_count=len(data), is_select_result=True, - row_format=row_format, + row_format="dict", ) row_count = self._resolve_count_result_rowcount(cursor, fallback=resolve_rowcount(cursor)) @@ -265,6 +327,14 @@ def dispatch_execute_script(self, cursor: "AdbcRawCursor", statement: "SQL") -> is_script_result=True, ) + def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "SyncRowStream[dict[str, Any]] | None": + """Return a native ADBC stream backed by ``fetch_record_batch()``.""" + _ = chunk_size + if not statement.returns_rows(): + return None + sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) + return SyncRowStream(AdbcSelectStreamSource(self, sql, prepared_parameters)) + # ───────────────────────────────────────────────────────────────────────────── # TRANSACTION MANAGEMENT # ───────────────────────────────────────────────────────────────────────────── diff --git a/sqlspec/adapters/adbc/type_converter.py b/sqlspec/adapters/adbc/type_converter.py index 21c3aef2f..e392690c2 100644 --- a/sqlspec/adapters/adbc/type_converter.py +++ b/sqlspec/adapters/adbc/type_converter.py @@ -7,63 +7,24 @@ from typing import Any -from sqlspec.core.type_converter import CachedOutputConverter from sqlspec.utils.serializers import to_json -__all__ = ("ADBC_SPECIAL_CHARS", "ADBCOutputConverter", "get_adbc_type_converter") +__all__ = ("ADBCOutputConverter", "get_adbc_type_converter") -ADBC_SPECIAL_CHARS: "frozenset[str]" = frozenset({"{", "[", "-", ":", "T", "."}) -# Native type support by dialect -_NATIVE_SUPPORT: "dict[str, list[str]]" = { - "postgres": ["uuid", "json", "interval", "pg_array"], - "postgresql": ["uuid", "json", "interval", "pg_array"], - "pgvector": ["uuid", "json", "interval", "pg_array"], - "paradedb": ["uuid", "json", "interval", "pg_array"], - "duckdb": ["uuid", "json"], - "bigquery": ["json"], - "sqlite": [], - "mysql": ["json"], - "snowflake": ["json"], -} - - -class ADBCOutputConverter(CachedOutputConverter): - """ADBC-specific output conversion with dialect awareness. - - Extends CachedOutputConverter with ADBC multi-backend functionality - including dialect-specific type handling for different database systems. - """ +class ADBCOutputConverter: + """ADBC-specific parameter conversion with dialect awareness.""" __slots__ = ("dialect",) - def __init__(self, dialect: str, cache_size: int = 5000) -> None: + def __init__(self, dialect: str) -> None: """Initialize with dialect-specific configuration. Args: dialect: Target database dialect (postgres, sqlite, duckdb, etc.) - cache_size: Maximum number of string values to cache (default: 5000) """ - super().__init__(special_chars=ADBC_SPECIAL_CHARS, cache_size=cache_size) self.dialect = dialect.lower() - def _convert_detected(self, value: str, detected_type: str) -> Any: - """Convert value with dialect-specific handling. - - Args: - value: String value to convert. - detected_type: Detected type name. - - Returns: - Converted value according to dialect requirements. - """ - try: - if self.dialect == "sqlite" and detected_type == "uuid": - return str(value) - return self.convert_value(value, detected_type) - except Exception: - return value - def convert_dict(self, value: "dict[str, Any]") -> Any: """Convert dictionary values with dialect-specific handling. @@ -77,42 +38,14 @@ def convert_dict(self, value: "dict[str, Any]") -> Any: return to_json(value) return value - def supports_native_type(self, type_name: str) -> bool: - """Check if dialect supports native handling of a type. - - Args: - type_name: Type name to check - - Returns: - True if dialect supports native handling, False otherwise. - """ - return type_name in _NATIVE_SUPPORT.get(self.dialect, []) - - def get_dialect_specific_converter(self, value: Any, target_type: str) -> Any: - """Apply dialect-specific conversion logic. - - Args: - value: Value to convert. - target_type: Target type for conversion. - - Returns: - Converted value according to dialect requirements. - """ - if self.dialect == "sqlite" and target_type == "uuid": - return str(value) - if self.dialect == "bigquery" and target_type == "uuid": - return str(self.convert_value(value, target_type)) - return self.convert_value(value, target_type) - -def get_adbc_type_converter(dialect: str, cache_size: int = 5000) -> ADBCOutputConverter: +def get_adbc_type_converter(dialect: str) -> ADBCOutputConverter: """Factory function to create dialect-specific ADBC type converter. Args: dialect: Database dialect name. - cache_size: Maximum number of string values to cache (default: 5000) Returns: Configured ADBCOutputConverter instance. """ - return ADBCOutputConverter(dialect, cache_size) + return ADBCOutputConverter(dialect) diff --git a/sqlspec/adapters/aiomysql/config.py b/sqlspec/adapters/aiomysql/config.py index e589eabea..2f01feb95 100644 --- a/sqlspec/adapters/aiomysql/config.py +++ b/sqlspec/adapters/aiomysql/config.py @@ -62,7 +62,6 @@ class AiomysqlConnectionParams(TypedDict): allow_local_infile: NotRequired[bool] echo: NotRequired[bool] local_infile: NotRequired[bool] - enable_local_infile: NotRequired[bool] ssl: NotRequired["SSLContext"] sql_mode: NotRequired[str] init_command: NotRequired[str] @@ -86,12 +85,12 @@ class AiomysqlPoolParams(AiomysqlConnectionParams): def _normalize_local_infile(connection_config: "Mapping[str, Any]", *, strip_consent_gate: bool) -> "dict[str, Any]": - """Normalize aiomysql local-infile aliases and SQLSpec's consent gate.""" + """Normalize aiomysql local-infile settings and SQLSpec's consent gate.""" config = dict(connection_config) - enable_local_infile = bool(config.get("enable_local_infile", False)) + config.pop("enable_local_infile", None) allow_local_infile = bool(config.get(_AIOMYSQL_LOCAL_INFILE_GATE, False)) - local_infile = bool(config.get("local_infile", False) or enable_local_infile) + local_infile = bool(config.get("local_infile", False)) if local_infile and not allow_local_infile: msg = ( "Aiomysql local_infile=True requires allow_local_infile=True because " @@ -99,7 +98,6 @@ def _normalize_local_infile(connection_config: "Mapping[str, Any]", *, strip_con ) raise ImproperConfigurationError(msg) config["local_infile"] = bool(local_infile and allow_local_infile) - config.pop("enable_local_infile", None) if strip_consent_gate: config.pop(_AIOMYSQL_LOCAL_INFILE_GATE, None) return config diff --git a/sqlspec/adapters/aiomysql/core.py b/sqlspec/adapters/aiomysql/core.py index fe0fc2ea9..63f4d9618 100644 --- a/sqlspec/adapters/aiomysql/core.py +++ b/sqlspec/adapters/aiomysql/core.py @@ -2,7 +2,7 @@ import contextlib from collections.abc import Callable, Sized -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast from sqlspec.core import DriverParameterProfile, ParameterStyle, StatementConfig, build_statement_config_from_profile from sqlspec.driver import rows_to_dicts @@ -39,6 +39,7 @@ "build_profile", "build_statement_config", "collect_rows", + "collect_stream_rows", "create_mapped_exception", "default_statement_config", "detect_json_columns", @@ -51,6 +52,7 @@ "normalize_lastrowid", "resolve_column_names", "resolve_many_rowcount", + "resolve_row_plan", "resolve_rowcount", ) @@ -177,15 +179,16 @@ def normalize_execute_parameters(parameters: Any) -> Any: class AiomysqlStreamSource: """Compiled async chunk source streaming dict rows from an aiomysql unbuffered ``SSCursor``.""" - __slots__ = ("_chunk_size", "_column_names", "_cursor", "_driver", "_parameters", "_sql") + __slots__ = ("_chunk_size", "_cursor", "_driver", "_json_type_codes", "_parameters", "_row_plan", "_sql") - def __init__(self, driver: Any, sql: str, parameters: Any, chunk_size: int) -> None: + def __init__(self, driver: Any, sql: str, parameters: Any, chunk_size: int, json_type_codes: "set[int]") -> None: self._driver = driver self._sql = sql self._parameters = parameters self._chunk_size = chunk_size self._cursor: Any = None - self._column_names: list[str] | None = None + self._json_type_codes = json_type_codes + self._row_plan: tuple[list[str], list[int] | None] | None = None async def start(self) -> None: from aiomysql import SSCursor @@ -195,6 +198,7 @@ async def start(self) -> None: cursor = await self._driver.connection.cursor(SSCursor) await cursor.execute(self._sql, normalize_execute_parameters(self._parameters)) self._cursor = cursor + self._row_plan = resolve_row_plan(self._cursor.description, self._json_type_codes) self._driver._check_pending_exception(handler) async def fetch_chunk(self) -> "list[dict[str, Any]]": @@ -205,9 +209,10 @@ async def fetch_chunk(self) -> "list[dict[str, Any]]": self._driver._check_pending_exception(handler) if not rows: return [] - if self._column_names is None: - self._column_names = [description[0] for description in self._cursor.description] - return rows_to_dicts(rows, self._column_names) + if self._row_plan is None: + self._row_plan = resolve_row_plan(self._cursor.description, self._json_type_codes) + deserializer = cast("Callable[[Any], Any]", self._driver.driver_features.get("json_deserializer", from_json)) + return collect_stream_rows(rows, self._row_plan, deserializer) async def close(self) -> None: cursor = self._cursor @@ -332,7 +337,7 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> "SQLSpe if error_code in _MYSQL_MIGRATION_ERROR_CODES: if logger is not None: - logger.warning("aiomysql MySQL expected migration error (ignoring): %s", error) + logger.warning("MySQL expected migration error (ignoring): %s", error) return True dispatch = _MYSQL_SQLSTATE_EXACT_DISPATCH.get(sqlstate) if sqlstate is not None else None @@ -388,23 +393,36 @@ def resolve_column_names(description: "Sequence[Any] | None") -> "list[str]": return [desc[0] for desc in description] -def detect_json_columns_from_description( +def resolve_row_plan( description: "Sequence[Any] | None", json_type_codes: "set[int]" -) -> "list[int]": - """Identify JSON column indexes from pre-fetched cursor description metadata.""" - if not description or not json_type_codes: - return [] +) -> "tuple[list[str], list[int] | None]": + """Resolve ordered column names and JSON column indexes in one pass.""" + if not description: + return [], None + + column_names: list[str] = [] + if not json_type_codes: + column_names.extend(column[0] for column in description) + return column_names, None json_indexes: list[int] = [] - append = json_indexes.append + append_json = json_indexes.append for index, column in enumerate(description): + column_names.append(column[0]) if isinstance(column, (tuple, list)): type_code = column[1] if len(column) > 1 else None else: type_code = column.type_code if isinstance(column, HasTypeCodeProtocol) else None if type_code in json_type_codes: - append(index) - return json_indexes + append_json(index) + return column_names, json_indexes or None + + +def detect_json_columns_from_description( + description: "Sequence[Any] | None", json_type_codes: "set[int]" +) -> "list[int]": + """Identify JSON column indexes from pre-fetched cursor description metadata.""" + return resolve_row_plan(description, json_type_codes)[1] or [] def detect_json_columns( @@ -515,46 +533,53 @@ def _deserialize_json_tuple_rows( return result -def collect_rows( +def collect_stream_rows( fetched_data: "Sequence[Any] | None", - description: "Sequence[Any] | None", - json_indexes: "list[int]", + row_plan: "tuple[list[str], list[int] | None]", deserializer: "Callable[[Any], Any]", *, - column_names: "list[str] | None" = None, logger: Any | None = None, -) -> "tuple[list[Any], list[str], str]": - """Collect aiomysql rows with JSON decoding, preserving raw format. +) -> "list[dict[str, Any]]": + """Materialize streamed rows as dicts and decode JSON columns in place.""" + column_names, json_indexes = row_plan + if not column_names or not fetched_data: + return [] - Args: - fetched_data: Rows returned from cursor.fetchall(). - description: Cursor description metadata. - json_indexes: Column indexes containing JSON values. - deserializer: JSON deserializer function. - column_names: Optional precomputed ordered column names. - logger: Optional logger for debug output. + rows = fetched_data if isinstance(fetched_data, list) else list(fetched_data) + if not rows: + return [] - Returns: - Tuple of (rows, column_names, row_format). - """ - if not description: + dict_rows = rows_to_dicts(rows, column_names) + if json_indexes: + dict_rows = _deserialize_json_dict_rows(column_names, dict_rows, json_indexes, deserializer, logger=logger) + return dict_rows + + +def collect_rows( + fetched_data: "Sequence[Any] | None", + row_plan: "tuple[list[str], list[int] | None]", + deserializer: "Callable[[Any], Any]", + *, + logger: Any | None = None, +) -> "tuple[list[Any], list[str], str]": + """Collect aiomysql rows with JSON decoding, preserving raw format.""" + column_names, json_indexes = row_plan + if not column_names: return [], [], "tuple" - resolved_column_names = resolve_column_names(description) if column_names is None else column_names if not fetched_data: - return [], resolved_column_names, "tuple" + return [], column_names, "tuple" first_row = fetched_data[0] if isinstance(first_row, dict): - if not json_indexes: - rows = fetched_data if isinstance(fetched_data, list) else list(fetched_data) - return rows, resolved_column_names, "dict" - rows = [dict(row) for row in fetched_data] - rows = _deserialize_json_dict_rows(resolved_column_names, rows, json_indexes, deserializer, logger=logger) - return rows, resolved_column_names, "dict" + rows = fetched_data if isinstance(fetched_data, list) else list(fetched_data) + if json_indexes: + rows = _deserialize_json_dict_rows(column_names, rows, json_indexes, deserializer, logger=logger) + return rows, column_names, "dict" + rows = fetched_data if isinstance(fetched_data, list) else list(fetched_data) if json_indexes: rows = _deserialize_json_tuple_rows(rows, json_indexes, deserializer, logger=logger) - return rows, resolved_column_names, "tuple" + return rows, column_names, "tuple" def resolve_rowcount(cursor: Any) -> int: diff --git a/sqlspec/adapters/aiomysql/driver.py b/sqlspec/adapters/aiomysql/driver.py index f450619cc..50717af8b 100644 --- a/sqlspec/adapters/aiomysql/driver.py +++ b/sqlspec/adapters/aiomysql/driver.py @@ -26,15 +26,14 @@ collect_rows, create_mapped_exception, default_statement_config, - detect_json_columns_from_description, driver_profile, encode_records_for_local_infile, format_identifier, normalize_execute_many_parameters, normalize_execute_parameters, normalize_lastrowid, - resolve_column_names, resolve_many_rowcount, + resolve_row_plan, resolve_rowcount, ) from sqlspec.adapters.aiomysql.data_dictionary import AiomysqlDataDictionary @@ -136,12 +135,9 @@ async def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionRes if statement.returns_rows(): fetched_data = await cursor.fetchall() description = cursor.description or None - column_names = resolve_column_names(description) - json_indexes = detect_json_columns_from_description(description, AIOMYSQL_JSON_TYPE_CODES) + row_plan = resolve_row_plan(description, AIOMYSQL_JSON_TYPE_CODES) deserializer = cast("Callable[[Any], Any]", self.driver_features.get("json_deserializer", from_json)) - rows, column_names, row_format = collect_rows( - fetched_data, description, json_indexes, deserializer, column_names=column_names, logger=logger - ) + rows, column_names, row_format = collect_rows(fetched_data, row_plan, deserializer, logger=logger) return self.create_execution_result( cursor, @@ -261,7 +257,9 @@ def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "AsyncRow if not statement.returns_rows(): return None sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - return AsyncRowStream(AiomysqlStreamSource(self, sql, prepared_parameters, chunk_size)) + return AsyncRowStream( + AiomysqlStreamSource(self, sql, prepared_parameters, chunk_size, AIOMYSQL_JSON_TYPE_CODES) + ) def handle_database_exceptions(self) -> "AiomysqlExceptionHandler": """Provide exception handling context manager. @@ -393,12 +391,9 @@ def data_dictionary(self) -> "AiomysqlDataDictionary": def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect aiomysql rows for the direct execution path.""" description = cursor.description or None - column_names = resolve_column_names(description) - json_indexes = detect_json_columns_from_description(description, AIOMYSQL_JSON_TYPE_CODES) + row_plan = resolve_row_plan(description, AIOMYSQL_JSON_TYPE_CODES) deserializer = cast("Callable[[Any], Any]", self.driver_features.get("json_deserializer", from_json)) - rows, column_names, _row_format = collect_rows( - fetched, description, json_indexes, deserializer, column_names=column_names, logger=logger - ) + rows, column_names, _row_format = collect_rows(fetched, row_plan, deserializer, logger=logger) return rows, column_names, len(rows) def resolve_rowcount(self, cursor: Any) -> int: diff --git a/sqlspec/adapters/aiosqlite/core.py b/sqlspec/adapters/aiosqlite/core.py index 544a78d5c..96b801d8b 100644 --- a/sqlspec/adapters/aiosqlite/core.py +++ b/sqlspec/adapters/aiosqlite/core.py @@ -30,7 +30,7 @@ from sqlspec.utils.type_guards import has_rowcount, has_sqlite_error if TYPE_CHECKING: - from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence + from collections.abc import Awaitable, Callable, Mapping, Sequence from sqlspec.adapters.aiosqlite._typing import AiosqliteConnection @@ -47,6 +47,8 @@ "create_mapped_exception", "default_statement_config", "driver_profile", + "execute_and_resolve_rowcount", + "execute_fetchall_with_description", "format_identifier", "normalize_execute_many_parameters", "normalize_execute_parameters", @@ -85,6 +87,32 @@ async def run_on_worker_thread( return await execute(function, *args, **kwargs) +def execute_fetchall_with_description( + connection: "AiosqliteConnection", sql: str, parameters: Any +) -> "tuple[list[Any], Any]": + """Execute a query and return raw rows plus description on the worker thread.""" + raw_connection = connection._conn # pyright: ignore[reportPrivateUsage] + cursor = raw_connection.execute(sql, normalize_execute_parameters(parameters)) + try: + return cursor.fetchall(), cursor.description + finally: + with contextlib.suppress(Exception): + cursor.close() + + +def execute_and_resolve_rowcount(connection: "AiosqliteConnection", sql: str, parameters: Any) -> int: + """Execute a statement and resolve rowcount on the worker thread.""" + raw_connection = connection._conn # pyright: ignore[reportPrivateUsage] + cursor = raw_connection.execute(sql, normalize_execute_parameters(parameters)) + try: + if has_rowcount(cursor) and isinstance(cursor.rowcount, int) and cursor.rowcount > 0: + return cursor.rowcount + return 0 + finally: + with contextlib.suppress(Exception): + cast("Any", cursor).close() + + def format_identifier(identifier: str) -> str: cleaned = identifier.strip() if not cleaned: @@ -100,9 +128,7 @@ def build_insert_statement(table: str, columns: "list[str]") -> str: return f"INSERT INTO {format_identifier(table)} ({column_clause}) VALUES ({placeholders})" -def collect_rows( - fetched_data: "Iterable[Any]", description: "Sequence[Any] | None" -) -> "tuple[list[Any], list[str], int]": +def collect_rows(fetched_data: "list[Any]", description: "Sequence[Any] | None") -> "tuple[list[Any], list[str], int]": """Collect aiosqlite result rows as raw tuples. Returns raw driver-native rows without dict conversion for lazy materialization. @@ -118,8 +144,7 @@ def collect_rows( return [], [], 0 column_names = [col[0] for col in description] - data = list(fetched_data) - return data, column_names, len(data) + return fetched_data, column_names, len(fetched_data) def resolve_rowcount(cursor: Any) -> int: @@ -406,7 +431,7 @@ def apply_driver_features( ) -> "tuple[StatementConfig, dict[str, Any]]": """Apply AIOSQLite driver feature defaults to statement config.""" features: dict[str, Any] = dict(driver_features) if driver_features else {} - features.setdefault("enable_custom_adapters", True) + features.setdefault("enable_custom_adapters", False) json_serializer = features.setdefault("json_serializer", to_json) json_deserializer = features.setdefault("json_deserializer", from_json) diff --git a/sqlspec/adapters/aiosqlite/driver.py b/sqlspec/adapters/aiosqlite/driver.py index 261c7a9ab..aec05a228 100644 --- a/sqlspec/adapters/aiosqlite/driver.py +++ b/sqlspec/adapters/aiosqlite/driver.py @@ -15,21 +15,31 @@ create_mapped_exception, default_statement_config, driver_profile, + execute_and_resolve_rowcount, + execute_fetchall_with_description, format_identifier, normalize_execute_many_parameters, normalize_execute_parameters, resolve_rowcount, + run_on_worker_thread, ) from sqlspec.adapters.aiosqlite.data_dictionary import AiosqliteDataDictionary -from sqlspec.core import ArrowResult, get_cache_config, register_driver_profile +from sqlspec.core import ArrowResult, ParameterStyle, TypedParameter, get_cache_config, register_driver_profile +from sqlspec.core.result import DMLResult from sqlspec.driver import AsyncDriverAdapterBase, AsyncRowStream, BaseAsyncExceptionHandler from sqlspec.exceptions import SQLSpecError +from sqlspec.utils.type_guards import resolve_row_format if TYPE_CHECKING: + from collections.abc import Sequence + from sqlspec.adapters.aiosqlite._typing import AiosqliteConnection - from sqlspec.core import SQL, StatementConfig + from sqlspec.builder import QueryBuilder + from sqlspec.core import SQL, SQLResult, Statement, StatementConfig, StatementFilter + from sqlspec.core.compiler import OperationType from sqlspec.driver import ExecutionResult from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry + from sqlspec.typing import StatementParameters __all__ = ( "AiosqliteCursor", @@ -88,14 +98,14 @@ def __init__( async def dispatch_execute(self, cursor: "AiosqliteRawCursor", statement: "SQL") -> "ExecutionResult": """Execute single SQL statement.""" sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - await cursor.execute(sql, normalize_execute_parameters(prepared_parameters)) + normalized_parameters = normalize_execute_parameters(prepared_parameters) if statement.returns_rows(): - fetched_data = await cursor.fetchall() - - # aiosqlite returns Iterable[Row], core helper expects Iterable[Any] - # Use cast to satisfy mypy and pyright - data, column_names, row_count = collect_rows(cast("list[Any]", fetched_data), cursor.description) + fetched_data, description = await run_on_worker_thread( + self.connection, execute_fetchall_with_description, self.connection, sql, normalized_parameters + ) + data, column_names, row_count = collect_rows(fetched_data, description) + row_format = resolve_row_format(data) return self.create_execution_result( cursor, @@ -103,10 +113,12 @@ async def dispatch_execute(self, cursor: "AiosqliteRawCursor", statement: "SQL") column_names=column_names, data_row_count=row_count, is_select_result=True, - row_format="tuple", + row_format=row_format, ) - affected_rows = resolve_rowcount(cursor) + affected_rows = await run_on_worker_thread( + self.connection, execute_and_resolve_rowcount, self.connection, sql, normalized_parameters + ) return self.create_execution_result(cursor, rowcount_override=affected_rows) async def dispatch_execute_many(self, cursor: "AiosqliteRawCursor", statement: "SQL") -> "ExecutionResult": @@ -135,6 +147,185 @@ async def dispatch_execute_script(self, cursor: "AiosqliteRawCursor", statement: last_cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True ) + async def _execute_cache_hit( + self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: Any + ) -> "SQLResult": + """Execute cached queries through the async cursor fast path.""" + prepared_params = self.prepare_driver_parameters(params, self.statement_config, prepared_statement=cached) + normalized_parameters = normalize_execute_parameters(prepared_params) + direct_statement: SQL | None = None + exc_handler = self.handle_database_exceptions() + result: SQLResult | None = None + try: + async with exc_handler, self.with_cursor(self.connection) as cursor: + await cursor.execute(cached.compiled_sql, normalized_parameters) + if cached.operation_profile.returns_rows: + fetched_data = cast("list[Any]", await cursor.fetchall()) + data, column_names, row_count = self.collect_rows(cursor, fetched_data) + execution_result = self.create_execution_result( + cursor, + selected_data=data, + column_names=column_names, + data_row_count=row_count, + is_select_result=True, + row_format=resolve_row_format(data), + ) + direct_statement = self._cached_statement( + sql, + params, + cached, + cast("tuple[Any, ...] | list[Any] | dict[str, Any]", prepared_params), + params_are_simple=True, + compiled_sql=cached.compiled_sql, + ) + result = self.build_statement_result(direct_statement, execution_result) + else: + affected_rows = self.resolve_rowcount(cursor) + result = DMLResult(cached.operation_type, affected_rows) + + self._check_pending_exception(exc_handler) + assert result is not None + return result + finally: + if direct_statement is not None: + self._release_pooled_statement(direct_statement) + + async def execute_many( + self, + statement: "SQL | Statement | QueryBuilder", + /, + parameters: "Sequence[StatementParameters]", + *filters: "StatementParameters | StatementFilter", + statement_config: "StatementConfig | None" = None, + **kwargs: Any, + ) -> "SQLResult": + """Execute many with an AIOSQLite thin path for simple qmark batches.""" + config = statement_config or self.statement_config + if ( + isinstance(statement, str) + and not filters + and not kwargs + and config is self.statement_config + and self.observability.is_idle + and self._can_use_execute_many_thin_path(statement, parameters, config) + ): + try: + cursor = await self.connection.executemany(statement, parameters) + except (aiosqlite.Error, sqlite3.Error) as exc: + raise create_mapped_exception(exc) from exc + + rowcount = cursor.rowcount + affected_rows = rowcount if isinstance(rowcount, int) and rowcount > 0 else 0 + operation = self._resolve_dml_operation_type(statement) + return DMLResult(operation, affected_rows) + return await super().execute_many(statement, parameters, *filters, statement_config=statement_config, **kwargs) + + def _can_use_execute_many_thin_path( + self, statement: str, parameters: "Sequence[StatementParameters]", config: "StatementConfig" + ) -> bool: + if type(parameters) is not list: + return False + if not parameters: + return False + if "?" not in statement: + return False + + parameter_config = config.parameter_config + if parameter_config.default_parameter_style is not ParameterStyle.QMARK: + return False + if ( + parameter_config.default_execution_parameter_style is not None + and parameter_config.default_execution_parameter_style is not ParameterStyle.QMARK + ): + return False + if parameter_config.ast_transformer is not None or parameter_config.output_transformer is not None: + return False + if parameter_config.needs_static_script_compilation: + return False + if config.output_transformer is not None or config.statement_transformers: + return False + + return self._thin_path_parameters_are_eligible(parameters, parameter_config.type_coercion_map) + + @staticmethod + def _thin_path_parameters_are_eligible( + parameters: "list[StatementParameters]", type_coercion_map: "dict[type, Any] | None" + ) -> bool: + first_sequence = AiosqliteDriver._as_sequence_parameter_set(parameters[0]) + if first_sequence is None: + return False + + first_type = type(first_sequence) + row_len = len(first_sequence) + coercion_map = type_coercion_map + has_type_coercion = bool(coercion_map) + + if row_len == 1: + if has_type_coercion and coercion_map is not None: + for param_set in parameters: + sequence = AiosqliteDriver._as_sequence_parameter_set(param_set) + if sequence is None or type(sequence) is not first_type: + return False + if len(sequence) != 1: + return False + value_type = type(sequence[0]) + if value_type is TypedParameter or value_type in coercion_map: + return False + return True + + for param_set in parameters: + sequence = AiosqliteDriver._as_sequence_parameter_set(param_set) + if sequence is None or type(sequence) is not first_type: + return False + if len(sequence) != 1: + return False + if type(sequence[0]) is TypedParameter: + return False + return True + + if has_type_coercion and coercion_map is not None: + for param_set in parameters: + sequence = AiosqliteDriver._as_sequence_parameter_set(param_set) + if sequence is None or type(sequence) is not first_type: + return False + if len(sequence) != row_len: + return False + for value in sequence: + value_type = type(value) + if value_type is TypedParameter or value_type in coercion_map: + return False + return True + + for param_set in parameters: + sequence = AiosqliteDriver._as_sequence_parameter_set(param_set) + if sequence is None or type(sequence) is not first_type: + return False + if len(sequence) != row_len: + return False + for value in sequence: + if type(value) is TypedParameter: + return False + return True + + @staticmethod + def _as_sequence_parameter_set(param_set: "StatementParameters") -> "list[Any] | tuple[Any, ...] | None": + if isinstance(param_set, list): + return param_set + if isinstance(param_set, tuple): + return param_set + return None + + @staticmethod + def _resolve_dml_operation_type(statement: str) -> "OperationType": + command_keyword = statement.lstrip().split(None, 1)[0].upper() if statement.strip() else "COMMAND" + if command_keyword == "INSERT": + return "INSERT" + if command_keyword == "UPDATE": + return "UPDATE" + if command_keyword == "DELETE": + return "DELETE" + return "COMMAND" + # ───────────────────────────────────────────────────────────────────────────── # TRANSACTION MANAGEMENT # ───────────────────────────────────────────────────────────────────────────── diff --git a/sqlspec/adapters/aiosqlite/pool.py b/sqlspec/adapters/aiosqlite/pool.py index df67785a6..47cf901ab 100644 --- a/sqlspec/adapters/aiosqlite/pool.py +++ b/sqlspec/adapters/aiosqlite/pool.py @@ -49,9 +49,15 @@ def _resolve_row_factory(row_factory: Any) -> Any: return row_factory +def _has_active_transaction(connection: "AiosqliteConnection") -> bool: + return bool(getattr(connection, "in_transaction", False)) + + async def _apply_runtime_setup(connection: "AiosqliteConnection", runtime_setup: "dict[str, Any]") -> None: - for pragma_name, pragma_value in runtime_setup.get("pragmas", ()): - await connection.execute(f"PRAGMA {pragma_name} = {pragma_value}") + pragmas = runtime_setup.get("pragmas", ()) + if pragmas: + pragma_script = "\n".join(f"PRAGMA {pragma_name} = {pragma_value};" for pragma_name, pragma_value in pragmas) + await connection.executescript(pragma_script) extensions = runtime_setup.get("extensions") if extensions: @@ -192,6 +198,8 @@ async def reset(self) -> None: """Reset connection to clean state.""" if self._closed: return + if not _has_active_transaction(self.connection): + return with suppress(Exception): await self.connection.rollback() @@ -200,8 +208,9 @@ async def close(self) -> None: if self._closed: return try: - with suppress(Exception): - await self.connection.rollback() + if _has_active_transaction(self.connection): + with suppress(Exception): + await self.connection.rollback() await self.connection.close() except Exception: # Note: No pool context available at connection level @@ -460,20 +469,23 @@ async def _create_connection(self) -> AiosqlitePoolConnection: is_memory_db = ":memory:" in database_path or "mode=memory" in database_path try: + pragma_lines: list[str] = [] if is_memory_db: - await connection.execute("PRAGMA journal_mode = MEMORY") - await connection.execute("PRAGMA synchronous = OFF") - await connection.execute("PRAGMA temp_store = MEMORY") - await connection.execute("PRAGMA cache_size = -16000") + pragma_lines.extend([ + "PRAGMA journal_mode = MEMORY;", + "PRAGMA synchronous = OFF;", + "PRAGMA temp_store = MEMORY;", + "PRAGMA cache_size = -16000;", + ]) else: - await connection.execute("PRAGMA journal_mode = WAL") - await connection.execute("PRAGMA synchronous = NORMAL") + pragma_lines.extend(["PRAGMA journal_mode = WAL;", "PRAGMA synchronous = NORMAL;"]) - await connection.execute("PRAGMA foreign_keys = ON") - await connection.execute("PRAGMA busy_timeout = 30000") + pragma_lines.extend(["PRAGMA foreign_keys = ON;", "PRAGMA busy_timeout = 30000;"]) if is_shared_cache and is_memory_db: - await connection.execute("PRAGMA read_uncommitted = ON") + pragma_lines.append("PRAGMA read_uncommitted = ON;") + + await connection.executescript("\n".join(pragma_lines)) await connection.commit() @@ -769,8 +781,9 @@ async def release(self, connection: AiosqlitePoolConnection) -> None: try: # Fast path: skip timeout wrapper for reset, just do the rollback directly # The rollback itself is fast for SQLite; timeout is overkill for hot path - with suppress(Exception): - await connection.connection.rollback() + if _has_active_transaction(connection.connection): + with suppress(Exception): + await connection.connection.rollback() connection.idle_since = time.time() # mark_as_idle inline self._queue.put_nowait(connection) except Exception as e: diff --git a/sqlspec/adapters/arrow_odbc/config.py b/sqlspec/adapters/arrow_odbc/config.py index e52b12b46..831e466dd 100644 --- a/sqlspec/adapters/arrow_odbc/config.py +++ b/sqlspec/adapters/arrow_odbc/config.py @@ -19,6 +19,7 @@ from sqlspec.extensions.events import EventRuntimeHints from sqlspec.protocols import SupportsCloseProtocol from sqlspec.utils.config_tools import normalize_connection_config +from sqlspec.utils.serializers import from_json, to_json if TYPE_CHECKING: from collections.abc import Callable @@ -70,6 +71,18 @@ class ArrowOdbcDriverFeatures(TypedDict): on_connection_create: "NotRequired[Callable[[ArrowOdbcConnection], None]]" +def _apply_json_serializer_override(statement_config: "StatementConfig", features: dict[str, Any]) -> "StatementConfig": + serializer = cast("Callable[[Any], str] | None", features.get("json_serializer")) + deserializer = cast("Callable[[str], Any] | None", features.get("json_deserializer")) + if serializer is to_json and deserializer is from_json: + return statement_config + return statement_config.replace( + parameter_config=statement_config.parameter_config.with_json_serializers( + serializer or to_json, deserializer=deserializer + ) + ) + + class ArrowOdbcConnectionContext(SyncPoolConnectionContext): """Context manager for arrow-odbc connections.""" @@ -121,6 +134,7 @@ class ArrowOdbcConfig(NoPoolSyncConfig[ArrowOdbcConnection, ArrowOdbcDriver]): supports_native_arrow_export: "ClassVar[bool]" = True supports_native_arrow_import: "ClassVar[bool]" = True supports_arrow_streaming: "ClassVar[bool]" = True + supports_native_row_streaming: "ClassVar[bool]" = True supports_native_parquet_export: "ClassVar[bool]" = False supports_native_parquet_import: "ClassVar[bool]" = False _connection_context_class: "ClassVar[type[ArrowOdbcConnectionContext]]" = ArrowOdbcConnectionContext @@ -230,7 +244,7 @@ def _close_arrow_odbc_connection(connection: "ArrowOdbcConnection") -> None: def _resolve_statement_config(features: dict[str, Any]) -> "StatementConfig": dialect = resolve_dialect_from_dbms_name(str(features.get("dbms_name") or features.get("connection_string") or "")) if dialect == "sqlite": - return default_statement_config + return _apply_json_serializer_override(default_statement_config, features) if dialect == "mssql": - return build_statement_config(dialect="tsql") - return build_statement_config(dialect=dialect) + return _apply_json_serializer_override(build_statement_config(dialect="tsql"), features) + return _apply_json_serializer_override(build_statement_config(dialect=dialect), features) diff --git a/sqlspec/adapters/arrow_odbc/driver.py b/sqlspec/adapters/arrow_odbc/driver.py index baee6d202..46a92253c 100644 --- a/sqlspec/adapters/arrow_odbc/driver.py +++ b/sqlspec/adapters/arrow_odbc/driver.py @@ -1,5 +1,6 @@ """arrow-odbc sync driver.""" +import contextlib import re from collections.abc import Iterable, Mapping from itertools import chain @@ -20,7 +21,7 @@ get_cache_config, register_driver_profile, ) -from sqlspec.driver import BaseSyncExceptionHandler, SyncDriverAdapterBase +from sqlspec.driver import BaseSyncExceptionHandler, SyncDriverAdapterBase, SyncRowStream from sqlspec.exceptions import ImproperConfigurationError, SQLSpecError from sqlspec.utils.module_loader import ensure_pyarrow from sqlspec.utils.text import quote_identifier, split_qualified_identifier @@ -53,6 +54,49 @@ def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "Ba return False +class ArrowOdbcStreamSource: + """Native Arrow ODBC chunk source backed by Arrow record batches.""" + + __slots__ = ("_chunk_size", "_driver", "_parameters", "_reader", "_sql") + + def __init__( + self, driver: "ArrowOdbcDriver", sql: str, parameters: "list[str | None] | None", chunk_size: int + ) -> None: + self._driver = driver + self._sql = sql + self._parameters = parameters + self._chunk_size = chunk_size + self._reader: Any = None + + def start(self) -> None: + handler = self._driver.handle_database_exceptions() + with handler: + reader = self._driver._read_arrow_batches(self._sql, self._parameters, self._chunk_size) + self._reader = iter(_to_pyarrow_reader(reader)) + self._driver._check_pending_exception(handler) + + def fetch_chunk(self) -> "list[dict[str, Any]]": + reader = self._reader + if reader is None: + return [] + while True: + try: + batch = next(reader) + except StopIteration: + return [] + rows = cast("list[dict[str, Any]]", batch.to_pylist()) + if rows: + return rows + + def close(self) -> None: + reader = self._reader + self._reader = None + close = getattr(reader, "close", None) + if callable(close): + with contextlib.suppress(Exception): + close() + + class ArrowOdbcDriver(SyncDriverAdapterBase): """Sync driver for generic ODBC connections with Arrow-native transfer.""" @@ -146,6 +190,15 @@ def dispatch_execute_script(self, cursor: "ArrowOdbcRawCursor", statement: "SQL" cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True ) + def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "SyncRowStream[dict[str, Any]] | None": + """Return a native Arrow ODBC row stream backed by record batches.""" + if not statement.returns_rows(): + return None + sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) + if self._dialect == "mssql": + sql, prepared_parameters = _inline_mssql_pagination_parameters(sql, prepared_parameters) + return SyncRowStream(ArrowOdbcStreamSource(self, sql, _odbc_parameters(prepared_parameters), chunk_size)) + def collect_rows(self, cursor: "ArrowOdbcRawCursor", fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": return fetched, [], len(fetched) diff --git a/sqlspec/adapters/asyncmy/core.py b/sqlspec/adapters/asyncmy/core.py index b27eff2b2..98eaff50d 100644 --- a/sqlspec/adapters/asyncmy/core.py +++ b/sqlspec/adapters/asyncmy/core.py @@ -2,7 +2,7 @@ import contextlib from collections.abc import Callable, Sized -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast from sqlspec.core import DriverParameterProfile, ParameterStyle, StatementConfig, build_statement_config_from_profile from sqlspec.driver import rows_to_dicts @@ -38,6 +38,7 @@ "build_profile", "build_statement_config", "collect_rows", + "collect_stream_rows", "create_mapped_exception", "default_statement_config", "detect_json_columns", @@ -49,6 +50,7 @@ "normalize_lastrowid", "resolve_column_names", "resolve_many_rowcount", + "resolve_row_plan", "resolve_rowcount", ) @@ -149,15 +151,16 @@ def normalize_execute_parameters(parameters: Any) -> Any: class AsyncmyStreamSource: """Compiled async chunk source streaming dict rows from an asyncmy unbuffered ``SSCursor``.""" - __slots__ = ("_chunk_size", "_column_names", "_cursor", "_driver", "_parameters", "_sql") + __slots__ = ("_chunk_size", "_cursor", "_driver", "_json_type_codes", "_parameters", "_row_plan", "_sql") - def __init__(self, driver: Any, sql: str, parameters: Any, chunk_size: int) -> None: + def __init__(self, driver: Any, sql: str, parameters: Any, chunk_size: int, json_type_codes: "set[int]") -> None: self._driver = driver self._sql = sql self._parameters = parameters self._chunk_size = chunk_size self._cursor: Any = None - self._column_names: list[str] | None = None + self._json_type_codes = json_type_codes + self._row_plan: tuple[list[str], list[int] | None] | None = None async def start(self) -> None: from asyncmy.cursors import SSCursor @@ -167,6 +170,7 @@ async def start(self) -> None: cursor = self._driver.connection.cursor(SSCursor) await cursor.execute(self._sql, normalize_execute_parameters(self._parameters)) self._cursor = cursor + self._row_plan = resolve_row_plan(self._cursor.description, self._json_type_codes) self._driver._check_pending_exception(handler) async def fetch_chunk(self) -> "list[dict[str, Any]]": @@ -177,9 +181,10 @@ async def fetch_chunk(self) -> "list[dict[str, Any]]": self._driver._check_pending_exception(handler) if not rows: return [] - if self._column_names is None: - self._column_names = [description[0] for description in self._cursor.description] - return rows_to_dicts(rows, self._column_names) + if self._row_plan is None: + self._row_plan = resolve_row_plan(self._cursor.description, self._json_type_codes) + deserializer = cast("Callable[[Any], Any]", self._driver.driver_features.get("json_deserializer", from_json)) + return collect_stream_rows(rows, self._row_plan, deserializer) async def close(self) -> None: cursor = self._cursor @@ -304,7 +309,7 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> "SQLSpe if error_code in _MYSQL_MIGRATION_ERROR_CODES: if logger is not None: - logger.warning("AsyncMy MySQL expected migration error (ignoring): %s", error) + logger.warning("MySQL expected migration error (ignoring): %s", error) return True dispatch = _MYSQL_SQLSTATE_EXACT_DISPATCH.get(sqlstate) if sqlstate is not None else None @@ -360,23 +365,36 @@ def resolve_column_names(description: "Sequence[Any] | None") -> "list[str]": return [desc[0] for desc in description] -def detect_json_columns_from_description( +def resolve_row_plan( description: "Sequence[Any] | None", json_type_codes: "set[int]" -) -> "list[int]": - """Identify JSON column indexes from pre-fetched cursor description metadata.""" - if not description or not json_type_codes: - return [] +) -> "tuple[list[str], list[int] | None]": + """Resolve ordered column names and JSON column indexes in one pass.""" + if not description: + return [], None + + column_names: list[str] = [] + if not json_type_codes: + column_names.extend(column[0] for column in description) + return column_names, None json_indexes: list[int] = [] - append = json_indexes.append + append_json = json_indexes.append for index, column in enumerate(description): + column_names.append(column[0]) if isinstance(column, (tuple, list)): type_code = column[1] if len(column) > 1 else None else: type_code = column.type_code if isinstance(column, HasTypeCodeProtocol) else None if type_code in json_type_codes: - append(index) - return json_indexes + append_json(index) + return column_names, json_indexes or None + + +def detect_json_columns_from_description( + description: "Sequence[Any] | None", json_type_codes: "set[int]" +) -> "list[int]": + """Identify JSON column indexes from pre-fetched cursor description metadata.""" + return resolve_row_plan(description, json_type_codes)[1] or [] def detect_json_columns( @@ -487,46 +505,53 @@ def _deserialize_json_tuple_rows( return result -def collect_rows( +def collect_stream_rows( fetched_data: "Sequence[Any] | None", - description: "Sequence[Any] | None", - json_indexes: "list[int]", + row_plan: "tuple[list[str], list[int] | None]", deserializer: "Callable[[Any], Any]", *, - column_names: "list[str] | None" = None, logger: Any | None = None, -) -> "tuple[list[Any], list[str], str]": - """Collect AsyncMy rows with JSON decoding, preserving raw format. +) -> "list[dict[str, Any]]": + """Materialize streamed rows as dicts and decode JSON columns in place.""" + column_names, json_indexes = row_plan + if not column_names or not fetched_data: + return [] - Args: - fetched_data: Rows returned from cursor.fetchall(). - description: Cursor description metadata. - json_indexes: Column indexes containing JSON values. - deserializer: JSON deserializer function. - column_names: Optional precomputed ordered column names. - logger: Optional logger for debug output. + rows = fetched_data if isinstance(fetched_data, list) else list(fetched_data) + if not rows: + return [] - Returns: - Tuple of (rows, column_names, row_format). - """ - if not description: + dict_rows = rows_to_dicts(rows, column_names) + if json_indexes: + dict_rows = _deserialize_json_dict_rows(column_names, dict_rows, json_indexes, deserializer, logger=logger) + return dict_rows + + +def collect_rows( + fetched_data: "Sequence[Any] | None", + row_plan: "tuple[list[str], list[int] | None]", + deserializer: "Callable[[Any], Any]", + *, + logger: Any | None = None, +) -> "tuple[list[Any], list[str], str]": + """Collect AsyncMy rows with JSON decoding, preserving raw format.""" + column_names, json_indexes = row_plan + if not column_names: return [], [], "tuple" - resolved_column_names = resolve_column_names(description) if column_names is None else column_names if not fetched_data: - return [], resolved_column_names, "tuple" + return [], column_names, "tuple" first_row = fetched_data[0] if isinstance(first_row, dict): - if not json_indexes: - rows = fetched_data if isinstance(fetched_data, list) else list(fetched_data) - return rows, resolved_column_names, "dict" - rows = [dict(row) for row in fetched_data] - rows = _deserialize_json_dict_rows(resolved_column_names, rows, json_indexes, deserializer, logger=logger) - return rows, resolved_column_names, "dict" + rows = fetched_data if isinstance(fetched_data, list) else list(fetched_data) + if json_indexes: + rows = _deserialize_json_dict_rows(column_names, rows, json_indexes, deserializer, logger=logger) + return rows, column_names, "dict" + rows = fetched_data if isinstance(fetched_data, list) else list(fetched_data) if json_indexes: rows = _deserialize_json_tuple_rows(rows, json_indexes, deserializer, logger=logger) - return rows, resolved_column_names, "tuple" + return rows, column_names, "tuple" def resolve_rowcount(cursor: Any) -> int: diff --git a/sqlspec/adapters/asyncmy/driver.py b/sqlspec/adapters/asyncmy/driver.py index 159b5b991..b04fa4d76 100644 --- a/sqlspec/adapters/asyncmy/driver.py +++ b/sqlspec/adapters/asyncmy/driver.py @@ -20,14 +20,13 @@ collect_rows, create_mapped_exception, default_statement_config, - detect_json_columns_from_description, driver_profile, format_identifier, normalize_execute_many_parameters, normalize_execute_parameters, normalize_lastrowid, - resolve_column_names, resolve_many_rowcount, + resolve_row_plan, resolve_rowcount, ) from sqlspec.adapters.asyncmy.data_dictionary import AsyncmyDataDictionary @@ -129,12 +128,9 @@ async def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionRes if statement.returns_rows(): fetched_data = await cursor.fetchall() description = cursor.description or None - column_names = resolve_column_names(description) - json_indexes = detect_json_columns_from_description(description, ASYNCMY_JSON_TYPE_CODES) + row_plan = resolve_row_plan(description, ASYNCMY_JSON_TYPE_CODES) deserializer = cast("Callable[[Any], Any]", self.driver_features.get("json_deserializer", from_json)) - rows, column_names, row_format = collect_rows( - fetched_data, description, json_indexes, deserializer, column_names=column_names, logger=logger - ) + rows, column_names, row_format = collect_rows(fetched_data, row_plan, deserializer, logger=logger) return self.create_execution_result( cursor, @@ -262,7 +258,7 @@ def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "AsyncRow if not statement.returns_rows(): return None sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - return AsyncRowStream(AsyncmyStreamSource(self, sql, prepared_parameters, chunk_size)) + return AsyncRowStream(AsyncmyStreamSource(self, sql, prepared_parameters, chunk_size, ASYNCMY_JSON_TYPE_CODES)) def handle_database_exceptions(self) -> "AsyncmyExceptionHandler": """Provide exception handling context manager. @@ -378,12 +374,9 @@ def data_dictionary(self) -> "AsyncmyDataDictionary": def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect asyncmy rows for the direct execution path.""" description = cursor.description or None - column_names = resolve_column_names(description) - json_indexes = detect_json_columns_from_description(description, ASYNCMY_JSON_TYPE_CODES) + row_plan = resolve_row_plan(description, ASYNCMY_JSON_TYPE_CODES) deserializer = cast("Callable[[Any], Any]", self.driver_features.get("json_deserializer", from_json)) - rows, column_names, _row_format = collect_rows( - fetched, description, json_indexes, deserializer, column_names=column_names, logger=logger - ) + rows, column_names, _row_format = collect_rows(fetched, row_plan, deserializer, logger=logger) return rows, column_names, len(rows) def resolve_rowcount(self, cursor: Any) -> int: diff --git a/sqlspec/adapters/bigquery/core.py b/sqlspec/adapters/bigquery/core.py index 7f4fbdd70..7e651fc6d 100644 --- a/sqlspec/adapters/bigquery/core.py +++ b/sqlspec/adapters/bigquery/core.py @@ -864,7 +864,7 @@ def fetch_chunk(self) -> "list[dict[str, Any]]": self._driver._check_pending_exception(handler) if page is None: return [] - rows = [dict(row.items()) for row in page] + rows = [dict(cast("Mapping[str, Any]", row)) for row in page] if rows: return rows @@ -1012,6 +1012,12 @@ def apply_driver_features( features: dict[str, Any] = dict(driver_features) if driver_features else {} features.setdefault("enable_uuid_conversion", True) features.setdefault("json_serializer", to_json) + if not features["enable_uuid_conversion"]: + type_coercion_map = dict(statement_config.parameter_config.type_coercion_map) + for uuid_type in build_uuid_coercions(): + type_coercion_map.pop(uuid_type, None) + parameter_config = statement_config.parameter_config.replace(type_coercion_map=type_coercion_map) + statement_config = statement_config.replace(parameter_config=parameter_config) return statement_config, features diff --git a/sqlspec/adapters/bigquery/type_converter.py b/sqlspec/adapters/bigquery/type_converter.py deleted file mode 100644 index 851b3ccb5..000000000 --- a/sqlspec/adapters/bigquery/type_converter.py +++ /dev/null @@ -1,96 +0,0 @@ -"""BigQuery-specific type conversion with native UUID support. - -Provides specialized type handling for BigQuery, including UUID support -for the native BigQuery driver and parameter creation. -""" - -from typing import Any -from uuid import UUID - -from typing_extensions import final - -from sqlspec.adapters.bigquery.core import _query_parameter_type -from sqlspec.core.type_converter import CachedOutputConverter, convert_uuid - -__all__ = ("BIGQUERY_SPECIAL_CHARS", "BigQueryOutputConverter") - -BIGQUERY_SPECIAL_CHARS: "frozenset[str]" = frozenset({"{", "[", "-", ":", "T", "."}) - - -@final -class BigQueryOutputConverter(CachedOutputConverter): - """BigQuery-specific output conversion with native UUID support. - - Extends CachedOutputConverter with BigQuery-specific functionality - including UUID handling and parameter creation for the native BigQuery driver. - """ - - __slots__ = ("_enable_uuid_conversion",) - - def __init__(self, cache_size: int = 5000, enable_uuid_conversion: bool = True) -> None: - """Initialize converter with BigQuery-specific options. - - Args: - cache_size: Maximum number of string values to cache (default: 5000) - enable_uuid_conversion: Enable automatic UUID string conversion (default: True) - """ - super().__init__(special_chars=BIGQUERY_SPECIAL_CHARS, cache_size=cache_size) - self._enable_uuid_conversion = enable_uuid_conversion - - def _convert_detected(self, value: str, detected_type: str) -> Any: - """Convert value with BigQuery-specific handling. - - Args: - value: String value to convert. - detected_type: Detected type name. - - Returns: - Converted value, respecting UUID conversion setting. - """ - try: - return self.convert_value(value, detected_type) - except Exception: - return value - - def create_parameter(self, name: str, value: Any) -> "Any | None": - """Create BigQuery parameter with proper type mapping. - - Args: - name: Parameter name. - value: Parameter value. - - Returns: - ScalarQueryParameter for native BigQuery driver, None if not available. - """ - try: - from google.cloud.bigquery import ScalarQueryParameter - except ImportError: - return None - - if self._enable_uuid_conversion: - if isinstance(value, UUID): - return ScalarQueryParameter(name, "STRING", str(value)) - - if isinstance(value, str): - detected_type = self.detect_type(value) - if detected_type == "uuid": - uuid_obj = convert_uuid(value) - return ScalarQueryParameter(name, "STRING", str(uuid_obj)) - - param_type, _ = _query_parameter_type(value) - param_type = param_type or "STRING" - return ScalarQueryParameter(name, param_type, value) - - def convert_bigquery_value(self, value: Any, column_type: str) -> Any: - """Convert BigQuery value based on column type. - - Args: - value: Value to convert. - column_type: BigQuery column type. - - Returns: - Converted value appropriate for the column type. - """ - if column_type == "STRING" and isinstance(value, str): - return self.convert(value) - return value diff --git a/sqlspec/adapters/cockroach_asyncpg/driver.py b/sqlspec/adapters/cockroach_asyncpg/driver.py index b8000a464..baa14d943 100644 --- a/sqlspec/adapters/cockroach_asyncpg/driver.py +++ b/sqlspec/adapters/cockroach_asyncpg/driver.py @@ -2,7 +2,7 @@ import asyncio import contextlib -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, TypeVar, cast from sqlspec.adapters.asyncpg.core import create_mapped_exception, driver_profile from sqlspec.adapters.asyncpg.driver import AsyncpgDriver @@ -20,7 +20,7 @@ from sqlspec.utils.type_guards import has_sqlstate if TYPE_CHECKING: - from collections.abc import Callable + from collections.abc import Awaitable, Callable from sqlspec.adapters.cockroach_asyncpg._typing import CockroachAsyncpgConnection from sqlspec.core import StatementConfig @@ -29,6 +29,7 @@ __all__ = ("CockroachAsyncpgDriver", "CockroachAsyncpgExceptionHandler", "CockroachAsyncpgSessionContext") logger = get_logger("sqlspec.adapters.cockroach_asyncpg") +_T = TypeVar("_T") class CockroachAsyncpgExceptionHandler(BaseAsyncExceptionHandler): @@ -66,21 +67,26 @@ def __init__( # Data dictionary is lazily initialized in property; use parent slot self._data_dictionary = None - async def _execute_with_retry(self, operation: "Callable[..., Any]", *args: Any) -> "ExecutionResult": - if not self._enable_retry: - return cast("ExecutionResult", await operation(*args)) + async def run_transaction_with_retry(self, operation: "Callable[[], Awaitable[_T]]") -> _T: + """Execute a full CockroachDB transaction callback with serialization retries.""" + if not self._enable_retry or self._connection_in_transaction(): + return await operation() - last_error: Exception | None = None + last_error: BaseException | None = None for attempt in range(self._retry_config.max_retries + 1): try: - return cast("ExecutionResult", await operation(*args)) + await self.begin() + result = await operation() + await self.commit() except Exception as exc: last_error = exc + with contextlib.suppress(Exception): + await self.rollback() if not is_retryable_error(exc) or attempt >= self._retry_config.max_retries: raise - with contextlib.suppress(Exception): - await self.connection.execute("ROLLBACK") + else: + return result delay = calculate_backoff_seconds(attempt, self._retry_config) if self._retry_config.enable_logging: logger.debug("CockroachDB retry %s/%s after %.3fs", attempt + 1, self._retry_config.max_retries, delay) @@ -112,19 +118,13 @@ async def _dispatch_execute_script_impl( return await AsyncpgDriver.dispatch_execute_script(self, cursor, statement) async def dispatch_execute(self, cursor: Any, statement: SQL) -> "ExecutionResult": - if not self._enable_retry: - return await self._dispatch_execute_impl(cursor, statement) - return await self._execute_with_retry(self._dispatch_execute_impl, cursor, statement) + return await self._dispatch_execute_impl(cursor, statement) async def dispatch_execute_many(self, cursor: Any, statement: SQL) -> "ExecutionResult": - if not self._enable_retry: - return await self._dispatch_execute_many_impl(cursor, statement) - return await self._execute_with_retry(self._dispatch_execute_many_impl, cursor, statement) + return await self._dispatch_execute_many_impl(cursor, statement) async def dispatch_execute_script(self, cursor: Any, statement: SQL) -> "ExecutionResult": - if not self._enable_retry: - return await self._dispatch_execute_script_impl(cursor, statement) - return await self._execute_with_retry(self._dispatch_execute_script_impl, cursor, statement) + return await self._dispatch_execute_script_impl(cursor, statement) def handle_database_exceptions(self) -> "CockroachAsyncpgExceptionHandler": # type: ignore[override] return CockroachAsyncpgExceptionHandler() diff --git a/sqlspec/adapters/cockroach_psycopg/driver.py b/sqlspec/adapters/cockroach_psycopg/driver.py index 4b12cb85d..aba93c4ff 100644 --- a/sqlspec/adapters/cockroach_psycopg/driver.py +++ b/sqlspec/adapters/cockroach_psycopg/driver.py @@ -3,7 +3,7 @@ import asyncio import contextlib import time -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, TypeVar, cast import psycopg @@ -33,7 +33,7 @@ from sqlspec.utils.type_guards import has_sqlstate if TYPE_CHECKING: - from collections.abc import Callable + from collections.abc import Awaitable, Callable from sqlspec.driver import ExecutionResult @@ -47,6 +47,7 @@ ) logger = get_logger("sqlspec.adapters.cockroach_psycopg") +_T = TypeVar("_T") class CockroachPsycopgSyncExceptionHandler(BaseSyncExceptionHandler): @@ -108,21 +109,26 @@ def __init__( # Data dictionary is lazily initialized in property; use parent slot self._data_dictionary = None - def _execute_with_retry(self, operation: "Callable[..., ExecutionResult]", *args: Any) -> "ExecutionResult": - if not self._enable_retry: - return operation(*args) + def run_transaction_with_retry(self, operation: "Callable[[], _T]") -> _T: + """Execute a full CockroachDB transaction callback with serialization retries.""" + if not self._enable_retry or self._connection_in_transaction(): + return operation() - last_error: Exception | None = None + last_error: BaseException | None = None for attempt in range(self._retry_config.max_retries + 1): try: - return operation(*args) + self.begin() + result = operation() + self.commit() except Exception as exc: last_error = exc + with contextlib.suppress(Exception): + self.rollback() if not is_retryable_error(exc) or attempt >= self._retry_config.max_retries: raise - with contextlib.suppress(Exception): - self.connection.rollback() + else: + return result delay = calculate_backoff_seconds(attempt, self._retry_config) if self._retry_config.enable_logging: logger.debug("CockroachDB retry %s/%s after %.3fs", attempt + 1, self._retry_config.max_retries, delay) @@ -150,13 +156,13 @@ def _dispatch_execute_script_impl(self, cursor: Any, statement: SQL) -> "Executi return PsycopgSyncDriver.dispatch_execute_script(self, cursor, statement) def dispatch_execute(self, cursor: Any, statement: SQL) -> "ExecutionResult": - return self._execute_with_retry(self._dispatch_execute_impl, cursor, statement) + return self._dispatch_execute_impl(cursor, statement) def dispatch_execute_many(self, cursor: Any, statement: SQL) -> "ExecutionResult": - return self._execute_with_retry(self._dispatch_execute_many_impl, cursor, statement) + return self._dispatch_execute_many_impl(cursor, statement) def dispatch_execute_script(self, cursor: Any, statement: SQL) -> "ExecutionResult": - return self._execute_with_retry(self._dispatch_execute_script_impl, cursor, statement) + return self._dispatch_execute_script_impl(cursor, statement) def handle_database_exceptions(self) -> "CockroachPsycopgSyncExceptionHandler": # type: ignore[override] return CockroachPsycopgSyncExceptionHandler() @@ -194,21 +200,26 @@ def __init__( # Data dictionary is lazily initialized in property; use parent slot self._data_dictionary = None - async def _execute_with_retry(self, operation: "Callable[..., Any]", *args: Any) -> "ExecutionResult": - if not self._enable_retry: - return cast("ExecutionResult", await operation(*args)) + async def run_transaction_with_retry(self, operation: "Callable[[], Awaitable[_T]]") -> _T: + """Execute a full CockroachDB transaction callback with serialization retries.""" + if not self._enable_retry or self._connection_in_transaction(): + return await operation() - last_error: Exception | None = None + last_error: BaseException | None = None for attempt in range(self._retry_config.max_retries + 1): try: - return cast("ExecutionResult", await operation(*args)) + await self.begin() + result = await operation() + await self.commit() except Exception as exc: last_error = exc + with contextlib.suppress(Exception): + await self.rollback() if not is_retryable_error(exc) or attempt >= self._retry_config.max_retries: raise - with contextlib.suppress(Exception): - await self.connection.rollback() + else: + return result delay = calculate_backoff_seconds(attempt, self._retry_config) if self._retry_config.enable_logging: logger.debug("CockroachDB retry %s/%s after %.3fs", attempt + 1, self._retry_config.max_retries, delay) @@ -236,13 +247,13 @@ async def _dispatch_execute_script_impl(self, cursor: Any, statement: SQL) -> "E return await PsycopgAsyncDriver.dispatch_execute_script(self, cursor, statement) async def dispatch_execute(self, cursor: Any, statement: SQL) -> "ExecutionResult": - return await self._execute_with_retry(self._dispatch_execute_impl, cursor, statement) + return await self._dispatch_execute_impl(cursor, statement) async def dispatch_execute_many(self, cursor: Any, statement: SQL) -> "ExecutionResult": - return await self._execute_with_retry(self._dispatch_execute_many_impl, cursor, statement) + return await self._dispatch_execute_many_impl(cursor, statement) async def dispatch_execute_script(self, cursor: Any, statement: SQL) -> "ExecutionResult": - return await self._execute_with_retry(self._dispatch_execute_script_impl, cursor, statement) + return await self._dispatch_execute_script_impl(cursor, statement) def handle_database_exceptions(self) -> "CockroachPsycopgAsyncExceptionHandler": # type: ignore[override] return CockroachPsycopgAsyncExceptionHandler() diff --git a/sqlspec/adapters/duckdb/core.py b/sqlspec/adapters/duckdb/core.py index d3eedc377..7c78cb9c0 100644 --- a/sqlspec/adapters/duckdb/core.py +++ b/sqlspec/adapters/duckdb/core.py @@ -4,7 +4,6 @@ from decimal import Decimal from typing import TYPE_CHECKING, Any, Final, cast -from sqlspec.adapters.duckdb.type_converter import DuckDBOutputConverter from sqlspec.core import DriverParameterProfile, ParameterStyle, StatementConfig, build_statement_config_from_profile from sqlspec.exceptions import ( CheckViolationError, @@ -173,8 +172,6 @@ def apply_driver_features( ) -> "tuple[StatementConfig, dict[str, Any]]": """Apply DuckDB-specific driver features to statement configuration.""" features: dict[str, Any] = dict(driver_features) if driver_features else {} - if not features: - return statement_config, features param_config = statement_config.parameter_config json_serializer = features.get("json_serializer") @@ -183,11 +180,10 @@ def apply_driver_features( cast("Callable[[Any], str]", json_serializer), tuple_strategy="tuple" ) - enable_uuid_conversion = features.get("enable_uuid_conversion", True) - if not enable_uuid_conversion: - type_converter = DuckDBOutputConverter(enable_uuid_conversion=enable_uuid_conversion) + if not features.get("enable_uuid_conversion", True): type_coercion_map = dict(param_config.type_coercion_map) - type_coercion_map[str] = type_converter.convert_if_detected + for uuid_type in build_uuid_coercions(): + type_coercion_map.pop(uuid_type, None) param_config = param_config.replace(type_coercion_map=type_coercion_map) if param_config is statement_config.parameter_config: diff --git a/sqlspec/adapters/duckdb/driver.py b/sqlspec/adapters/duckdb/driver.py index d088912f1..551022c11 100644 --- a/sqlspec/adapters/duckdb/driver.py +++ b/sqlspec/adapters/duckdb/driver.py @@ -2,9 +2,10 @@ import contextlib from typing import TYPE_CHECKING, Any, cast -from uuid import uuid4 +from uuid import UUID, uuid4 import duckdb +from sqlglot import exp from sqlspec.adapters.duckdb._typing import DuckDBCursor, DuckDBSessionContext from sqlspec.adapters.duckdb.core import ( @@ -24,6 +25,7 @@ get_cache_config, register_driver_profile, ) +from sqlspec.core.result import DMLResult from sqlspec.driver import BaseSyncExceptionHandler, SyncDriverAdapterBase from sqlspec.exceptions import SQLSpecError from sqlspec.utils.logging import get_logger @@ -31,9 +33,11 @@ from sqlspec.utils.text import quote_identifier if TYPE_CHECKING: + from collections.abc import Sequence + from sqlspec.adapters.duckdb._typing import DuckDBConnection from sqlspec.builder import QueryBuilder - from sqlspec.core import ArrowResult, Statement, StatementFilter + from sqlspec.core import ArrowResult, SQLResult, Statement, StatementFilter from sqlspec.driver import ExecutionResult from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry from sqlspec.typing import ArrowReturnFormat, StatementParameters @@ -116,9 +120,10 @@ def dispatch_execute(self, cursor: "DuckDBConnection", statement: SQL) -> "Execu is_select_like = statement.returns_rows() or self._should_force_select(statement, cursor) if is_select_like: - fetched_data = cursor.fetchall() - data, column_names = collect_rows(cast("list[Any] | None", fetched_data), cursor.description) - row_format = "dict" if data and isinstance(data[0], dict) else "tuple" + arrow_table = cursor.fetch_arrow_table() + data = arrow_table.to_pylist() + _restore_uuid_columns(data, cursor.description) + column_names = list(arrow_table.column_names) return self.create_execution_result( cursor, @@ -126,7 +131,7 @@ def dispatch_execute(self, cursor: "DuckDBConnection", statement: SQL) -> "Execu column_names=column_names, data_row_count=len(data), is_select_result=True, - row_format=row_format, + row_format="dict", ) row_count = resolve_rowcount(cursor) @@ -158,6 +163,28 @@ def dispatch_execute_many(self, cursor: "DuckDBConnection", statement: SQL) -> " return self.create_execution_result(cursor, rowcount_override=row_count, is_many_result=True) + def execute_many( + self, + statement: "SQL | Statement | QueryBuilder", + /, + parameters: "Sequence[StatementParameters]", + *filters: "StatementParameters | StatementFilter", + statement_config: "StatementConfig | None" = None, + **kwargs: Any, + ) -> "SQLResult": + """Execute many with a DuckDB bulk insert fast path for simple INSERT batches.""" + config = statement_config or self.statement_config + if isinstance(statement, str) and not filters and not kwargs and config is self.statement_config: + prepared_statement = SQL(statement, tuple(parameters), statement_config=config, is_many=True) + _, prepared_parameters = self._compiled_sql(prepared_statement, config) + processed_state = prepared_statement.get_processed_state() + parsed_expression = processed_state.parsed_expression + if isinstance(parsed_expression, exp.Insert) and not parsed_expression.args.get("returning"): + bulk_result = self._execute_bulk_insert_many(parsed_expression, prepared_parameters) + if bulk_result is not None: + return bulk_result + return super().execute_many(statement, parameters, *filters, statement_config=statement_config, **kwargs) + def dispatch_execute_script(self, cursor: "DuckDBConnection", statement: SQL) -> "ExecutionResult": """Execute SQL script with statement splitting and parameter handling. @@ -438,6 +465,64 @@ def collect_rows(self, cursor: "DuckDBConnection", fetched: "list[Any]") -> "tup data, column_names = collect_rows(cast("list[Any] | None", fetched), cursor.description) return data, column_names, len(data) + def _execute_bulk_insert_many(self, expression: exp.Insert, prepared_parameters: Any) -> "DMLResult | None": + """Execute a batch INSERT via Arrow registration when the payload is simple.""" + if not isinstance(prepared_parameters, list) or not prepared_parameters: + return None + if not isinstance(expression.this, exp.Schema): + return None + + table_expr = expression.this.this + if not isinstance(table_expr, exp.Table): + return None + + if table_expr.alias: + return None + + rows = prepared_parameters + column_names = [column.name for column in expression.this.expressions] + arrow_table = self._build_arrow_table(rows, column_names) + if arrow_table is None: + return None + + target_table = table_expr.sql(dialect="duckdb") + temp_view = f"_sqlspec_batch_{uuid4().hex}" + self.connection.register(temp_view, arrow_table) + try: + self.connection.execute(f"INSERT INTO {target_table} SELECT * FROM {temp_view}") + finally: + with contextlib.suppress(Exception): + self.connection.unregister(temp_view) + + return DMLResult("INSERT", len(rows)) + + @staticmethod + def _build_arrow_table(rows: "list[Any]", column_names: "list[str]") -> Any | None: + """Build a pyarrow table from batch rows when they share a stable shape.""" + if not rows: + return None + first_row = rows[0] + + if isinstance(first_row, dict): + keys = column_names or list(first_row.keys()) + if any(not isinstance(row, dict) for row in rows): + return None + import pyarrow as pa + + return pa.table({key: [row.get(key) for row in rows] for key in keys}) + + if isinstance(first_row, (list, tuple)): + values = list(first_row) + if not column_names: + column_names = [f"col_{index}" for index in range(len(values))] + if any(not isinstance(row, (list, tuple)) or len(row) != len(column_names) for row in rows): + return None + import pyarrow as pa + + return pa.table({name: [row[index] for row in rows] for index, name in enumerate(column_names)}) + + return None + def resolve_rowcount(self, cursor: "DuckDBConnection") -> int: """Resolve rowcount from DuckDB cursor for the direct execution path.""" return resolve_rowcount(cursor) @@ -453,6 +538,20 @@ def _connection_in_transaction(self) -> bool: return False +def _restore_uuid_columns(rows: "list[dict[str, Any]]", description: "list[Any] | None") -> None: + """Restore DuckDB UUID columns after Arrow materialization.""" + if not rows or not description: + return + uuid_columns = [column[0] for column in description if len(column) > 1 and str(column[1]).upper() == "UUID"] + if not uuid_columns: + return + for row in rows: + for column in uuid_columns: + value = row.get(column) + if isinstance(value, str): + row[column] = UUID(value) + + def _resolve_duckdb_inserted_rows(result: object) -> int: fetchall = getattr(result, "fetchall", None) if not callable(fetchall): diff --git a/sqlspec/adapters/duckdb/type_converter.py b/sqlspec/adapters/duckdb/type_converter.py deleted file mode 100644 index 19ae28501..000000000 --- a/sqlspec/adapters/duckdb/type_converter.py +++ /dev/null @@ -1,85 +0,0 @@ -"""DuckDB-specific type conversion with native UUID support. - -Provides specialized type handling for DuckDB, including native UUID -support and standardized datetime formatting. -""" - -from datetime import datetime -from typing import Any, Final -from uuid import UUID - -from typing_extensions import final - -from sqlspec.core.type_converter import CachedOutputConverter, convert_uuid, format_datetime_rfc3339 - -__all__ = ("DUCKDB_SPECIAL_CHARS", "DuckDBOutputConverter") - -DUCKDB_SPECIAL_CHARS: Final[frozenset[str]] = frozenset({"-", ":", "T", ".", "[", "{"}) - - -@final -class DuckDBOutputConverter(CachedOutputConverter): - """DuckDB-specific output conversion with native UUID support. - - Extends CachedOutputConverter with DuckDB-specific functionality - including native UUID handling and standardized datetime formatting. - """ - - __slots__ = ("_enable_uuid_conversion",) - - def __init__(self, cache_size: int = 5000, enable_uuid_conversion: bool = True) -> None: - """Initialize converter with DuckDB-specific options. - - Args: - cache_size: Maximum number of string values to cache (default: 5000) - enable_uuid_conversion: Enable automatic UUID string conversion (default: True) - """ - super().__init__(special_chars=DUCKDB_SPECIAL_CHARS, cache_size=cache_size) - self._enable_uuid_conversion = enable_uuid_conversion - - def _convert_detected(self, value: str, detected_type: str) -> Any: - """Convert value with DuckDB-specific UUID handling. - - Args: - value: String value to convert. - detected_type: Detected type name. - - Returns: - Converted value, respecting UUID conversion setting. - """ - if detected_type == "uuid" and not self._enable_uuid_conversion: - return value - try: - return self.convert_value(value, detected_type) - except Exception: - return value - - def handle_uuid(self, value: Any) -> Any: - """Handle UUID conversion for DuckDB. - - Args: - value: Value that might be a UUID. - - Returns: - UUID object if value is UUID-like and conversion enabled, original value otherwise. - """ - if isinstance(value, UUID): - return value - - if isinstance(value, str) and self._enable_uuid_conversion: - detected_type = self.detect_type(value) - if detected_type == "uuid": - return convert_uuid(value) - - return value - - def format_datetime(self, dt: datetime) -> str: - """Standardized datetime formatting for DuckDB. - - Args: - dt: datetime object to format. - - Returns: - RFC 3339 formatted datetime string. - """ - return format_datetime_rfc3339(dt) diff --git a/sqlspec/adapters/mssql_python/config.py b/sqlspec/adapters/mssql_python/config.py index 0c8580104..736200aec 100644 --- a/sqlspec/adapters/mssql_python/config.py +++ b/sqlspec/adapters/mssql_python/config.py @@ -18,6 +18,7 @@ from sqlspec.driver._async import AsyncPoolConnectionContext, AsyncPoolSessionFactory from sqlspec.driver._sync import SyncPoolConnectionContext, SyncPoolSessionFactory from sqlspec.utils.config_tools import normalize_connection_config +from sqlspec.utils.serializers import from_json, to_json if TYPE_CHECKING: from collections.abc import Callable @@ -96,6 +97,20 @@ class MssqlPythonPoolParams(MssqlPythonConnectionParams): pool_enabled: NotRequired[bool] +def _apply_json_serializer_override(statement_config: Any, features_dict: dict[str, Any]) -> Any: + serializer = cast("Callable[[Any], str] | None", features_dict.get("json_serializer")) + deserializer = cast("Callable[[str], Any] | None", features_dict.get("json_deserializer")) + if statement_config is not default_statement_config: + return statement_config + if serializer is to_json and deserializer is from_json: + return statement_config + return statement_config.replace( + parameter_config=statement_config.parameter_config.with_json_serializers( + serializer or to_json, deserializer=deserializer + ) + ) + + class MssqlPythonDriverFeatures(TypedDict): """mssql-python driver feature flags.""" @@ -208,6 +223,7 @@ class MssqlPythonConfig(SyncDatabaseConfig[MssqlPythonConnection, MssqlPythonCon supports_native_arrow_export: "ClassVar[bool]" = True supports_native_arrow_import: "ClassVar[bool]" = True supports_arrow_streaming: "ClassVar[bool]" = True + supports_native_row_streaming: "ClassVar[bool]" = True supports_native_parquet_export: "ClassVar[bool]" = False supports_native_parquet_import: "ClassVar[bool]" = False _connection_context_class: "ClassVar[type[MssqlPythonConnectionContext]]" = MssqlPythonConnectionContext @@ -234,12 +250,13 @@ def __init__( connection_config, driver_features ) self._user_connection_hook = user_connection_hook + statement_config = _apply_json_serializer_override(statement_config or default_statement_config, features_dict) super().__init__( bind_key=bind_key, connection_config=normalized, connection_instance=connection_instance, migration_config=migration_config, - statement_config=statement_config or default_statement_config, + statement_config=statement_config, driver_features=features_dict, extension_config=extension_config, observability_config=observability_config, @@ -287,6 +304,7 @@ class MssqlPythonAsyncConfig( supports_native_arrow_export: "ClassVar[bool]" = True supports_native_arrow_import: "ClassVar[bool]" = True supports_arrow_streaming: "ClassVar[bool]" = True + supports_native_row_streaming: "ClassVar[bool]" = True supports_native_parquet_export: "ClassVar[bool]" = False supports_native_parquet_import: "ClassVar[bool]" = False _connection_context_class: "ClassVar[type[MssqlPythonAsyncConnectionContext]]" = MssqlPythonAsyncConnectionContext @@ -313,12 +331,13 @@ def __init__( connection_config, driver_features ) self._user_connection_hook = user_connection_hook + statement_config = _apply_json_serializer_override(statement_config or default_statement_config, features_dict) super().__init__( bind_key=bind_key, connection_config=normalized, connection_instance=connection_instance, migration_config=migration_config, - statement_config=statement_config or default_statement_config, + statement_config=statement_config, driver_features=features_dict, extension_config=extension_config, observability_config=observability_config, diff --git a/sqlspec/adapters/mssql_python/driver.py b/sqlspec/adapters/mssql_python/driver.py index b2c66a12d..58351c6ee 100644 --- a/sqlspec/adapters/mssql_python/driver.py +++ b/sqlspec/adapters/mssql_python/driver.py @@ -1,6 +1,7 @@ """mssql-python sync and async drivers.""" import asyncio +import contextlib from typing import TYPE_CHECKING, Any, TypedDict, cast from typing_extensions import NotRequired @@ -24,9 +25,12 @@ ) from sqlspec.driver import ( AsyncDriverAdapterBase, + AsyncRowStream, BaseAsyncExceptionHandler, BaseSyncExceptionHandler, SyncDriverAdapterBase, + SyncRowStream, + rows_to_dicts, ) from sqlspec.exceptions import SQLSpecError from sqlspec.utils.arrow_helpers import arrow_reader_with_deferred_close @@ -62,6 +66,7 @@ def _quote_mssql_table(table: str) -> str: logger = get_logger("sqlspec.adapters.mssql_python") _MSSQL_ERROR = cast("type[BaseException]", getattr(MSSQL_PYTHON_MODULE, "Error", Exception)) +_COLUMN_CACHE_MAX_SIZE = 256 class MssqlPythonBulkCopyResult(TypedDict): @@ -100,10 +105,116 @@ def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "Ba return False +class MssqlPythonStreamSource: + """Native mssql-python chunk source backed by ``cursor.fetchmany()``.""" + + __slots__ = ("_chunk_size", "_column_names", "_cursor_manager", "_driver", "_parameters", "_sql") + + def __init__(self, driver: "MssqlPythonDriver", sql: str, parameters: Any, chunk_size: int) -> None: + self._driver = driver + self._sql = sql + self._parameters = parameters + self._chunk_size = chunk_size + self._cursor_manager: MssqlPythonCursor | None = None + self._column_names: list[str] | None = None + + def start(self) -> None: + cursor_manager = self._driver.with_cursor(self._driver.connection) + try: + cursor = cursor_manager.__enter__() + handler = self._driver.handle_database_exceptions() + with handler: + _execute_cursor(cursor, self._sql, self._parameters) + self._driver._check_pending_exception(handler) + except BaseException: + with contextlib.suppress(Exception): + cursor_manager.__exit__(None, None, None) + raise + self._cursor_manager = cursor_manager + + def fetch_chunk(self) -> "list[dict[str, Any]]": + cursor_manager = self._cursor_manager + if cursor_manager is None or cursor_manager.cursor is None: + return [] + cursor = cursor_manager.cursor + handler = self._driver.handle_database_exceptions() + rows: Any = [] + with handler: + rows = cursor.fetchmany(self._chunk_size) + self._driver._check_pending_exception(handler) + if not rows: + return [] + column_names = self._column_names + if column_names is None: + column_names = _resolve_column_names(cursor.description, self._driver._column_name_cache) + self._column_names = column_names + return rows_to_dicts(rows, column_names) + + def close(self) -> None: + cursor_manager = self._cursor_manager + self._cursor_manager = None + if cursor_manager is not None: + with contextlib.suppress(Exception): + cursor_manager.__exit__(None, None, None) + + +class MssqlPythonAsyncStreamSource: + """Async mssql-python chunk source backed by ``cursor.fetchmany()``.""" + + __slots__ = ("_chunk_size", "_column_names", "_cursor_manager", "_driver", "_parameters", "_sql") + + def __init__(self, driver: "MssqlPythonAsyncDriver", sql: str, parameters: Any, chunk_size: int) -> None: + self._driver = driver + self._sql = sql + self._parameters = parameters + self._chunk_size = chunk_size + self._cursor_manager: MssqlPythonAsyncCursor | None = None + self._column_names: list[str] | None = None + + async def start(self) -> None: + cursor_manager = self._driver.with_cursor(self._driver.connection) + try: + cursor = await cursor_manager.__aenter__() + handler = self._driver.handle_database_exceptions() + async with handler: + await asyncio.to_thread(_execute_cursor, cursor, self._sql, self._parameters) + self._driver._check_pending_exception(handler) + except BaseException: + with contextlib.suppress(Exception): + await cursor_manager.__aexit__(None, None, None) + raise + self._cursor_manager = cursor_manager + + async def fetch_chunk(self) -> "list[dict[str, Any]]": + cursor_manager = self._cursor_manager + if cursor_manager is None or cursor_manager.cursor is None: + return [] + cursor = cursor_manager.cursor + handler = self._driver.handle_database_exceptions() + rows: Any = [] + async with handler: + rows = await asyncio.to_thread(cursor.fetchmany, self._chunk_size) + self._driver._check_pending_exception(handler) + if not rows: + return [] + column_names = self._column_names + if column_names is None: + column_names = _resolve_column_names(cursor.description, self._driver._column_name_cache) + self._column_names = column_names + return rows_to_dicts(rows, column_names) + + async def close(self) -> None: + cursor_manager = self._cursor_manager + self._cursor_manager = None + if cursor_manager is not None: + with contextlib.suppress(Exception): + await cursor_manager.__aexit__(None, None, None) + + class MssqlPythonDriver(SyncDriverAdapterBase): """mssql-python sync driver.""" - __slots__ = ("_data_dictionary",) + __slots__ = ("_column_name_cache", "_data_dictionary") dialect = "tsql" def __init__( @@ -118,6 +229,7 @@ def __init__( ) super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features) self._data_dictionary: MssqlPythonSyncDataDictionary | None = None + self._column_name_cache: dict[int, tuple[Any, list[str]]] = {} @property def data_dictionary(self) -> "MssqlPythonSyncDataDictionary": @@ -131,7 +243,7 @@ def dispatch_execute(self, cursor: "MssqlPythonRawCursor", statement: "SQL") -> if statement.returns_rows(): fetched = cursor.fetchall() - column_names = [desc[0] for desc in (cursor.description or [])] + column_names = _resolve_column_names(cursor.description, self._column_name_cache) return self.create_execution_result( cursor, selected_data=fetched, @@ -160,7 +272,7 @@ def dispatch_execute_script(self, cursor: "MssqlPythonRawCursor", statement: "SQ ) def collect_rows(self, cursor: "MssqlPythonRawCursor", fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": - column_names = [desc[0] for desc in (cursor.description or [])] + column_names = _resolve_column_names(cursor.description, self._column_name_cache) return fetched, column_names, len(fetched) def resolve_rowcount(self, cursor: "MssqlPythonRawCursor") -> int: @@ -192,6 +304,13 @@ def with_cursor(self, connection: "MssqlPythonConnection") -> "MssqlPythonCursor def handle_database_exceptions(self) -> "MssqlPythonExceptionHandler": return MssqlPythonExceptionHandler() + def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "SyncRowStream[dict[str, Any]] | None": + """Return a native mssql-python row stream backed by ``fetchmany()``.""" + if not statement.returns_rows(): + return None + sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) + return SyncRowStream(MssqlPythonStreamSource(self, sql, prepared_parameters, chunk_size)) + def create_savepoint(self, name: str) -> None: self.execute_script(f"SAVE TRANSACTION {name}") @@ -361,7 +480,7 @@ def load_from_storage( class MssqlPythonAsyncDriver(AsyncDriverAdapterBase): """Async wrapper around mssql-python's sync DB-API via asyncio.to_thread.""" - __slots__ = ("_data_dictionary",) + __slots__ = ("_column_name_cache", "_data_dictionary") dialect = "tsql" def __init__( @@ -376,6 +495,7 @@ def __init__( ) super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features) self._data_dictionary: MssqlPythonAsyncDataDictionary | None = None + self._column_name_cache: dict[int, tuple[Any, list[str]]] = {} @property def data_dictionary(self) -> "MssqlPythonAsyncDataDictionary": @@ -389,7 +509,7 @@ async def dispatch_execute(self, cursor: "MssqlPythonRawCursor", statement: "SQL if statement.returns_rows(): fetched = await asyncio.to_thread(cursor.fetchall) - column_names = [desc[0] for desc in (cursor.description or [])] + column_names = _resolve_column_names(cursor.description, self._column_name_cache) return self.create_execution_result( cursor, selected_data=fetched, @@ -418,7 +538,7 @@ async def dispatch_execute_script(self, cursor: "MssqlPythonRawCursor", statemen ) def collect_rows(self, cursor: "MssqlPythonRawCursor", fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": - column_names = [desc[0] for desc in (cursor.description or [])] + column_names = _resolve_column_names(cursor.description, self._column_name_cache) return fetched, column_names, len(fetched) def resolve_rowcount(self, cursor: "MssqlPythonRawCursor") -> int: @@ -450,6 +570,13 @@ def with_cursor(self, connection: "MssqlPythonConnection") -> "MssqlPythonAsyncC def handle_database_exceptions(self) -> "MssqlPythonAsyncExceptionHandler": return MssqlPythonAsyncExceptionHandler() + def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "AsyncRowStream[dict[str, Any]] | None": + """Return a native async mssql-python row stream backed by ``fetchmany()``.""" + if not statement.returns_rows(): + return None + sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) + return AsyncRowStream(MssqlPythonAsyncStreamSource(self, sql, prepared_parameters, chunk_size)) + async def create_savepoint(self, name: str) -> None: await self.execute_script(f"SAVE TRANSACTION {name}") @@ -634,6 +761,20 @@ def _cursor_rowcount(cursor: "MssqlPythonRawCursor") -> int: return rowcount if isinstance(rowcount, int) and rowcount > 0 else 0 +def _resolve_column_names(description: Any, cache: "dict[int, tuple[Any, list[str]]]") -> list[str]: + if not description: + return [] + cache_key = id(description) + cached = cache.get(cache_key) + if cached is not None and cached[0] is description: + return cached[1] + column_names = [desc[0] for desc in description] + if len(cache) >= _COLUMN_CACHE_MAX_SIZE: + cache.pop(next(iter(cache))) + cache[cache_key] = (description, column_names) + return column_names + + def _cursor_arrow_reader( cursor: "MssqlPythonRawCursor", arrow_kwargs: "dict[str, int]" ) -> "ArrowRecordBatchReader | None": diff --git a/sqlspec/adapters/mysqlconnector/config.py b/sqlspec/adapters/mysqlconnector/config.py index e1f9596ac..ad54330c5 100644 --- a/sqlspec/adapters/mysqlconnector/config.py +++ b/sqlspec/adapters/mysqlconnector/config.py @@ -1,7 +1,7 @@ """MysqlConnector database configuration.""" import contextlib -from collections.abc import Callable +from collections.abc import Callable, Mapping from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypedDict, cast from weakref import WeakSet @@ -194,6 +194,14 @@ class MysqlConnectorDriverFeatures(TypedDict): enable_local_infile_bulk_load: NotRequired[bool] +def _normalize_local_infile(connection_config: "Mapping[str, Any] | None") -> "dict[str, Any]": + """Normalize mysql-connector local-infile consent.""" + config = normalize_connection_config(connection_config) + config.pop("local_infile", None) + config["allow_local_infile"] = bool(config.get("allow_local_infile", False)) + return config + + class MysqlConnectorSyncConnectionContext(SyncPoolConnectionContext): """Context manager for mysql-connector sync connections.""" @@ -301,10 +309,9 @@ def __init__( observability_config: "ObservabilityConfig | None" = None, **kwargs: Any, ) -> None: - connection_config = normalize_connection_config(connection_config) - connection_config.setdefault("host", "127.0.0.1") + connection_config = _normalize_local_infile(connection_config) + connection_config.setdefault("host", "localhost") connection_config.setdefault("port", 3306) - connection_config.setdefault("allow_local_infile", False) statement_config = statement_config or default_statement_config statement_config, driver_features = apply_driver_features(statement_config, driver_features) @@ -422,10 +429,9 @@ def __init__( observability_config: "ObservabilityConfig | None" = None, **kwargs: Any, ) -> None: - self.connection_config = normalize_connection_config(connection_config) - self.connection_config.setdefault("host", "127.0.0.1") + self.connection_config = _normalize_local_infile(connection_config) + self.connection_config.setdefault("host", "localhost") self.connection_config.setdefault("port", 3306) - self.connection_config.setdefault("allow_local_infile", False) statement_config = statement_config or default_statement_config statement_config, driver_features = apply_driver_features(statement_config, driver_features) diff --git a/sqlspec/adapters/mysqlconnector/core.py b/sqlspec/adapters/mysqlconnector/core.py index 777089085..eccfe2d56 100644 --- a/sqlspec/adapters/mysqlconnector/core.py +++ b/sqlspec/adapters/mysqlconnector/core.py @@ -6,7 +6,7 @@ import contextlib from collections.abc import Callable, Sized -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast from sqlspec.core import DriverParameterProfile, ParameterStyle, StatementConfig, build_statement_config_from_profile from sqlspec.driver import rows_to_dicts @@ -44,6 +44,7 @@ "build_profile", "build_statement_config", "collect_rows", + "collect_stream_rows", "create_mapped_exception", "default_statement_config", "detect_json_columns", @@ -56,6 +57,7 @@ "normalize_lastrowid", "resolve_column_names", "resolve_many_rowcount", + "resolve_row_plan", "resolve_rowcount", ) @@ -175,22 +177,43 @@ def normalize_execute_parameters(parameters: Any) -> Any: class MysqlConnectorSyncStreamSource: """Compiled chunk source streaming dict rows from an unbuffered mysql-connector cursor.""" - __slots__ = ("_chunk_size", "_column_names", "_cursor", "_driver", "_parameters", "_sql") + __slots__ = ( + "_chunk_size", + "_cursor", + "_cursor_options", + "_driver", + "_json_type_codes", + "_parameters", + "_row_plan", + "_sql", + ) - def __init__(self, driver: Any, sql: str, parameters: Any, chunk_size: int) -> None: + def __init__( + self, + driver: Any, + sql: str, + parameters: Any, + chunk_size: int, + json_type_codes: "set[int]", + cursor_options: "dict[str, Any] | None" = None, + ) -> None: self._driver = driver self._sql = sql self._parameters = parameters self._chunk_size = chunk_size self._cursor: Any = None - self._column_names: list[str] | None = None + self._json_type_codes = json_type_codes + self._cursor_options = dict(cursor_options or {}) + self._cursor_options["buffered"] = False + self._row_plan: tuple[list[str], list[int] | None] | None = None def start(self) -> None: handler = self._driver.handle_database_exceptions() with handler: - cursor = self._driver.connection.cursor(buffered=False) + cursor = self._driver.connection.cursor(**self._cursor_options) cursor.execute(self._sql, normalize_execute_parameters(self._parameters)) self._cursor = cursor + self._row_plan = resolve_row_plan(self._cursor.description, self._json_type_codes) self._driver._check_pending_exception(handler) def fetch_chunk(self) -> "list[dict[str, Any]]": @@ -201,9 +224,10 @@ def fetch_chunk(self) -> "list[dict[str, Any]]": self._driver._check_pending_exception(handler) if not rows: return [] - if self._column_names is None: - self._column_names = [description[0] for description in self._cursor.description] - return rows_to_dicts(rows, self._column_names) + if self._row_plan is None: + self._row_plan = resolve_row_plan(self._cursor.description, self._json_type_codes) + deserializer = cast("Callable[[Any], Any]", self._driver.driver_features.get("json_deserializer", from_json)) + return collect_stream_rows(rows, self._row_plan, deserializer) def close(self) -> None: cursor = self._cursor @@ -218,22 +242,43 @@ def close(self) -> None: class MysqlConnectorAsyncStreamSource: """Compiled async chunk source streaming dict rows from an unbuffered mysql-connector cursor.""" - __slots__ = ("_chunk_size", "_column_names", "_cursor", "_driver", "_parameters", "_sql") + __slots__ = ( + "_chunk_size", + "_cursor", + "_cursor_options", + "_driver", + "_json_type_codes", + "_parameters", + "_row_plan", + "_sql", + ) - def __init__(self, driver: Any, sql: str, parameters: Any, chunk_size: int) -> None: + def __init__( + self, + driver: Any, + sql: str, + parameters: Any, + chunk_size: int, + json_type_codes: "set[int]", + cursor_options: "dict[str, Any] | None" = None, + ) -> None: self._driver = driver self._sql = sql self._parameters = parameters self._chunk_size = chunk_size self._cursor: Any = None - self._column_names: list[str] | None = None + self._json_type_codes = json_type_codes + self._cursor_options = dict(cursor_options or {}) + self._cursor_options["buffered"] = False + self._row_plan: tuple[list[str], list[int] | None] | None = None async def start(self) -> None: handler = self._driver.handle_database_exceptions() async with handler: - cursor = await self._driver.connection.cursor(buffered=False) + cursor = await self._driver.connection.cursor(**self._cursor_options) await cursor.execute(self._sql, normalize_execute_parameters(self._parameters)) self._cursor = cursor + self._row_plan = resolve_row_plan(self._cursor.description, self._json_type_codes) self._driver._check_pending_exception(handler) async def fetch_chunk(self) -> "list[dict[str, Any]]": @@ -244,9 +289,10 @@ async def fetch_chunk(self) -> "list[dict[str, Any]]": self._driver._check_pending_exception(handler) if not rows: return [] - if self._column_names is None: - self._column_names = [description[0] for description in self._cursor.description] - return rows_to_dicts(rows, self._column_names) + if self._row_plan is None: + self._row_plan = resolve_row_plan(self._cursor.description, self._json_type_codes) + deserializer = cast("Callable[[Any], Any]", self._driver.driver_features.get("json_deserializer", from_json)) + return collect_stream_rows(rows, self._row_plan, deserializer) async def close(self) -> None: cursor = self._cursor @@ -366,7 +412,7 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> "SQLSpe if error_code in _MYSQL_MIGRATION_ERROR_CODES: if logger is not None: - logger.warning("MysqlConnector expected migration error (ignoring): %s", error) + logger.warning("MySQL expected migration error (ignoring): %s", error) return True dispatch = _MYSQL_SQLSTATE_EXACT_DISPATCH.get(sqlstate) if sqlstate is not None else None @@ -422,23 +468,36 @@ def resolve_column_names(description: "Sequence[Any] | None") -> "list[str]": return [desc[0] for desc in description] -def detect_json_columns_from_description( +def resolve_row_plan( description: "Sequence[Any] | None", json_type_codes: "set[int]" -) -> "list[int]": - """Identify JSON column indexes from pre-fetched cursor description metadata.""" - if not description or not json_type_codes: - return [] +) -> "tuple[list[str], list[int] | None]": + """Resolve ordered column names and JSON column indexes in one pass.""" + if not description: + return [], None + + column_names: list[str] = [] + if not json_type_codes: + column_names.extend(column[0] for column in description) + return column_names, None json_indexes: list[int] = [] - append = json_indexes.append + append_json = json_indexes.append for index, column in enumerate(description): + column_names.append(column[0]) if isinstance(column, (tuple, list)): type_code = column[1] if len(column) > 1 else None else: type_code = column.type_code if isinstance(column, HasTypeCodeProtocol) else None if type_code in json_type_codes: - append(index) - return json_indexes + append_json(index) + return column_names, json_indexes or None + + +def detect_json_columns_from_description( + description: "Sequence[Any] | None", json_type_codes: "set[int]" +) -> "list[int]": + """Identify JSON column indexes from pre-fetched cursor description metadata.""" + return resolve_row_plan(description, json_type_codes)[1] or [] def detect_json_columns( @@ -518,38 +577,52 @@ def _deserialize_json_tuple_rows( return result +def collect_stream_rows( + fetched_data: "Sequence[Any] | None", + row_plan: "tuple[list[str], list[int] | None]", + deserializer: "Callable[[Any], Any]", + *, + logger: Any | None = None, +) -> "list[dict[str, Any]]": + """Materialize streamed rows as dicts and decode JSON columns in place.""" + column_names, json_indexes = row_plan + if not column_names or not fetched_data: + return [] + + rows = fetched_data if isinstance(fetched_data, list) else list(fetched_data) + if not rows: + return [] + + dict_rows = rows_to_dicts(rows, column_names) + if json_indexes: + dict_rows = _deserialize_json_dict_rows(column_names, dict_rows, json_indexes, deserializer, logger=logger) + return dict_rows + + def collect_rows( fetched_data: "Sequence[Any] | None", - description: "Sequence[Any] | None", - json_indexes: "list[int]", + row_plan: "tuple[list[str], list[int] | None]", deserializer: "Callable[[Any], Any]", *, - column_names: "list[str] | None" = None, logger: Any | None = None, ) -> "tuple[list[Any], list[str], str]": - """Collect mysql-connector rows with JSON decoding, preserving raw format. - - Returns: - Tuple of (rows, column_names, row_format). - """ - if not description: + """Collect mysql-connector rows with JSON decoding, preserving raw format.""" + column_names, json_indexes = row_plan + if not column_names: return [], [], "tuple" - resolved_column_names = resolve_column_names(description) if column_names is None else column_names if not fetched_data: - return [], resolved_column_names, "tuple" + return [], column_names, "tuple" first_row = fetched_data[0] if isinstance(first_row, dict): - if not json_indexes: - rows = fetched_data if isinstance(fetched_data, list) else list(fetched_data) - return rows, resolved_column_names, "dict" - rows = [dict(row) for row in fetched_data] - rows = _deserialize_json_dict_rows(resolved_column_names, rows, json_indexes, deserializer, logger=logger) - return rows, resolved_column_names, "dict" + rows = fetched_data if isinstance(fetched_data, list) else list(fetched_data) + if json_indexes: + rows = _deserialize_json_dict_rows(column_names, rows, json_indexes, deserializer, logger=logger) + return rows, column_names, "dict" rows = fetched_data if isinstance(fetched_data, list) else list(fetched_data) if json_indexes: rows = _deserialize_json_tuple_rows(rows, json_indexes, deserializer, logger=logger) - return rows, resolved_column_names, "tuple" + return rows, column_names, "tuple" def resolve_rowcount(cursor: Any) -> int: diff --git a/sqlspec/adapters/mysqlconnector/driver.py b/sqlspec/adapters/mysqlconnector/driver.py index 4a962304b..933582bce 100644 --- a/sqlspec/adapters/mysqlconnector/driver.py +++ b/sqlspec/adapters/mysqlconnector/driver.py @@ -25,15 +25,14 @@ collect_rows, create_mapped_exception, default_statement_config, - detect_json_columns_from_description, driver_profile, encode_records_for_local_infile, format_identifier, normalize_execute_many_parameters, normalize_execute_parameters, normalize_lastrowid, - resolve_column_names, resolve_many_rowcount, + resolve_row_plan, resolve_rowcount, ) from sqlspec.adapters.mysqlconnector.data_dictionary import ( @@ -123,12 +122,9 @@ def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionResult": if statement.returns_rows() or getattr(cursor, "with_rows", False): fetched_data = cursor.fetchall() description = cursor.description or None - column_names = resolve_column_names(description) - json_indexes = detect_json_columns_from_description(description, MYSQLCONNECTOR_JSON_TYPE_CODES) + row_plan = resolve_row_plan(description, MYSQLCONNECTOR_JSON_TYPE_CODES) deserializer = cast("Callable[[Any], Any]", self.driver_features.get("json_deserializer", from_json)) - rows, column_names, row_format = collect_rows( - fetched_data, description, json_indexes, deserializer, column_names=column_names, logger=logger - ) + rows, column_names, row_format = collect_rows(fetched_data, row_plan, deserializer, logger=logger) return self.create_execution_result( cursor, @@ -203,7 +199,12 @@ def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "SyncRowS if not statement.returns_rows(): return None sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - return SyncRowStream(MysqlConnectorSyncStreamSource(self, sql, prepared_parameters, chunk_size)) + cursor_options = cast("dict[str, Any]", self.driver_features.get("cursor_options") or {}) + return SyncRowStream( + MysqlConnectorSyncStreamSource( + self, sql, prepared_parameters, chunk_size, MYSQLCONNECTOR_JSON_TYPE_CODES, cursor_options + ) + ) def handle_database_exceptions(self) -> "MysqlConnectorSyncExceptionHandler": return MysqlConnectorSyncExceptionHandler() @@ -305,12 +306,9 @@ def data_dictionary(self) -> "MysqlConnectorSyncDataDictionary": def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect mysql-connector sync rows for the direct execution path.""" description = cursor.description or None - column_names = resolve_column_names(description) - json_indexes = detect_json_columns_from_description(description, MYSQLCONNECTOR_JSON_TYPE_CODES) + row_plan = resolve_row_plan(description, MYSQLCONNECTOR_JSON_TYPE_CODES) deserializer = cast("Callable[[Any], Any]", self.driver_features.get("json_deserializer", from_json)) - rows, column_names, _row_format = collect_rows( - fetched, description, json_indexes, deserializer, column_names=column_names, logger=logger - ) + rows, column_names, _row_format = collect_rows(fetched, row_plan, deserializer, logger=logger) return rows, column_names, len(rows) def resolve_rowcount(self, cursor: Any) -> int: @@ -377,12 +375,9 @@ async def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionRes if statement.returns_rows() or getattr(cursor, "with_rows", False): fetched_data = await cursor.fetchall() description = cursor.description or None - column_names = resolve_column_names(description) - json_indexes = detect_json_columns_from_description(description, MYSQLCONNECTOR_JSON_TYPE_CODES) + row_plan = resolve_row_plan(description, MYSQLCONNECTOR_JSON_TYPE_CODES) deserializer = cast("Callable[[Any], Any]", self.driver_features.get("json_deserializer", from_json)) - rows, column_names, row_format = collect_rows( - fetched_data, description, json_indexes, deserializer, column_names=column_names, logger=logger - ) + rows, column_names, row_format = collect_rows(fetched_data, row_plan, deserializer, logger=logger) return self.create_execution_result( cursor, @@ -457,7 +452,12 @@ def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "AsyncRow if not statement.returns_rows(): return None sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - return AsyncRowStream(MysqlConnectorAsyncStreamSource(self, sql, prepared_parameters, chunk_size)) + cursor_options = cast("dict[str, Any]", self.driver_features.get("cursor_options") or {}) + return AsyncRowStream( + MysqlConnectorAsyncStreamSource( + self, sql, prepared_parameters, chunk_size, MYSQLCONNECTOR_JSON_TYPE_CODES, cursor_options + ) + ) def handle_database_exceptions(self) -> "MysqlConnectorAsyncExceptionHandler": return MysqlConnectorAsyncExceptionHandler() @@ -561,12 +561,9 @@ def data_dictionary(self) -> "MysqlConnectorAsyncDataDictionary": def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect mysql-connector async rows for the direct execution path.""" description = cursor.description or None - column_names = resolve_column_names(description) - json_indexes = detect_json_columns_from_description(description, MYSQLCONNECTOR_JSON_TYPE_CODES) + row_plan = resolve_row_plan(description, MYSQLCONNECTOR_JSON_TYPE_CODES) deserializer = cast("Callable[[Any], Any]", self.driver_features.get("json_deserializer", from_json)) - rows, column_names, _row_format = collect_rows( - fetched, description, json_indexes, deserializer, column_names=column_names, logger=logger - ) + rows, column_names, _row_format = collect_rows(fetched, row_plan, deserializer, logger=logger) return rows, column_names, len(rows) def resolve_rowcount(self, cursor: Any) -> int: diff --git a/sqlspec/adapters/oracledb/_json_handlers.py b/sqlspec/adapters/oracledb/_json_handlers.py index 31e27187a..a8028d329 100644 --- a/sqlspec/adapters/oracledb/_json_handlers.py +++ b/sqlspec/adapters/oracledb/_json_handlers.py @@ -17,9 +17,12 @@ Routing matrix (output): * ``DB_TYPE_JSON``: passthrough (python-oracledb already returns ``dict``). -* ``DB_TYPE_BLOB`` with ``JSON`` in column ``type_name``: parse via +* ``FetchInfo.is_oson`` + ``DB_TYPE_BLOB``: decode through + ``connection.decode_oson``. This branch precedes ``is_json`` because OSON + metadata can report both flags. +* ``FetchInfo.is_json`` + ``DB_TYPE_BLOB``: fetch as bytes and parse via ``json_converter_out_blob``. -* ``DB_TYPE_CLOB`` with ``JSON`` in column ``type_name``: parse via +* ``FetchInfo.is_json`` + ``DB_TYPE_CLOB`` or string types: parse via ``json_converter_out_clob``. Handlers chain to any pre-existing ``inputtypehandler`` / ``outputtypehandler`` @@ -31,7 +34,9 @@ from functools import partial from typing import TYPE_CHECKING, Any, cast -from sqlspec.adapters.oracledb._typing import DB_TYPE_BLOB, DB_TYPE_CLOB, DB_TYPE_JSON +from oracledb import DB_TYPE_CHAR, DB_TYPE_NCHAR, DB_TYPE_NVARCHAR, DB_TYPE_VARCHAR + +from sqlspec.adapters.oracledb._typing import DB_TYPE_BLOB, DB_TYPE_CLOB, DB_TYPE_JSON, DB_TYPE_LONG, DB_TYPE_LONG_RAW from sqlspec.utils.serializers import from_json, to_json if TYPE_CHECKING: @@ -44,19 +49,18 @@ "json_converter_in_clob", "json_converter_out_blob", "json_converter_out_clob", + "json_converter_out_oson", "json_input_type_handler", "json_output_type_handler", "register_json_handlers", ) - -_JSON_TYPE_NAME_MARKER = "JSON" - # Server-version thresholds for JSON binding strategy selection. # 21c+ supports DB_TYPE_JSON (binary OSON); 19c-20c uses BLOB CHECK (... IS JSON); # pre-19c uses CLOB CHECK (... IS JSON). _NATIVE_JSON_MIN_MAJOR = 21 _BLOB_IS_JSON_MIN_MAJOR = 19 +_JSON_STRING_TYPE_CODES = (DB_TYPE_VARCHAR, DB_TYPE_CHAR, DB_TYPE_NVARCHAR, DB_TYPE_NCHAR) def json_converter_in_clob(value: Any) -> str: @@ -83,6 +87,13 @@ def json_converter_out_blob(value: "bytes | None") -> Any: return from_json(value) +def json_converter_out_oson(connection: "Connection | AsyncConnection", value: "bytes | None") -> Any: + """Decode OSON bytes from a BLOB read back into a Python value.""" + if value is None: + return None + return connection.decode_oson(value) + + def _is_json_payload(value: Any) -> bool: """Return True if the value should be claimed by the JSON input handler. @@ -129,14 +140,21 @@ def _output_type_handler(cursor: "Cursor | AsyncCursor", metadata: Any) -> Any: if type_code is DB_TYPE_JSON: return None - type_name = (getattr(metadata, "type_name", "") or "").upper() - if _JSON_TYPE_NAME_MARKER not in type_name: - return None + if getattr(metadata, "is_oson", False) and type_code is DB_TYPE_BLOB: + return cursor.var( + DB_TYPE_LONG_RAW, + arraysize=cursor.arraysize, + outconverter=partial(json_converter_out_oson, cursor.connection), + ) + + if getattr(metadata, "is_json", False): + if type_code is DB_TYPE_BLOB: + return cursor.var(DB_TYPE_LONG_RAW, arraysize=cursor.arraysize, outconverter=json_converter_out_blob) + if type_code is DB_TYPE_CLOB: + return cursor.var(DB_TYPE_LONG, arraysize=cursor.arraysize, outconverter=json_converter_out_clob) + if type_code in _JSON_STRING_TYPE_CODES: + return cursor.var(type_code, arraysize=cursor.arraysize, outconverter=json_converter_out_clob) - if type_code is DB_TYPE_BLOB: - return cursor.var(DB_TYPE_BLOB, arraysize=cursor.arraysize, outconverter=json_converter_out_blob) - if type_code is DB_TYPE_CLOB: - return cursor.var(DB_TYPE_CLOB, arraysize=cursor.arraysize, outconverter=json_converter_out_clob) return None diff --git a/sqlspec/adapters/oracledb/_typing.py b/sqlspec/adapters/oracledb/_typing.py index 55ece2adf..e6c296c25 100644 --- a/sqlspec/adapters/oracledb/_typing.py +++ b/sqlspec/adapters/oracledb/_typing.py @@ -12,6 +12,8 @@ DB_TYPE_BLOB, DB_TYPE_CLOB, DB_TYPE_JSON, + DB_TYPE_LONG, + DB_TYPE_LONG_RAW, DB_TYPE_RAW, AsyncConnection, AsyncCursor, @@ -60,6 +62,8 @@ "DB_TYPE_BLOB", "DB_TYPE_CLOB", "DB_TYPE_JSON", + "DB_TYPE_LONG", + "DB_TYPE_LONG_RAW", "DB_TYPE_RAW", "DB_TYPE_VECTOR", "SPARSE_VECTOR_TYPE", diff --git a/sqlspec/adapters/oracledb/core.py b/sqlspec/adapters/oracledb/core.py index 2483ffc5a..134f13ffc 100644 --- a/sqlspec/adapters/oracledb/core.py +++ b/sqlspec/adapters/oracledb/core.py @@ -53,6 +53,7 @@ class _OracleSyncStreamDriver(Protocol): connection: "OracleSyncConnection" + driver_features: dict[str, Any] def handle_database_exceptions(self) -> "SyncExceptionHandler": ... @@ -63,6 +64,7 @@ def _resolve_row_metadata(self, description: object) -> tuple[list[str], bool]: class _OracleAsyncStreamDriver(Protocol): connection: "OracleAsyncConnection" + driver_features: dict[str, Any] def handle_database_exceptions(self) -> "AsyncExceptionHandler": ... @@ -219,6 +221,9 @@ def build_fetch_kwargs(driver_features: "dict[str, Any]") -> "dict[str, object]" def build_arrow_fetch_kwargs(driver_features: "dict[str, Any]") -> "dict[str, object]": """Build Oracle Arrow/DataFrame fetch keyword arguments.""" fetch_kwargs: dict[str, object] = {} + fetch_lobs = driver_features.get("fetch_lobs") + if fetch_lobs is not None: + fetch_kwargs["fetch_lobs"] = fetch_lobs fetch_decimals = driver_features.get("fetch_decimals") if fetch_decimals is not None: fetch_kwargs["fetch_decimals"] = fetch_decimals @@ -537,11 +542,17 @@ def resolve_rowcount(cursor: Any) -> int: def apply_driver_features( statement_config: "StatementConfig", driver_features: "Mapping[str, Any] | None" ) -> "tuple[StatementConfig, dict[str, Any]]": - """Apply OracleDB driver feature defaults.""" + """Apply OracleDB driver feature defaults. + + Defaults ``fetch_lobs`` to ``False`` so direct LOB reads under the 1 GB + fetch ceiling return text/bytes by default. Locators remain available for + OUT binds, PL/SQL, and explicit ``fetch_lobs=True``. + """ features: dict[str, Any] = dict(driver_features) if driver_features else {} features.setdefault("enable_numpy_vectors", NUMPY_INSTALLED) features.setdefault("enable_lowercase_column_names", True) features.setdefault("enable_uuid_binary", True) + features.setdefault("fetch_lobs", False) if "vector_return_format" not in features: if not NUMPY_INSTALLED: features["vector_return_format"] = "list" @@ -640,7 +651,7 @@ def resolve_row_metadata( class OracleSyncStreamSource: """Compiled chunk source streaming dict rows from an oracledb cursor via ``fetchmany``.""" - __slots__ = ("_chunk_size", "_column_names", "_cursor", "_driver", "_parameters", "_sql") + __slots__ = ("_chunk_size", "_column_names", "_cursor", "_driver", "_fetch_lobs", "_parameters", "_sql") def __init__( self, @@ -648,11 +659,13 @@ def __init__( sql: str, parameters: "list[object] | tuple[object, ...] | dict[object, object] | None", chunk_size: int, + fetch_lobs: bool | None = None, ) -> None: self._driver = driver self._sql = sql self._parameters = parameters self._chunk_size = chunk_size + self._fetch_lobs = fetch_lobs self._cursor: OracleSyncRawCursor | None = None self._column_names: list[str] | None = None @@ -662,7 +675,10 @@ def start(self) -> None: cursor = self._driver.connection.cursor() cursor.arraysize = self._chunk_size cursor.prefetchrows = self._chunk_size - cursor.execute(self._sql, self._parameters or {}) + fetch_kwargs = build_fetch_kwargs(self._driver.driver_features) + if self._fetch_lobs is not None: + fetch_kwargs["fetch_lobs"] = self._fetch_lobs + cast("Any", cursor).execute(self._sql, self._parameters or {}, **fetch_kwargs) self._cursor = cursor self._driver._check_pending_exception(handler) @@ -694,7 +710,7 @@ def close(self) -> None: class OracleAsyncStreamSource: """Compiled async chunk source streaming dict rows from an oracledb cursor via ``fetchmany``.""" - __slots__ = ("_chunk_size", "_column_names", "_cursor", "_driver", "_parameters", "_sql") + __slots__ = ("_chunk_size", "_column_names", "_cursor", "_driver", "_fetch_lobs", "_parameters", "_sql") def __init__( self, @@ -702,11 +718,13 @@ def __init__( sql: str, parameters: "list[object] | tuple[object, ...] | dict[object, object] | None", chunk_size: int, + fetch_lobs: bool | None = None, ) -> None: self._driver = driver self._sql = sql self._parameters = parameters self._chunk_size = chunk_size + self._fetch_lobs = fetch_lobs self._cursor: OracleAsyncRawCursor | None = None self._column_names: list[str] | None = None @@ -716,7 +734,10 @@ async def start(self) -> None: cursor = self._driver.connection.cursor() cursor.arraysize = self._chunk_size cursor.prefetchrows = self._chunk_size - await cursor.execute(self._sql, self._parameters or {}) + fetch_kwargs = build_fetch_kwargs(self._driver.driver_features) + if self._fetch_lobs is not None: + fetch_kwargs["fetch_lobs"] = self._fetch_lobs + await cast("Any", cursor).execute(self._sql, self._parameters or {}, **fetch_kwargs) self._cursor = cursor self._driver._check_pending_exception(handler) @@ -759,8 +780,9 @@ def _row_requires_lob_coercion(row: "tuple[Any, ...]") -> bool: def _coerce_sync_row_values(row: "tuple[Any, ...]") -> "tuple[Any, ...]": """Coerce LOB handles to concrete values for synchronous execution. - Processes each value in the row, reading LOB objects and applying - type detection for JSON values stored in CLOBs. + Processes each value in the row, reading locator objects returned for + OUT binds, PL/SQL, or explicit ``fetch_lobs=True``. Direct-fetched + ``str``/``bytes`` values from the default path pass through unchanged. Args: row: Tuple of column values from database fetch. @@ -783,12 +805,8 @@ def _coerce_sync_row_values(row: "tuple[Any, ...]") -> "tuple[Any, ...]": if coerced_values is not None: coerced_values.append(value) continue - if isinstance(processed_value, str): - processed_value = TYPE_CONVERTER.convert_if_detected(processed_value) if coerced_values is None: - if processed_value is value: - continue coerced_values = list(row[:index]) coerced_values.append(processed_value) continue @@ -804,8 +822,9 @@ def _coerce_sync_row_values(row: "tuple[Any, ...]") -> "tuple[Any, ...]": async def _coerce_async_row_values(row: "tuple[Any, ...]") -> "tuple[Any, ...]": """Coerce LOB handles to concrete values for asynchronous execution. - Processes each value in the row, reading LOB objects asynchronously - and applying type detection for JSON values stored in CLOBs. + Processes each value in the row, reading locator objects returned for + OUT binds, PL/SQL, or explicit ``fetch_lobs=True``. Direct-fetched + ``str``/``bytes`` values from the default path pass through unchanged. Args: row: Tuple of column values from database fetch. @@ -828,8 +847,6 @@ async def _coerce_async_row_values(row: "tuple[Any, ...]") -> "tuple[Any, ...]": if coerced_values is not None: coerced_values.append(value) continue - if isinstance(processed_value, str): - processed_value = TYPE_CONVERTER.convert_if_detected(processed_value) if coerced_values is None: if processed_value is value: @@ -856,9 +873,11 @@ def collect_sync_rows( ) -> "tuple[list[tuple[Any, ...]], list[str]]": """Collect OracleDB sync rows as tuples with normalized column names. - LOB coercion is still applied to each row. The raw coerced tuples are - returned instead of dicts so that ``SQLResult`` can handle lazy dict - materialization based on ``row_format``. + LOB locators are still read for OUT binds, PL/SQL, and explicit + ``fetch_lobs=True``. Direct-fetched LOB values under the 1 GB ceiling + arrive as plain strings/bytes and pass through unchanged. The raw coerced + tuples are returned instead of dicts so that ``SQLResult`` can handle lazy + dict materialization based on ``row_format``. Args: fetched_data: Rows returned from cursor.fetchall(). @@ -905,9 +924,11 @@ async def collect_async_rows( ) -> "tuple[list[tuple[Any, ...]], list[str]]": """Collect OracleDB async rows as tuples with normalized column names. - LOB coercion is still applied to each row. The raw coerced tuples are - returned instead of dicts so that ``SQLResult`` can handle lazy dict - materialization based on ``row_format``. + LOB locators are still read for OUT binds, PL/SQL, and explicit + ``fetch_lobs=True``. Direct-fetched LOB values under the 1 GB ceiling + arrive as plain strings/bytes and pass through unchanged. The raw coerced + tuples are returned instead of dicts so that ``SQLResult`` can handle lazy + dict materialization based on ``row_format``. Args: fetched_data: Rows returned from cursor.fetchall(). diff --git a/sqlspec/adapters/oracledb/driver.py b/sqlspec/adapters/oracledb/driver.py index 1b6b457fd..bfbcf8b4c 100644 --- a/sqlspec/adapters/oracledb/driver.py +++ b/sqlspec/adapters/oracledb/driver.py @@ -1,7 +1,7 @@ """Oracle Driver""" import logging -from typing import TYPE_CHECKING, Any, NamedTuple, cast +from typing import TYPE_CHECKING, Any, NamedTuple, cast, overload from sqlspec.adapters.oracledb._typing import ( DB_TYPE_BLOB, @@ -67,7 +67,7 @@ from sqlspec.utils.logging import get_logger, log_with_context from sqlspec.utils.module_loader import ensure_pyarrow from sqlspec.utils.text import normalize_identifier, quote_identifier, split_qualified_identifier -from sqlspec.utils.type_guards import has_pipeline_capability +from sqlspec.utils.type_guards import has_pipeline_capability, is_async_readable if TYPE_CHECKING: from collections.abc import Sequence @@ -79,7 +79,7 @@ from sqlspec.data_dictionary import VersionInfo from sqlspec.driver import ExecutionResult from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry - from sqlspec.typing import ArrowRecordBatch, ArrowReturnFormat, ArrowSchema, StatementParameters + from sqlspec.typing import ArrowRecordBatch, ArrowReturnFormat, ArrowSchema, SchemaT, StatementParameters __all__ = ( "OracleAsyncDriver", @@ -104,6 +104,24 @@ def _resolve_direct_path_target(connection: Any, table: str) -> tuple[str, str]: return ".".join(parts[:-1]), parts[-1] +def _row_has_async_readable(row: Any) -> bool: + if isinstance(row, dict): + return any(is_async_readable(value) for value in row.values()) + try: + values = row if isinstance(row, (list, tuple)) else tuple(row) + except TypeError: + values = (row,) + return any(is_async_readable(value) for value in values) + + +def _retry_arrow_without_fetch_lobs(exc: TypeError, fetch_kwargs: "dict[str, object]") -> "dict[str, object] | None": + if "fetch_lobs" not in fetch_kwargs or "fetch_lobs" not in str(exc): + return None + retry_kwargs = dict(fetch_kwargs) + retry_kwargs.pop("fetch_lobs", None) + return retry_kwargs + + # Oracle SQL-context byte thresholds (4000 / 2000) live in driver_features so users # on MAX_STRING_SIZE=EXTENDED databases can override them; defaults are wired in # core.apply_driver_features and read at the dispatch_execute call sites below. @@ -524,13 +542,70 @@ def with_cursor(self, connection: OracleSyncConnection) -> OracleSyncCursor: prefetchrows=self.driver_features.get("prefetchrows"), ) - def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "SyncRowStream[dict[str, Any]] | None": + @overload + def select_stream( + self, + statement: "SQL | Statement | QueryBuilder", + /, + *parameters: "StatementParameters | StatementFilter", + schema_type: "type[SchemaT]", + statement_config: "StatementConfig | None" = None, + chunk_size: int = 1000, + native_only: bool = False, + **kwargs: Any, + ) -> "SyncRowStream[SchemaT]": ... + + @overload + def select_stream( + self, + statement: "SQL | Statement | QueryBuilder", + /, + *parameters: "StatementParameters | StatementFilter", + schema_type: None = None, + statement_config: "StatementConfig | None" = None, + chunk_size: int = 1000, + native_only: bool = False, + **kwargs: Any, + ) -> "SyncRowStream[dict[str, Any]]": ... + + def select_stream( + self, + statement: "SQL | Statement | QueryBuilder", + /, + *parameters: "StatementParameters | StatementFilter", + schema_type: "type[SchemaT] | None" = None, + statement_config: "StatementConfig | None" = None, + chunk_size: int = 1000, + native_only: bool = False, + **kwargs: Any, + ) -> "SyncRowStream[SchemaT] | SyncRowStream[dict[str, Any]]": + """Execute a query and stream rows in chunks with Oracle fetch tuning.""" + fetch_lobs = cast("bool | None", kwargs.pop("fetch_lobs", None)) + if chunk_size < 1: + msg = "chunk_size must be greater than or equal to 1" + raise ValueError(msg) + config = statement_config or self.statement_config + sql_statement = self.prepare_statement(statement, parameters, statement_config=config, kwargs=kwargs) + stream = self.dispatch_select_stream(sql_statement, chunk_size, fetch_lobs=fetch_lobs) + if stream is not None: + return stream._with_schema_type(schema_type) + return super().select_stream( + sql_statement, + schema_type=schema_type, + statement_config=config, + chunk_size=chunk_size, + native_only=native_only, + ) + + def dispatch_select_stream( + self, statement: "SQL", chunk_size: int, fetch_lobs: bool | None = None + ) -> "SyncRowStream[dict[str, Any]] | None": """Return a native oracledb row stream backed by chunked ``fetchmany``.""" if not statement.returns_rows(): return None sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) stream_parameters = cast("list[object] | tuple[object, ...] | dict[object, object] | None", prepared_parameters) - return SyncRowStream(OracleSyncStreamSource(self, sql, stream_parameters, chunk_size)) + return SyncRowStream(OracleSyncStreamSource(self, sql, stream_parameters, chunk_size, fetch_lobs=fetch_lobs)) def handle_database_exceptions(self) -> "OracleSyncExceptionHandler": """Handle database-specific exceptions and wrap them appropriately.""" @@ -794,12 +869,18 @@ def _execute_arrow_dataframe(self, sql: str, parameters: "Any", batch_size: int return execute_df(sql, params, arraysize=batch_size or 1000) except TypeError: return execute_df(sql, params) - return self.connection.fetch_df_all( - statement=sql, - parameters=params, - arraysize=batch_size or 1000, - **build_arrow_fetch_kwargs(self.driver_features), - ) + fetch_kwargs = build_arrow_fetch_kwargs(self.driver_features) + try: + return self.connection.fetch_df_all( + statement=sql, parameters=params, arraysize=batch_size or 1000, **fetch_kwargs + ) + except TypeError as exc: + retry_kwargs = _retry_arrow_without_fetch_lobs(exc, fetch_kwargs) + if retry_kwargs is None: + raise + return self.connection.fetch_df_all( + statement=sql, parameters=params, arraysize=batch_size or 1000, **retry_kwargs + ) def _fetch_arrow_record_batches( self, @@ -813,9 +894,19 @@ def _fetch_arrow_record_batches( params = parameters if parameters is not None else [] record_batches: list[ArrowRecordBatch] = [] batch_schema: ArrowSchema | None = None - for oracle_df in self.connection.fetch_df_batches( - statement=sql, parameters=params, size=batch_size or 1000, **build_arrow_fetch_kwargs(self.driver_features) - ): + fetch_kwargs = build_arrow_fetch_kwargs(self.driver_features) + try: + oracle_dfs = self.connection.fetch_df_batches( + statement=sql, parameters=params, size=batch_size or 1000, **fetch_kwargs + ) + except TypeError as exc: + retry_kwargs = _retry_arrow_without_fetch_lobs(exc, fetch_kwargs) + if retry_kwargs is None: + raise + oracle_dfs = self.connection.fetch_df_batches( + statement=sql, parameters=params, size=batch_size or 1000, **retry_kwargs + ) + for oracle_df in oracle_dfs: batch_table = pa.table(oracle_df) column_names = normalize_column_names(batch_table.column_names, self.driver_features) if column_names != batch_table.column_names: @@ -1104,13 +1195,70 @@ def with_cursor(self, connection: OracleAsyncConnection) -> OracleAsyncCursor: prefetchrows=self.driver_features.get("prefetchrows"), ) - def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "AsyncRowStream[dict[str, Any]] | None": + @overload + def select_stream( + self, + statement: "SQL | Statement | QueryBuilder", + /, + *parameters: "StatementParameters | StatementFilter", + schema_type: "type[SchemaT]", + statement_config: "StatementConfig | None" = None, + chunk_size: int = 1000, + native_only: bool = False, + **kwargs: Any, + ) -> "AsyncRowStream[SchemaT]": ... + + @overload + def select_stream( + self, + statement: "SQL | Statement | QueryBuilder", + /, + *parameters: "StatementParameters | StatementFilter", + schema_type: None = None, + statement_config: "StatementConfig | None" = None, + chunk_size: int = 1000, + native_only: bool = False, + **kwargs: Any, + ) -> "AsyncRowStream[dict[str, Any]]": ... + + def select_stream( + self, + statement: "SQL | Statement | QueryBuilder", + /, + *parameters: "StatementParameters | StatementFilter", + schema_type: "type[SchemaT] | None" = None, + statement_config: "StatementConfig | None" = None, + chunk_size: int = 1000, + native_only: bool = False, + **kwargs: Any, + ) -> "AsyncRowStream[SchemaT] | AsyncRowStream[dict[str, Any]]": + """Execute a query and stream rows in chunks with Oracle fetch tuning.""" + fetch_lobs = cast("bool | None", kwargs.pop("fetch_lobs", None)) + if chunk_size < 1: + msg = "chunk_size must be greater than or equal to 1" + raise ValueError(msg) + config = statement_config or self.statement_config + sql_statement = self.prepare_statement(statement, parameters, statement_config=config, kwargs=kwargs) + stream = self.dispatch_select_stream(sql_statement, chunk_size, fetch_lobs=fetch_lobs) + if stream is not None: + return stream._with_schema_type(schema_type) + return super().select_stream( + sql_statement, + schema_type=schema_type, + statement_config=config, + chunk_size=chunk_size, + native_only=native_only, + ) + + def dispatch_select_stream( + self, statement: "SQL", chunk_size: int, fetch_lobs: bool | None = None + ) -> "AsyncRowStream[dict[str, Any]] | None": """Return a native oracledb row stream backed by chunked ``fetchmany``.""" if not statement.returns_rows(): return None sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) stream_parameters = cast("list[object] | tuple[object, ...] | dict[object, object] | None", prepared_parameters) - return AsyncRowStream(OracleAsyncStreamSource(self, sql, stream_parameters, chunk_size)) + return AsyncRowStream(OracleAsyncStreamSource(self, sql, stream_parameters, chunk_size, fetch_lobs=fetch_lobs)) def handle_database_exceptions(self) -> "OracleAsyncExceptionHandler": """Handle database-specific exceptions and wrap them appropriately.""" @@ -1336,9 +1484,12 @@ def data_dictionary(self) -> "OracledbAsyncDataDictionary": def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect Oracle async rows for the direct execution path. - Uses synchronous LOB coercion. For async LOB coercion, the standard - dispatch path via collect_async_rows is used instead. + Falls back to the standard async dispatch path when rows contain async + LOB locators, because those must be read with ``collect_async_rows``. """ + if fetched and any(_row_has_async_readable(row) for row in fetched): + msg = "Oracle async LOB locators require async row collection" + raise NotImplementedError(msg) column_names, requires_lob_coercion = self._resolve_row_metadata(cursor.description) data, column_names = collect_sync_rows( cast("list[Any] | None", fetched), @@ -1382,12 +1533,18 @@ async def _execute_arrow_dataframe(self, sql: str, parameters: "Any", batch_size return await execute_df(sql, params, arraysize=batch_size or 1000) except TypeError: return await execute_df(sql, params) - return await self.connection.fetch_df_all( - statement=sql, - parameters=params, - arraysize=batch_size or 1000, - **build_arrow_fetch_kwargs(self.driver_features), - ) + fetch_kwargs = build_arrow_fetch_kwargs(self.driver_features) + try: + return await self.connection.fetch_df_all( + statement=sql, parameters=params, arraysize=batch_size or 1000, **fetch_kwargs + ) + except TypeError as exc: + retry_kwargs = _retry_arrow_without_fetch_lobs(exc, fetch_kwargs) + if retry_kwargs is None: + raise + return await self.connection.fetch_df_all( + statement=sql, parameters=params, arraysize=batch_size or 1000, **retry_kwargs + ) async def _fetch_arrow_record_batches( self, @@ -1401,9 +1558,19 @@ async def _fetch_arrow_record_batches( params = parameters if parameters is not None else [] record_batches: list[ArrowRecordBatch] = [] batch_schema: ArrowSchema | None = None - async for oracle_df in self.connection.fetch_df_batches( - statement=sql, parameters=params, size=batch_size or 1000, **build_arrow_fetch_kwargs(self.driver_features) - ): + fetch_kwargs = build_arrow_fetch_kwargs(self.driver_features) + try: + oracle_dfs = self.connection.fetch_df_batches( + statement=sql, parameters=params, size=batch_size or 1000, **fetch_kwargs + ) + except TypeError as exc: + retry_kwargs = _retry_arrow_without_fetch_lobs(exc, fetch_kwargs) + if retry_kwargs is None: + raise + oracle_dfs = self.connection.fetch_df_batches( + statement=sql, parameters=params, size=batch_size or 1000, **retry_kwargs + ) + async for oracle_df in oracle_dfs: batch_table = pa.table(oracle_df) column_names = normalize_column_names(batch_table.column_names, self.driver_features) if column_names != batch_table.column_names: diff --git a/sqlspec/adapters/oracledb/type_converter.py b/sqlspec/adapters/oracledb/type_converter.py index 11c61e4af..188bce10d 100644 --- a/sqlspec/adapters/oracledb/type_converter.py +++ b/sqlspec/adapters/oracledb/type_converter.py @@ -1,65 +1,27 @@ """Oracle-specific type conversion with LOB optimization. Provides specialized type handling for Oracle databases, including -efficient LOB (Large Object) processing and JSON storage detection. +efficient LOB (Large Object) processing and vector helpers. """ import array -import re -from typing import Any, Final +from typing import Any -from sqlspec.core.type_converter import CachedOutputConverter from sqlspec.typing import NUMPY_INSTALLED from sqlspec.utils.sync_tools import ensure_async_ from sqlspec.utils.type_guards import is_readable -__all__ = ("ORACLE_JSON_STORAGE_REGEX", "ORACLE_SPECIAL_CHARS", "OracleOutputConverter") +__all__ = ("OracleOutputConverter",) -ORACLE_JSON_STORAGE_REGEX: Final[re.Pattern[str]] = re.compile( - r"^(?:" - r"(?PJSON)|" - r"(?PBLOB.*OSON)|" - r"(?PBLOB.*JSON)|" - r"(?PCLOB.*JSON)" - r")$", - re.IGNORECASE, -) -ORACLE_SPECIAL_CHARS: Final[frozenset[str]] = frozenset({"{", "[", "-", ":", "T", "."}) - - -class OracleOutputConverter(CachedOutputConverter): +class OracleOutputConverter: """Oracle-specific output conversion with LOB optimization. - Extends CachedOutputConverter with Oracle-specific functionality - including streaming LOB support and JSON storage type detection. + Provides streaming LOB support and NumPy vector helpers. """ __slots__ = () - def __init__(self, cache_size: int = 5000) -> None: - """Initialize converter with Oracle-specific options. - - Args: - cache_size: Maximum number of string values to cache (default: 5000) - """ - super().__init__(special_chars=ORACLE_SPECIAL_CHARS, cache_size=cache_size) - - def _convert_detected(self, value: str, detected_type: str) -> Any: - """Convert value with Oracle-specific handling. - - Args: - value: String value to convert. - detected_type: Detected type name. - - Returns: - Converted value, or original on failure. - """ - try: - return self.convert_value(value, detected_type) - except Exception: - return value - async def process_lob(self, value: Any) -> Any: """Process Oracle LOB objects efficiently. diff --git a/sqlspec/adapters/psqlpy/core.py b/sqlspec/adapters/psqlpy/core.py index 140ad8241..394204cbf 100644 --- a/sqlspec/adapters/psqlpy/core.py +++ b/sqlspec/adapters/psqlpy/core.py @@ -46,7 +46,7 @@ if TYPE_CHECKING: from collections.abc import Callable, Mapping - from sqlspec.core import SQL, ParameterStyleConfig, StatementConfig + from sqlspec.core import ParameterStyleConfig, StatementConfig __all__ = ( "PsqlpyStreamSource", @@ -228,11 +228,16 @@ def _parameter_config(base_config: "ParameterStyleConfig") -> "ParameterStyleCon return base_config.replace(type_coercion_map=updated_type_map) -def build_statement_config(*, json_serializer: "Callable[[Any], str] | None" = None) -> "StatementConfig": +def build_statement_config( + *, json_serializer: "Callable[[Any], str] | None" = None, json_deserializer: "Callable[[str], Any] | None" = None +) -> "StatementConfig": """Construct the psqlpy statement configuration with optional JSON codecs.""" serializer = json_serializer or to_json + deserializer = json_deserializer or from_json profile = driver_profile - base_config = build_statement_config_from_profile(profile, json_serializer=serializer) + base_config = build_statement_config_from_profile( + profile, json_serializer=serializer, json_deserializer=deserializer + ) parameter_config = _parameter_config(base_config.parameter_config) return base_config.replace(parameter_config=parameter_config) @@ -258,12 +263,16 @@ def apply_driver_features( """Apply psqlpy driver feature defaults to statement config.""" features: dict[str, Any] = dict(driver_features) if driver_features else {} serializer = features.get("json_serializer", to_json) + deserializer = features.get("json_deserializer", from_json) features.setdefault("json_serializer", serializer) + features.setdefault("json_deserializer", deserializer) features.setdefault("enable_cast_detection", True) features.setdefault("enable_pgvector", PGVECTOR_INSTALLED) features.setdefault("enable_paradedb", True) - base_config = build_statement_config_from_profile(driver_profile, json_serializer=serializer) + base_config = build_statement_config_from_profile( + driver_profile, json_serializer=serializer, json_deserializer=deserializer + ) parameter_config = _parameter_config(base_config.parameter_config) statement_config = statement_config.replace(parameter_config=parameter_config) @@ -485,11 +494,19 @@ def extract_rows_affected(result: Any) -> int: return -1 -def get_parameter_casts(statement: "SQL") -> "dict[int, str]": +def get_parameter_casts(statement: Any) -> "dict[int, str]": """Get parameter cast metadata from compiled statements.""" - processed_state = statement.get_processed_state() - if processed_state is not Empty: - return processed_state.parameter_casts or {} + processed_state = getattr(statement, "get_processed_state", None) + if callable(processed_state): + state = cast("Any", processed_state()) + if state is not Empty: + return state.parameter_casts or {} + parameter_casts = getattr(statement, "parameter_casts", None) + if parameter_casts: + return cast("dict[int, str]", parameter_casts) + processed_state = getattr(statement, "processed_state", None) + if processed_state is not None and processed_state is not Empty: + return cast("dict[int, str]", getattr(processed_state, "parameter_casts", {}) or {}) return {} diff --git a/sqlspec/adapters/psqlpy/driver.py b/sqlspec/adapters/psqlpy/driver.py index 53c4f34ed..64bfb5155 100644 --- a/sqlspec/adapters/psqlpy/driver.py +++ b/sqlspec/adapters/psqlpy/driver.py @@ -26,7 +26,6 @@ split_schema_and_table, ) from sqlspec.adapters.psqlpy.data_dictionary import PsqlpyDataDictionary -from sqlspec.adapters.psqlpy.type_converter import PostgreSQLOutputConverter from sqlspec.core import SQL, StatementConfig, get_cache_config, register_driver_profile from sqlspec.driver import AsyncDriverAdapterBase, AsyncRowStream, BaseAsyncExceptionHandler from sqlspec.exceptions import SQLSpecError @@ -37,7 +36,7 @@ from collections.abc import Mapping, Sequence from sqlspec.adapters.psqlpy._typing import PsqlpyConnection - from sqlspec.core import ArrowResult + from sqlspec.core import ArrowResult, SQLResult from sqlspec.driver import ExecutionResult from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry @@ -45,8 +44,6 @@ logger = get_logger("sqlspec.adapters.psqlpy") -_type_converter = PostgreSQLOutputConverter() - class PsqlpyExceptionHandler(BaseAsyncExceptionHandler): """Async context manager for handling psqlpy database exceptions. @@ -110,12 +107,7 @@ async def dispatch_execute(self, cursor: "PsqlpyConnection", statement: SQL) -> ExecutionResult with execution metadata """ sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - - driver_parameters = prepared_parameters - operation_type = statement.operation_type - should_coerce = operation_type != "SELECT" - effective_parameters = coerce_numeric_for_write(driver_parameters) if should_coerce else driver_parameters - params = cast("Sequence[Any] | Mapping[str, Any] | None", effective_parameters) or [] + params = cast("Sequence[Any] | Mapping[str, Any] | None", prepared_parameters) or [] if statement.returns_rows(): query_result = await cursor.fetch(sql, params) @@ -150,9 +142,7 @@ async def dispatch_execute_many(self, cursor: "PsqlpyConnection", statement: SQL if not prepared_parameters: return self.create_execution_result(cursor, rowcount_override=0, is_many_result=True) - operation_type = statement.operation_type - should_coerce = operation_type != "SELECT" - formatted_parameters = format_execute_many_parameters(prepared_parameters, coerce_numeric=should_coerce) + formatted_parameters = format_execute_many_parameters(prepared_parameters, coerce_numeric=False) await cursor.execute_many(sql, formatted_parameters) @@ -192,6 +182,21 @@ async def dispatch_execute_script(self, cursor: "PsqlpyConnection", statement: S last_result, statement_count=len(statements), successful_statements=successful_count, is_script_result=True ) + async def _execute_cache_hit( + self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: Any + ) -> "SQLResult": + """Execute cached psqlpy queries with cast-aware parameter preparation.""" + prepared_params = self.prepare_driver_parameters(params, self.statement_config, prepared_statement=cached) + direct_statement = self._cached_statement( + sql, + params, + cached, + cast("tuple[Any, ...] | list[Any] | dict[str, Any]", prepared_params), + params_are_simple=True, + compiled_sql=cached.compiled_sql, + ) + return await self._execute_cached_statement(direct_statement) + # ───────────────────────────────────────────────────────────────────────────── # TRANSACTION MANAGEMENT # ───────────────────────────────────────────────────────────────────────────── @@ -405,6 +410,13 @@ def prepare_driver_parameters( else: prepared = super().prepare_driver_parameters(parameters, statement_config, is_many, prepared_statement) + if not is_many: + operation_type = getattr(prepared_statement, "operation_type", None) + if operation_type != "SELECT": + prepared = coerce_numeric_for_write(prepared) + elif getattr(prepared_statement, "operation_type", None) != "SELECT": + prepared = coerce_numeric_for_write(prepared) + if not is_many and isinstance(prepared, list): prepared = tuple(prepared) diff --git a/sqlspec/adapters/psqlpy/type_converter.py b/sqlspec/adapters/psqlpy/type_converter.py index 50efa9f01..a29f4902a 100644 --- a/sqlspec/adapters/psqlpy/type_converter.py +++ b/sqlspec/adapters/psqlpy/type_converter.py @@ -1,94 +1,17 @@ -"""PostgreSQL-specific type conversion for psqlpy adapter. +"""PostgreSQL-specific helpers for the psqlpy adapter. -Combines output conversion (database results → Python) and input conversion -(Python params → PostgreSQL format) in a single module. Designed for mypyc -compilation with no nested functions. - -Output conversion handles: - - PostgreSQL-specific types like intervals and arrays - - Standard type detection (UUID, JSON, datetime, etc.) - -Input conversion handles: - - pgvector type handlers (placeholder for future support) +This module preserves the ``register_pgvector`` placeholder used by the +driver configuration layer. """ -import re -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING -from sqlspec.core.type_converter import CachedOutputConverter from sqlspec.typing import PGVECTOR_INSTALLED -from sqlspec.utils.logging import get_logger if TYPE_CHECKING: from psqlpy import Connection -__all__ = ("PG_SPECIAL_CHARS", "PG_SPECIFIC_REGEX", "PostgreSQLOutputConverter", "register_pgvector") - -logger = get_logger(__name__) - -PG_SPECIFIC_REGEX: Final[re.Pattern[str]] = re.compile( - r"^(?:" - r"(?P(?:(?:\d+\s+(?:year|month|day|hour|minute|second)s?\s*)+)|(?:P(?:\d+Y)?(?:\d+M)?(?:\d+D)?(?:T(?:\d+H)?(?:\d+M)?(?:\d+(?:\.\d+)?S)?)?))|" - r"(?P\{(?:[^{}]+|\{[^{}]*\})*\})" - r")$", - re.IGNORECASE, -) - -PG_SPECIAL_CHARS: Final[frozenset[str]] = frozenset({"{", "-", ":", "T", ".", "P", "[", "Y", "M", "D", "H", "S"}) - - -class PostgreSQLOutputConverter(CachedOutputConverter): - """PostgreSQL-specific output conversion with interval and array support. - - Extends CachedOutputConverter with PostgreSQL-specific functionality - for interval and array type handling. - """ - - __slots__ = () - - def __init__(self, cache_size: int = 5000) -> None: - """Initialize converter with PostgreSQL-specific options. - - Args: - cache_size: Maximum number of string values to cache (default: 5000) - """ - super().__init__(special_chars=PG_SPECIAL_CHARS, cache_size=cache_size) - - def _convert_detected(self, value: str, detected_type: str) -> Any: - """Convert value with PostgreSQL-specific handling. - - Args: - value: String value to convert. - detected_type: Detected type name. - - Returns: - Converted value or original for PostgreSQL-specific types. - """ - if detected_type in {"interval", "pg_array"}: - return value - try: - return self.convert_value(value, detected_type) - except Exception: - return value - - def detect_type(self, value: str) -> "str | None": - """Detect types including PostgreSQL-specific types. - - Args: - value: String value to analyze. - - Returns: - Type name if detected, None otherwise. - """ - detected_type = super().detect_type(value) - if detected_type: - return detected_type - - match = PG_SPECIFIC_REGEX.match(value) - if match: - return match.lastgroup - - return None +__all__ = ("register_pgvector",) def register_pgvector(connection: "Connection") -> None: diff --git a/sqlspec/adapters/psycopg/core.py b/sqlspec/adapters/psycopg/core.py index 02ef8f124..99c707d75 100644 --- a/sqlspec/adapters/psycopg/core.py +++ b/sqlspec/adapters/psycopg/core.py @@ -38,10 +38,10 @@ map_sqlstate_to_exception, ) from sqlspec.typing import PGVECTOR_INSTALLED -from sqlspec.utils.serializers import to_json +from sqlspec.utils.serializers import from_json, to_json from sqlspec.utils.text import split_qualified_identifier from sqlspec.utils.type_converters import build_json_list_converter, build_json_tuple_converter, build_uuid_coercions -from sqlspec.utils.type_guards import has_rowcount, has_sqlstate +from sqlspec.utils.type_guards import has_rowcount, has_sqlstate, resolve_row_format # Module-level lazy import for psycopg errors (mypyc optimization) try: @@ -152,18 +152,23 @@ def _custom_type_coercions() -> "dict[type, Callable[[Any], Any]]": } -def _parameter_config(profile: "DriverParameterProfile", serializer: "Callable[[Any], str]") -> "ParameterStyleConfig": +def _parameter_config( + profile: "DriverParameterProfile", serializer: "Callable[[Any], str]", deserializer: "Callable[[str], Any]" +) -> "ParameterStyleConfig": """Construct parameter configuration with shared JSON serializer support. Args: profile: Driver parameter profile to extend. serializer: JSON serializer for parameter coercion. + deserializer: JSON deserializer for result coercion. Returns: ParameterStyleConfig with updated type coercions. """ - base_config = build_statement_config_from_profile(profile, json_serializer=serializer).parameter_config + base_config = build_statement_config_from_profile( + profile, json_serializer=serializer, json_deserializer=deserializer + ).parameter_config updated_type_map = dict(base_config.type_coercion_map) updated_type_map[list] = build_json_list_converter(serializer) @@ -195,12 +200,17 @@ def build_profile() -> "DriverParameterProfile": driver_profile = build_profile() -def build_statement_config(*, json_serializer: "Callable[[Any], str] | None" = None) -> "StatementConfig": +def build_statement_config( + *, json_serializer: "Callable[[Any], str] | None" = None, json_deserializer: "Callable[[str], Any] | None" = None +) -> "StatementConfig": """Construct the psycopg statement configuration with optional JSON codecs.""" serializer = json_serializer or to_json + deserializer = json_deserializer or from_json profile = driver_profile - parameter_config = _parameter_config(profile, serializer) - base_config = build_statement_config_from_profile(profile, json_serializer=serializer) + parameter_config = _parameter_config(profile, serializer, deserializer) + base_config = build_statement_config_from_profile( + profile, json_serializer=serializer, json_deserializer=deserializer + ) return base_config.replace(parameter_config=parameter_config) @@ -213,11 +223,13 @@ def apply_driver_features( """Apply psycopg driver feature defaults to statement config.""" features: dict[str, Any] = dict(driver_features) if driver_features else {} serializer = features.get("json_serializer", to_json) + deserializer = features.get("json_deserializer", from_json) features.setdefault("json_serializer", serializer) + features.setdefault("json_deserializer", deserializer) features.setdefault("enable_pgvector", PGVECTOR_INSTALLED) features.setdefault("enable_paradedb", True) - parameter_config = _parameter_config(driver_profile, serializer) + parameter_config = _parameter_config(driver_profile, serializer, deserializer) statement_config = statement_config.replace(parameter_config=parameter_config) return statement_config, features @@ -455,6 +467,7 @@ def build_pipeline_execution_result( fetched_data, column_names = collect_rows(fetched_data, cursor.description) else: column_names = column_name_resolver(cursor.description) + row_format = resolve_row_format(fetched_data) return ExecutionResult( cursor_result=cursor, rowcount_override=None, @@ -467,7 +480,7 @@ def build_pipeline_execution_result( is_script_result=False, is_select_result=True, is_many_result=False, - row_format="tuple", + row_format=row_format, last_inserted_id=None, ) @@ -508,6 +521,7 @@ async def build_async_pipeline_execution_result( fetched_data, column_names = collect_rows(fetched_data, cursor.description) else: column_names = column_name_resolver(cursor.description) + row_format = resolve_row_format(fetched_data) return ExecutionResult( cursor_result=cursor, rowcount_override=None, @@ -520,7 +534,7 @@ async def build_async_pipeline_execution_result( is_script_result=False, is_select_result=True, is_many_result=False, - row_format="tuple", + row_format=row_format, last_inserted_id=None, ) diff --git a/sqlspec/adapters/psycopg/driver.py b/sqlspec/adapters/psycopg/driver.py index b0973bb90..794a673ef 100644 --- a/sqlspec/adapters/psycopg/driver.py +++ b/sqlspec/adapters/psycopg/driver.py @@ -62,7 +62,7 @@ from sqlspec.exceptions import SQLSpecError, StackExecutionError from sqlspec.utils.logging import get_logger from sqlspec.utils.text import normalize_identifier, quote_identifier -from sqlspec.utils.type_guards import is_readable +from sqlspec.utils.type_guards import is_readable, resolve_row_format if TYPE_CHECKING: from sqlspec.adapters.psycopg._typing import PsycopgPipelineDriver @@ -83,7 +83,6 @@ ) logger = get_logger("sqlspec.adapters.psycopg") -COLUMN_CACHE_MAX_SIZE = 256 class PsycopgPipelineMixin: @@ -156,7 +155,7 @@ class PsycopgSyncDriver(PsycopgPipelineMixin, SyncDriverAdapterBase): bulk data transfer, and PostgreSQL-specific error handling. """ - __slots__ = ("_column_name_cache", "_data_dictionary") + __slots__ = ("_data_dictionary",) dialect = "postgres" def __init__( @@ -172,7 +171,6 @@ def __init__( super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features) self._data_dictionary: PsycopgSyncDataDictionary | None = None - self._column_name_cache: dict[int, tuple[Any, list[str]]] = {} # ───────────────────────────────────────────────────────────────────────────── # CORE DISPATCH METHODS @@ -196,6 +194,7 @@ def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionResult": fetched_data = cursor.fetchall() data = cast("list[Any] | None", fetched_data) or [] column_names = self._resolve_column_names(cursor.description) + row_format = resolve_row_format(data) return self.create_execution_result( cursor, @@ -203,7 +202,7 @@ def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionResult": column_names=column_names, data_row_count=len(data), is_select_result=True, - row_format="tuple", + row_format=row_format, ) affected_rows = resolve_rowcount(cursor) @@ -584,21 +583,10 @@ def data_dictionary(self) -> "PsycopgSyncDataDictionary": # ───────────────────────────────────────────────────────────────────────────── def _resolve_column_names(self, description: Any) -> list[str]: - """Resolve and cache psycopg column names for hot row materialization paths.""" + """Resolve psycopg column names for row materialization paths.""" if not description: return [] - - cache_key = id(description) - cached = self._column_name_cache.get(cache_key) - if cached is not None and cached[0] is description: - return cached[1] - - column_names = [col.name for col in description] - - if len(self._column_name_cache) >= COLUMN_CACHE_MAX_SIZE: - self._column_name_cache.pop(next(iter(self._column_name_cache))) - self._column_name_cache[cache_key] = (description, column_names) - return column_names + return [col.name for col in description] def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect psycopg sync rows for the direct execution path.""" @@ -648,7 +636,7 @@ class PsycopgAsyncDriver(PsycopgPipelineMixin, AsyncDriverAdapterBase): and async pub/sub support. """ - __slots__ = ("_column_name_cache", "_data_dictionary") + __slots__ = ("_data_dictionary",) dialect = "postgres" def __init__( @@ -664,7 +652,6 @@ def __init__( super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features) self._data_dictionary: PsycopgAsyncDataDictionary | None = None - self._column_name_cache: dict[int, tuple[Any, list[str]]] = {} # ───────────────────────────────────────────────────────────────────────────── # CORE DISPATCH METHODS @@ -688,6 +675,7 @@ async def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionRes fetched_data = await cursor.fetchall() data = cast("list[Any] | None", fetched_data) or [] column_names = self._resolve_column_names(cursor.description) + row_format = resolve_row_format(data) return self.create_execution_result( cursor, @@ -695,7 +683,7 @@ async def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionRes column_names=column_names, data_row_count=len(data), is_select_result=True, - row_format="tuple", + row_format=row_format, ) affected_rows = resolve_rowcount(cursor) @@ -1088,21 +1076,10 @@ def data_dictionary(self) -> "PsycopgAsyncDataDictionary": # ───────────────────────────────────────────────────────────────────────────── def _resolve_column_names(self, description: Any) -> list[str]: - """Resolve and cache psycopg column names for hot row materialization paths.""" + """Resolve psycopg column names for row materialization paths.""" if not description: return [] - - cache_key = id(description) - cached = self._column_name_cache.get(cache_key) - if cached is not None and cached[0] is description: - return cached[1] - - column_names = [col.name for col in description] - - if len(self._column_name_cache) >= COLUMN_CACHE_MAX_SIZE: - self._column_name_cache.pop(next(iter(self._column_name_cache))) - self._column_name_cache[cache_key] = (description, column_names) - return column_names + return [col.name for col in description] def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect psycopg async rows for the direct execution path.""" diff --git a/sqlspec/adapters/pymssql/config.py b/sqlspec/adapters/pymssql/config.py index f99050a06..9487ee3e4 100644 --- a/sqlspec/adapters/pymssql/config.py +++ b/sqlspec/adapters/pymssql/config.py @@ -96,7 +96,7 @@ class PymssqlConfig(SyncDatabaseConfig[PymssqlConnection, PymssqlConnectionPool, supports_native_arrow_import: "ClassVar[bool]" = False supports_native_parquet_export: "ClassVar[bool]" = False supports_native_parquet_import: "ClassVar[bool]" = False - supports_native_row_streaming: "ClassVar[bool]" = False + supports_native_row_streaming: "ClassVar[bool]" = True _connection_context_class: "ClassVar[type[PymssqlConnectionContext]]" = PymssqlConnectionContext _session_factory_class: "ClassVar[type[_PymssqlSessionConnectionHandler]]" = _PymssqlSessionConnectionHandler _session_context_class: "ClassVar[type[PymssqlSessionContext]]" = PymssqlSessionContext diff --git a/sqlspec/adapters/pymssql/core.py b/sqlspec/adapters/pymssql/core.py index 7f8dfefe7..53e655795 100644 --- a/sqlspec/adapters/pymssql/core.py +++ b/sqlspec/adapters/pymssql/core.py @@ -45,6 +45,7 @@ ) _ERROR_NUMBER_PATTERN: Final[re.Pattern[str]] = re.compile(r"\(([-]?\d+)\)") +_COLUMN_CACHE_MAX_SIZE: Final[int] = 256 _ERROR_CODE_MAPPING: Final[dict[int, tuple[type[SQLSpecError], str]]] = { 2601: (UniqueViolationError, "unique constraint violation"), 2627: (UniqueViolationError, "unique constraint violation"), @@ -180,18 +181,32 @@ def create_mapped_exception(error: Exception, *, logger: "Logger | None" = None) return SQLSpecError(f"SQL Server database error. Original error: {error}") -def resolve_column_names(description: "Sequence[Any] | None") -> "list[str]": +def resolve_column_names( + description: "Sequence[Any] | None", column_name_cache: "dict[int, tuple[Any, list[str]]] | None" = None +) -> "list[str]": """Resolve ordered column names from cursor metadata.""" if not description: return [] - return [desc[0] for desc in description] + if column_name_cache is None: + return [desc[0] for desc in description] + cache_key = id(description) + cached = column_name_cache.get(cache_key) + if cached is not None and cached[0] is description: + return cached[1] + column_names = [desc[0] for desc in description] + if len(column_name_cache) >= _COLUMN_CACHE_MAX_SIZE: + column_name_cache.pop(next(iter(column_name_cache))) + column_name_cache[cache_key] = (description, column_names) + return column_names def collect_rows( - fetched_data: "Sequence[Any] | None", description: "Sequence[Any] | None" + fetched_data: "Sequence[Any] | None", + description: "Sequence[Any] | None", + column_name_cache: "dict[int, tuple[Any, list[str]]] | None" = None, ) -> "tuple[list[Any], list[str], str]": """Collect pymssql rows, preserving tuple row shape.""" - column_names = resolve_column_names(description) + column_names = resolve_column_names(description, column_name_cache) if not fetched_data: return [], column_names, "tuple" return list(fetched_data), column_names, "tuple" diff --git a/sqlspec/adapters/pymssql/driver.py b/sqlspec/adapters/pymssql/driver.py index b58438820..9b4dec567 100644 --- a/sqlspec/adapters/pymssql/driver.py +++ b/sqlspec/adapters/pymssql/driver.py @@ -1,5 +1,6 @@ """pymssql SQL Server driver implementation.""" +import contextlib from collections.abc import Sized from typing import TYPE_CHECKING, Any, cast @@ -23,7 +24,13 @@ ) from sqlspec.adapters.pymssql.data_dictionary import PymssqlSyncDataDictionary from sqlspec.core import SQL, StatementConfig, get_cache_config, register_driver_profile -from sqlspec.driver import BaseSyncExceptionHandler, ExecutionResult, SyncDriverAdapterBase +from sqlspec.driver import ( + BaseSyncExceptionHandler, + ExecutionResult, + SyncDriverAdapterBase, + SyncRowStream, + rows_to_dicts, +) from sqlspec.exceptions import SQLSpecError from sqlspec.utils.logging import get_logger @@ -57,10 +64,63 @@ def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "Ba return False +class PymssqlStreamSource: + """Native pymssql chunk source backed by ``cursor.fetchmany()``.""" + + __slots__ = ("_chunk_size", "_column_names", "_cursor_manager", "_driver", "_parameters", "_sql") + + def __init__(self, driver: "PymssqlDriver", sql: str, parameters: Any, chunk_size: int) -> None: + self._driver = driver + self._sql = sql + self._parameters = parameters + self._chunk_size = chunk_size + self._cursor_manager: PymssqlCursor | None = None + self._column_names: list[str] | None = None + + def start(self) -> None: + cursor_manager = self._driver.with_cursor(self._driver.connection) + try: + cursor = cursor_manager.__enter__() + handler = self._driver.handle_database_exceptions() + with handler: + cursor.execute(self._sql, normalize_execute_parameters(self._parameters)) + self._driver._check_pending_exception(handler) + except BaseException: + with contextlib.suppress(Exception): + cursor_manager.__exit__(None, None, None) + raise + self._cursor_manager = cursor_manager + + def fetch_chunk(self) -> "list[dict[str, Any]]": + cursor_manager = self._cursor_manager + if cursor_manager is None or cursor_manager.cursor is None: + return [] + cursor = cursor_manager.cursor + handler = self._driver.handle_database_exceptions() + rows: Any = [] + with handler: + rows = cursor.fetchmany(self._chunk_size) + self._driver._check_pending_exception(handler) + if not rows: + return [] + column_names = self._column_names + if column_names is None: + column_names = resolve_column_names(cursor.description or None, self._driver._column_name_cache) + self._column_names = column_names + return rows_to_dicts(rows, column_names) + + def close(self) -> None: + cursor_manager = self._cursor_manager + self._cursor_manager = None + if cursor_manager is not None: + with contextlib.suppress(Exception): + cursor_manager.__exit__(None, None, None) + + class PymssqlDriver(SyncDriverAdapterBase): """SQL Server database driver using pymssql.""" - __slots__ = ("_data_dictionary",) + __slots__ = ("_column_name_cache", "_data_dictionary") dialect = "tsql" def __init__( @@ -76,6 +136,7 @@ def __init__( super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features) self._data_dictionary: PymssqlSyncDataDictionary | None = None + self._column_name_cache: dict[int, tuple[Any, list[str]]] = {} def dispatch_execute(self, cursor: "PymssqlRawCursor", statement: "SQL") -> "ExecutionResult": sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) @@ -84,7 +145,7 @@ def dispatch_execute(self, cursor: "PymssqlRawCursor", statement: "SQL") -> "Exe if statement.returns_rows(): fetched_data = cursor.fetchall() description = cursor.description or None - rows, column_names, row_format = collect_rows(fetched_data, description) + rows, column_names, row_format = collect_rows(fetched_data, description, self._column_name_cache) return self.create_execution_result( cursor, selected_data=rows, @@ -149,6 +210,13 @@ def with_cursor(self, connection: "PymssqlConnection") -> "PymssqlCursor": def handle_database_exceptions(self) -> "PymssqlExceptionHandler": return PymssqlExceptionHandler() + def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "SyncRowStream[dict[str, Any]] | None": + """Return a native pymssql row stream backed by ``fetchmany()``.""" + if not statement.returns_rows(): + return None + sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) + return SyncRowStream(PymssqlStreamSource(self, sql, prepared_parameters, chunk_size)) + def create_savepoint(self, name: str) -> None: self.execute_script(f"SAVE TRANSACTION {name}") @@ -165,7 +233,7 @@ def data_dictionary(self) -> "PymssqlSyncDataDictionary": return self._data_dictionary def collect_rows(self, cursor: "PymssqlRawCursor", fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": - column_names = resolve_column_names(cursor.description or None) + column_names = resolve_column_names(cursor.description or None, self._column_name_cache) return fetched, column_names, len(fetched) def resolve_rowcount(self, cursor: "PymssqlRawCursor") -> int: diff --git a/sqlspec/adapters/pymssql/type_converter.py b/sqlspec/adapters/pymssql/type_converter.py deleted file mode 100644 index 3e15b0545..000000000 --- a/sqlspec/adapters/pymssql/type_converter.py +++ /dev/null @@ -1,92 +0,0 @@ -"""Type converters for pymssql parameter binding.""" - -from typing import TYPE_CHECKING, Any, Final, cast -from uuid import UUID - -from sqlspec.utils.module_loader import ensure_pyarrow -from sqlspec.utils.serializers import from_json, to_json - -if TYPE_CHECKING: - from collections.abc import Callable - - import pyarrow as pa - -__all__ = ("PymssqlTypeConverter", "mssql_type_to_arrow") - -_MSSQL_ARROW_TYPE_SPECS: Final[dict[str, tuple[str, tuple[Any, ...], dict[str, Any]]]] = { - "bit": ("bool_", (), {}), - "tinyint": ("uint8", (), {}), - "smallint": ("int16", (), {}), - "int": ("int32", (), {}), - "bigint": ("int64", (), {}), - "float": ("float64", (), {}), - "real": ("float32", (), {}), - "smallmoney": ("decimal128", (10, 4), {}), - "money": ("decimal128", (19, 4), {}), - "date": ("date32", (), {}), - "datetime": ("timestamp", ("ms",), {}), - "datetime2": ("timestamp", ("us",), {}), - "smalldatetime": ("timestamp", ("s",), {}), - "datetimeoffset": ("timestamp", ("us",), {"tz": "UTC"}), - "uniqueidentifier": ("string", (), {}), - "xml": ("string", (), {}), - "image": ("binary", (), {}), - "binary": ("binary", (), {}), - "varbinary": ("binary", (), {}), - "timestamp": ("binary", (), {}), - "rowversion": ("binary", (), {}), - "char": ("string", (), {}), - "varchar": ("string", (), {}), - "nchar": ("string", (), {}), - "nvarchar": ("string", (), {}), - "text": ("string", (), {}), - "ntext": ("string", (), {}), -} - - -class PymssqlTypeConverter: - """Utility converter for explicit pymssql value coercion. - - The driver pipeline builds its internal coercions directly from feature - flags. This class is a public utility for callers that need per-value bind - or result coercion outside the driver execution path. - """ - - __slots__ = ("_json_deserializer", "_json_serializer") - - def __init__( - self, json_serializer: "Callable[[Any], str]" = to_json, json_deserializer: "Callable[[str], Any]" = from_json - ) -> None: - self._json_serializer = json_serializer - self._json_deserializer = json_deserializer - - def coerce_bind_value(self, value: "Any") -> "Any": - """Coerce Python values before pymssql parameter binding.""" - if isinstance(value, (dict, list)): - return self._json_serializer(value) - if isinstance(value, UUID): - return value - return value - - def coerce_read_value(self, value: "Any") -> "Any": - """Coerce pymssql result values after fetching.""" - return value - - -def mssql_type_to_arrow(sql_type: str, *, precision: int | None = None, scale: int | None = None) -> "pa.DataType": - """Resolve a T-SQL type name to an Arrow data type.""" - normalized_type = sql_type.lower().split("(", 1)[0].strip() - if normalized_type in {"decimal", "numeric"} and precision is not None and scale is not None: - return _arrow_type("decimal128", (precision, scale)) - spec = _MSSQL_ARROW_TYPE_SPECS.get(normalized_type) - if spec is None: - return _arrow_type("string") - name, args, kwargs = spec - return _arrow_type(name, args, kwargs) - - -def _arrow_type(name: str, args: tuple[Any, ...] = (), kwargs: dict[str, Any] | None = None) -> "pa.DataType": - ensure_pyarrow() - import pyarrow as pa - - return cast("pa.DataType", getattr(pa, name)(*args, **(kwargs or {}))) diff --git a/sqlspec/adapters/pymysql/core.py b/sqlspec/adapters/pymysql/core.py index 725634687..4848114fa 100644 --- a/sqlspec/adapters/pymysql/core.py +++ b/sqlspec/adapters/pymysql/core.py @@ -6,7 +6,7 @@ import contextlib from collections.abc import Callable, Sized -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast from sqlspec.core import DriverParameterProfile, ParameterStyle, StatementConfig, build_statement_config_from_profile from sqlspec.driver import rows_to_dicts @@ -43,6 +43,7 @@ "build_profile", "build_statement_config", "collect_rows", + "collect_stream_rows", "create_mapped_exception", "default_statement_config", "detect_json_columns", @@ -55,6 +56,7 @@ "normalize_lastrowid", "resolve_column_names", "resolve_many_rowcount", + "resolve_row_plan", "resolve_rowcount", ) @@ -174,15 +176,16 @@ def normalize_execute_parameters(parameters: Any) -> Any: class PymysqlStreamSource: """Compiled chunk source streaming dict rows from a PyMySQL unbuffered ``SSCursor``.""" - __slots__ = ("_chunk_size", "_column_names", "_cursor", "_driver", "_parameters", "_sql") + __slots__ = ("_chunk_size", "_cursor", "_driver", "_json_type_codes", "_parameters", "_row_plan", "_sql") - def __init__(self, driver: Any, sql: str, parameters: Any, chunk_size: int) -> None: + def __init__(self, driver: Any, sql: str, parameters: Any, chunk_size: int, json_type_codes: "set[int]") -> None: self._driver = driver self._sql = sql self._parameters = parameters self._chunk_size = chunk_size self._cursor: Any = None - self._column_names: list[str] | None = None + self._json_type_codes = json_type_codes + self._row_plan: tuple[list[str], list[int] | None] | None = None def start(self) -> None: from pymysql.cursors import SSCursor @@ -192,6 +195,7 @@ def start(self) -> None: cursor = self._driver.connection.cursor(SSCursor) cursor.execute(self._sql, normalize_execute_parameters(self._parameters)) self._cursor = cursor + self._row_plan = resolve_row_plan(self._cursor.description, self._json_type_codes) self._driver._check_pending_exception(handler) def fetch_chunk(self) -> "list[dict[str, Any]]": @@ -202,9 +206,10 @@ def fetch_chunk(self) -> "list[dict[str, Any]]": self._driver._check_pending_exception(handler) if not rows: return [] - if self._column_names is None: - self._column_names = [description[0] for description in self._cursor.description] - return rows_to_dicts(rows, self._column_names) + if self._row_plan is None: + self._row_plan = resolve_row_plan(self._cursor.description, self._json_type_codes) + deserializer = cast("Callable[[Any], Any]", self._driver.driver_features.get("json_deserializer", from_json)) + return collect_stream_rows(rows, self._row_plan, deserializer) def close(self) -> None: cursor = self._cursor @@ -320,7 +325,7 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> "SQLSpe if error_code in _MYSQL_MIGRATION_ERROR_CODES: if logger is not None: - logger.warning("PyMySQL expected migration error (ignoring): %s", error) + logger.warning("MySQL expected migration error (ignoring): %s", error) return True dispatch = _MYSQL_SQLSTATE_EXACT_DISPATCH.get(sqlstate) if sqlstate is not None else None @@ -376,23 +381,36 @@ def resolve_column_names(description: "Sequence[Any] | None") -> "list[str]": return [desc[0] for desc in description] -def detect_json_columns_from_description( +def resolve_row_plan( description: "Sequence[Any] | None", json_type_codes: "set[int]" -) -> "list[int]": - """Identify JSON column indexes from pre-fetched cursor description metadata.""" - if not description or not json_type_codes: - return [] +) -> "tuple[list[str], list[int] | None]": + """Resolve ordered column names and JSON column indexes in one pass.""" + if not description: + return [], None + + column_names: list[str] = [] + if not json_type_codes: + column_names.extend(column[0] for column in description) + return column_names, None json_indexes: list[int] = [] - append = json_indexes.append + append_json = json_indexes.append for index, column in enumerate(description): + column_names.append(column[0]) if isinstance(column, (tuple, list)): type_code = column[1] if len(column) > 1 else None else: type_code = column.type_code if isinstance(column, HasTypeCodeProtocol) else None if type_code in json_type_codes: - append(index) - return json_indexes + append_json(index) + return column_names, json_indexes or None + + +def detect_json_columns_from_description( + description: "Sequence[Any] | None", json_type_codes: "set[int]" +) -> "list[int]": + """Identify JSON column indexes from pre-fetched cursor description metadata.""" + return resolve_row_plan(description, json_type_codes)[1] or [] def detect_json_columns( @@ -472,38 +490,53 @@ def _deserialize_json_tuple_rows( return result +def collect_stream_rows( + fetched_data: "Sequence[Any] | None", + row_plan: "tuple[list[str], list[int] | None]", + deserializer: "Callable[[Any], Any]", + *, + logger: Any | None = None, +) -> "list[dict[str, Any]]": + """Materialize streamed rows as dicts and decode JSON columns in place.""" + column_names, json_indexes = row_plan + if not column_names or not fetched_data: + return [] + + rows = fetched_data if isinstance(fetched_data, list) else list(fetched_data) + if not rows: + return [] + + dict_rows = rows_to_dicts(rows, column_names) + if json_indexes: + dict_rows = _deserialize_json_dict_rows(column_names, dict_rows, json_indexes, deserializer, logger=logger) + return dict_rows + + def collect_rows( fetched_data: "Sequence[Any] | None", - description: "Sequence[Any] | None", - json_indexes: "list[int]", + row_plan: "tuple[list[str], list[int] | None]", deserializer: "Callable[[Any], Any]", *, - column_names: "list[str] | None" = None, logger: Any | None = None, ) -> "tuple[list[Any], list[str], str]": - """Collect PyMySQL rows with JSON decoding, preserving raw format. - - Returns: - Tuple of (rows, column_names, row_format). - """ - if not description: + """Collect PyMySQL rows with JSON decoding, preserving raw format.""" + column_names, json_indexes = row_plan + if not column_names: return [], [], "tuple" - resolved_column_names = resolve_column_names(description) if column_names is None else column_names if not fetched_data: - return [], resolved_column_names, "tuple" + return [], column_names, "tuple" first_row = fetched_data[0] if isinstance(first_row, dict): - if not json_indexes: - rows = fetched_data if isinstance(fetched_data, list) else list(fetched_data) - return rows, resolved_column_names, "dict" - rows = [dict(row) for row in fetched_data] - rows = _deserialize_json_dict_rows(resolved_column_names, rows, json_indexes, deserializer, logger=logger) - return rows, resolved_column_names, "dict" + rows = fetched_data if isinstance(fetched_data, list) else list(fetched_data) + if json_indexes: + rows = _deserialize_json_dict_rows(column_names, rows, json_indexes, deserializer, logger=logger) + return rows, column_names, "dict" + rows = fetched_data if isinstance(fetched_data, list) else list(fetched_data) if json_indexes: rows = _deserialize_json_tuple_rows(rows, json_indexes, deserializer, logger=logger) - return rows, resolved_column_names, "tuple" + return rows, column_names, "tuple" def resolve_rowcount(cursor: Any) -> int: diff --git a/sqlspec/adapters/pymysql/driver.py b/sqlspec/adapters/pymysql/driver.py index 463ff900a..61a3797a4 100644 --- a/sqlspec/adapters/pymysql/driver.py +++ b/sqlspec/adapters/pymysql/driver.py @@ -19,15 +19,14 @@ collect_rows, create_mapped_exception, default_statement_config, - detect_json_columns_from_description, driver_profile, encode_records_for_local_infile, format_identifier, normalize_execute_many_parameters, normalize_execute_parameters, normalize_lastrowid, - resolve_column_names, resolve_many_rowcount, + resolve_row_plan, resolve_rowcount, ) from sqlspec.adapters.pymysql.data_dictionary import PyMysqlDataDictionary @@ -98,12 +97,9 @@ def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionResult": if statement.returns_rows(): fetched_data = cursor.fetchall() description = cursor.description or None - column_names = resolve_column_names(description) - json_indexes = detect_json_columns_from_description(description, PYMYSQL_JSON_TYPE_CODES) + row_plan = resolve_row_plan(description, PYMYSQL_JSON_TYPE_CODES) deserializer = cast("Callable[[Any], Any]", self.driver_features.get("json_deserializer", from_json)) - rows, column_names, row_format = collect_rows( - fetched_data, description, json_indexes, deserializer, column_names=column_names, logger=logger - ) + rows, column_names, row_format = collect_rows(fetched_data, row_plan, deserializer, logger=logger) return self.create_execution_result( cursor, @@ -177,7 +173,7 @@ def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "SyncRowS if not statement.returns_rows(): return None sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - return SyncRowStream(PymysqlStreamSource(self, sql, prepared_parameters, chunk_size)) + return SyncRowStream(PymysqlStreamSource(self, sql, prepared_parameters, chunk_size, PYMYSQL_JSON_TYPE_CODES)) def handle_database_exceptions(self) -> "PyMysqlExceptionHandler": return PyMysqlExceptionHandler() @@ -279,12 +275,9 @@ def data_dictionary(self) -> "PyMysqlDataDictionary": def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect PyMySQL rows for the direct execution path.""" description = cursor.description or None - column_names = resolve_column_names(description) - json_indexes = detect_json_columns_from_description(description, PYMYSQL_JSON_TYPE_CODES) + row_plan = resolve_row_plan(description, PYMYSQL_JSON_TYPE_CODES) deserializer = cast("Callable[[Any], Any]", self.driver_features.get("json_deserializer", from_json)) - rows, column_names, _row_format = collect_rows( - fetched, description, json_indexes, deserializer, column_names=column_names, logger=logger - ) + rows, column_names, _row_format = collect_rows(fetched, row_plan, deserializer, logger=logger) return rows, column_names, len(rows) def resolve_rowcount(self, cursor: Any) -> int: diff --git a/sqlspec/adapters/spanner/core.py b/sqlspec/adapters/spanner/core.py index 89d1f8264..12bb4d2fe 100644 --- a/sqlspec/adapters/spanner/core.py +++ b/sqlspec/adapters/spanner/core.py @@ -1,8 +1,11 @@ """Spanner adapter compiled helpers.""" -from typing import TYPE_CHECKING, Any +from functools import partial +from typing import TYPE_CHECKING, Any, cast from google.api_core import exceptions as api_exceptions +from google.cloud.spanner_v1.data_types import JsonObject +from google.cloud.spanner_v1.types.type import TypeCode from sqlspec.adapters.spanner.type_converter import coerce_params_for_spanner, infer_spanner_param_types from sqlspec.core import DriverParameterProfile, ParameterStyle, StatementConfig, build_statement_config_from_profile @@ -37,6 +40,7 @@ "driver_profile", "infer_param_types", "resolve_column_names", + "resolve_row_plan", "supports_batch_update", "supports_write", ) @@ -144,6 +148,53 @@ def resolve_column_names(fields: "Sequence[Any] | None", cache: "dict[int, tuple return column_names +def _convert_json_row_value(value: Any, *, json_deserializer: "Callable[[str], Any]") -> Any: + """Convert a native Spanner JSON cell using the configured deserializer.""" + if isinstance(value, JsonObject): + json_value = cast("Any", value).serialize() + elif isinstance(value, str): + json_value = value + else: + return value + + try: + return json_deserializer(json_value) + except (TypeError, ValueError): + return value + + +def resolve_row_plan( + fields: "Sequence[Any] | None", + cache: "dict[int, tuple[Any, list[str], tuple[tuple[int, Any], ...] | None]]", + *, + json_deserializer: "Callable[[str], Any]", +) -> "tuple[list[str], tuple[tuple[int, Any], ...] | None]": + """Resolve column names and metadata-driven converters for a Spanner result.""" + if not fields: + return [], None + + cache_key = id(fields) + cached = cache.get(cache_key) + if cached is not None and cached[0] is fields: + return cached[1], cached[2] + + column_names = [getattr(field, "name", "") for field in fields] + plan_entries: list[tuple[int, Any]] = [] + append_plan_entry = plan_entries.append + convert_json = partial(_convert_json_row_value, json_deserializer=json_deserializer) + for index, field in enumerate(fields): + field_type = getattr(field, "type_", None) + type_code = getattr(field_type, "code", None) + if type_code == TypeCode.JSON: + append_plan_entry((index, convert_json)) + + plan = tuple(plan_entries) if plan_entries else None + if len(cache) >= COLUMN_CACHE_MAX_SIZE: + cache.pop(next(iter(cache))) + cache[cache_key] = (fields, column_names, plan) + return column_names, plan + + def coerce_params( params: "dict[str, Any] | list[Any] | tuple[Any, ...] | None", *, @@ -156,36 +207,39 @@ def coerce_params( def collect_rows( - rows: "Sequence[Any]", fields: "Sequence[Any]", converter: Any, *, column_names: "list[str] | None" = None -) -> "tuple[list[tuple[Any, ...]], list[str]]": + rows: "list[Any]", + fields: "Sequence[Any]", + *, + column_names: "list[str] | None" = None, + column_plan: "tuple[tuple[int, Any], ...] | None" = None, +) -> "tuple[list[Any], list[str]]": """Collect Spanner rows as tuples with type conversion applied. - Type conversion is still applied to each row value. The raw converted - tuples are returned instead of dicts so that ``SQLResult`` can handle - lazy dict materialization based on ``row_format``. + When a metadata-driven column plan is unavailable, the original row list + is returned unchanged. When a plan exists, only the planned columns are + converted and the rows are copied once. Args: rows: Rows from result set. fields: Result set fields metadata. - converter: Type converter for row values. column_names: Optional precomputed column names. + column_plan: Optional zero-copy column conversion plan. Returns: Tuple of (rows, column_names). """ resolved_column_names = column_names if column_names is not None else [field.name for field in fields] - num_columns = len(resolved_column_names) - string_converter = converter.convert_if_detected - data: list[tuple[Any, ...]] = [] + if column_plan is None and all(type(row) is tuple for row in rows): + return rows, resolved_column_names + if column_plan is None: + return [tuple(row) for row in rows], resolved_column_names + + data: list[Any] = [] append = data.append for row in rows: - converted_row: list[Any] = [] - append_value = converted_row.append - for index in range(num_columns): - value = row[index] - if isinstance(value, str): - value = string_converter(value) - append_value(value) + converted_row = list(row) + for index, converter in column_plan: + converted_row[index] = converter(converted_row[index]) append(tuple(converted_row)) return data, resolved_column_names diff --git a/sqlspec/adapters/spanner/driver.py b/sqlspec/adapters/spanner/driver.py index 6c41cbed9..fb673dc7c 100644 --- a/sqlspec/adapters/spanner/driver.py +++ b/sqlspec/adapters/spanner/driver.py @@ -1,7 +1,9 @@ """Spanner driver implementation.""" +import contextlib from collections.abc import Iterator -from typing import TYPE_CHECKING, Any, Protocol, cast +from itertools import islice +from typing import TYPE_CHECKING, Any, Protocol, cast, overload import sqlglot as _sqlglot from sqlglot import exp as _sqlglot_exp @@ -21,14 +23,19 @@ default_statement_config, driver_profile, infer_param_types, - resolve_column_names, + resolve_row_plan, supports_batch_update, supports_write, ) from sqlspec.adapters.spanner.data_dictionary import SpannerDataDictionary -from sqlspec.adapters.spanner.type_converter import SpannerOutputConverter from sqlspec.core import StatementConfig, create_arrow_result, register_driver_profile -from sqlspec.driver import BaseSyncExceptionHandler, ExecutionResult, SyncDriverAdapterBase +from sqlspec.driver import ( + BaseSyncExceptionHandler, + ExecutionResult, + SyncDriverAdapterBase, + SyncRowStream, + rows_to_dicts, +) from sqlspec.exceptions import SQLConversionError from sqlspec.utils.serializers import from_json @@ -51,7 +58,7 @@ from sqlspec.core import ArrowResult, SQLResult, Statement, StatementFilter from sqlspec.core.statement import SQL from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry - from sqlspec.typing import ArrowReturnFormat, StatementParameters + from sqlspec.typing import ArrowReturnFormat, SchemaT, StatementParameters __all__ = ( "SpannerDataDictionary", @@ -140,11 +147,106 @@ def __init__( self.timeout = timeout +class _SpannerSelectStreamSource: + """Native chunk source for Spanner SELECT streaming.""" + + __slots__ = ( + "_chunk_size", + "_column_names", + "_column_plan", + "_driver", + "_execute_kwargs", + "_param_types", + "_params", + "_result_set", + "_row_iterator", + "_sql", + ) + + def __init__( + self, + driver: "SpannerSyncDriver", + sql: str, + params: "dict[str, Any] | None", + param_types: "dict[str, Any]", + chunk_size: int, + execute_kwargs: "dict[str, Any]", + ) -> None: + self._driver = driver + self._sql = sql + self._params = params + self._param_types = param_types + self._chunk_size = chunk_size + self._execute_kwargs = execute_kwargs + self._column_names: list[str] | None = None + self._column_plan: tuple[tuple[int, Any], ...] | None = None + self._result_set: _SpannerResultSetProtocol | None = None + self._row_iterator: Iterator[Any] | None = None + + def start(self) -> None: + handler = self._driver.handle_database_exceptions() + with handler: + result_set = self._driver.connection.execute_sql( + self._sql, params=self._params, param_types=self._param_types, **self._execute_kwargs + ) + self._result_set = result_set + self._row_iterator = iter(result_set) + self._driver._check_pending_exception(handler) + + def fetch_chunk(self) -> "list[dict[str, Any]]": + result_set = self._result_set + row_iterator = self._row_iterator + column_names = self._column_names + if result_set is None or row_iterator is None: + return [] + + handler = self._driver.handle_database_exceptions() + rows: list[Any] = [] + with handler: + rows = list(islice(row_iterator, self._chunk_size)) + self._driver._check_pending_exception(handler) + if not rows: + return [] + + if column_names is None: + try: + metadata = result_set.metadata + row_type = metadata.row_type + fields = row_type.fields + except AttributeError: + msg = "Result set metadata not available." + raise SQLConversionError(msg) + if not fields: + msg = "Result set metadata not available." + raise SQLConversionError(msg) + column_names, column_plan = self._driver._resolve_row_plan(fields) + self._column_names = column_names + self._column_plan = column_plan + + converted_rows, resolved_column_names = collect_rows( + rows, (), column_names=column_names, column_plan=self._column_plan + ) + self._column_names = resolved_column_names + return rows_to_dicts(converted_rows, resolved_column_names) + + def close(self) -> None: + result_set = self._result_set + if result_set is not None: + close = getattr(result_set, "close", None) + if callable(close): + with contextlib.suppress(Exception): + close() + self._result_set = None + self._row_iterator = None + self._column_names = None + self._column_plan = None + + class SpannerSyncDriver(SyncDriverAdapterBase): """Synchronous Spanner driver operating on Snapshot or Transaction contexts.""" dialect: "DialectType" = "spanner" - __slots__ = ("_column_name_cache", "_data_dictionary", "_pending_execute_options", "_type_converter") + __slots__ = ("_data_dictionary", "_pending_execute_options", "_row_plan_cache") def __init__( self, @@ -157,15 +259,9 @@ def __init__( statement_config = default_statement_config super().__init__(connection=connection, statement_config=statement_config, driver_features=features) - - json_deserializer = features.get("json_deserializer") - self._type_converter = SpannerOutputConverter( - enable_uuid_conversion=features.get("enable_uuid_conversion", True), - json_deserializer=cast("Callable[[str], Any]", json_deserializer or from_json), - ) - self._column_name_cache: dict[int, tuple[Any, list[str]]] = {} self._data_dictionary: SpannerDataDictionary | None = None self._pending_execute_options: _PerCallExecuteOptions | None = None + self._row_plan_cache: dict[int, tuple[Any, list[str], tuple[tuple[int, Any], ...] | None]] = {} # ───────────────────────────────────────────────────────────────────────────── # CORE DISPATCH METHODS - The Execution Engine @@ -191,8 +287,8 @@ def dispatch_execute(self, cursor: "SpannerConnection", statement: "SQL") -> Exe if not fields: msg = "Result set metadata not available." raise SQLConversionError(msg) - column_names = self._resolve_column_names(fields) - data, column_names = collect_rows(rows, fields, self._type_converter, column_names=column_names) + column_names, column_plan = self._resolve_row_plan(fields) + data, column_names = collect_rows(rows, fields, column_names=column_names, column_plan=column_plan) return self.create_execution_result( cursor, selected_data=data, @@ -210,6 +306,19 @@ def dispatch_execute(self, cursor: "SpannerConnection", statement: "SQL") -> Exe raise SQLConversionError(_READ_ONLY_SNAPSHOT_ERROR_MESSAGE) + def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "SyncRowStream[dict[str, Any]] | None": + if not statement.returns_rows(): + return None + sql, params = self._compiled_sql(statement, self.statement_config) + params = cast("dict[str, Any] | None", params) + coerced_params = self._coerce_params(params) + param_types_map = self._infer_param_types(coerced_params) + return SyncRowStream( + _SpannerSelectStreamSource( + self, sql, coerced_params, param_types_map, chunk_size, self._execute_kwargs(for_read=True) + ) + ) + def dispatch_execute_many(self, cursor: "SpannerConnection", statement: "SQL") -> ExecutionResult: if not supports_batch_update(cursor): msg = "execute_many requires a Transaction context" @@ -366,6 +475,70 @@ def execute_script( finally: self._pending_execute_options = previous_options + @overload + def select_stream( + self, + statement: "SQL | Statement | QueryBuilder", + /, + *parameters: "StatementParameters | StatementFilter", + schema_type: "type[SchemaT]", + statement_config: "StatementConfig | None" = None, + chunk_size: int = 1000, + native_only: bool = False, + **kwargs: Any, + ) -> "SyncRowStream[SchemaT]": ... + + @overload + def select_stream( + self, + statement: "SQL | Statement | QueryBuilder", + /, + *parameters: "StatementParameters | StatementFilter", + schema_type: None = None, + statement_config: "StatementConfig | None" = None, + chunk_size: int = 1000, + native_only: bool = False, + **kwargs: Any, + ) -> "SyncRowStream[dict[str, Any]]": ... + + def select_stream( + self, + statement: "SQL | Statement | QueryBuilder", + /, + *parameters: "StatementParameters | StatementFilter", + schema_type: "type[SchemaT] | None" = None, + statement_config: "StatementConfig | None" = None, + chunk_size: int = 1000, + native_only: bool = False, + **kwargs: Any, + ) -> "SyncRowStream[SchemaT] | SyncRowStream[dict[str, Any]]": + """Execute a query and stream rows with optional Spanner per-call options.""" + execute_options = self._pop_execute_options(kwargs) + if execute_options is None: + return super().select_stream( + statement, + *parameters, + schema_type=schema_type, + statement_config=statement_config, + chunk_size=chunk_size, + native_only=native_only, + **kwargs, + ) + previous_options = self._pending_execute_options + self._pending_execute_options = execute_options + try: + return super().select_stream( + statement, + *parameters, + schema_type=schema_type, + statement_config=statement_config, + chunk_size=chunk_size, + native_only=native_only, + **kwargs, + ) + finally: + self._pending_execute_options = previous_options + def _execute_kwargs(self, *, for_read: bool = False) -> dict[str, Any]: kwargs = {key: self.driver_features[key] for key in ("retry", "timeout") if key in self.driver_features} request_options = self.driver_features.get("request_options") @@ -580,8 +753,9 @@ def _coerce_params(self, params: "dict[str, Any] | list[Any] | tuple[Any, ...] | def _infer_param_types(self, params: "dict[str, Any] | list[Any] | tuple[Any, ...] | None") -> "dict[str, Any]": return infer_param_types(params) - def _resolve_column_names(self, fields: Any) -> list[str]: - return resolve_column_names(fields, self._column_name_cache) + def _resolve_row_plan(self, fields: Any) -> "tuple[list[str], tuple[tuple[int, Any], ...] | None]": + json_deserializer = cast("Callable[[str], Any]", self.driver_features.get("json_deserializer", from_json)) + return resolve_row_plan(fields, self._row_plan_cache, json_deserializer=json_deserializer) register_driver_profile("spanner", driver_profile) diff --git a/sqlspec/adapters/spanner/type_converter.py b/sqlspec/adapters/spanner/type_converter.py index 2dd18906a..025cbe99c 100644 --- a/sqlspec/adapters/spanner/type_converter.py +++ b/sqlspec/adapters/spanner/type_converter.py @@ -21,12 +21,8 @@ from typing import TYPE_CHECKING, Any, cast from uuid import UUID -from typing_extensions import final - from sqlspec.core import TypedParameter -from sqlspec.core.type_converter import CachedOutputConverter, convert_uuid from sqlspec.utils.module_loader import import_optional_attr -from sqlspec.utils.serializers import from_json from sqlspec.utils.type_converters import should_json_encode_sequence _UUID_TYPES: "tuple[type[Any], ...]" = (UUID,) @@ -40,8 +36,6 @@ from sqlspec.protocols import SpannerParamTypesProtocol __all__ = ( - "SPANNER_SPECIAL_CHARS", - "SpannerOutputConverter", "bytes_to_spanner", "coerce_params_for_spanner", "infer_spanner_param_types", @@ -51,7 +45,6 @@ "uuid_to_spanner", ) -SPANNER_SPECIAL_CHARS: "frozenset[str]" = frozenset({"{", "[", "-", ":", "T", "."}) UUID_BYTE_LENGTH: int = 16 _SPANNER_PARAM_TYPES: "SpannerParamTypesProtocol | None" = None _JSON_OBJECT_TYPE: "type[Any] | None" = None @@ -75,77 +68,6 @@ def _get_json_object_type() -> "type[Any]": return _JSON_OBJECT_TYPE -@final -class SpannerOutputConverter(CachedOutputConverter): - """Spanner-specific output conversion with UUID and JSON support. - - Extends CachedOutputConverter with Spanner-specific functionality - including UUID bytes handling and JSON deserialization. - """ - - __slots__ = ("_enable_uuid_conversion", "_json_deserializer") - - def __init__( - self, - cache_size: int = 5000, - enable_uuid_conversion: bool = True, - json_deserializer: "Callable[[str], Any] | None" = None, - ) -> None: - """Initialize converter with Spanner-specific options. - - Args: - cache_size: Maximum number of string values to cache (default: 5000) - enable_uuid_conversion: Enable automatic UUID conversion (default: True) - json_deserializer: Custom JSON deserializer (default: from_json) - """ - super().__init__(special_chars=SPANNER_SPECIAL_CHARS, cache_size=cache_size) - self._enable_uuid_conversion = enable_uuid_conversion - self._json_deserializer = json_deserializer if json_deserializer is not None else from_json - - def _convert_detected(self, value: str, detected_type: str) -> Any: - """Convert value with Spanner-specific handling. - - Args: - value: String value to convert. - detected_type: Detected type name. - - Returns: - Converted value according to Spanner requirements. - """ - if detected_type == "uuid": - if not self._enable_uuid_conversion: - return value - try: - return convert_uuid(value) - except ValueError: - return value - if detected_type == "json": - try: - return self._json_deserializer(value) - except (ValueError, TypeError): - return value - return value - - def convert(self, value: Any) -> Any: - """Convert value with Spanner-specific byte UUID handling. - - Args: - value: Value to potentially convert. - - Returns: - Converted value or original value. - """ - if self._enable_uuid_conversion and isinstance(value, bytes) and len(value) == UUID_BYTE_LENGTH: - try: - return UUID(bytes=value) - except ValueError: - return value - - if not isinstance(value, str): - return value - return self._convert_cache(value) - - def _json_param_type() -> Any: """Get Spanner JSON param type with fallback to STRING. diff --git a/sqlspec/adapters/sqlite/core.py b/sqlspec/adapters/sqlite/core.py index c1ef32519..a71646b2b 100644 --- a/sqlspec/adapters/sqlite/core.py +++ b/sqlspec/adapters/sqlite/core.py @@ -417,7 +417,7 @@ def apply_driver_features( ) -> "tuple[StatementConfig, dict[str, Any]]": """Apply SQLite driver feature defaults to statement config.""" features: dict[str, Any] = dict(driver_features) if driver_features else {} - features.setdefault("enable_custom_adapters", True) + features.setdefault("enable_custom_adapters", False) json_serializer = features.setdefault("json_serializer", to_json) json_deserializer = features.setdefault("json_deserializer", from_json) diff --git a/sqlspec/adapters/sqlite/driver.py b/sqlspec/adapters/sqlite/driver.py index 9f1701a0e..b2226a667 100644 --- a/sqlspec/adapters/sqlite/driver.py +++ b/sqlspec/adapters/sqlite/driver.py @@ -21,6 +21,7 @@ from sqlspec.core.result import DMLResult from sqlspec.driver import BaseSyncExceptionHandler, SyncDriverAdapterBase, SyncRowStream from sqlspec.exceptions import SQLSpecError +from sqlspec.utils.type_guards import resolve_row_format if TYPE_CHECKING: from collections.abc import Sequence @@ -110,6 +111,7 @@ def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionResult": if statement.returns_rows(): fetched_data = cursor.fetchall() data, column_names, row_count = collect_rows(fetched_data, cursor.description) + row_format = resolve_row_format(data) return self.create_execution_result( cursor, @@ -117,7 +119,7 @@ def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionResult": column_names=column_names, data_row_count=row_count, is_select_result=True, - row_format="tuple", + row_format=row_format, ) affected_rows = resolve_rowcount(cursor) @@ -223,13 +225,14 @@ def _execute_cache_hit( if column_names is None: description = cursor.description column_names = [col[0] for col in description] if description else [] + row_format = resolve_row_format(fetched_data) execution_result = self.create_execution_result( cursor, selected_data=fetched_data, column_names=column_names, data_row_count=len(fetched_data), is_select_result=True, - row_format="tuple", + row_format=row_format, ) direct_statement = self._cached_statement( sql, params, cached, params, params_are_simple=True, compiled_sql=cached.compiled_sql diff --git a/sqlspec/adapters/sqlite/pool.py b/sqlspec/adapters/sqlite/pool.py index 627e9eb9c..07b3915ea 100644 --- a/sqlspec/adapters/sqlite/pool.py +++ b/sqlspec/adapters/sqlite/pool.py @@ -45,8 +45,10 @@ def _load_extensions(connection: SqliteConnection, extensions: "list[str]") -> N def _apply_runtime_setup(connection: SqliteConnection, runtime_setup: "dict[str, Any]") -> None: - for pragma_name, pragma_value in runtime_setup.get("pragmas", ()): - connection.execute(f"PRAGMA {pragma_name} = {pragma_value}") + pragmas = runtime_setup.get("pragmas", ()) + if pragmas: + pragma_script = "\n".join(f"PRAGMA {pragma_name} = {pragma_value};" for pragma_name, pragma_value in pragmas) + connection.executescript(pragma_script) extensions = runtime_setup.get("extensions") if extensions: diff --git a/sqlspec/core/__init__.py b/sqlspec/core/__init__.py index c264055f1..b7be2f8b8 100644 --- a/sqlspec/core/__init__.py +++ b/sqlspec/core/__init__.py @@ -230,11 +230,7 @@ get_default_parameter_config, ) from sqlspec.core.type_converter import ( - DEFAULT_CACHE_SIZE, - DEFAULT_SPECIAL_CHARS, BaseInputConverter, - BaseTypeConverter, - CachedOutputConverter, convert_decimal, convert_iso_date, convert_iso_datetime, @@ -247,8 +243,6 @@ from sqlspec.exceptions import StackExecutionError __all__ = ( - "DEFAULT_CACHE_SIZE", - "DEFAULT_SPECIAL_CHARS", "DRIVER_PARAMETER_PROFILES", "EXECUTE_MANY_MIN_ROWS", "PARAMETER_REGEX", @@ -256,13 +250,11 @@ "AnyCollectionFilter", "ArrowResult", "BaseInputConverter", - "BaseTypeConverter", "BeforeAfterFilter", "BooleanFilter", "CacheConfig", "CacheKey", "CacheStats", - "CachedOutputConverter", "CachedStatement", "ChoicesFilter", "CompiledSQL", diff --git a/sqlspec/core/result/_base.py b/sqlspec/core/result/_base.py index 86a577c96..79341e7bb 100644 --- a/sqlspec/core/result/_base.py +++ b/sqlspec/core/result/_base.py @@ -285,6 +285,7 @@ def _get_rows(self) -> "list[dict[str, Any]]": if not raw: self._materialized_dicts = [] return self._materialized_dicts + self._assert_row_format_fidelity(raw) fmt = self._row_format if fmt == "dict": self._materialized_dicts = raw @@ -305,6 +306,27 @@ def _get_rows(self) -> "list[dict[str, Any]]": self._materialized_dicts = [dict(record) for record in raw] return self._materialized_dicts + def _assert_row_format_fidelity(self, rows: "list[Any]") -> None: + """Debug-only row-format shape check for first-row fidelity.""" + if not __debug__ or not rows: + return + + first_row = rows[0] + row_format = self._row_format + if row_format == "dict": + assert isinstance(first_row, dict), "row_format='dict' requires dict rows" + return + if row_format == "tuple": + assert isinstance(first_row, tuple), "row_format='tuple' requires tuple rows" + return + if row_format == "record": + try: + keys = first_row.keys + except AttributeError as exc: + msg = "row_format='record' requires dict-like rows" + raise AssertionError(msg) from exc + assert callable(keys), "row_format='record' requires dict-like rows" + @property def raw_data(self) -> "tuple[list[Any], list[str]]": """Zero-copy access to raw driver-native rows and column names. @@ -378,6 +400,7 @@ def _get_schema_rows(self, schema_type: "type[SchemaT]", rows: "list[dict[str, A if cached_rows is not None: return cast("list[SchemaT]", cached_rows) converted_rows = cast("list[SchemaT]", to_schema(rows, schema_type=schema_type)) + self._materialized_dicts = None if cache is None: self._schema_rows_cache = {schema_type: converted_rows} else: @@ -391,6 +414,7 @@ def _get_schema_row(self, schema_type: "type[SchemaT]", row: "dict[str, Any]") - if cached_row is not None: return cast("SchemaT", cached_row) converted_row = to_schema(row, schema_type=schema_type) + self._materialized_dicts = None if row_cache is None: self._schema_row_cache = {schema_type: converted_row} else: diff --git a/sqlspec/core/type_converter.py b/sqlspec/core/type_converter.py index 657ea4d67..1c69d4db8 100644 --- a/sqlspec/core/type_converter.py +++ b/sqlspec/core/type_converter.py @@ -1,22 +1,15 @@ -"""Base classes and detection for adapter type conversion.""" +"""Base classes and scalar helpers for adapter type conversion.""" import json -import re -from collections.abc import Callable from datetime import date, datetime, time, timezone from decimal import Decimal -from functools import lru_cache -from typing import Any, Final +from typing import Any from uuid import UUID from mypy_extensions import mypyc_attr __all__ = ( - "DEFAULT_CACHE_SIZE", - "DEFAULT_SPECIAL_CHARS", "BaseInputConverter", - "BaseTypeConverter", - "CachedOutputConverter", "convert_decimal", "convert_iso_date", "convert_iso_datetime", @@ -27,26 +20,8 @@ "parse_datetime_rfc3339", ) -DEFAULT_SPECIAL_CHARS: Final[frozenset[str]] = frozenset({"{", "[", "-", ":", "T", "."}) -DEFAULT_CACHE_SIZE: Final[int] = 5000 -DEFAULT_DETECTION_CHARS: Final[frozenset[str]] = frozenset({"{", "[", "-", ":", "T"}) - -SPECIAL_TYPE_REGEX: Final[re.Pattern[str]] = re.compile( - r"^(?:" - r"(?P[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12})|" - r"(?P\d{4}-\d{2}-\d{2}[T ]\d{2}:\d{2}:\d{2}(?:\.\d+)?(?:Z|[+-]\d{2}:?\d{2})?)|" - r"(?P\d{4}-\d{2}-\d{2})|" - r"(?P\d{2}:\d{2}:\d{2}(?:\.\d+)?)|" - r"(?P[\[{].*[\]}])|" - r"(?P(?:\d{1,3}\.){3}\d{1,3})|" - r"(?P(?:[0-9a-f]{1,4}:){7}[0-9a-f]{1,4})|" - r"(?P(?:[0-9a-f]{2}:){5}[0-9a-f]{2})" - r")$", - re.IGNORECASE | re.DOTALL, -) - -def convert_uuid(value: str) -> "UUID": +def convert_uuid(value: str) -> UUID: """Convert UUID string to UUID object. Args: @@ -152,163 +127,6 @@ def parse_datetime_rfc3339(dt_str: str) -> "datetime": return convert_iso_datetime(dt_str) -_TYPE_CONVERTERS: Final[dict[str, Callable[[str], Any]]] = { - "uuid": convert_uuid, - "iso_datetime": convert_iso_datetime, - "iso_date": convert_iso_date, - "iso_time": convert_iso_time, - "json": convert_json, -} - - -class _CachedConverter: - __slots__ = ("_cached", "_converter", "_special_chars") - - def __init__(self, converter: "CachedOutputConverter", special_chars: "frozenset[str]", cache_size: int) -> None: - self._converter = converter - self._special_chars = special_chars - self._cached = lru_cache(maxsize=cache_size)(self._convert) - - def _convert(self, value: str) -> "Any": - if not value or not any(c in value for c in self._special_chars): - return value - detected_type = self._converter.detect_type(value) - if detected_type: - return self._converter._convert_detected(value, detected_type) # pyright: ignore[reportPrivateUsage] - return value - - def __call__(self, value: str) -> "Any": - return self._cached(value) - - -def _make_cached_converter( - converter: "CachedOutputConverter", special_chars: "frozenset[str]", cache_size: int -) -> "Callable[[str], Any]": - """Create a cached conversion function for an output converter. - - Args: - converter: The output converter instance to use for type detection/conversion. - special_chars: Characters that trigger type detection. - cache_size: Maximum entries in the LRU cache. - - Returns: - A cached function that converts string values. - """ - - return _CachedConverter(converter, special_chars, cache_size) - - -@mypyc_attr(allow_interpreted_subclasses=True) -class BaseTypeConverter: - """Universal type detection and conversion for all adapters.""" - - __slots__ = () - - def detect_type(self, value: str) -> str | None: - """Detect special types from string values. - - Args: - value: String value to analyze. - - Returns: - Type name if detected, None otherwise. - """ - if not isinstance(value, str): # pyright: ignore - return None - if not value: - return None - - match = SPECIAL_TYPE_REGEX.match(value) - if not match: - return None - - return match.lastgroup - - def convert_value(self, value: str, detected_type: str) -> "Any": - """Convert string value to appropriate Python type. - - Args: - value: String value to convert. - detected_type: Detected type name. - - Returns: - Converted value in appropriate Python type. - """ - converter = _TYPE_CONVERTERS.get(detected_type) - if converter: - return converter(value) - return value - - def convert_if_detected(self, value: "Any") -> "Any": - """Convert value only if special type detected, else return original. - - Args: - value: Value to potentially convert. - - Returns: - Converted value if special type detected, original value otherwise. - """ - if not isinstance(value, str): - return value - - if not any(c in value for c in DEFAULT_DETECTION_CHARS): - return value - - detected_type = self.detect_type(value) - if detected_type: - try: - return self.convert_value(value, detected_type) - except Exception: - return value - return value - - -@mypyc_attr(allow_interpreted_subclasses=True) -class CachedOutputConverter(BaseTypeConverter): - """Base class for converting database results to Python types.""" - - __slots__ = ("_convert_cache", "_special_chars") - - def __init__(self, special_chars: "frozenset[str] | None" = None, cache_size: int = DEFAULT_CACHE_SIZE) -> None: - """Initialize converter with caching. - - Args: - special_chars: Characters that trigger type detection. - cache_size: Maximum entries in LRU cache. - """ - super().__init__() - self._special_chars = special_chars if special_chars is not None else DEFAULT_SPECIAL_CHARS - self._convert_cache = _make_cached_converter(self, self._special_chars, cache_size) - - def _convert_detected(self, value: str, detected_type: str) -> "Any": - """Convert value with detected type. Override for adapter-specific logic. - - Args: - value: String value to convert. - detected_type: Detected type name from detect_type(). - - Returns: - Converted value, or original value on conversion failure. - """ - try: - return self.convert_value(value, detected_type) - except Exception: - return value - - def convert(self, value: "Any") -> "Any": - """Convert value using cached detection and conversion. - - Args: - value: Value to potentially convert. - - Returns: - Converted value if string with special type, original otherwise. - """ - if not isinstance(value, str): - return value - return self._convert_cache(value) - - @mypyc_attr(allow_interpreted_subclasses=True) class BaseInputConverter: """Base class for converting Python params to database format.""" @@ -325,14 +143,3 @@ def convert_params(self, params: "dict[str, Any] | None") -> "dict[str, Any] | N Converted parameters dictionary, or None if input was None. """ return params - - def convert_value(self, value: "Any") -> "Any": - """Convert a single parameter value. - - Args: - value: Value to convert. - - Returns: - Converted value. - """ - return value diff --git a/sqlspec/typing.py b/sqlspec/typing.py index 89dd8e8a1..5bffb2980 100644 --- a/sqlspec/typing.py +++ b/sqlspec/typing.py @@ -52,6 +52,7 @@ FailFast, Gauge, Histogram, + MsgspecValidationError, NumpyArray, PandasDataFrame, PolarsDataFrame, @@ -120,6 +121,7 @@ "FailFast", "Gauge", "Histogram", + "MsgspecValidationError", "NumpyArray", "PandasDataFrame", "PolarsDataFrame", diff --git a/sqlspec/utils/schema.py b/sqlspec/utils/schema.py index 922c628d5..215e02e94 100644 --- a/sqlspec/utils/schema.py +++ b/sqlspec/utils/schema.py @@ -16,6 +16,7 @@ from sqlspec.typing import ( CATTRS_INSTALLED, NUMPY_INSTALLED, + MsgspecValidationError, SchemaT, attrs_asdict, cattrs_structure, @@ -200,14 +201,14 @@ def _coerce_foreign_key_item(data: Any, schema_type: Any) -> Any: def _convert_typed_dict(data: Any, schema_type: Any) -> Any: """Convert data to TypedDict.""" - return [item for item in data if is_dict(item)] if isinstance(data, list) else data + return data def _convert_dataclass(data: Any, schema_type: Any) -> Any: """Convert data to dataclass.""" if isinstance(data, list): - return [schema_type(**dict(item)) if is_dict(item) else item for item in data] - return schema_type(**dict(data)) if is_dict(data) else (schema_type(**data) if isinstance(data, dict) else data) + return [schema_type(**item) if is_dict(item) else item for item in data] + return schema_type(**data) if is_dict(data) else data class _IsTypePredicate: @@ -379,15 +380,19 @@ def _convert_msgspec(data: Any, schema_type: Any) -> Any: except Exception as e: logger.debug("Field name transformation failed for msgspec schema: %s", e) - if NUMPY_INSTALLED: - transformed_data = _convert_numpy_recursive(transformed_data) + target_type = list[schema_type] if isinstance(transformed_data, Sequence) else schema_type - return convert( - obj=transformed_data, - type=(list[schema_type] if isinstance(transformed_data, Sequence) else schema_type), - from_attributes=True, - dec_hook=_DEFAULT_MSGSPEC_DESERIALIZER, - ) + try: + return convert( + obj=transformed_data, type=target_type, from_attributes=True, dec_hook=_DEFAULT_MSGSPEC_DESERIALIZER + ) + except MsgspecValidationError: + if not NUMPY_INSTALLED: + raise + transformed_data = _convert_numpy_recursive(transformed_data) + return convert( + obj=transformed_data, type=target_type, from_attributes=True, dec_hook=_DEFAULT_MSGSPEC_DESERIALIZER + ) def _convert_pydantic(data: Any, schema_type: Any) -> Any: @@ -406,8 +411,8 @@ def _convert_attrs(data: Any, schema_type: Any) -> Any: return cattrs_structure(structured, schema_type) if isinstance(data, list): - return [schema_type(**dict(item)) if is_dict(item) else schema_type(**attrs_asdict(item)) for item in data] - return schema_type(**dict(data)) if is_dict(data) else data + return [schema_type(**item) if is_dict(item) else schema_type(**attrs_asdict(item)) for item in data] + return schema_type(**data) if is_dict(data) else data # Cache for schema converters - maps type directly to converter callable (or None if unsupported) diff --git a/sqlspec/utils/serializers/_schema.py b/sqlspec/utils/serializers/_schema.py index 6f2a8fe96..b6fc3a344 100644 --- a/sqlspec/utils/serializers/_schema.py +++ b/sqlspec/utils/serializers/_schema.py @@ -2,11 +2,13 @@ import os from collections import OrderedDict +from dataclasses import Field +from dataclasses import fields as dataclasses_fields from functools import partial from threading import RLock from typing import TYPE_CHECKING, Any, Final, cast -from sqlspec.typing import MSGSPEC_INSTALLED, UNSET, ArrowReturnFormat, attrs_asdict, msgspec_fields +from sqlspec.typing import MSGSPEC_INSTALLED, UNSET, ArrowReturnFormat, Empty, attrs_asdict, msgspec_fields from sqlspec.utils.arrow_helpers import convert_dict_to_arrow from sqlspec.utils.type_guards import ( dataclass_to_dict, @@ -114,6 +116,7 @@ def to_arrow( _SERIALIZER_LOCK: RLock = RLock() _SCHEMA_SERIALIZERS: "OrderedDict[tuple[type[Any] | None, bool, bool], SchemaSerializer]" = OrderedDict() _SERIALIZER_METRICS = _SerializerCacheMetrics() +_DATACLASS_FIELDS_CACHE: "dict[type[Any], tuple[Field[Any], ...]]" = {} def _make_serializer_key(sample: Any, exclude_unset: bool, wire_format: bool) -> "tuple[type[Any] | None, bool, bool]": @@ -126,6 +129,15 @@ def _dump_identity_dict(value: Any) -> "dict[str, Any]": return cast("dict[str, Any]", value) +def _dataclass_fields(schema_type: type[Any]) -> "tuple[Field[Any], ...]": + cached = _DATACLASS_FIELDS_CACHE.get(schema_type) + if cached is not None: + return cached + fields = dataclasses_fields(schema_type) + _DATACLASS_FIELDS_CACHE[schema_type] = fields + return fields + + def _msgspec_field_pairs(schema_type: type[Any], *, wire_format: bool) -> "tuple[tuple[str, str], ...]": if not MSGSPEC_INSTALLED: msg = "msgspec is required to serialize msgspec.Struct values" @@ -148,8 +160,17 @@ def _dump_msgspec_struct( } -def _dump_dataclass(value: Any, *, exclude_unset: bool) -> "dict[str, Any]": - return dataclass_to_dict(value, exclude_empty=exclude_unset) +def _dump_dataclass(value: Any, *, dataclass_fields: "tuple[Field[Any], ...]", exclude_unset: bool) -> "dict[str, Any]": + result: dict[str, Any] = {} + for field in dataclass_fields: + field_value = value.__getattribute__(field.name) + if exclude_unset and field_value is Empty: + continue + if is_dataclass_instance(field_value): + result[field.name] = dataclass_to_dict(field_value, exclude_empty=exclude_unset) + else: + result[field.name] = field_value + return result def _dump_pydantic(value: Any, *, exclude_unset: bool) -> "dict[str, Any]": @@ -176,11 +197,19 @@ def _msgspec_dump_function(sample: Any, *, exclude_unset: bool, wire_format: boo ) +def _dataclass_dump_function(sample: Any, *, exclude_unset: bool) -> "Callable[[Any], dict[str, Any]]": + field_tuple = _dataclass_fields(type(sample)) + return cast( + "Callable[[Any], dict[str, Any]]", + partial(_dump_dataclass, dataclass_fields=field_tuple, exclude_unset=exclude_unset), + ) + + def _dump_function(sample: Any, exclude_unset: bool, wire_format: bool) -> "Callable[[Any], dict[str, Any]]": if sample is None or isinstance(sample, dict): return _dump_identity_dict if is_dataclass_instance(sample): - return cast("Callable[[Any], dict[str, Any]]", partial(_dump_dataclass, exclude_unset=exclude_unset)) + return _dataclass_dump_function(sample, exclude_unset=exclude_unset) if is_pydantic_model(sample): return cast("Callable[[Any], dict[str, Any]]", partial(_dump_pydantic, exclude_unset=exclude_unset)) if is_msgspec_struct(sample): @@ -248,6 +277,7 @@ def reset_serializer_cache() -> None: """Clear cached serializer pipelines.""" with _SERIALIZER_LOCK: _SCHEMA_SERIALIZERS.clear() + _DATACLASS_FIELDS_CACHE.clear() _SERIALIZER_METRICS.reset() diff --git a/sqlspec/utils/type_guards.py b/sqlspec/utils/type_guards.py index 558a99dc3..3ae7be4e3 100644 --- a/sqlspec/utils/type_guards.py +++ b/sqlspec/utils/type_guards.py @@ -10,7 +10,6 @@ from dataclasses import Field from dataclasses import fields as dataclasses_fields from dataclasses import is_dataclass as dataclasses_is_dataclass -from functools import lru_cache from typing import TYPE_CHECKING, Any, cast from sqlglot import exp @@ -183,6 +182,7 @@ "is_string_literal", "is_typed_dict", "is_typed_parameter", + "resolve_row_format", "supports_arrow_results", "supports_async_delete", "supports_async_read_bytes", @@ -468,6 +468,19 @@ def is_mapping_like(obj: Any) -> "TypeGuard[MappingLikeProtocol]": return False +def resolve_row_format(rows: "Sequence[Any] | None", *, default: str = "tuple") -> str: + """Resolve the raw row format from the first row in a result set.""" + if not rows: + return default + + first_row = rows[0] + if is_dict_row(first_row): + return "dict" + if is_mapping_like(first_row): + return "record" + return default + + def has_asdict_method(obj: Any) -> "TypeGuard[HasAsDictProtocol]": """Check if an object has _asdict() method. @@ -685,7 +698,6 @@ def is_msgspec_struct_without_field(obj: Any, field_name: str) -> "TypeGuard[Str return all(field.name != field_name for field in fields) -@lru_cache(maxsize=500) def _detect_rename_pattern(field_name: str, encode_name: str) -> "str | None": """Detect the rename pattern by comparing field name transformations. @@ -696,15 +708,26 @@ def _detect_rename_pattern(field_name: str, encode_name: str) -> "str | None": Returns: The detected rename pattern ("camel", "kebab", "pascal") or None """ + key = (field_name, encode_name) + if key in _MSGSPEC_RENAME_PATTERN_CACHE: + return _MSGSPEC_RENAME_PATTERN_CACHE[key] + + result: str | None if encode_name == camelize(field_name) and encode_name != field_name: - return "camel" + result = "camel" + elif encode_name == kebabize(field_name) and encode_name != field_name: + result = "kebab" + elif encode_name == pascalize(field_name) and encode_name != field_name: + result = "pascal" + else: + result = None - if encode_name == kebabize(field_name) and encode_name != field_name: - return "kebab" + _MSGSPEC_RENAME_PATTERN_CACHE[key] = result + return result - if encode_name == pascalize(field_name) and encode_name != field_name: - return "pascal" - return None + +_MSGSPEC_RENAME_CONFIG_CACHE: "dict[object, str | None]" = {} +_MSGSPEC_RENAME_PATTERN_CACHE: "dict[tuple[str, str], str | None]" = {} def get_msgspec_rename_config(schema_type: type) -> "str | None": @@ -721,22 +744,27 @@ def get_msgspec_rename_config(schema_type: type) -> "str | None": The rename configuration value ("camel", "kebab", "pascal", etc.) if detected, None if no rename configuration exists or if not a msgspec struct. """ - if not MSGSPEC_INSTALLED: - return None + if schema_type in _MSGSPEC_RENAME_CONFIG_CACHE: + return _MSGSPEC_RENAME_CONFIG_CACHE[schema_type] - if not is_msgspec_struct(schema_type): + if not MSGSPEC_INSTALLED or not is_msgspec_struct(schema_type): + _MSGSPEC_RENAME_CONFIG_CACHE[schema_type] = None return None from msgspec import structs fields: tuple[Any, ...] = structs.fields(cast("Any", schema_type)) if not fields: + _MSGSPEC_RENAME_CONFIG_CACHE[schema_type] = None return None for field in fields: if field.name != field.encode_name: - return _detect_rename_pattern(field.name, field.encode_name) + rename_config = _detect_rename_pattern(field.name, field.encode_name) + _MSGSPEC_RENAME_CONFIG_CACHE[schema_type] = rename_config + return rename_config + _MSGSPEC_RENAME_CONFIG_CACHE[schema_type] = None return None diff --git a/tests/integration/adapters/aiomysql/test_local_infile_bulk_load.py b/tests/integration/adapters/aiomysql/test_local_infile_bulk_load.py index 967d248a7..fe1755eb1 100644 --- a/tests/integration/adapters/aiomysql/test_local_infile_bulk_load.py +++ b/tests/integration/adapters/aiomysql/test_local_infile_bulk_load.py @@ -22,7 +22,7 @@ async def aiomysql_infile_config(mysql_service: "MySQLService") -> "AsyncGenerat "db": mysql_service.db, "autocommit": True, "allow_local_infile": True, - "enable_local_infile": True, + "local_infile": True, }, driver_features={"enable_local_infile_bulk_load": True}, ) diff --git a/tests/integration/adapters/contracts/_cases.py b/tests/integration/adapters/contracts/_cases.py index 220f9edb4..3005a50de 100644 --- a/tests/integration/adapters/contracts/_cases.py +++ b/tests/integration/adapters/contracts/_cases.py @@ -299,8 +299,10 @@ class DriverCaseContext: supports_storage_bridge=True, supports_native_bulk_ingest=True, supports_load_from_records=True, + supports_native_row_streaming=True, supports_exception_translation=False, execute_rowcount_policy="unavailable", + stream_chunk_policy="advisory", extra_assertions=("param_codecs:adbc_sqlite",), ), DriverCase( @@ -321,7 +323,9 @@ class DriverCaseContext: supports_native_bulk_ingest=True, supports_load_from_records=True, supports_vector=True, + supports_native_row_streaming=True, execute_rowcount_policy="unavailable", + stream_chunk_policy="advisory", extra_assertions=("param_codecs:adbc_duckdb",), ), DriverCase( @@ -343,7 +347,9 @@ class DriverCaseContext: supports_load_from_records=True, supports_for_update=True, supports_for_share=True, + supports_native_row_streaming=True, execute_rowcount_policy="unavailable", + stream_chunk_policy="advisory", unsupported_explain_reason="ADBC PostgreSQL EXPLAIN is incompatible with the driver's COPY result transfer", extra_assertions=("arrow_specifics:adbc_select_to_arrow_error", "param_codecs:adbc_postgres"), ), @@ -369,6 +375,7 @@ class DriverCaseContext: supports_connection_instance=True, supports_merge=True, supports_vector=True, + supports_lob=True, supports_lowercase_columns=True, supports_uuid_feature=True, supports_native_row_streaming=True, @@ -379,6 +386,7 @@ class DriverCaseContext: "driver_features:oracle_json_native", "driver_features:oracle_batch_errors", "driver_features:oracle_plsql", + "oracle_lob_fetch:matrix", "streaming_native:oracledb", ), ), @@ -396,8 +404,10 @@ class DriverCaseContext: supports_native_bulk_ingest=True, supports_execute_many=False, supports_load_from_records=False, + supports_native_row_streaming=True, supports_exception_translation=False, execute_rowcount_policy="unavailable", + supports_stream_reopen_after_partial_iteration=False, ), DriverCase( id="bigquery-sync", @@ -719,6 +729,7 @@ class DriverCaseContext: supports_merge=True, supports_merge_bulk=True, supports_vector=True, + supports_lob=True, supports_lowercase_columns=True, supports_uuid_feature=True, supports_native_row_streaming=True, @@ -731,6 +742,7 @@ class DriverCaseContext: "driver_features:oracle_batch_errors", "driver_features:oracle_plsql", "statement_stack:oracle_native", + "oracle_lob_fetch:matrix", "streaming_native:oracledb", ), ), @@ -745,6 +757,7 @@ class DriverCaseContext: "sync", integration_status="deferred", reason="No active integration fixture exists for mssql_python.", + supports_native_row_streaming=True, ), DriverCase( "pymssql-sync", @@ -757,6 +770,7 @@ class DriverCaseContext: supports_execute_many=True, supports_migrations=True, supports_data_dictionary=True, + supports_native_row_streaming=True, ), DriverCase( "spanner-sync", diff --git a/tests/integration/adapters/contracts/_driver_type_system.py b/tests/integration/adapters/contracts/_driver_type_system.py new file mode 100644 index 000000000..d7b977d4b --- /dev/null +++ b/tests/integration/adapters/contracts/_driver_type_system.py @@ -0,0 +1,552 @@ +"""Driver-type-system registry and source-equality helpers for contract tests.""" + +import ast +import difflib +import importlib +import inspect +import textwrap +from dataclasses import dataclass +from typing import Any, cast + +__all__ = ( + "DRIVER_FEATURE_CONSUMED_KEYS", + "DRIVER_FEATURE_CONTRACT_ENABLED", + "DRIVER_FEATURE_CONTRACT_GROUPS", + "DRIVER_FEATURE_TYPED_DICTS", + "SOURCE_EQUIVALENCE_ALLOWLIST", + "SOURCE_EQUIVALENCE_CASES", + "DriverFeatureContractGroup", + "SourceEquivalenceCase", + "assert_source_equivalent", + "get_driver_feature_consumed_keys", + "get_driver_feature_keys", + "normalize_source", +) + + +@dataclass(frozen=True) +class DriverFeatureContractGroup: + """Registry entry for a disabled-first contract group.""" + + description: str + enabled: bool = False + enabled_adapters: tuple[str, ...] = () + + +@dataclass(frozen=True) +class SourceEquivalenceCase: + """One source-equality tripwire across adapter helper implementations.""" + + case_id: str + group: str + symbol: str + module_paths: tuple[str, ...] + allowlist_key: str + + +DRIVER_FEATURE_TYPED_DICTS: dict[str, tuple[str, str]] = { + "adbc": ("sqlspec.adapters.adbc.config", "AdbcDriverFeatures"), + "aiomysql": ("sqlspec.adapters.aiomysql.config", "AiomysqlDriverFeatures"), + "aiosqlite": ("sqlspec.adapters.aiosqlite.config", "AiosqliteDriverFeatures"), + "arrow_odbc": ("sqlspec.adapters.arrow_odbc.config", "ArrowOdbcDriverFeatures"), + "asyncmy": ("sqlspec.adapters.asyncmy.config", "AsyncmyDriverFeatures"), + "asyncpg": ("sqlspec.adapters.asyncpg.config", "AsyncpgDriverFeatures"), + "bigquery": ("sqlspec.adapters.bigquery.config", "BigQueryDriverFeatures"), + "cockroach_asyncpg": ("sqlspec.adapters.cockroach_asyncpg.config", "CockroachAsyncpgDriverFeatures"), + "cockroach_psycopg": ("sqlspec.adapters.cockroach_psycopg.config", "CockroachPsycopgDriverFeatures"), + "duckdb": ("sqlspec.adapters.duckdb.config", "DuckDBDriverFeatures"), + "mssql_python": ("sqlspec.adapters.mssql_python.config", "MssqlPythonDriverFeatures"), + "mysqlconnector": ("sqlspec.adapters.mysqlconnector.config", "MysqlConnectorDriverFeatures"), + "oracledb": ("sqlspec.adapters.oracledb.config", "OracleDriverFeatures"), + "psqlpy": ("sqlspec.adapters.psqlpy.config", "PsqlpyDriverFeatures"), + "psycopg": ("sqlspec.adapters.psycopg.config", "PsycopgDriverFeatures"), + "pymssql": ("sqlspec.adapters.pymssql.config", "PymssqlDriverFeatures"), + "pymysql": ("sqlspec.adapters.pymysql.config", "PyMysqlDriverFeatures"), + "spanner": ("sqlspec.adapters.spanner.config", "SpannerDriverFeatures"), + "sqlite": ("sqlspec.adapters.sqlite.config", "SqliteDriverFeatures"), +} + +_MIGRATED_DRIVER_FEATURE_ADAPTERS = ( + "aiomysql", + "aiosqlite", + "arrow_odbc", + "asyncmy", + "bigquery", + "duckdb", + "mssql_python", + "mysqlconnector", + "psqlpy", + "psycopg", + "pymssql", + "pymysql", + "spanner", + "sqlite", +) + +# Downstream migrations populate the consumed-key registry adapter by adapter. +DRIVER_FEATURE_CONSUMED_KEYS: dict[str, tuple[str, ...]] = dict.fromkeys(DRIVER_FEATURE_TYPED_DICTS, ()) +DRIVER_FEATURE_CONSUMED_KEYS.update({ + "aiomysql": ( + "json_serializer", + "json_deserializer", + "on_connection_create", + "enable_events", + "events_backend", + "enable_local_infile_bulk_load", + ), + "aiosqlite": ( + "enable_custom_adapters", + "json_serializer", + "json_deserializer", + "on_connection_create", + "enable_events", + "events_backend", + "custom_functions", + "custom_collations", + "custom_aggregates", + "authorizer_callback", + "trace_callback", + "progress_handler", + "progress_handler_interval", + "row_factory", + "text_factory", + "pragmas", + "extensions", + ), + "arrow_odbc": ( + "chunk_size", + "max_bytes_per_batch", + "max_text_size", + "max_binary_size", + "fetch_concurrently", + "query_timeout_sec", + "connection_string", + "dbms_name", + "json_serializer", + "json_deserializer", + "enable_events", + "on_connection_create", + ), + "asyncmy": ("json_serializer", "json_deserializer", "on_connection_create", "enable_events", "events_backend"), + "bigquery": ( + "connection_instance", + "on_connection_create", + "json_serializer", + "enable_uuid_conversion", + "enable_events", + "events_backend", + "job_retry_deadline", + "job_result_timeout", + "use_query_and_wait", + "request_timeout", + "query_page_size", + "query_max_results", + "enable_storage_write_api", + ), + "duckdb": ( + "extensions", + "secrets", + "on_connection_create", + "json_serializer", + "enable_uuid_conversion", + "extension_flags", + "enable_events", + "events_backend", + ), + "mssql_python": ("use_pool", "json_serializer", "json_deserializer", "on_connection_create", "enable_events"), + "mysqlconnector": ( + "json_serializer", + "json_deserializer", + "on_connection_create", + "enable_events", + "events_backend", + "cursor_options", + "enable_local_infile_bulk_load", + ), + "psqlpy": ( + "enable_cast_detection", + "enable_pgvector", + "enable_paradedb", + "json_serializer", + "json_deserializer", + "on_connection_create", + "enable_events", + "events_backend", + ), + "psycopg": ( + "enable_pgvector", + "enable_paradedb", + "json_serializer", + "json_deserializer", + "on_connection_create", + "enable_events", + "events_backend", + "enable_alloydb", + "alloydb_instance_uri", + "enable_alloydb_iam_auth", + "alloydb_ip_type", + ), + "pymssql": ("json_serializer", "json_deserializer", "on_connection_create", "enable_events", "events_backend"), + "pymysql": ( + "json_serializer", + "json_deserializer", + "on_connection_create", + "enable_events", + "events_backend", + "enable_local_infile_bulk_load", + "enable_cloud_sql", + "cloud_sql_instance", + "cloud_sql_enable_iam_auth", + "cloud_sql_ip_type", + ), + "spanner": ( + "enable_uuid_conversion", + "json_serializer", + "json_deserializer", + "retry", + "timeout", + "request_options", + "directed_read_options", + "session_labels", + "enable_events", + "events_backend", + "enable_batch_write_api", + ), + "sqlite": ( + "enable_custom_adapters", + "json_serializer", + "json_deserializer", + "on_connection_create", + "enable_events", + "events_backend", + "custom_functions", + "custom_collations", + "custom_aggregates", + "authorizer_callback", + "trace_callback", + "progress_handler", + "progress_handler_interval", + "row_factory", + "text_factory", + "pragmas", + "extensions", + ), +}) + +DRIVER_FEATURE_CONTRACT_GROUPS: dict[str, DriverFeatureContractGroup] = { + "feature_honesty": DriverFeatureContractGroup( + description="TypedDict keys must be consumed by the owning adapter implementation.", + enabled=True, + enabled_adapters=_MIGRATED_DRIVER_FEATURE_ADAPTERS, + ), + "enable_false_semantics": DriverFeatureContractGroup( + description="enable_*=False must be inert and not leak process-global state.", + enabled=True, + enabled_adapters=("aiosqlite", "sqlite"), + ), + "parity": DriverFeatureContractGroup( + description="Buffered, streamed, and Arrow values must stay equal on canonical typed fixtures.", + enabled=True, + enabled_adapters=( + "adbc", + "aiomysql", + "arrow_odbc", + "asyncmy", + "bigquery", + "duckdb", + "mysqlconnector", + "pymysql", + "spanner", + ), + ), + "row_format": DriverFeatureContractGroup( + description="The first materialized row shape must match the row_format tag.", + enabled=True, + enabled_adapters=("aiosqlite", "psycopg", "sqlite"), + ), +} + +DRIVER_FEATURE_CONTRACT_ENABLED: dict[str, bool] = { + group_name: group.enabled for group_name, group in DRIVER_FEATURE_CONTRACT_GROUPS.items() +} + +SOURCE_EQUIVALENCE_CASES: tuple[SourceEquivalenceCase, ...] = ( + SourceEquivalenceCase( + case_id="mysql_four_way_format_identifier", + group="mysql_four_way", + symbol="format_identifier", + module_paths=( + "sqlspec.adapters.aiomysql.core", + "sqlspec.adapters.asyncmy.core", + "sqlspec.adapters.mysqlconnector.core", + "sqlspec.adapters.pymysql.core", + ), + allowlist_key="mysql_four_way:format_identifier", + ), + SourceEquivalenceCase( + case_id="mysql_four_way_build_insert_statement", + group="mysql_four_way", + symbol="build_insert_statement", + module_paths=( + "sqlspec.adapters.aiomysql.core", + "sqlspec.adapters.asyncmy.core", + "sqlspec.adapters.mysqlconnector.core", + "sqlspec.adapters.pymysql.core", + ), + allowlist_key="mysql_four_way:build_insert_statement", + ), + SourceEquivalenceCase( + case_id="mysql_four_way_normalize_execute_parameters", + group="mysql_four_way", + symbol="normalize_execute_parameters", + module_paths=( + "sqlspec.adapters.aiomysql.core", + "sqlspec.adapters.asyncmy.core", + "sqlspec.adapters.mysqlconnector.core", + "sqlspec.adapters.pymysql.core", + ), + allowlist_key="mysql_four_way:normalize_execute_parameters", + ), + SourceEquivalenceCase( + case_id="mysql_four_way_normalize_execute_many_parameters", + group="mysql_four_way", + symbol="normalize_execute_many_parameters", + module_paths=( + "sqlspec.adapters.aiomysql.core", + "sqlspec.adapters.asyncmy.core", + "sqlspec.adapters.mysqlconnector.core", + "sqlspec.adapters.pymysql.core", + ), + allowlist_key="mysql_four_way:normalize_execute_many_parameters", + ), + SourceEquivalenceCase( + case_id="mysql_four_way_resolve_column_names", + group="mysql_four_way", + symbol="resolve_column_names", + module_paths=( + "sqlspec.adapters.aiomysql.core", + "sqlspec.adapters.asyncmy.core", + "sqlspec.adapters.mysqlconnector.core", + "sqlspec.adapters.pymysql.core", + ), + allowlist_key="mysql_four_way:resolve_column_names", + ), + SourceEquivalenceCase( + case_id="mysql_four_way_detect_json_columns_from_description", + group="mysql_four_way", + symbol="detect_json_columns_from_description", + module_paths=( + "sqlspec.adapters.aiomysql.core", + "sqlspec.adapters.asyncmy.core", + "sqlspec.adapters.mysqlconnector.core", + "sqlspec.adapters.pymysql.core", + ), + allowlist_key="mysql_four_way:detect_json_columns_from_description", + ), + SourceEquivalenceCase( + case_id="mysql_four_way_detect_json_columns", + group="mysql_four_way", + symbol="detect_json_columns", + module_paths=( + "sqlspec.adapters.aiomysql.core", + "sqlspec.adapters.asyncmy.core", + "sqlspec.adapters.mysqlconnector.core", + "sqlspec.adapters.pymysql.core", + ), + allowlist_key="mysql_four_way:detect_json_columns", + ), + SourceEquivalenceCase( + case_id="mysql_four_way_resolve_row_plan", + group="mysql_four_way", + symbol="resolve_row_plan", + module_paths=( + "sqlspec.adapters.aiomysql.core", + "sqlspec.adapters.asyncmy.core", + "sqlspec.adapters.mysqlconnector.core", + "sqlspec.adapters.pymysql.core", + ), + allowlist_key="mysql_four_way:resolve_row_plan", + ), + SourceEquivalenceCase( + case_id="mysql_four_way_collect_rows", + group="mysql_four_way", + symbol="collect_rows", + module_paths=( + "sqlspec.adapters.aiomysql.core", + "sqlspec.adapters.asyncmy.core", + "sqlspec.adapters.mysqlconnector.core", + "sqlspec.adapters.pymysql.core", + ), + allowlist_key="mysql_four_way:collect_rows", + ), + SourceEquivalenceCase( + case_id="mysql_four_way_collect_stream_rows", + group="mysql_four_way", + symbol="collect_stream_rows", + module_paths=( + "sqlspec.adapters.aiomysql.core", + "sqlspec.adapters.asyncmy.core", + "sqlspec.adapters.mysqlconnector.core", + "sqlspec.adapters.pymysql.core", + ), + allowlist_key="mysql_four_way:collect_stream_rows", + ), + SourceEquivalenceCase( + case_id="mysql_four_way_create_mapped_exception", + group="mysql_four_way", + symbol="create_mapped_exception", + module_paths=( + "sqlspec.adapters.aiomysql.core", + "sqlspec.adapters.asyncmy.core", + "sqlspec.adapters.mysqlconnector.core", + "sqlspec.adapters.pymysql.core", + ), + allowlist_key="mysql_four_way:create_mapped_exception", + ), + SourceEquivalenceCase( + case_id="sqlite_pair_apply_driver_features", + group="sqlite_pair", + symbol="apply_driver_features", + module_paths=("sqlspec.adapters.sqlite.core", "sqlspec.adapters.aiosqlite.core"), + allowlist_key="sqlite_pair:apply_driver_features", + ), + SourceEquivalenceCase( + case_id="sqlite_pair_build_insert_statement", + group="sqlite_pair", + symbol="build_insert_statement", + module_paths=("sqlspec.adapters.sqlite.core", "sqlspec.adapters.aiosqlite.core"), + allowlist_key="sqlite_pair:build_insert_statement", + ), + SourceEquivalenceCase( + case_id="sqlite_pair_normalize_execute_parameters", + group="sqlite_pair", + symbol="normalize_execute_parameters", + module_paths=("sqlspec.adapters.sqlite.core", "sqlspec.adapters.aiosqlite.core"), + allowlist_key="sqlite_pair:normalize_execute_parameters", + ), + SourceEquivalenceCase( + case_id="sqlite_pair_normalize_execute_many_parameters", + group="sqlite_pair", + symbol="normalize_execute_many_parameters", + module_paths=("sqlspec.adapters.sqlite.core", "sqlspec.adapters.aiosqlite.core"), + allowlist_key="sqlite_pair:normalize_execute_many_parameters", + ), + SourceEquivalenceCase( + case_id="mssql_pair_create_mapped_exception", + group="mssql_pair", + symbol="create_mapped_exception", + module_paths=("sqlspec.adapters.mssql_python.core", "sqlspec.adapters.pymssql.core"), + allowlist_key="mssql_pair:create_mapped_exception", + ), +) + +SOURCE_EQUIVALENCE_ALLOWLIST: dict[str, frozenset[str]] = { + case.allowlist_key: frozenset() for case in SOURCE_EQUIVALENCE_CASES +} +# mysqlconnector intentionally resolves error codes from errno instead of args[0]. +SOURCE_EQUIVALENCE_ALLOWLIST["mysql_four_way:create_mapped_exception"] = frozenset({ + "sqlspec.adapters.mysqlconnector.core" +}) + + +def _load_symbol(module_path: str, symbol_name: str) -> Any: + try: + module = importlib.import_module(module_path) + return getattr(module, symbol_name) + except (AttributeError, ImportError) as exc: + msg = f"Unable to load {symbol_name} from {module_path}: {exc}" + raise AssertionError(msg) from exc + + +def get_driver_feature_keys(adapter: str) -> tuple[str, ...]: + """Return the canonical DriverFeatures keys declared by an adapter TypedDict.""" + module_path, typed_dict_name = DRIVER_FEATURE_TYPED_DICTS[adapter] + typed_dict = _load_symbol(module_path, typed_dict_name) + annotations = getattr(typed_dict, "__annotations__", None) + if not isinstance(annotations, dict): + msg = f"{typed_dict_name} in {module_path} does not expose TypedDict annotations" + raise AssertionError(msg) + return tuple(annotations) + + +def get_driver_feature_consumed_keys(adapter: str) -> tuple[str, ...]: + """Return the registry-tracked keys consumed by an adapter.""" + return DRIVER_FEATURE_CONSUMED_KEYS[adapter] + + +class _DocstringStripper(ast.NodeTransformer): + def _strip_body(self, body: list[ast.stmt]) -> list[ast.stmt]: + if body and isinstance(body[0], ast.Expr): + value = body[0].value + if isinstance(value, ast.Constant) and isinstance(value.value, str): + return body[1:] + return body + + def visit_Module(self, node: ast.Module) -> ast.Module: + node.body = self._strip_body(node.body) + return cast("ast.Module", self.generic_visit(node)) + + def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef: + node.body = self._strip_body(node.body) + return cast("ast.FunctionDef", self.generic_visit(node)) + + def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> ast.AsyncFunctionDef: + node.body = self._strip_body(node.body) + return cast("ast.AsyncFunctionDef", self.generic_visit(node)) + + def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef: + node.body = self._strip_body(node.body) + return cast("ast.ClassDef", self.generic_visit(node)) + + +def normalize_source(source: str) -> str: + """Normalize Python source so source-equality checks ignore docstrings and formatting.""" + normalized = textwrap.dedent(source).strip() + tree = ast.parse(normalized) + stripped = _DocstringStripper().visit(tree) + ast.fix_missing_locations(stripped) + return ast.unparse(stripped).strip() + + +def _load_symbol_source(module_path: str, symbol_name: str) -> str: + symbol = _load_symbol(module_path, symbol_name) + try: + return inspect.getsource(symbol) + except OSError as exc: + msg = f"Unable to read source for {symbol_name} from {module_path}: {exc}" + raise AssertionError(msg) from exc + + +def assert_source_equivalent(case: SourceEquivalenceCase) -> None: + """Assert that every unallowlisted module exports the same normalized source.""" + allowed_modules = SOURCE_EQUIVALENCE_ALLOWLIST.get(case.allowlist_key, frozenset()) + normalized_sources: dict[str, str] = {} + for module_path in case.module_paths: + if module_path in allowed_modules: + continue + normalized_sources[module_path] = normalize_source(_load_symbol_source(module_path, case.symbol)) + + if len(normalized_sources) < 2: + msg = f"{case.case_id} does not have enough unallowlisted modules to compare" + raise AssertionError(msg) + + baseline_module, baseline_source = next(iter(normalized_sources.items())) + mismatches: list[str] = [] + for module_path, source in normalized_sources.items(): + if source == baseline_source: + continue + diff = "\n".join( + difflib.unified_diff( + baseline_source.splitlines(), + source.splitlines(), + fromfile=baseline_module, + tofile=module_path, + lineterm="", + ) + ) + mismatches.append(diff) + + if mismatches: + details = "\n\n".join(mismatches) + msg = f"{case.case_id} source mismatch for {case.symbol}:\n{details}" + raise AssertionError(msg) diff --git a/tests/integration/adapters/contracts/behaviors.py b/tests/integration/adapters/contracts/behaviors.py index a59f9e3af..322ec0a8d 100644 --- a/tests/integration/adapters/contracts/behaviors.py +++ b/tests/integration/adapters/contracts/behaviors.py @@ -4,7 +4,7 @@ import inspect import json import math -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Iterator, Mapping from decimal import Decimal from typing import TYPE_CHECKING, Any, Literal, Protocol, cast from uuid import NAMESPACE_DNS, UUID, uuid1, uuid4, uuid5 @@ -376,16 +376,22 @@ def _stream_seed_rows(count: int) -> "list[tuple[str, int, None]]": return [(f"row-{i:06d}", i, None) for i in range(count)] -def _seed_stream_table(driver: SyncContractDriver, table: ContractTable, count: int) -> None: +def _stream_contract_rows(count: int) -> tuple[ContractRow, ...]: + return tuple(ContractRow(name, value, note) for name, value, note in _stream_seed_rows(count)) + + +def _seed_stream_table(driver: SyncContractDriver, table: ContractTable, count: int, case: DriverCase) -> None: driver.execute(table.delete_sql) - driver.execute_many(table.insert_qmark_sql, _stream_seed_rows(count)) driver.commit() + _seed_sync(driver, _stream_contract_rows(count), table, case) -async def _seed_stream_table_async(driver: AsyncContractDriver, table: ContractTable, count: int) -> None: +async def _seed_stream_table_async( + driver: AsyncContractDriver, table: ContractTable, count: int, case: DriverCase +) -> None: await driver.execute(table.delete_sql) - await driver.execute_many(table.insert_qmark_sql, _stream_seed_rows(count)) await driver.commit() + await _seed_async(driver, _stream_contract_rows(count), table, case) def assert_sync_streaming_unsupported_contract(driver: object, case: DriverCase) -> None: @@ -431,7 +437,7 @@ def assert_sync_streaming_contract(driver: object, case: DriverCase) -> None: sync_driver = cast("SyncContractDriver", driver) table = case.table count = case.streaming_row_count - _seed_stream_table(sync_driver, table, count) + _seed_stream_table(sync_driver, table, count, case) stream = sync_driver.select_stream(table.select_ordered_sql, chunk_size=_STREAM_CHUNK_SIZE) iterator = iter(stream) @@ -469,7 +475,7 @@ async def assert_async_streaming_contract(driver: object, case: DriverCase) -> N async_driver = cast("AsyncContractDriver", driver) table = case.table count = case.streaming_row_count - await _seed_stream_table_async(async_driver, table, count) + await _seed_stream_table_async(async_driver, table, count, case) stream = async_driver.select_stream(table.select_ordered_sql, chunk_size=_STREAM_CHUNK_SIZE) iterator = aiter(stream) @@ -3691,6 +3697,326 @@ async def _oracle_json_native_async(driver: object, case: DriverCase) -> None: register_async_extra_assertion("driver_features:oracle_json_native", DRIVER_FEATURES_SCOPE, _oracle_json_native_async) +ORACLE_LOB_FETCH_SCOPE = "oracle_lob_fetch" +_ORACLE_LOB_JSON_PAYLOAD: dict[str, object] = {"integer": 42, "fraction": 1.25, "nested": {"label": "metadata-driven"}} +_ORACLE_OSON_PAYLOAD: dict[str, object] = {"kind": "oson", "tags": ["contract", "oracle"]} +_ORACLE_PLAIN_JSON_TEXT = to_json({"looks": "json", "but": "plain-clob"}) +_ORACLE_LONG_CLOB_TEXT = "long oracle clob " + ("x" * 5000) +_ORACLE_PLAIN_BLOB_BYTES = b"plain-oracle-blob" + + +def _oracle_lob_table(case: DriverCase, suffix: str) -> str: + return _pc_table(case, f"lob_{suffix}") + + +def _oracle_lob_select_sql(table: str) -> str: + return f"SELECT native_json, json_clob, json_blob, plain_clob, plain_blob, long_clob FROM {table} WHERE id = 1" + + +@contextlib.contextmanager +def _oracle_driver_feature(driver: object, key: str, value: object) -> "Iterator[None]": + features = cast("Any", driver).driver_features + missing = object() + previous = features.get(key, missing) + features[key] = value + try: + yield + finally: + if previous is missing: + features.pop(key, None) + else: + features[key] = previous + + +def _json_dict(value: object) -> dict[str, Any]: + if isinstance(value, memoryview): + return cast("dict[str, Any]", from_json(value.tobytes())) + if isinstance(value, bytes | bytearray): + return cast("dict[str, Any]", from_json(bytes(value))) + if isinstance(value, str): + return cast("dict[str, Any]", from_json(value)) + return cast("dict[str, Any]", value) + + +def _bytes_value(value: object) -> bytes: + if isinstance(value, memoryview): + return value.tobytes() + if isinstance(value, bytearray): + return bytes(value) + return cast("bytes", value) + + +def _assert_oracle_json_lanes(row: dict[str, object]) -> None: + native_json = _json_dict(row["native_json"]) + json_clob = _json_dict(row["json_clob"]) + json_blob = _json_dict(row["json_blob"]) + + assert native_json["nested"] == _ORACLE_LOB_JSON_PAYLOAD["nested"] + assert isinstance(native_json["integer"], Decimal) + assert isinstance(native_json["fraction"], Decimal) + assert json_clob == _ORACLE_LOB_JSON_PAYLOAD + assert json_blob == _ORACLE_LOB_JSON_PAYLOAD + assert type(json_clob["integer"]) is int + assert type(json_clob["fraction"]) is float + assert type(json_blob["integer"]) is int + assert type(json_blob["fraction"]) is float + + +def _assert_oracle_lob_materialized_row(row: Mapping[str, object]) -> None: + normalized = _lower_keys(dict(row)) + _assert_oracle_json_lanes(normalized) + assert normalized["plain_clob"] == _ORACLE_PLAIN_JSON_TEXT + assert _bytes_value(normalized["plain_blob"]) == _ORACLE_PLAIN_BLOB_BYTES + assert normalized["long_clob"] == _ORACLE_LONG_CLOB_TEXT + + +def _assert_oracle_lob_arrow_row(row: Mapping[str, object]) -> None: + normalized = _lower_keys(dict(row)) + assert normalized["plain_clob"] == _ORACLE_PLAIN_JSON_TEXT + assert _bytes_value(normalized["plain_blob"]) == _ORACLE_PLAIN_BLOB_BYTES + assert normalized["long_clob"] == _ORACLE_LONG_CLOB_TEXT + + +def _read_sync_lob(value: object) -> object: + read = getattr(value, "read", None) + assert callable(read) + return read() + + +async def _read_async_lob(value: object) -> object: + read = getattr(value, "read", None) + assert callable(read) + result = read() + if inspect.isawaitable(result): + result = await result + return result + + +def _seed_oracle_lob_fetch_sync(driver: "SyncContractDriver", case: DriverCase, table: str) -> None: + from sqlspec.adapters.oracledb import OracleBlob, OracleClob + + _sync_drop_table(driver, table) + driver.execute_script( + f""" + CREATE TABLE {table} ( + id NUMBER PRIMARY KEY, + native_json JSON, + json_clob CLOB CHECK (json_clob IS JSON), + json_blob BLOB CHECK (json_blob IS JSON), + plain_clob CLOB, + plain_blob BLOB, + long_clob CLOB + ) + """ + ) + driver.execute( + f""" + INSERT INTO {table} + (id, native_json, json_clob, json_blob, plain_clob, plain_blob, long_clob) + VALUES + (:id, :native_json, :json_clob, :json_blob, :plain_clob, :plain_blob, :long_clob) + """, + { + "id": 1, + "native_json": _ORACLE_LOB_JSON_PAYLOAD, + "json_clob": OracleClob(to_json(_ORACLE_LOB_JSON_PAYLOAD)), + "json_blob": OracleBlob(to_json(_ORACLE_LOB_JSON_PAYLOAD, as_bytes=True)), + "plain_clob": OracleClob(_ORACLE_PLAIN_JSON_TEXT), + "plain_blob": OracleBlob(_ORACLE_PLAIN_BLOB_BYTES), + "long_clob": _ORACLE_LONG_CLOB_TEXT, + }, + ) + driver.commit() + + +async def _seed_oracle_lob_fetch_async(driver: "AsyncContractDriver", case: DriverCase, table: str) -> None: + from sqlspec.adapters.oracledb import OracleBlob, OracleClob + + await _async_drop_table(driver, table) + await driver.execute_script( + f""" + CREATE TABLE {table} ( + id NUMBER PRIMARY KEY, + native_json JSON, + json_clob CLOB CHECK (json_clob IS JSON), + json_blob BLOB CHECK (json_blob IS JSON), + plain_clob CLOB, + plain_blob BLOB, + long_clob CLOB + ) + """ + ) + await driver.execute( + f""" + INSERT INTO {table} + (id, native_json, json_clob, json_blob, plain_clob, plain_blob, long_clob) + VALUES + (:id, :native_json, :json_clob, :json_blob, :plain_clob, :plain_blob, :long_clob) + """, + { + "id": 1, + "native_json": _ORACLE_LOB_JSON_PAYLOAD, + "json_clob": OracleClob(to_json(_ORACLE_LOB_JSON_PAYLOAD)), + "json_blob": OracleBlob(to_json(_ORACLE_LOB_JSON_PAYLOAD, as_bytes=True)), + "plain_clob": OracleClob(_ORACLE_PLAIN_JSON_TEXT), + "plain_blob": OracleBlob(_ORACLE_PLAIN_BLOB_BYTES), + "long_clob": _ORACLE_LONG_CLOB_TEXT, + }, + ) + await driver.commit() + + +def _assert_oracle_lob_fetch_true_stream_sync(driver: "SyncContractDriver", table: str) -> None: + with driver.select_stream( + f"SELECT plain_clob, plain_blob FROM {table} WHERE id = 1", chunk_size=1, fetch_lobs=True + ) as stream: + row = _lower_keys(next(iter(stream))) + assert _read_sync_lob(row["plain_clob"]) == _ORACLE_PLAIN_JSON_TEXT + assert _read_sync_lob(row["plain_blob"]) == _ORACLE_PLAIN_BLOB_BYTES + + +async def _assert_oracle_lob_fetch_true_stream_async(driver: "AsyncContractDriver", table: str) -> None: + async with driver.select_stream( + f"SELECT plain_clob, plain_blob FROM {table} WHERE id = 1", chunk_size=1, fetch_lobs=True + ) as stream: + row = _lower_keys(await anext(aiter(stream))) + assert await _read_async_lob(row["plain_clob"]) == _ORACLE_PLAIN_JSON_TEXT + assert await _read_async_lob(row["plain_blob"]) == _ORACLE_PLAIN_BLOB_BYTES + + +def _assert_oracle_lob_oson_sync(driver: "SyncContractDriver", case: DriverCase) -> None: + from sqlspec.adapters.oracledb import OracleBlob + + connection = cast("Any", driver).connection + encode_oson = getattr(connection, "encode_oson", None) + if not callable(encode_oson): + return + table = _oracle_lob_table(case, "oson") + _sync_drop_table(driver, table) + try: + try: + driver.execute_script( + f"CREATE TABLE {table} (id NUMBER PRIMARY KEY, payload BLOB CHECK (payload IS JSON FORMAT OSON))" + ) + except Exception as exc: + if "ORA-" in str(exc): + return + raise + driver.execute( + f"INSERT INTO {table} (id, payload) VALUES (:id, :payload)", + {"id": 1, "payload": OracleBlob(cast("bytes", encode_oson(_ORACLE_OSON_PAYLOAD)))}, + ) + row = _lower_keys(driver.select_one(f"SELECT payload FROM {table} WHERE id = 1")) + assert _json_dict(row["payload"]) == _ORACLE_OSON_PAYLOAD + finally: + _sync_drop_table(driver, table) + + +async def _assert_oracle_lob_oson_async(driver: "AsyncContractDriver", case: DriverCase) -> None: + from sqlspec.adapters.oracledb import OracleBlob + + connection = cast("Any", driver).connection + encode_oson = getattr(connection, "encode_oson", None) + if not callable(encode_oson): + return + encoded = encode_oson(_ORACLE_OSON_PAYLOAD) + if inspect.isawaitable(encoded): + encoded = await encoded + table = _oracle_lob_table(case, "oson") + await _async_drop_table(driver, table) + try: + try: + await driver.execute_script( + f"CREATE TABLE {table} (id NUMBER PRIMARY KEY, payload BLOB CHECK (payload IS JSON FORMAT OSON))" + ) + except Exception as exc: + if "ORA-" in str(exc): + return + raise + await driver.execute( + f"INSERT INTO {table} (id, payload) VALUES (:id, :payload)", + {"id": 1, "payload": OracleBlob(cast("bytes", encoded))}, + ) + row = _lower_keys(await driver.select_one(f"SELECT payload FROM {table} WHERE id = 1")) + assert _json_dict(row["payload"]) == _ORACLE_OSON_PAYLOAD + finally: + await _async_drop_table(driver, table) + + +def _oracle_lob_fetch_matrix(driver: object, case: DriverCase) -> None: + sync_driver = cast("SyncContractDriver", driver) + table = _oracle_lob_table(case, "matrix") + _seed_oracle_lob_fetch_sync(sync_driver, case, table) + select_sql = _oracle_lob_select_sql(table) + try: + _assert_oracle_lob_materialized_row(sync_driver.select_one(select_sql)) + with _oracle_driver_feature(sync_driver, "fetch_lobs", True): + _assert_oracle_lob_materialized_row(sync_driver.select_one(select_sql)) + with sync_driver.select_stream(select_sql, chunk_size=1) as stream: + _assert_oracle_lob_materialized_row(next(iter(stream))) + _assert_oracle_lob_fetch_true_stream_sync(sync_driver, table) + _assert_oracle_lob_arrow_row( + sync_driver.select_to_arrow( + f"SELECT plain_clob, plain_blob, long_clob FROM {table} WHERE id = 1" + ).data.to_pylist()[0] + ) + with _oracle_driver_feature(sync_driver, "fetch_lobs", True): + _assert_oracle_lob_arrow_row( + sync_driver.select_to_arrow( + f"SELECT plain_clob, plain_blob, long_clob FROM {table} WHERE id = 1" + ).data.to_pylist()[0] + ) + _assert_oracle_lob_oson_sync(sync_driver, case) + finally: + _sync_drop_table(sync_driver, table) + + +async def _oracle_lob_fetch_matrix_async(driver: object, case: DriverCase) -> None: + async_driver = cast("AsyncContractDriver", driver) + table = _oracle_lob_table(case, "matrix") + await _seed_oracle_lob_fetch_async(async_driver, case, table) + select_sql = _oracle_lob_select_sql(table) + try: + _assert_oracle_lob_materialized_row(await async_driver.select_one(select_sql)) + with _oracle_driver_feature(async_driver, "fetch_lobs", True): + _assert_oracle_lob_materialized_row(await async_driver.select_one(select_sql)) + async with async_driver.select_stream(select_sql, chunk_size=1) as stream: + _assert_oracle_lob_materialized_row(await anext(aiter(stream))) + await _assert_oracle_lob_fetch_true_stream_async(async_driver, table) + _assert_oracle_lob_arrow_row( + ( + await async_driver.select_to_arrow( + f"SELECT plain_clob, plain_blob, long_clob FROM {table} WHERE id = 1" + ) + ).data.to_pylist()[0] + ) + with _oracle_driver_feature(async_driver, "fetch_lobs", True): + _assert_oracle_lob_arrow_row( + ( + await async_driver.select_to_arrow( + f"SELECT plain_clob, plain_blob, long_clob FROM {table} WHERE id = 1" + ) + ).data.to_pylist()[0] + ) + await _assert_oracle_lob_oson_async(async_driver, case) + finally: + await _async_drop_table(async_driver, table) + + +register_sync_extra_assertion("oracle_lob_fetch:matrix", ORACLE_LOB_FETCH_SCOPE, _oracle_lob_fetch_matrix) +register_async_extra_assertion("oracle_lob_fetch:matrix", ORACLE_LOB_FETCH_SCOPE, _oracle_lob_fetch_matrix_async) + + +def assert_sync_oracle_lob_fetch_contract(driver: object, case: DriverCase) -> None: + """Run Oracle's shared LOB fetch matrix proof.""" + dispatch_sync_extra_assertions(driver, case, ORACLE_LOB_FETCH_SCOPE) + + +async def assert_async_oracle_lob_fetch_contract(driver: object, case: DriverCase) -> None: + """Run Oracle's shared async LOB fetch matrix proof.""" + await dispatch_async_extra_assertions(driver, case, ORACLE_LOB_FETCH_SCOPE) + + def _bigquery_sql_features(driver: object, case: DriverCase) -> None: """Fold BigQuery dialect scalar SQL: math/string/conditional/specific functions + ARRAY/STRUCT. @@ -4308,6 +4634,188 @@ def assert_sync_custom_type_adapters_contract(make_config: SyncConfigFactory, ca disabled.close_pool() +async def assert_async_custom_type_adapters_contract(make_config: AsyncConfigFactory, case: DriverCase) -> None: + """Async mirror of the sqlite custom-adapter enable-false semantics contract.""" + import sqlite3 + + table = _pc_table(case, "adapters") + test_dict = {"key": "value", "count": 42} + test_list = [1, 2, 3, "four"] + + enabled = make_config( + driver_features={"enable_custom_adapters": True}, connection_overrides={"detect_types": sqlite3.PARSE_DECLTYPES} + ) + try: + async with enabled.provide_session() as session: + await session.execute_script(f"DROP TABLE IF EXISTS {table}") + await session.execute(f"CREATE TABLE {table} (id INTEGER, data JSON, items JSON)") + await session.execute( + f"INSERT INTO {table} (id, data, items) VALUES (?, ?, ?)", + (1, json.dumps(test_dict), json.dumps(test_list)), + ) + await session.commit() + row = await session.select_one(f"SELECT data, items FROM {table} WHERE id = 1") + assert row["data"] == test_dict + assert row["items"] == test_list + finally: + await enabled.close_pool() + + disabled = make_config() + try: + async with disabled.provide_session() as session: + await session.execute_script(f"DROP TABLE IF EXISTS {table}") + await session.execute(f"CREATE TABLE {table} (id INTEGER, data TEXT)") + await session.execute(f"INSERT INTO {table} (id, data) VALUES (?, ?)", (1, json.dumps(test_dict))) + await session.commit() + row = await session.select_one(f"SELECT data FROM {table} WHERE id = 1") + assert isinstance(row["data"], str) + assert json.loads(row["data"]) == test_dict + finally: + await disabled.close_pool() + + +_DRIVER_FEATURE_PARITY_ROWS = ( + ContractRow("feature-a", 1, "note-a"), + ContractRow("feature-b", 2, None), + ContractRow("feature-c", 3, "note-c"), +) + + +def _expected_feature_rows(rows: tuple[ContractRow, ...] = _DRIVER_FEATURE_PARITY_ROWS) -> list[dict[str, object]]: + return [{"name": row.name, "value": row.value, "note": row.note} for row in rows] + + +def _assert_row_format_matches_data(result: SQLResult, expected_rows: list[dict[str, object]]) -> None: + assert result.get_data() == expected_rows + assert result.data is not None + if not result.data: + return + row_format = result._row_format # pyright: ignore[reportPrivateUsage] + first = result.data[0] + if isinstance(first, dict): + assert row_format == "dict" + return + if isinstance(first, Mapping): + assert row_format in {"dict", "record"} + return + if hasattr(first, "keys"): + assert row_format == "record" + return + assert row_format == "tuple" + + +def _reset_contract_table_sync(driver: SyncContractDriver, table: ContractTable) -> None: + with contextlib.suppress(Exception): + driver.execute_script(f"DROP TABLE IF EXISTS {table.name}") + driver.execute_script(table.create_sql) + driver.commit() + + +async def _reset_contract_table_async(driver: AsyncContractDriver, table: ContractTable) -> None: + with contextlib.suppress(Exception): + await driver.execute_script(f"DROP TABLE IF EXISTS {table.name}") + await driver.execute_script(table.create_sql) + await driver.commit() + + +def _row_format_config_kwargs(case: DriverCase) -> dict[str, Any]: + if case.adapter == "psycopg": + from psycopg.rows import dict_row + + return {"connection_overrides": {"row_factory": dict_row}} + if case.adapter == "sqlite": + return {"driver_features": {"row_factory": "dict"}} + return {} + + +def _async_row_format_config_kwargs(case: DriverCase) -> dict[str, Any]: + if case.adapter == "psycopg": + from psycopg.rows import dict_row + + return {"connection_overrides": {"row_factory": dict_row}} + if case.adapter == "aiosqlite": + return {"driver_features": {"row_factory": "dict"}} + return {} + + +def assert_sync_driver_feature_parity_contract(driver: object, case: DriverCase) -> None: + """Assert buffered, native row-stream, and Arrow rows agree for one canonical fixture.""" + if not (case.supports_native_row_streaming or case.supports_arrow): + pytest.skip(f"{case.adapter} has no row-stream or Arrow feature path") + sync_driver = cast("SyncContractDriver", driver) + table = case.table + sync_driver.execute(table.delete_sql) + sync_driver.commit() + _seed_sync(sync_driver, _DRIVER_FEATURE_PARITY_ROWS, table, case) + expected = _expected_feature_rows() + + buffered = sync_driver.execute(table.select_ordered_sql).get_data() + assert buffered == expected + + if case.supports_native_row_streaming: + with sync_driver.select_stream(table.select_ordered_sql, chunk_size=1) as stream: + assert list(stream) == expected + + if case.supports_arrow: + assert sync_driver.select_to_arrow(table.select_ordered_sql).data.to_pylist() == expected + + +async def assert_async_driver_feature_parity_contract(driver: object, case: DriverCase) -> None: + """Async mirror of the buffered/native-stream/Arrow parity contract.""" + if not (case.supports_native_row_streaming or case.supports_arrow): + pytest.skip(f"{case.adapter} has no row-stream or Arrow feature path") + async_driver = cast("AsyncContractDriver", driver) + table = case.table + await async_driver.execute(table.delete_sql) + await async_driver.commit() + await _seed_async(async_driver, _DRIVER_FEATURE_PARITY_ROWS, table, case) + expected = _expected_feature_rows() + + buffered = (await async_driver.execute(table.select_ordered_sql)).get_data() + assert buffered == expected + + if case.supports_native_row_streaming: + async with async_driver.select_stream(table.select_ordered_sql, chunk_size=1) as stream: + assert [row async for row in stream] == expected + + if case.supports_arrow: + assert (await async_driver.select_to_arrow(table.select_ordered_sql)).data.to_pylist() == expected + + +def assert_sync_driver_feature_row_format_contract(make_config: SyncConfigFactory | None, case: DriverCase) -> None: + """Assert configured row factories are tagged consistently with their materialized data.""" + if make_config is None: + pytest.skip(f"{case.adapter} has no config factory for row-format verification") + config = make_config(**_row_format_config_kwargs(case)) + try: + with config.provide_session() as session: + sync_driver = session + _reset_contract_table_sync(sync_driver, case.table) + _seed_sync(sync_driver, _DRIVER_FEATURE_PARITY_ROWS, case.table, case) + result = assert_sql_result(sync_driver.execute(case.table.select_ordered_sql)) + _assert_row_format_matches_data(result, _expected_feature_rows()) + finally: + config.close_pool() + + +async def assert_async_driver_feature_row_format_contract( + make_config: AsyncConfigFactory | None, case: DriverCase +) -> None: + """Async mirror of the configured row-factory row-format contract.""" + if make_config is None: + pytest.skip(f"{case.adapter} has no config factory for row-format verification") + config = make_config(**_async_row_format_config_kwargs(case)) + try: + async with config.provide_session() as session: + async_driver = session + await _reset_contract_table_async(async_driver, case.table) + await _seed_async(async_driver, _DRIVER_FEATURE_PARITY_ROWS, case.table, case) + result = assert_sql_result(await async_driver.execute(case.table.select_ordered_sql)) + _assert_row_format_matches_data(result, _expected_feature_rows()) + finally: + await config.close_pool() + + def assert_sync_statement_input_contract(driver: object, case: DriverCase, input_case: StatementInputCase) -> None: """Assert sync drivers return equivalent rows for one statement input shape.""" sync_driver = cast("SyncContractDriver", driver) diff --git a/tests/integration/adapters/contracts/conftest.py b/tests/integration/adapters/contracts/conftest.py index ab00aa4a4..3e69731bd 100644 --- a/tests/integration/adapters/contracts/conftest.py +++ b/tests/integration/adapters/contracts/conftest.py @@ -1160,9 +1160,12 @@ def make( *, pooled: bool = False, driver_features: "AiosqliteDriverFeatures | None" = None, + connection_overrides: "dict[str, Any] | None" = None, connection_instance: object | None = None, ) -> AiosqliteConfig: - connection_config = _lifecycle_connection_config(str(tmp_path / "lifecycle_aiosqlite.db"), pooled=pooled) + connection_config = _lifecycle_connection_config( + str(tmp_path / "lifecycle_aiosqlite.db"), pooled=pooled, connection_overrides=connection_overrides + ) extra: dict[str, Any] = {} if driver_features is not None: extra["driver_features"] = driver_features @@ -1248,9 +1251,12 @@ def make( *, pooled: bool = False, driver_features: "PsycopgDriverFeatures | None" = None, + connection_overrides: "dict[str, Any] | None" = None, connection_instance: object | None = None, ) -> PsycopgSyncConfig: connection_config: dict[str, Any] = {"conninfo": _postgres_conninfo(postgres_service), "autocommit": True} + if connection_overrides: + connection_config.update(connection_overrides) if pooled: connection_config.update({"min_size": 2, "max_size": 5}) extra: dict[str, Any] = {} @@ -1271,9 +1277,12 @@ def make( *, pooled: bool = False, driver_features: "PsycopgDriverFeatures | None" = None, + connection_overrides: "dict[str, Any] | None" = None, connection_instance: object | None = None, ) -> PsycopgAsyncConfig: connection_config: dict[str, Any] = {"conninfo": _postgres_conninfo(postgres_service), "autocommit": True} + if connection_overrides: + connection_config.update(connection_overrides) if pooled: connection_config.update({"min_size": 2, "max_size": 5}) extra: dict[str, Any] = {} diff --git a/tests/integration/adapters/contracts/test_driver_features_contract.py b/tests/integration/adapters/contracts/test_driver_features_contract.py index 5fbf75099..3e2ea1b5e 100644 --- a/tests/integration/adapters/contracts/test_driver_features_contract.py +++ b/tests/integration/adapters/contracts/test_driver_features_contract.py @@ -1,12 +1,187 @@ -"""Shared driver-feature contracts: folded driver-specific behaviors (COPY, SET-variable, native types).""" - -from tests.integration.adapters.contracts._cases import DriverCaseContext, get_driver_case +"""Shared driver-feature contracts: registry scaffolds, source tripwires, and folded driver-specific behaviors.""" + +import importlib +from collections.abc import Sequence +from typing import cast + +import pytest +from _pytest.mark.structures import ParameterSet + +from tests.integration.adapters.contracts._cases import ( + ASYNC_DRIVER_CASES, + DRIVER_CASES, + SYNC_DRIVER_CASES, + DriverCase, + DriverCaseContext, + get_driver_case, +) +from tests.integration.adapters.contracts._driver_type_system import ( + DRIVER_FEATURE_CONTRACT_GROUPS, + SOURCE_EQUIVALENCE_CASES, + SourceEquivalenceCase, + assert_source_equivalent, + get_driver_feature_consumed_keys, + get_driver_feature_keys, +) from tests.integration.adapters.contracts.behaviors import ( + AsyncConfigFactory, + SyncConfigFactory, + assert_async_custom_type_adapters_contract, + assert_async_driver_feature_parity_contract, + assert_async_driver_feature_row_format_contract, assert_async_driver_features_contract, + assert_sync_custom_type_adapters_contract, + assert_sync_driver_feature_parity_contract, + assert_sync_driver_feature_row_format_contract, assert_sync_driver_features_contract, validate_extra_assertions, ) +_DRIVER_TYPE_SYSTEM_MODULE = "tests.integration.adapters.contracts._driver_type_system" + + +def test_driver_type_system_registry_module_exports_contract_groups() -> None: + """The driver-type-system registry module is importable and exposes the expected contract groups.""" + module = importlib.import_module(_DRIVER_TYPE_SYSTEM_MODULE) + + assert hasattr(module, "DRIVER_FEATURE_CONTRACT_GROUPS") + assert hasattr(module, "DRIVER_FEATURE_CONTRACT_ENABLED") + assert hasattr(module, "DRIVER_FEATURE_CONSUMED_KEYS") + assert hasattr(module, "SOURCE_EQUIVALENCE_CASES") + assert hasattr(module, "SOURCE_EQUIVALENCE_ALLOWLIST") + assert set(module.DRIVER_FEATURE_CONTRACT_GROUPS) == { + "feature_honesty", + "enable_false_semantics", + "parity", + "row_format", + } + assert module.DRIVER_FEATURE_CONTRACT_ENABLED == { + group_name: group.enabled for group_name, group in module.DRIVER_FEATURE_CONTRACT_GROUPS.items() + } + assert set(module.DRIVER_FEATURE_CONSUMED_KEYS) == {case.adapter for case in DRIVER_CASES} + assert {case.group for case in module.SOURCE_EQUIVALENCE_CASES} == {"mysql_four_way", "sqlite_pair", "mssql_pair"} + + +def _disabled_first_params(group_name: str, cases: Sequence[DriverCase]) -> tuple[ParameterSet, ...]: + group = DRIVER_FEATURE_CONTRACT_GROUPS[group_name] + params: list[ParameterSet] = [] + for case in cases: + marks = list(case.marks) + if not group.enabled or case.adapter not in group.enabled_adapters: + marks.append( + pytest.mark.skip( + reason=( + f"{group_name} scaffold is disabled until the adapter type-system registry is populated " + f"for {case.adapter}" + ) + ) + ) + params.append(pytest.param(case, id=case.id, marks=tuple(marks))) + return tuple(params) + + +def _assert_driver_feature_honesty(adapter: str) -> None: + declared = set(get_driver_feature_keys(adapter)) + consumed = set(get_driver_feature_consumed_keys(adapter)) + missing = sorted(declared - consumed) + + assert consumed, f"{adapter} has no consumed DriverFeatures registry entry" + assert not missing, f"{adapter} declares unconsumed DriverFeatures keys: {missing!r}" + + +@pytest.mark.parametrize("adapter", DRIVER_FEATURE_CONTRACT_GROUPS["feature_honesty"].enabled_adapters) +def test_driver_feature_honesty_registry_covers_enabled_adapters(adapter: str) -> None: + """Every feature-honesty-enabled adapter has an explicit consumed-key registry entry.""" + _assert_driver_feature_honesty(adapter) + + +@pytest.mark.parametrize("sync_driver_case", _disabled_first_params("feature_honesty", SYNC_DRIVER_CASES)) +def test_sync_driver_feature_honesty_contract(sync_driver_case: DriverCase) -> None: + """Sync DriverFeatures TypedDict keys are tracked as consumed after adapter migration.""" + _assert_driver_feature_honesty(sync_driver_case.adapter) + + +@pytest.mark.parametrize("async_driver_case", _disabled_first_params("feature_honesty", ASYNC_DRIVER_CASES)) +async def test_async_driver_feature_honesty_contract(async_driver_case: DriverCase) -> None: + """Async DriverFeatures TypedDict keys are tracked as consumed after adapter migration.""" + _assert_driver_feature_honesty(async_driver_case.adapter) + + +@pytest.mark.parametrize( + "sync_driver_case", + _disabled_first_params("enable_false_semantics", (get_driver_case("sqlite-sync"),)), + indirect=True, +) +def test_sync_driver_feature_semantics_contract(sync_driver_case: DriverCaseContext) -> None: + """Sync enable-false semantics stays inert unless the feature is explicitly enabled.""" + assert sync_driver_case.make_config is not None + assert_sync_custom_type_adapters_contract( + cast("SyncConfigFactory", sync_driver_case.make_config), sync_driver_case.case + ) + + +@pytest.mark.parametrize( + "async_driver_case", + _disabled_first_params("enable_false_semantics", (get_driver_case("aiosqlite-async"),)), + indirect=True, +) +async def test_async_driver_feature_semantics_contract(async_driver_case: DriverCaseContext) -> None: + """Async enable-false semantics stays inert unless the feature is explicitly enabled.""" + assert async_driver_case.make_config is not None + await assert_async_custom_type_adapters_contract( + cast("AsyncConfigFactory", async_driver_case.make_config), async_driver_case.case + ) + + +@pytest.mark.parametrize( + "sync_driver_case", + _disabled_first_params( + "parity", + tuple( + case for case in SYNC_DRIVER_CASES if case.supports_arrow_streaming or case.supports_native_row_streaming + ), + ), + indirect=True, +) +def test_sync_driver_feature_parity_contract(sync_driver_case: DriverCaseContext) -> None: + """Sync buffered, streamed, and Arrow values stay equal for a canonical fixture.""" + assert_sync_driver_feature_parity_contract(sync_driver_case.driver, sync_driver_case.case) + + +@pytest.mark.parametrize( + "async_driver_case", + _disabled_first_params( + "parity", + tuple( + case for case in ASYNC_DRIVER_CASES if case.supports_arrow_streaming or case.supports_native_row_streaming + ), + ), + indirect=True, +) +async def test_async_driver_feature_parity_contract(async_driver_case: DriverCaseContext) -> None: + """Async buffered, streamed, and Arrow values stay equal for a canonical fixture.""" + await assert_async_driver_feature_parity_contract(async_driver_case.driver, async_driver_case.case) + + +@pytest.mark.parametrize("sync_driver_case", _disabled_first_params("row_format", SYNC_DRIVER_CASES), indirect=True) +def test_sync_driver_feature_row_format_contract(sync_driver_case: DriverCaseContext) -> None: + """Sync row-format tags match the actual configured row shape.""" + make_config = cast("SyncConfigFactory | None", sync_driver_case.make_config) + assert_sync_driver_feature_row_format_contract(make_config, sync_driver_case.case) + + +@pytest.mark.parametrize("async_driver_case", _disabled_first_params("row_format", ASYNC_DRIVER_CASES), indirect=True) +async def test_async_driver_feature_row_format_contract(async_driver_case: DriverCaseContext) -> None: + """Async row-format tags match the actual configured row shape.""" + make_config = cast("AsyncConfigFactory | None", async_driver_case.make_config) + await assert_async_driver_feature_row_format_contract(make_config, async_driver_case.case) + + +@pytest.mark.parametrize("source_equivalence_case", SOURCE_EQUIVALENCE_CASES, ids=lambda case: case.case_id) +def test_driver_feature_source_equivalence(source_equivalence_case: SourceEquivalenceCase) -> None: + """Duplicate helper implementations stay source-equivalent across adapter families.""" + assert_source_equivalent(source_equivalence_case) + def test_sync_driver_features_contract(sync_driver_case: DriverCaseContext) -> None: """Sync drivers run folded driver-feature proofs opted in via extra_assertions.""" @@ -47,6 +222,17 @@ def test_oracle_contract_declares_plsql_assertions() -> None: validate_extra_assertions(async_case) +def test_oracle_contract_declares_lob_fetch_matrix_assertions() -> None: + """Oracle LOB fetch behavior stays in the shared contract matrix.""" + sync_case = get_driver_case("oracledb-sync") + async_case = get_driver_case("oracledb-async") + + assert "oracle_lob_fetch:matrix" in sync_case.extra_assertions + assert "oracle_lob_fetch:matrix" in async_case.extra_assertions + validate_extra_assertions(sync_case) + validate_extra_assertions(async_case) + + def test_spanner_contract_status_defers_session_controls_explicitly() -> None: """Spanner remains deferred in the shared matrix until a safe active contract is available.""" case = get_driver_case("spanner-sync") diff --git a/tests/integration/adapters/contracts/test_oracle_lob_fetch_contract.py b/tests/integration/adapters/contracts/test_oracle_lob_fetch_contract.py new file mode 100644 index 000000000..dc192ba32 --- /dev/null +++ b/tests/integration/adapters/contracts/test_oracle_lob_fetch_contract.py @@ -0,0 +1,27 @@ +"""Shared Oracle LOB fetch-path contracts.""" + +import pytest +from _pytest.mark.structures import ParameterSet + +from tests.integration.adapters.contracts._cases import DriverCaseContext, get_driver_case +from tests.integration.adapters.contracts.behaviors import ( + assert_async_oracle_lob_fetch_contract, + assert_sync_oracle_lob_fetch_contract, +) + + +def _case_param(case_id: str) -> ParameterSet: + case = get_driver_case(case_id) + return pytest.param(case, id=case.id, marks=case.marks) + + +@pytest.mark.parametrize("sync_driver_case", (_case_param("oracledb-sync"),), indirect=True) +def test_sync_oracle_lob_fetch_contract(sync_driver_case: DriverCaseContext) -> None: + """Sync Oracle proves metadata-driven LOB/JSON fetches through the shared matrix.""" + assert_sync_oracle_lob_fetch_contract(sync_driver_case.driver, sync_driver_case.case) + + +@pytest.mark.parametrize("async_driver_case", (_case_param("oracledb-async"),), indirect=True) +async def test_async_oracle_lob_fetch_contract(async_driver_case: DriverCaseContext) -> None: + """Async Oracle proves metadata-driven LOB/JSON fetches through the shared matrix.""" + await assert_async_oracle_lob_fetch_contract(async_driver_case.driver, async_driver_case.case) diff --git a/tests/integration/adapters/oracledb/test_bench_scenarios.py b/tests/integration/adapters/oracledb/test_bench_scenarios.py new file mode 100644 index 000000000..62415997c --- /dev/null +++ b/tests/integration/adapters/oracledb/test_bench_scenarios.py @@ -0,0 +1,92 @@ +"""Service-backed smoke tests for Oracle benchmark scenarios.""" + +import json +from pathlib import Path + +import pytest +from click.testing import CliRunner +from pytest_databases.docker.oracle import OracleService +from tools.scripts import bench + +pytestmark = pytest.mark.xdist_group("oracle") + + +ORACLE_LOB_SCENARIOS = { + ("raw", "oracle", "lob_fetch_1k"): "raw_oracle_lob_fetch_1k", + ("raw", "oracle", "lob_fetch_100k"): "raw_oracle_lob_fetch_100k", + ("sqlspec", "oracle", "lob_fetch_1k"): "sqlspec_oracle_lob_fetch_1k", + ("sqlspec", "oracle", "lob_fetch_100k"): "sqlspec_oracle_lob_fetch_100k", + ("sqlspec_fetch_lobs_true", "oracle", "lob_fetch_1k"): "sqlspec_oracle_lob_fetch_fetch_lobs_true_1k", + ("sqlspec_fetch_lobs_true", "oracle", "lob_fetch_100k"): "sqlspec_oracle_lob_fetch_fetch_lobs_true_100k", + ("sqlspec_async", "oracle", "lob_fetch_1k"): "sqlspec_oracle_lob_fetch_async_1k", + ("sqlspec_async", "oracle", "lob_fetch_100k"): "sqlspec_oracle_lob_fetch_async_100k", + ("sqlspec_async_fetch_lobs_true", "oracle", "lob_fetch_1k"): "sqlspec_oracle_lob_fetch_async_fetch_lobs_true_1k", + ( + "sqlspec_async_fetch_lobs_true", + "oracle", + "lob_fetch_100k", + ): "sqlspec_oracle_lob_fetch_async_fetch_lobs_true_100k", +} + + +def _configure_oracle_bench_env(monkeypatch: pytest.MonkeyPatch, oracle_service: OracleService) -> None: + monkeypatch.setenv("SQLSPEC_BENCH_ORACLE_HOST", oracle_service.host) + monkeypatch.setenv("SQLSPEC_BENCH_ORACLE_PORT", str(oracle_service.port)) + monkeypatch.setenv("SQLSPEC_BENCH_ORACLE_SERVICE_NAME", oracle_service.service_name) + monkeypatch.setenv("SQLSPEC_BENCH_ORACLE_USER", oracle_service.user) + monkeypatch.setenv("SQLSPEC_BENCH_ORACLE_PASSWORD", oracle_service.password) + monkeypatch.setattr(bench, "ORACLE_LOB_ROWS", 2) + monkeypatch.setattr(bench, "ORACLE_LOB_PAYLOAD_SIZES", {"1k": 128, "100k": 256}) + + +def test_oracle_lob_benchmark_scenarios_are_registered() -> None: + for key, function_name in ORACLE_LOB_SCENARIOS.items(): + assert bench.SCENARIO_REGISTRY[key] is getattr(bench, function_name) + assert function_name in bench.__all__ + + +def test_oracle_lob_benchmark_scenarios_run_with_service_config( + monkeypatch: pytest.MonkeyPatch, oracle_23ai_service: OracleService +) -> None: + _configure_oracle_bench_env(monkeypatch, oracle_23ai_service) + + errors: list[str] = [] + results = bench.run_extended_benchmark("oracle", errors, iterations=1, warmup=0) + + assert errors == [] + assert {(result["library"], result["driver"], result["scenario"]) for result in results} == { + (bench._benchmark_library_label(library), driver, scenario) + for library, driver, scenario in ORACLE_LOB_SCENARIOS + } + assert len(results) == len(ORACLE_LOB_SCENARIOS) + + +def test_oracle_extended_cli_runs_without_core_missing_scenario_errors( + monkeypatch: pytest.MonkeyPatch, oracle_23ai_service: OracleService, tmp_path: Path +) -> None: + _configure_oracle_bench_env(monkeypatch, oracle_23ai_service) + json_output = tmp_path / "oracle-bench.json" + + result = CliRunner().invoke( + bench.main, + ["--driver", "oracle", "--extended", "--iterations", "1", "--warmup", "0", "--json-output", str(json_output)], + ) + + assert result.exit_code == 0 + assert "No implementation" not in result.output + payload = json.loads(json_output.read_text()) + assert {result["scenario"] for result in payload["results"]} == {"lob_fetch_1k", "lob_fetch_100k"} + + +def test_oracle_lob_benchmark_connection_config_requires_service_safe_env(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ( + "SQLSPEC_BENCH_ORACLE_HOST", + "SQLSPEC_BENCH_ORACLE_PORT", + "SQLSPEC_BENCH_ORACLE_SERVICE_NAME", + "SQLSPEC_BENCH_ORACLE_USER", + "SQLSPEC_BENCH_ORACLE_PASSWORD", + ): + monkeypatch.delenv(name, raising=False) + + with pytest.raises(RuntimeError, match="SQLSPEC_BENCH_ORACLE_HOST"): + bench._oracle_connection_config_from_env() diff --git a/tests/integration/adapters/oracledb/test_lob_fetch_parity.py b/tests/integration/adapters/oracledb/test_lob_fetch_parity.py new file mode 100644 index 000000000..790905361 --- /dev/null +++ b/tests/integration/adapters/oracledb/test_lob_fetch_parity.py @@ -0,0 +1,103 @@ +"""Oracle LOB fetch parity for streaming and Arrow paths.""" + +import pytest + +from sqlspec.adapters.oracledb import OracleAsyncDriver, OracleBlob, OracleClob, OracleSyncDriver + +pytestmark = pytest.mark.xdist_group("oracle") + +_CLOB_TEXT = "streamed oracle clob payload" +_BLOB_BYTES = b"streamed-oracle-blob" + + +def _drop_sync(driver: "OracleSyncDriver", table_name: str) -> None: + driver.execute_script( + f"BEGIN EXECUTE IMMEDIATE 'DROP TABLE {table_name}'; " + "EXCEPTION WHEN OTHERS THEN IF SQLCODE != -942 THEN RAISE; END IF; END;" + ) + + +async def _drop_async(driver: "OracleAsyncDriver", table_name: str) -> None: + await driver.execute_script( + f"BEGIN EXECUTE IMMEDIATE 'DROP TABLE {table_name}'; " + "EXCEPTION WHEN OTHERS THEN IF SQLCODE != -942 THEN RAISE; END IF; END;" + ) + + +def test_sync_stream_returns_concrete_lobs_by_default(oracle_sync_session: "OracleSyncDriver") -> None: + """Native sync streaming should inherit the default ``fetch_lobs=False`` behavior.""" + table_name = "lob_fetch_stream_sync" + _drop_sync(oracle_sync_session, table_name) + oracle_sync_session.execute_script(f"CREATE TABLE {table_name} (id NUMBER PRIMARY KEY, content CLOB, data BLOB)") + oracle_sync_session.execute( + f"INSERT INTO {table_name} (id, content, data) VALUES (:id, :content, :data)", + {"id": 1, "content": OracleClob(_CLOB_TEXT), "data": OracleBlob(_BLOB_BYTES)}, + ) + + try: + rows = list(oracle_sync_session.select_stream(f"SELECT content, data FROM {table_name}", chunk_size=1)) + + assert rows == [{"content": _CLOB_TEXT, "data": _BLOB_BYTES}] + finally: + _drop_sync(oracle_sync_session, table_name) + + +def test_sync_stream_fetch_lobs_true_keeps_locator(oracle_sync_session: "OracleSyncDriver") -> None: + """Per-call ``fetch_lobs=True`` keeps the Oracle locator escape hatch.""" + table_name = "lob_fetch_stream_locator_sync" + _drop_sync(oracle_sync_session, table_name) + oracle_sync_session.execute_script(f"CREATE TABLE {table_name} (id NUMBER PRIMARY KEY, content CLOB)") + oracle_sync_session.execute( + f"INSERT INTO {table_name} (id, content) VALUES (:id, :content)", {"id": 1, "content": OracleClob(_CLOB_TEXT)} + ) + + try: + with oracle_sync_session.select_stream( + f"SELECT content FROM {table_name}", chunk_size=1, fetch_lobs=True + ) as stream: + row = next(stream) + locator = row["content"] + assert callable(locator.read) + assert locator.read() == _CLOB_TEXT + finally: + _drop_sync(oracle_sync_session, table_name) + + +async def test_async_stream_returns_concrete_lobs_by_default(oracle_async_session: "OracleAsyncDriver") -> None: + """Native async streaming should inherit the default ``fetch_lobs=False`` behavior.""" + table_name = "lob_fetch_stream_async" + await _drop_async(oracle_async_session, table_name) + await oracle_async_session.execute_script( + f"CREATE TABLE {table_name} (id NUMBER PRIMARY KEY, content CLOB, data BLOB)" + ) + await oracle_async_session.execute( + f"INSERT INTO {table_name} (id, content, data) VALUES (:id, :content, :data)", + {"id": 1, "content": OracleClob(_CLOB_TEXT), "data": OracleBlob(_BLOB_BYTES)}, + ) + + try: + rows: list[dict[str, object]] = [ + row + async for row in oracle_async_session.select_stream(f"SELECT content, data FROM {table_name}", chunk_size=1) + ] + + assert rows == [{"content": _CLOB_TEXT, "data": _BLOB_BYTES}] + finally: + await _drop_async(oracle_async_session, table_name) + + +async def test_arrow_export_returns_concrete_clob_values(oracle_async_session: "OracleAsyncDriver") -> None: + """Oracle Arrow export should forward ``fetch_lobs=False`` and return CLOB text.""" + table_name = "lob_fetch_arrow_async" + await _drop_async(oracle_async_session, table_name) + await oracle_async_session.execute_script(f"CREATE TABLE {table_name} (id NUMBER PRIMARY KEY, content CLOB)") + await oracle_async_session.execute( + f"INSERT INTO {table_name} (id, content) VALUES (:id, :content)", {"id": 1, "content": OracleClob(_CLOB_TEXT)} + ) + + try: + result = await oracle_async_session.select_to_arrow(f"SELECT content FROM {table_name}") + + assert result.get_data().to_pydict() == {"content": [_CLOB_TEXT]} + finally: + await _drop_async(oracle_async_session, table_name) diff --git a/tests/integration/adapters/oracledb/test_msgspec_clob.py b/tests/integration/adapters/oracledb/test_msgspec_clob.py index 47ef6d27f..ba20f79d0 100644 --- a/tests/integration/adapters/oracledb/test_msgspec_clob.py +++ b/tests/integration/adapters/oracledb/test_msgspec_clob.py @@ -45,6 +45,13 @@ class JsonDocumentRecord(msgspec.Struct): metadata: dict +class ClobTextDocumentRecord(msgspec.Struct): + """Document with unconstrained CLOB text.""" + + id: int + metadata: str + + class BinaryDocumentRecord(msgspec.Struct): """Document with BLOB binary data.""" @@ -206,12 +213,8 @@ def test_oracle_sync_mixed_clob_varchar2_msgspec(oracle_sync_session: OracleSync ) -async def test_oracle_async_json_in_clob_detection(oracle_async_session: OracleAsyncDriver) -> None: - """Test JSON detection in CLOB with msgspec hydration. - - Verifies that JSON stored in CLOB is automatically detected and parsed, - then successfully hydrated into msgspec struct. - """ +async def test_oracle_async_json_in_clob_returns_str(oracle_async_session: OracleAsyncDriver) -> None: + """Test unconstrained CLOB JSON hydrates as text.""" await oracle_async_session.execute_script( "BEGIN EXECUTE IMMEDIATE 'DROP TABLE test_json_clob_oracledb_async'; EXCEPTION WHEN OTHERS THEN IF SQLCODE != -942 THEN RAISE; END IF; END;" ) @@ -232,6 +235,40 @@ async def test_oracle_async_json_in_clob_detection(oracle_async_session: OracleA "SELECT id, metadata FROM test_json_clob_oracledb_async WHERE id = :1", (1,) ) + hydrated = result.get_first(schema_type=ClobTextDocumentRecord) + assert hydrated is not None + assert isinstance(hydrated, ClobTextDocumentRecord) + assert hydrated.id == 1 + assert isinstance(hydrated.metadata, str) + assert hydrated.metadata == json_text + + await oracle_async_session.execute_script( + "BEGIN EXECUTE IMMEDIATE 'DROP TABLE test_json_clob_oracledb_async'; EXCEPTION WHEN OTHERS THEN NULL; END;" + ) + + +async def test_oracle_async_json_in_is_json_clob_detection(oracle_async_session: OracleAsyncDriver) -> None: + """Test IS JSON CLOB hydration keeps Ch2 JSON parsing.""" + await oracle_async_session.execute_script( + "BEGIN EXECUTE IMMEDIATE 'DROP TABLE test_json_is_json_clob_oracledb_async'; EXCEPTION WHEN OTHERS THEN IF SQLCODE != -942 THEN RAISE; END IF; END;" + ) + + await oracle_async_session.execute_script(""" + CREATE TABLE test_json_is_json_clob_oracledb_async ( + id NUMBER PRIMARY KEY, + metadata CLOB CHECK (metadata IS JSON) + ) + """) + + json_text = json.dumps(LARGE_JSON_CONTENT) + await oracle_async_session.execute( + "INSERT INTO test_json_is_json_clob_oracledb_async (id, metadata) VALUES (:1, :2)", (1, json_text) + ) + + result = await oracle_async_session.execute( + "SELECT id, metadata FROM test_json_is_json_clob_oracledb_async WHERE id = :1", (1,) + ) + hydrated = result.get_first(schema_type=JsonDocumentRecord) assert hydrated is not None assert isinstance(hydrated, JsonDocumentRecord) @@ -242,16 +279,12 @@ async def test_oracle_async_json_in_clob_detection(oracle_async_session: OracleA assert len(hydrated.metadata["nested"]["data"]) == 5000 await oracle_async_session.execute_script( - "BEGIN EXECUTE IMMEDIATE 'DROP TABLE test_json_clob_oracledb_async'; EXCEPTION WHEN OTHERS THEN NULL; END;" + "BEGIN EXECUTE IMMEDIATE 'DROP TABLE test_json_is_json_clob_oracledb_async'; EXCEPTION WHEN OTHERS THEN NULL; END;" ) -def test_oracle_sync_json_in_clob_detection(oracle_sync_session: OracleSyncDriver) -> None: - """Test JSON detection in CLOB with msgspec hydration (sync). - - Verifies that JSON stored in CLOB is automatically detected and parsed, - then successfully hydrated into msgspec struct. - """ +def test_oracle_sync_json_in_clob_returns_str(oracle_sync_session: OracleSyncDriver) -> None: + """Test unconstrained CLOB JSON hydrates as text (sync).""" oracle_sync_session.execute_script( "BEGIN EXECUTE IMMEDIATE 'DROP TABLE test_json_clob_sync'; EXCEPTION WHEN OTHERS THEN IF SQLCODE != -942 THEN RAISE; END IF; END;" ) @@ -268,6 +301,38 @@ def test_oracle_sync_json_in_clob_detection(oracle_sync_session: OracleSyncDrive result = oracle_sync_session.execute("SELECT id, metadata FROM test_json_clob_sync WHERE id = :1", (1,)) + hydrated = result.get_first(schema_type=ClobTextDocumentRecord) + assert hydrated is not None + assert isinstance(hydrated, ClobTextDocumentRecord) + assert hydrated.id == 1 + assert isinstance(hydrated.metadata, str) + assert hydrated.metadata == json_text + + oracle_sync_session.execute_script( + "BEGIN EXECUTE IMMEDIATE 'DROP TABLE test_json_clob_sync'; EXCEPTION WHEN OTHERS THEN NULL; END;" + ) + + +def test_oracle_sync_json_in_is_json_clob_detection(oracle_sync_session: OracleSyncDriver) -> None: + """Test IS JSON CLOB hydration keeps Ch2 JSON parsing (sync).""" + oracle_sync_session.execute_script( + "BEGIN EXECUTE IMMEDIATE 'DROP TABLE test_json_is_json_clob_sync'; EXCEPTION WHEN OTHERS THEN IF SQLCODE != -942 THEN RAISE; END IF; END;" + ) + + oracle_sync_session.execute_script(""" + CREATE TABLE test_json_is_json_clob_sync ( + id NUMBER PRIMARY KEY, + metadata CLOB CHECK (metadata IS JSON) + ) + """) + + json_text = json.dumps(LARGE_JSON_CONTENT) + oracle_sync_session.execute( + "INSERT INTO test_json_is_json_clob_sync (id, metadata) VALUES (:1, :2)", (1, json_text) + ) + + result = oracle_sync_session.execute("SELECT id, metadata FROM test_json_is_json_clob_sync WHERE id = :1", (1,)) + hydrated = result.get_first(schema_type=JsonDocumentRecord) assert hydrated is not None assert isinstance(hydrated, JsonDocumentRecord) @@ -278,7 +343,7 @@ def test_oracle_sync_json_in_clob_detection(oracle_sync_session: OracleSyncDrive assert len(hydrated.metadata["nested"]["data"]) == 5000 oracle_sync_session.execute_script( - "BEGIN EXECUTE IMMEDIATE 'DROP TABLE test_json_clob_sync'; EXCEPTION WHEN OTHERS THEN NULL; END;" + "BEGIN EXECUTE IMMEDIATE 'DROP TABLE test_json_is_json_clob_sync'; EXCEPTION WHEN OTHERS THEN NULL; END;" ) diff --git a/tests/integration/adapters/sqlite/test_driver_features.py b/tests/integration/adapters/sqlite/test_driver_features.py index 4972bddc2..4c20493c2 100644 --- a/tests/integration/adapters/sqlite/test_driver_features.py +++ b/tests/integration/adapters/sqlite/test_driver_features.py @@ -8,10 +8,10 @@ @pytest.mark.sqlite -def test_driver_features_enabled_by_default() -> None: - """Test that driver features are enabled by default for stdlib types.""" +def test_custom_adapters_disabled_by_default() -> None: + """Test that process-global custom adapters are opt-in.""" config = SqliteConfig(connection_config={"database": ":memory:"}) - assert config.driver_features.get("enable_custom_adapters") is True + assert config.driver_features.get("enable_custom_adapters") is False @pytest.mark.sqlite diff --git a/tests/unit/adapters/test_adbc/test_arrow_streaming.py b/tests/unit/adapters/test_adbc/test_arrow_streaming.py index c4ce41afe..31cf31343 100644 --- a/tests/unit/adapters/test_adbc/test_arrow_streaming.py +++ b/tests/unit/adapters/test_adbc/test_arrow_streaming.py @@ -70,6 +70,18 @@ def test_select_to_arrow_reader_uses_fetch_record_batch_and_defers_cursor_close( assert connection.cursor_obj.closed is True +def test_select_stream_uses_record_batch_reader_without_eager_table_materialization() -> None: + connection = _AdbcStreamingConnection() + driver = AdbcDriver(cast("AdbcConnection", connection), dialect="sqlite") + + with driver.select_stream("SELECT 1 AS x", native_only=True, chunk_size=2) as stream: + rows = list(stream) + + assert rows == [{"x": 1}, {"x": 2}, {"x": 3}] + assert connection.cursor_obj.executed == [("SELECT 1 AS x", [])] + assert connection.cursor_obj.closed is True + + def test_load_from_arrow_passes_record_batch_reader_to_adbc_ingest() -> None: connection = _AdbcStreamingConnection() driver = AdbcDriver( diff --git a/tests/unit/adapters/test_adbc/test_driver.py b/tests/unit/adapters/test_adbc/test_driver.py index 76c0e7e7f..094996d7e 100644 --- a/tests/unit/adapters/test_adbc/test_driver.py +++ b/tests/unit/adapters/test_adbc/test_driver.py @@ -1,7 +1,9 @@ """Unit tests for ADBC driver execute path optimizations.""" +from types import SimpleNamespace from unittest.mock import MagicMock +import pyarrow as pa import pytest from sqlspec.adapters.adbc.driver import AdbcDriver, AdbcExceptionHandler @@ -58,12 +60,35 @@ def _unexpected_prepare(*_: object, **__: object) -> object: cursor = MagicMock() cursor.rowcount = 1 - result = driver.dispatch_execute(cursor, statement) # type: ignore[protected-access] + result = driver.dispatch_execute(cursor, statement) # type: ignore[arg-type, protected-access] cursor.execute.assert_called_once_with("UPDATE users SET name = $2 WHERE id = $1", parameters=prepared_parameters) assert result.rowcount_override == 1 +def test_dispatch_execute_uses_fetch_arrow_table_and_dict_rows(monkeypatch: pytest.MonkeyPatch) -> None: + connection = MagicMock() + connection.adbc_get_info.return_value = {"vendor_name": "sqlite", "driver_name": "sqlite"} + driver = AdbcDriver(connection) + + table = pa.table({"id": [1]}) + cursor = SimpleNamespace( + description=[("id",)], + execute=lambda *_args, **_kwargs: None, + fetch_arrow_table=lambda: table, + fetchall=lambda: (_ for _ in ()).throw(AssertionError("fetchall should not be used")), + ) + monkeypatch.setattr(AdbcDriver, "_compiled_sql", lambda self, statement, config: ("SELECT 1", None)) + + statement = driver.prepare_statement("SELECT 1", statement_config=driver.statement_config) + result = driver.dispatch_execute(cursor, statement) # type: ignore[arg-type, protected-access] + + assert result.selected_data == [{"id": 1}] + assert result.column_names == ["id"] + assert result.data_row_count == 1 + assert result.row_format == "dict" + + def test_prepare_driver_parameters_postgres_applies_cast_aware_preparation_when_casts_present( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/unit/adapters/test_aiomysql/test_config.py b/tests/unit/adapters/test_aiomysql/test_config.py index a6b2acf9b..68d089049 100644 --- a/tests/unit/adapters/test_aiomysql/test_config.py +++ b/tests/unit/adapters/test_aiomysql/test_config.py @@ -62,7 +62,7 @@ def test_aiomysql_connection_kwargs_normalize_cursor_alias_and_omit_pool_only_ke "minsize": 2, "maxsize": 8, "pool_recycle": 30, - "enable_local_infile": True, + "local_infile": True, "allow_local_infile": True, } ) @@ -168,8 +168,8 @@ async def fake_create_pool(**kwargs: object) -> object: config = AiomysqlConfig( connection_config={ "cursor_class": AiomysqlDictCursor, - "enable_local_infile": True, "allow_local_infile": True, + "local_infile": True, "minsize": 2, "maxsize": 8, "pool_recycle": 30, diff --git a/tests/unit/adapters/test_aiomysql/test_local_infile_bulk_load.py b/tests/unit/adapters/test_aiomysql/test_local_infile_bulk_load.py index e626c8f01..1195c8d0a 100644 --- a/tests/unit/adapters/test_aiomysql/test_local_infile_bulk_load.py +++ b/tests/unit/adapters/test_aiomysql/test_local_infile_bulk_load.py @@ -31,9 +31,9 @@ def test_config_gate_raises_when_local_infile_disabled() -> None: AiomysqlConfig(driver_features={"enable_local_infile_bulk_load": True}) -def test_config_gate_allows_when_enable_local_infile_set() -> None: +def test_config_gate_allows_when_local_infile_set() -> None: config = AiomysqlConfig( - connection_config={"enable_local_infile": True, "allow_local_infile": True}, + connection_config={"local_infile": True, "allow_local_infile": True}, driver_features={"enable_local_infile_bulk_load": True}, ) assert config.driver_features["enable_local_infile_bulk_load"] is True diff --git a/tests/unit/adapters/test_aiosqlite/test_config.py b/tests/unit/adapters/test_aiosqlite/test_config.py index bb35555c1..b0d72493a 100644 --- a/tests/unit/adapters/test_aiosqlite/test_config.py +++ b/tests/unit/adapters/test_aiosqlite/test_config.py @@ -216,4 +216,4 @@ def test_runtime_keys_removed_from_driver_features() -> None: assert "row_factory" not in config.driver_features assert "pragmas" not in config.driver_features assert "extensions" not in config.driver_features - assert config.driver_features["enable_custom_adapters"] is True + assert config.driver_features["enable_custom_adapters"] is False diff --git a/tests/unit/adapters/test_aiosqlite/test_driver.py b/tests/unit/adapters/test_aiosqlite/test_driver.py index 426786b1b..deaad0c36 100644 --- a/tests/unit/adapters/test_aiosqlite/test_driver.py +++ b/tests/unit/adapters/test_aiosqlite/test_driver.py @@ -1,9 +1,9 @@ import asyncio import sqlite3 import time -from collections.abc import Iterator +from collections.abc import Callable, Iterator from threading import Thread -from typing import Any +from typing import Any, cast import aiosqlite import pytest @@ -15,8 +15,10 @@ default_statement_config, driver_profile, ) +from sqlspec.adapters.aiosqlite.driver import AiosqliteDriver from sqlspec.adapters.aiosqlite.pool import AiosqliteConnectionPool -from sqlspec.core import ParameterStyle +from sqlspec.core import SQL, ParameterStyle +from sqlspec.core.result import DMLResult async def test_cursor_lifecycle_cursor_closed_after_normal_exit() -> None: @@ -70,13 +72,89 @@ async def test_cursor_lifecycle_aexit_does_not_suppress_exception() -> None: await conn.close() +class _MappingRow: + def __init__(self, data: dict[str, object]) -> None: + self._data = data + + def keys(self) -> object: + return self._data.keys() + + def __getitem__(self, key: str) -> object: + return self._data[key] + + +class _SelectCursor: + def __init__(self) -> None: + self.description = [("id",), ("name",)] + self.rowcount = 2 + self.closed = False + + def fetchall(self) -> list[_MappingRow]: + return [_MappingRow({"id": 1, "name": "alice"})] + + def close(self) -> None: + self.closed = True + + +class _WorkerConnection: + def __init__(self) -> None: + self._conn = self + self.cursor = _SelectCursor() + self.executemany_calls: list[tuple[str, object]] = [] + + def execute(self, sql: str, parameters: object) -> _SelectCursor: + _ = sql, parameters + return self.cursor + + async def executemany(self, sql: str, parameters: object) -> _SelectCursor: + self.executemany_calls.append((sql, parameters)) + return self.cursor + + +async def test_dispatch_execute_detects_record_row_format(monkeypatch: pytest.MonkeyPatch) -> None: + async def _run_on_worker_thread( + _connection: object, function: Callable[..., object], *args: object, **kwargs: object + ) -> object: + return function(*args, **kwargs) + + monkeypatch.setattr("sqlspec.adapters.aiosqlite.driver.run_on_worker_thread", _run_on_worker_thread) + monkeypatch.setattr(AiosqliteDriver, "_compiled_sql", lambda *_args, **_kwargs: ("SELECT id, name FROM users", [])) + + driver = AiosqliteDriver(connection=cast("Any", _WorkerConnection()), statement_config=default_statement_config) + + result = await driver.dispatch_execute( + cast("Any", object()), SQL("SELECT id, name FROM users", statement_config=default_statement_config) + ) + + assert result.row_format == "record" + selected_data = result.selected_data + assert selected_data is not None + assert dict(selected_data[0]) == {"id": 1, "name": "alice"} + + +async def test_execute_many_uses_thin_qmark_path() -> None: + connection = _WorkerConnection() + driver = AiosqliteDriver(connection=cast("Any", connection), statement_config=default_statement_config) + + result = await driver.execute_many("INSERT INTO test_thin_path (value) VALUES (?)", [("a",), ("b",)]) + + assert isinstance(result, DMLResult) + assert result.operation_type == "INSERT" + assert result.rows_affected == 2 + assert connection.executemany_calls == [("INSERT INTO test_thin_path (value) VALUES (?)", [("a",), ("b",)])] + + class _FakeConnection: def __init__(self) -> None: self.executed: list[str] = [] + self.scripted: list[str] = [] async def execute(self, sql: str) -> None: self.executed.append(sql) + async def executescript(self, script: str) -> None: + self.scripted.append(script) + async def commit(self) -> None: return None diff --git a/tests/unit/adapters/test_aiosqlite/test_type_converter.py b/tests/unit/adapters/test_aiosqlite/test_type_converter.py index 2018d5d76..592d119c0 100644 --- a/tests/unit/adapters/test_aiosqlite/test_type_converter.py +++ b/tests/unit/adapters/test_aiosqlite/test_type_converter.py @@ -3,6 +3,8 @@ import json from unittest.mock import MagicMock, patch +import pytest + def test_import_is_from_aiosqlite_not_sqlite() -> None: from sqlspec.adapters.aiosqlite.type_converter import register_type_handlers @@ -91,6 +93,18 @@ def test_unregister_type_handlers_is_noop() -> None: unregister_type_handlers() +@pytest.mark.anyio +async def test_aiosqlite_config_does_not_register_type_handlers_by_default() -> None: + import sqlspec.adapters.aiosqlite.config as config_mod + + with patch.object(config_mod, "register_type_handlers") as mock_register: + config = config_mod.AiosqliteConfig(connection_config={"database": ":memory:"}) + pool = await config._create_pool() # pyright: ignore[reportPrivateUsage] + await pool.close() + + mock_register.assert_not_called() + + def test_aiosqlite_config_uses_aiosqlite_type_converter() -> None: import sqlspec.adapters.aiosqlite.config as config_mod import sqlspec.adapters.sqlite.type_converter as sqlite_type_converter diff --git a/tests/unit/adapters/test_arrow_odbc/test_driver.py b/tests/unit/adapters/test_arrow_odbc/test_driver.py index 102fa54bd..7cd8e77dd 100644 --- a/tests/unit/adapters/test_arrow_odbc/test_driver.py +++ b/tests/unit/adapters/test_arrow_odbc/test_driver.py @@ -69,6 +69,23 @@ def close(self) -> None: self.closed = True +def test_arrow_odbc_config_applies_driver_feature_json_serializer_to_statement_config() -> None: + """Custom JSON serializers should reach the statement parameter config.""" + + def serializer(value: object) -> str: + return f"json:{value!r}" + + config = ArrowOdbcConfig( + connection_config={"connection_string": "Driver={ODBC Driver 18 for SQL Server};"}, + driver_features={"json_serializer": serializer}, + ) + + parameter_config = config.statement_config.parameter_config + assert parameter_config.json_serializer is serializer + assert parameter_config.type_coercion_map[dict] is serializer + assert parameter_config.type_coercion_map[list] is serializer + + class NoCloseConnection: """Connection stub matching arrow-odbc 10.4's no-close public surface.""" @@ -242,6 +259,25 @@ def fail_reader_to_table(_reader: object) -> pa.Table: assert result.get_data().read_all().to_pydict() == {"x": [1, 2, 3]} +def test_arrow_odbc_select_stream_uses_native_batches(monkeypatch: pytest.MonkeyPatch) -> None: + """select_stream should stream Arrow batches instead of materializing the whole table.""" + connection = FakeConnection() + driver = ArrowOdbcDriver(cast("ArrowOdbcConnection", connection), driver_features={"chunk_size": 2}) + + def fail_reader_to_table(_reader: object) -> pa.Table: + msg = "select_stream should not materialize via _reader_to_table" + raise AssertionError(msg) + + monkeypatch.setattr("sqlspec.adapters.arrow_odbc.driver._reader_to_table", fail_reader_to_table) + + with driver.select_stream("SELECT 1 AS x", native_only=True, chunk_size=2) as stream: + rows = list(stream) + + assert rows == [{"x": 1}, {"x": 2}, {"x": 3}] + assert connection.read_calls[0]["query"] == "SELECT 1 AS x" + assert connection.read_calls[0]["batch_size"] == 2 + + def test_arrow_odbc_select_to_arrow_raises_mapped_driver_exception(monkeypatch: pytest.MonkeyPatch) -> None: """Native read failures should use SQLSpec's deferred exception mapping.""" monkeypatch.setattr("sqlspec.adapters.arrow_odbc.driver.ArrowOdbcError", FakeOdbcError) diff --git a/tests/unit/adapters/test_async_adapters.py b/tests/unit/adapters/test_async_adapters.py index 07b1032f9..b672475fa 100644 --- a/tests/unit/adapters/test_async_adapters.py +++ b/tests/unit/adapters/test_async_adapters.py @@ -332,7 +332,7 @@ async def test_async_driver_execute_script_method(aiosqlite_async_driver: Aiosql ("method_name", "call_args"), [ pytest.param("execute", ("SELECT * FROM users WHERE id = ?", 1), id="execute"), - pytest.param("execute_many", ("INSERT INTO users (name) VALUES (?)", [["alice"]]), id="execute_many"), + pytest.param("execute_many", (SQL("INSERT INTO users (name) VALUES (?)"), [["alice"]]), id="execute_many"), pytest.param("execute_script", ("INSERT INTO users (name) VALUES ('alice');",), id="execute_script"), ], ) diff --git a/tests/unit/adapters/test_bigquery/test_config.py b/tests/unit/adapters/test_bigquery/test_config.py index 74b57145d..769c84125 100644 --- a/tests/unit/adapters/test_bigquery/test_config.py +++ b/tests/unit/adapters/test_bigquery/test_config.py @@ -1,12 +1,13 @@ """BigQuery configuration tests covering statement config builders.""" +import uuid from typing import Any, cast from google.cloud.bigquery import LoadJobConfig, QueryJobConfig from pytest import MonkeyPatch from sqlspec.adapters.bigquery.config import BigQueryConfig, BigQueryConnectionParams, BigQueryDriverFeatures -from sqlspec.adapters.bigquery.core import build_statement_config +from sqlspec.adapters.bigquery.core import apply_driver_features, build_statement_config class _RecordingBigQueryClient: @@ -40,6 +41,19 @@ def serializer(_: object) -> str: assert config.driver_features["json_serializer"] is serializer +def test_bigquery_driver_features_honor_uuid_conversion_flag() -> None: + """The BigQuery UUID conversion flag should control UUID parameter coercion.""" + enabled_config, enabled_features = apply_driver_features(build_statement_config(), {"enable_uuid_conversion": True}) + disabled_config, disabled_features = apply_driver_features( + build_statement_config(), {"enable_uuid_conversion": False} + ) + + assert enabled_features["enable_uuid_conversion"] is True + assert disabled_features["enable_uuid_conversion"] is False + assert uuid.UUID in enabled_config.parameter_config.type_coercion_map + assert uuid.UUID not in disabled_config.parameter_config.type_coercion_map + + def test_bigquery_config_wires_query_timeout_ms_to_default_job_config() -> None: """query_timeout_ms on connection_config reaches QueryJobConfig.job_timeout_ms (existing behaviour).""" config = BigQueryConfig(connection_config={"project": "p", "dataset_id": "d", "query_timeout_ms": 12345}) diff --git a/tests/unit/adapters/test_bigquery/test_core.py b/tests/unit/adapters/test_bigquery/test_core.py index b1cd4e09e..522dcffc0 100644 --- a/tests/unit/adapters/test_bigquery/test_core.py +++ b/tests/unit/adapters/test_bigquery/test_core.py @@ -1,11 +1,12 @@ """Unit tests for BigQuery core performance helpers.""" -from collections.abc import Iterable +import importlib +from collections.abc import ItemsView, Iterable, Iterator, Mapping from contextlib import nullcontext from types import SimpleNamespace from typing import Any, cast -from uuid import UUID +import pytest from google.cloud.bigquery import QueryJobConfig from sqlspec.adapters.bigquery.core import ( @@ -21,7 +22,6 @@ run_query_job, ) from sqlspec.adapters.bigquery.driver import BigQueryDriver -from sqlspec.adapters.bigquery.type_converter import BigQueryOutputConverter from sqlspec.core import ParameterStyle from sqlspec.utils.serializers import to_json @@ -73,20 +73,35 @@ def result(self, **kwargs: Any) -> None: self.result_calls.append(kwargs) -class _RecordingRow: +class _RecordingRow(Mapping[str, object]): def __init__(self, values: dict[str, object]) -> None: self._values = values - def items(self) -> Iterable[tuple[str, object]]: + def __iter__(self) -> Iterator[str]: + return iter(self._values) + + def __len__(self) -> int: + return len(self._values) + + def __getitem__(self, key: str) -> object: + return self._values[key] + + def items(self) -> ItemsView[str, object]: return self._values.items() +class _MappingRow(_RecordingRow): + def items(self) -> ItemsView[str, object]: + msg = "fetch_chunk should use dict(row), not row.items()" + raise AssertionError(msg) + + def _schema_field(name: str) -> SimpleNamespace: return SimpleNamespace(name=name) class _RecordingStreamJob: - def __init__(self, pages: Iterable[Iterable[_RecordingRow]] | None = None) -> None: + def __init__(self, pages: Iterable[Iterable[Mapping[str, object]]] | None = None) -> None: self.result_kwargs: dict[str, Any] = {} self._pages = iter(pages or ()) @@ -96,7 +111,7 @@ def result(self, **kwargs: Any) -> SimpleNamespace: class _RecordingStreamDriver: - def __init__(self, api_base_url: str, pages: Iterable[Iterable[_RecordingRow]] | None = None) -> None: + def __init__(self, api_base_url: str, pages: Iterable[Iterable[Mapping[str, object]]] | None = None) -> None: self.connection = SimpleNamespace(_connection=SimpleNamespace(API_BASE_URL=api_base_url)) self._job_retry: Any = build_retry(2.0) self._job_retry_deadline = 2.0 @@ -376,24 +391,14 @@ def test_stream_source_skips_empty_pages_before_rows() -> None: assert source.fetch_chunk() == [] -def test_type_converter_bigquery_output_converter_is_final() -> None: - assert getattr(BigQueryOutputConverter, "__final__", False) is True - - -def test_type_converter_bigquery_output_converter_instantiates() -> None: - converter = BigQueryOutputConverter() - assert isinstance(converter, BigQueryOutputConverter) - +def test_stream_source_converts_rows_without_items_copying() -> None: + driver = _RecordingStreamDriver("https://bigquery.googleapis.com", pages=((_MappingRow({"value": 1}),),)) + source = BigQueryStreamSource(cast(Any, driver), "SELECT 1", None, 100) + source.start() -def test_type_converter_bigquery_output_converter_converts_uuid_string() -> None: - converter = BigQueryOutputConverter(enable_uuid_conversion=True) - uuid_str = "550e8400-e29b-41d4-a716-446655440000" - result = converter.convert(uuid_str) - assert isinstance(result, UUID) - assert str(result) == uuid_str + assert source.fetch_chunk() == [{"value": 1}] -def test_type_converter_bigquery_output_converter_convert_detected_fallback() -> None: - converter = BigQueryOutputConverter() - result = converter._convert_detected("hello", "unknown_type") - assert result == "hello" +def test_bigquery_type_converter_module_is_gone() -> None: + with pytest.raises(ModuleNotFoundError): + importlib.import_module("sqlspec.adapters.bigquery.type_converter") diff --git a/tests/unit/adapters/test_cockroach_asyncpg/test_core.py b/tests/unit/adapters/test_cockroach_asyncpg/test_core.py index 7c270c160..15446df73 100644 --- a/tests/unit/adapters/test_cockroach_asyncpg/test_core.py +++ b/tests/unit/adapters/test_cockroach_asyncpg/test_core.py @@ -194,6 +194,48 @@ async def _dispatch_execute_script_impl(self, cursor: object, statement: SQL) -> return _execution_result("script") +class _RetryableCockroachError(Exception): + sqlstate = "40001" + + +class _RetryConnection: + def __init__(self) -> None: + self.commands: list[str] = [] + self.in_transaction = False + + async def execute(self, sql: str) -> None: + self.commands.append(sql) + if sql == "BEGIN": + self.in_transaction = True + elif sql in {"COMMIT", "ROLLBACK"}: + self.in_transaction = False + + def is_in_transaction(self) -> bool: + return self.in_transaction + + +class _RetryingCockroachAsyncpgDriver(CockroachAsyncpgDriver): + def __init__(self, connection: _RetryConnection) -> None: + super().__init__( + connection=cast("CockroachAsyncpgConnection", connection), + driver_features={ + "enable_auto_retry": True, + "enable_retry_logging": False, + "max_retries": 1, + "retry_delay_base_ms": 0.0, + "retry_delay_max_ms": 0.0, + }, + ) + self.calls = 0 + + async def _dispatch_execute_impl(self, cursor: object, statement: SQL) -> ExecutionResult: + _ = cursor, statement + self.calls += 1 + if self.calls == 1: + raise _RetryableCockroachError("restart transaction") + return _execution_result("retried") + + @pytest.mark.anyio async def test_driver_cockroach_asyncpg_non_retry_execute_many_uses_impl_wrapper() -> None: driver = _RecordingCockroachAsyncpgDriver() @@ -208,3 +250,44 @@ async def test_driver_cockroach_asyncpg_non_retry_execute_script_uses_impl_wrapp result = await driver.dispatch_execute_script(object(), SQL("SELECT 1")) assert result.cursor_result == "script" assert driver.calls == ["script"] + + +@pytest.mark.anyio +async def test_driver_cockroach_asyncpg_dispatch_execute_does_not_retry_inside_statement() -> None: + connection = _RetryConnection() + driver = _RetryingCockroachAsyncpgDriver(connection) + + with pytest.raises(_RetryableCockroachError): + await driver.dispatch_execute(object(), SQL("SELECT 1")) + + assert driver.calls == 1 + assert connection.commands == [] + + +@pytest.mark.anyio +async def test_driver_cockroach_asyncpg_transaction_retry_replays_whole_callback() -> None: + connection = _RetryConnection() + driver = CockroachAsyncpgDriver( + connection=cast("CockroachAsyncpgConnection", connection), + driver_features={ + "enable_auto_retry": True, + "enable_retry_logging": False, + "max_retries": 1, + "retry_delay_base_ms": 0.0, + "retry_delay_max_ms": 0.0, + }, + ) + calls = 0 + + async def operation() -> str: + nonlocal calls + calls += 1 + if calls == 1: + raise _RetryableCockroachError("restart transaction") + return "ok" + + result = await driver.run_transaction_with_retry(operation) + + assert result == "ok" + assert calls == 2 + assert connection.commands == ["BEGIN", "ROLLBACK", "BEGIN", "COMMIT"] diff --git a/tests/unit/adapters/test_cockroach_psycopg/test_core.py b/tests/unit/adapters/test_cockroach_psycopg/test_core.py index 622826450..3eef222cc 100644 --- a/tests/unit/adapters/test_cockroach_psycopg/test_core.py +++ b/tests/unit/adapters/test_cockroach_psycopg/test_core.py @@ -218,6 +218,91 @@ async def _dispatch_execute_script_impl(self, cursor: object, statement: SQL) -> return _execution_result("script") +class _RetryableCockroachError(Exception): + sqlstate = "40001" + + +class _SyncRetryConnection: + def __init__(self) -> None: + self.autocommit = True + self.commits = 0 + self.rollbacks = 0 + + def commit(self) -> None: + self.commits += 1 + + def rollback(self) -> None: + self.rollbacks += 1 + + +class _AsyncRetryConnection: + def __init__(self) -> None: + self.autocommit = True + self.commits = 0 + self.rollbacks = 0 + self.autocommit_values: list[bool] = [] + + async def set_autocommit(self, value: bool) -> None: + self.autocommit = value + self.autocommit_values.append(value) + + async def commit(self) -> None: + self.commits += 1 + + async def rollback(self) -> None: + self.rollbacks += 1 + + +class _RetryingCockroachPsycopgSyncDriver(CockroachPsycopgSyncDriver): + def __init__(self, connection: _SyncRetryConnection) -> None: + super().__init__( + connection=cast("CockroachSyncConnection", connection), + driver_features={ + "enable_auto_retry": True, + "enable_retry_logging": False, + "max_retries": 1, + "retry_delay_base_ms": 0.0, + "retry_delay_max_ms": 0.0, + }, + ) + self.calls = 0 + + def _connection_in_transaction(self) -> bool: + return False + + def _dispatch_execute_impl(self, cursor: object, statement: SQL) -> ExecutionResult: + _ = cursor, statement + self.calls += 1 + if self.calls == 1: + raise _RetryableCockroachError("restart transaction") + return _execution_result("retried") + + +class _RetryingCockroachPsycopgAsyncDriver(CockroachPsycopgAsyncDriver): + def __init__(self, connection: _AsyncRetryConnection) -> None: + super().__init__( + connection=cast("CockroachAsyncConnection", connection), + driver_features={ + "enable_auto_retry": True, + "enable_retry_logging": False, + "max_retries": 1, + "retry_delay_base_ms": 0.0, + "retry_delay_max_ms": 0.0, + }, + ) + self.calls = 0 + + def _connection_in_transaction(self) -> bool: + return False + + async def _dispatch_execute_impl(self, cursor: object, statement: SQL) -> ExecutionResult: + _ = cursor, statement + self.calls += 1 + if self.calls == 1: + raise _RetryableCockroachError("restart transaction") + return _execution_result("retried") + + def test_driver_cockroach_psycopg_sync_non_retry_execute_many_uses_impl_wrapper() -> None: driver = _RecordingCockroachPsycopgSyncDriver() result = driver.dispatch_execute_many(object(), SQL("SELECT 1")) @@ -248,6 +333,70 @@ async def test_driver_cockroach_psycopg_async_non_retry_execute_script_uses_impl assert driver.calls == ["script"] +def test_driver_cockroach_psycopg_sync_dispatch_execute_does_not_retry_inside_statement() -> None: + connection = _SyncRetryConnection() + driver = _RetryingCockroachPsycopgSyncDriver(connection) + + with pytest.raises(_RetryableCockroachError): + driver.dispatch_execute(object(), SQL("SELECT 1")) + + assert driver.calls == 1 + assert connection.rollbacks == 0 + + +def test_driver_cockroach_psycopg_sync_transaction_retry_replays_whole_callback() -> None: + connection = _SyncRetryConnection() + driver = _RetryingCockroachPsycopgSyncDriver(connection) + calls = 0 + + def operation() -> str: + nonlocal calls + calls += 1 + if calls == 1: + raise _RetryableCockroachError("restart transaction") + return "ok" + + result = driver.run_transaction_with_retry(operation) + + assert result == "ok" + assert calls == 2 + assert connection.rollbacks == 1 + assert connection.commits == 1 + + +@pytest.mark.anyio +async def test_driver_cockroach_psycopg_async_dispatch_execute_does_not_retry_inside_statement() -> None: + connection = _AsyncRetryConnection() + driver = _RetryingCockroachPsycopgAsyncDriver(connection) + + with pytest.raises(_RetryableCockroachError): + await driver.dispatch_execute(object(), SQL("SELECT 1")) + + assert driver.calls == 1 + assert connection.rollbacks == 0 + + +@pytest.mark.anyio +async def test_driver_cockroach_psycopg_async_transaction_retry_replays_whole_callback() -> None: + connection = _AsyncRetryConnection() + driver = _RetryingCockroachPsycopgAsyncDriver(connection) + calls = 0 + + async def operation() -> str: + nonlocal calls + calls += 1 + if calls == 1: + raise _RetryableCockroachError("restart transaction") + return "ok" + + result = await driver.run_transaction_with_retry(operation) + + assert result == "ok" + assert calls == 2 + assert connection.rollbacks == 1 + assert connection.commits == 1 + + def test_driver_cockroach_psycopg_sync_data_dictionary_uses_parent_slot() -> None: driver = CockroachPsycopgSyncDriver(connection=_sync_connection()) assert isinstance(driver.data_dictionary, CockroachPsycopgSyncDataDictionary) diff --git a/tests/unit/adapters/test_contract_arrow_odbc_harness.py b/tests/unit/adapters/test_contract_arrow_odbc_harness.py index edf32ee83..bede08ce5 100644 --- a/tests/unit/adapters/test_contract_arrow_odbc_harness.py +++ b/tests/unit/adapters/test_contract_arrow_odbc_harness.py @@ -56,6 +56,7 @@ def test_arrow_odbc_sync_case_is_active_sql_server_bulk_only() -> None: assert not case.supports_execute_many assert not case.supports_load_from_records assert case.execute_rowcount_policy == "unavailable" + assert not case.supports_stream_reopen_after_partial_iteration def test_seed_sync_uses_arrow_bulk_ingest_when_execute_many_is_disabled() -> None: diff --git a/tests/unit/adapters/test_duckdb/test_arrow_streaming.py b/tests/unit/adapters/test_duckdb/test_arrow_streaming.py index e3e1b8f04..8e19f6bdf 100644 --- a/tests/unit/adapters/test_duckdb/test_arrow_streaming.py +++ b/tests/unit/adapters/test_duckdb/test_arrow_streaming.py @@ -2,12 +2,14 @@ from collections.abc import Iterator from contextlib import contextmanager +from uuid import UUID import pyarrow as pa import pytest from sqlspec.adapters.duckdb import DuckDBConfig from sqlspec.adapters.duckdb.driver import DuckDBDriver +from sqlspec.core.result import DMLResult from sqlspec.exceptions import MissingDependencyError @@ -55,7 +57,39 @@ def test_load_from_arrow_registers_record_batch_reader() -> None: assert job.telemetry["rows_processed"] == 2 result = driver.execute("SELECT id, name FROM reader_target ORDER BY id") - assert result.data == [(10, "ten"), (11, "eleven")] + assert result.get_data() == [{"id": 10, "name": "ten"}, {"id": 11, "name": "eleven"}] + + +def test_execute_many_uses_bulk_insert_fast_path() -> None: + with _seed_driver() as driver: + driver.execute("CREATE OR REPLACE TABLE bulk_target (id INTEGER, name VARCHAR)") + + result = driver.execute_many( + "INSERT INTO bulk_target (id, name) VALUES (?, ?)", [(1, "one"), (2, "two"), (3, "three")] + ) + + assert isinstance(result, DMLResult) + assert result.operation_type == "INSERT" + assert result.rows_affected == 3 + assert driver.execute("SELECT id, name FROM bulk_target ORDER BY id").get_data() == [ + {"id": 1, "name": "one"}, + {"id": 2, "name": "two"}, + {"id": 3, "name": "three"}, + ] + + +def test_select_arrow_path_restores_uuid_columns_only() -> None: + uuid_value = "550e8400-e29b-41d4-a716-446655440000" + with _seed_driver() as driver: + driver.execute("CREATE OR REPLACE TABLE uuid_target (id UUID, text_id VARCHAR)") + driver.execute("INSERT INTO uuid_target VALUES (?, ?)", uuid_value, uuid_value) + + row = driver.select_one("SELECT id, text_id FROM uuid_target") + + assert isinstance(row["id"], UUID) + assert str(row["id"]) == uuid_value + assert row["text_id"] == uuid_value + assert isinstance(row["text_id"], str) def test_missing_pyarrow_raises_cleanly(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/tests/unit/adapters/test_duckdb/test_config.py b/tests/unit/adapters/test_duckdb/test_config.py index 309c6c3fd..03556fc6a 100644 --- a/tests/unit/adapters/test_duckdb/test_config.py +++ b/tests/unit/adapters/test_duckdb/test_config.py @@ -2,6 +2,7 @@ from typing import Any, get_type_hints from unittest.mock import patch +from uuid import UUID import duckdb from typing_extensions import NotRequired @@ -44,6 +45,23 @@ def serializer(_: object) -> str: assert parameter_config.json_serializer is serializer +def test_duckdb_config_honors_uuid_conversion_flag() -> None: + """Driver features should control native UUID parameter coercion without string sniffing.""" + + uuid_value = UUID("550e8400-e29b-41d4-a716-446655440000") + + enabled_config = DuckDBConfig(driver_features={"enable_uuid_conversion": True}) + disabled_config = DuckDBConfig(driver_features={"enable_uuid_conversion": False}) + + enabled_map = enabled_config.statement_config.parameter_config.type_coercion_map + disabled_map = disabled_config.statement_config.parameter_config.type_coercion_map + + assert str not in enabled_map + assert str not in disabled_map + assert enabled_map[UUID](uuid_value) == str(uuid_value) + assert UUID not in disabled_map + + def test_connection_param_types_match_current_duckdb_settings() -> None: """Typed config keys should match DuckDB 1.5 setting types.""" diff --git a/tests/unit/adapters/test_duckdb/test_type_converter.py b/tests/unit/adapters/test_duckdb/test_type_converter.py index 3778a9e09..122984964 100644 --- a/tests/unit/adapters/test_duckdb/test_type_converter.py +++ b/tests/unit/adapters/test_duckdb/test_type_converter.py @@ -1,7 +1,6 @@ # pyright: reportArgumentType=false, reportIncompatibleMethodOverride=false -"""Unit tests for DuckDB type converter UUID conversion control.""" +"""Unit tests for DuckDB Arrow paths.""" -import uuid from contextlib import contextmanager from typing import TYPE_CHECKING, Any, cast @@ -9,7 +8,6 @@ from sqlspec.adapters.duckdb import default_statement_config from sqlspec.adapters.duckdb.driver import DuckDBDriver -from sqlspec.adapters.duckdb.type_converter import DuckDBOutputConverter if TYPE_CHECKING: from sqlspec.adapters.duckdb._typing import DuckDBConnection @@ -26,71 +24,6 @@ } -def test_uuid_conversion_enabled_by_default() -> None: - """Test that UUID conversion is enabled by default.""" - converter = DuckDBOutputConverter() - uuid_str = "550e8400-e29b-41d4-a716-446655440000" - result = converter.handle_uuid(uuid_str) - assert isinstance(result, uuid.UUID) - assert str(result) == uuid_str - - -def test_uuid_conversion_can_be_disabled() -> None: - """Test that UUID conversion can be disabled.""" - converter = DuckDBOutputConverter(enable_uuid_conversion=False) - uuid_str = "550e8400-e29b-41d4-a716-446655440000" - result = converter.handle_uuid(uuid_str) - assert isinstance(result, str) - assert result == uuid_str - - -def test_uuid_objects_pass_through_regardless_of_flag() -> None: - """Test that UUID objects pass through unchanged regardless of conversion flag.""" - converter_enabled = DuckDBOutputConverter(enable_uuid_conversion=True) - converter_disabled = DuckDBOutputConverter(enable_uuid_conversion=False) - uuid_obj = uuid.UUID("550e8400-e29b-41d4-a716-446655440000") - result_enabled = converter_enabled.handle_uuid(uuid_obj) - result_disabled = converter_disabled.handle_uuid(uuid_obj) - assert result_enabled is uuid_obj - assert result_disabled is uuid_obj - - -def test_convert_respects_uuid_flag() -> None: - """Test that convert respects UUID conversion flag.""" - converter_enabled = DuckDBOutputConverter(enable_uuid_conversion=True) - converter_disabled = DuckDBOutputConverter(enable_uuid_conversion=False) - uuid_str = "550e8400-e29b-41d4-a716-446655440000" - result_enabled = converter_enabled.convert(uuid_str) - result_disabled = converter_disabled.convert(uuid_str) - assert isinstance(result_enabled, uuid.UUID) - assert isinstance(result_disabled, str) - assert result_disabled == uuid_str - - -def test_non_uuid_strings_unaffected_by_flag() -> None: - """Test that non-UUID strings are unaffected by the conversion flag.""" - converter_enabled = DuckDBOutputConverter(enable_uuid_conversion=True) - converter_disabled = DuckDBOutputConverter(enable_uuid_conversion=False) - regular_str = "just a regular string" - result_enabled = converter_enabled.convert(regular_str) - result_disabled = converter_disabled.convert(regular_str) - assert result_enabled == regular_str - assert result_disabled == regular_str - - -def test_datetime_conversion_unaffected_by_uuid_flag() -> None: - """Test that datetime conversion works regardless of UUID flag.""" - converter_enabled = DuckDBOutputConverter(enable_uuid_conversion=True) - converter_disabled = DuckDBOutputConverter(enable_uuid_conversion=False) - datetime_str = "2024-01-15T10:30:00" - result_enabled = converter_enabled.convert(datetime_str) - result_disabled = converter_disabled.convert(datetime_str) - from datetime import datetime - - assert isinstance(result_enabled, datetime) - assert isinstance(result_disabled, datetime) - - class _ArrowCursor: def __init__(self) -> None: self.executed: list[tuple[str, object]] = [] diff --git a/tests/unit/adapters/test_mssql_python/test_arrow.py b/tests/unit/adapters/test_mssql_python/test_arrow.py index 892114e8e..bb9f8d85a 100644 --- a/tests/unit/adapters/test_mssql_python/test_arrow.py +++ b/tests/unit/adapters/test_mssql_python/test_arrow.py @@ -18,17 +18,24 @@ class ArrowCursor: """Cursor stub for Arrow and BulkCopy tests.""" - description = (("x",),) - def __init__(self) -> None: self.closed = False self.executed: list[tuple[str, object]] = [] self.bulkcopy_calls: list[tuple[str, list[object], dict[str, object]]] = [] + self.description = (("id",), ("name",)) + self.fetchmany_sizes: list[int] = [] self.rowcount = 0 + self.rows = [(1, "Ada"), (2, "Grace"), (3, "Linus")] def execute(self, sql: str, parameters: object = None) -> None: self.executed.append((sql, parameters)) + def fetchmany(self, size: int) -> list[tuple[int, str]]: + self.fetchmany_sizes.append(size) + chunk = self.rows[:size] + self.rows = self.rows[size:] + return chunk + def arrow(self, batch_size: int = 8192) -> object: import pyarrow as pa @@ -150,6 +157,20 @@ def test_select_to_arrow_reader_uses_arrow_reader_and_defers_cursor_close() -> N assert connection.cursor_obj.closed is True +def test_select_stream_uses_fetchmany_chunks() -> None: + """The sync driver should stream rows with cursor.fetchmany().""" + connection = ArrowConnection() + driver = MssqlPythonDriver(cast("MssqlPythonConnection", connection)) + + with driver.select_stream("SELECT id, name FROM dbo.users", native_only=True, chunk_size=2) as stream: + rows = list(stream) + + assert rows == [{"id": 1, "name": "Ada"}, {"id": 2, "name": "Grace"}, {"id": 3, "name": "Linus"}] + assert connection.cursor_obj.executed == [("SELECT id, name FROM dbo.users", [])] + assert connection.cursor_obj.fetchmany_sizes == [2, 2, 2] + assert connection.cursor_obj.closed is True + + def test_bulk_copy_forwards_options_to_cursor_bulkcopy() -> None: """BulkCopy options should be forwarded without SQLSpec-side rewriting.""" connection = ArrowConnection() @@ -296,6 +317,30 @@ async def fake_to_thread(func: Callable[..., object], /, *args: object, **kwargs assert "arrow_reader" in calls +@pytest.mark.anyio +async def test_async_select_stream_offloads_fetchmany(monkeypatch: pytest.MonkeyPatch) -> None: + """The async driver should stream rows with cursor.fetchmany() via to_thread.""" + calls: list[str] = [] + + async def fake_to_thread(func: Callable[..., object], /, *args: object, **kwargs: object) -> object: + calls.append(getattr(func, "__name__", repr(func))) + return func(*args, **kwargs) + + monkeypatch.setattr(driver_module.asyncio, "to_thread", fake_to_thread) + + connection = ArrowConnection() + driver = MssqlPythonAsyncDriver(cast("MssqlPythonConnection", connection)) + + async with driver.select_stream("SELECT id, name FROM dbo.users", native_only=True, chunk_size=2) as stream: + rows = [row async for row in stream] + + assert rows == [{"id": 1, "name": "Ada"}, {"id": 2, "name": "Grace"}, {"id": 3, "name": "Linus"}] + assert connection.cursor_obj.fetchmany_sizes == [2, 2, 2] + assert "_execute_cursor" in calls + assert "fetchmany" in calls + assert "close" in calls + + @pytest.mark.anyio async def test_async_bulk_copy_offloads_bulkcopy(monkeypatch: pytest.MonkeyPatch) -> None: """The async driver should offload cursor.bulkcopy().""" diff --git a/tests/unit/adapters/test_mssql_python/test_config.py b/tests/unit/adapters/test_mssql_python/test_config.py index d51395cf4..b6c1236a4 100644 --- a/tests/unit/adapters/test_mssql_python/test_config.py +++ b/tests/unit/adapters/test_mssql_python/test_config.py @@ -15,6 +15,22 @@ from sqlspec.adapters.mssql_python.pool import MssqlPythonConnectionPool +def test_config_applies_driver_feature_json_serializer_to_statement_config() -> None: + """Custom JSON serializers should reach the statement parameter config.""" + + def serializer(value: object) -> str: + return f"json:{value!r}" + + config = MssqlPythonConfig( + connection_config={"server": "localhost"}, driver_features={"json_serializer": serializer} + ) + + parameter_config = config.statement_config.parameter_config + assert parameter_config.json_serializer is serializer + assert parameter_config.type_coercion_map[dict] is serializer + assert parameter_config.type_coercion_map[list] is serializer + + class DummyConnection: """Minimal connection object for pool lifecycle assertions.""" diff --git a/tests/unit/adapters/test_mysql_family_core.py b/tests/unit/adapters/test_mysql_family_core.py index 4d9faf51b..adbe94bc4 100644 --- a/tests/unit/adapters/test_mysql_family_core.py +++ b/tests/unit/adapters/test_mysql_family_core.py @@ -40,16 +40,28 @@ def test_resolve_many_rowcount_falls_back_to_parameter_count(module_path: str) - @pytest.mark.parametrize("module_path", MODULE_PATHS) -def test_collect_rows_accepts_precomputed_column_names(module_path: str) -> None: +def test_resolve_row_plan_combines_column_names_and_json_indexes(module_path: str) -> None: + core = importlib.import_module(module_path) + description = [("id", 3), ("payload", 245), ("name", 253)] + + column_names, json_indexes = core.resolve_row_plan(description, {245}) + + assert column_names == ["id", "payload", "name"] + assert json_indexes == [1] + + +@pytest.mark.parametrize("module_path", MODULE_PATHS) +def test_collect_rows_accepts_precomputed_row_plan(module_path: str) -> None: core = importlib.import_module(module_path) description = [("id", 3), ("payload", 245)] - column_names = ["id", "payload"] + row_plan = core.resolve_row_plan(description, {245}) rows = [{"id": 1, "payload": json.dumps({"a": 1})}] - (data, resolved_column_names, row_format) = core.collect_rows( - rows, description, [1], json.loads, column_names=column_names - ) - assert resolved_column_names is column_names + + data, resolved_column_names, row_format = core.collect_rows(rows, row_plan, json.loads) + + assert resolved_column_names is row_plan[0] assert row_format == "dict" + assert data is rows assert data[0]["payload"] == {"a": 1} @@ -65,16 +77,28 @@ def test_detect_json_columns_uses_provided_description(module_path: str) -> None def test_collect_rows_without_json_reuses_input_list(module_path: str) -> None: core = importlib.import_module(module_path) description = [("id", 3), ("name", 253)] - column_names = ["id", "name"] + row_plan = core.resolve_row_plan(description, set()) rows = [{"id": 1, "name": "a"}] - (data, resolved_column_names, row_format) = core.collect_rows( - rows, description, [], json.loads, column_names=column_names - ) + + data, resolved_column_names, row_format = core.collect_rows(rows, row_plan, json.loads) + assert data is rows - assert resolved_column_names is column_names + assert resolved_column_names is row_plan[0] assert row_format == "dict" +@pytest.mark.parametrize("module_path", MODULE_PATHS) +def test_collect_stream_rows_decodes_json_using_row_plan(module_path: str) -> None: + core = importlib.import_module(module_path) + description = [("id", 3), ("payload", 245)] + row_plan = core.resolve_row_plan(description, {245}) + rows = [(1, json.dumps({"a": 1}))] + + data = core.collect_stream_rows(rows, row_plan, json.loads) + + assert data == [{"id": 1, "payload": {"a": 1}}] + + def test_asyncmy_uses_canonical_private_json_helper_names() -> None: core = importlib.import_module("sqlspec.adapters.asyncmy.core") assert hasattr(core, "_deserialize_json_dict_rows") @@ -107,7 +131,8 @@ def deserialize( monkeypatch.setattr(core, "_deserialize_json_dict_rows", deserialize) rows = [{"payload": json.dumps({"ok": True})}] - (data, _columns, row_format) = core.collect_rows(rows, [("payload", 245)], [0], json.loads) + row_plan = core.resolve_row_plan([("payload", 245)], {245}) + (data, _columns, row_format) = core.collect_rows(rows, row_plan, json.loads) assert row_format == "dict" assert data == rows assert calls == [(["payload"], rows, [0])] @@ -126,7 +151,8 @@ def deserialize( monkeypatch.setattr(core, "_deserialize_json_tuple_rows", deserialize) rows = [(json.dumps({"ok": True}),)] - (data, _columns, row_format) = core.collect_rows(rows, [("payload", 245)], [0], json.loads) + row_plan = core.resolve_row_plan([("payload", 245)], {245}) + (data, _columns, row_format) = core.collect_rows(rows, row_plan, json.loads) assert row_format == "tuple" assert data == rows assert calls == [(rows, [0])] diff --git a/tests/unit/adapters/test_mysqlconnector/test_config.py b/tests/unit/adapters/test_mysqlconnector/test_config.py index 6e16990d2..65c1170cd 100644 --- a/tests/unit/adapters/test_mysqlconnector/test_config.py +++ b/tests/unit/adapters/test_mysqlconnector/test_config.py @@ -1,5 +1,6 @@ """Unit tests for mysql-connector configuration modernization.""" +import contextlib from typing import TYPE_CHECKING, Any from unittest.mock import AsyncMock, MagicMock @@ -17,6 +18,7 @@ MysqlConnectorSyncConfig, MysqlConnectorSyncConnectionParams, ) +from sqlspec.adapters.mysqlconnector.core import MysqlConnectorAsyncStreamSource, MysqlConnectorSyncStreamSource from sqlspec.adapters.mysqlconnector.driver import MysqlConnectorAsyncDriver, MysqlConnectorSyncDriver if TYPE_CHECKING: @@ -27,7 +29,7 @@ def test_sync_config_uses_connector_python_host_default_and_disables_local_infil """SQLSpec should preserve the driver host default and close the local infile gate.""" config = MysqlConnectorSyncConfig() - assert config.connection_config["host"] == "127.0.0.1" + assert config.connection_config["host"] == "localhost" assert config.connection_config["port"] == 3306 assert config.connection_config["allow_local_infile"] is False @@ -36,16 +38,18 @@ def test_async_config_uses_connector_python_host_default_and_disables_local_infi """The async config should use the same connection defaults as sync.""" config = MysqlConnectorAsyncConfig() - assert config.connection_config["host"] == "127.0.0.1" + assert config.connection_config["host"] == "localhost" assert config.connection_config["port"] == 3306 assert config.connection_config["allow_local_infile"] is False -def test_local_infile_must_be_enabled_explicitly() -> None: - """Explicit local infile opt-in should be preserved for callers that need it.""" - config = MysqlConnectorSyncConfig(connection_config={"allow_local_infile": True}) +@pytest.mark.parametrize("config_cls", [MysqlConnectorSyncConfig, MysqlConnectorAsyncConfig]) +def test_local_infile_uses_connector_python_security_gate(config_cls: type[Any]) -> None: + """LOAD DATA LOCAL INFILE should use mysql-connector's native consent gate.""" + config = config_cls(connection_config={"allow_local_infile": True}) assert config.connection_config["allow_local_infile"] is True + assert "local_infile" not in config.connection_config def test_sync_connection_params_type_accepts_modern_connector_options() -> None: @@ -195,6 +199,92 @@ def test_async_driver_with_cursor_uses_driver_cursor_options() -> None: assert cursor_context.cursor_options == {"prepared": True} +def test_sync_dispatch_select_stream_uses_driver_cursor_options(monkeypatch: pytest.MonkeyPatch) -> None: + """Sync stream execution should reuse cursor options and force unbuffered cursors.""" + connection = MagicMock() + cursor = MagicMock() + cursor.description = [] + connection.cursor.return_value = cursor + driver = MysqlConnectorSyncDriver(connection=connection, driver_features={"cursor_options": {"prepared": True}}) + statement = MagicMock() + statement.returns_rows.return_value = True + + monkeypatch.setattr(MysqlConnectorSyncDriver, "_compiled_sql", lambda self, stmt, config: ("SELECT 1", None)) + monkeypatch.setattr(MysqlConnectorSyncDriver, "handle_database_exceptions", lambda self: contextlib.nullcontext()) + monkeypatch.setattr(MysqlConnectorSyncDriver, "_check_pending_exception", lambda self, handler: None) + stream = driver.dispatch_select_stream(statement, 10) + + assert stream is not None + source = stream._source + assert isinstance(source, MysqlConnectorSyncStreamSource) + assert source._cursor_options == {"prepared": True, "buffered": False} + + source.start() + + connection.cursor.assert_called_once_with(prepared=True, buffered=False) + + +@pytest.mark.anyio +async def test_async_dispatch_select_stream_uses_driver_cursor_options(monkeypatch: pytest.MonkeyPatch) -> None: + """Async stream execution should reuse cursor options and force unbuffered cursors.""" + connection = AsyncMock() + cursor = AsyncMock() + cursor.description = [] + connection.cursor.return_value = cursor + driver = MysqlConnectorAsyncDriver(connection=connection, driver_features={"cursor_options": {"prepared": True}}) + statement = MagicMock() + statement.returns_rows.return_value = True + + monkeypatch.setattr(MysqlConnectorAsyncDriver, "_compiled_sql", lambda self, stmt, config: ("SELECT 1", None)) + monkeypatch.setattr(MysqlConnectorAsyncDriver, "handle_database_exceptions", lambda self: _AsyncNullContext()) + monkeypatch.setattr(MysqlConnectorAsyncDriver, "_check_pending_exception", lambda self, handler: None) + stream = driver.dispatch_select_stream(statement, 10) + + assert stream is not None + source = stream._source + assert isinstance(source, MysqlConnectorAsyncStreamSource) + assert source._cursor_options == {"prepared": True, "buffered": False} + + await source.start() + + connection.cursor.assert_awaited_once_with(prepared=True, buffered=False) + + +def test_sync_create_connection_forwards_local_infile_gate(monkeypatch: pytest.MonkeyPatch) -> None: + """Connection creation should forward mysql-connector's local infile gate.""" + calls: list[dict[str, Any]] = [] + connection = MagicMock() + + def connect(**kwargs: Any) -> MagicMock: + calls.append(kwargs) + return connection + + monkeypatch.setattr("sqlspec.adapters.mysqlconnector.config.mysql.connector.connect", connect) + MysqlConnectorSyncConfig(connection_config={"allow_local_infile": True}).create_connection() + + assert calls[0]["allow_local_infile"] is True + assert "local_infile" not in calls[0] + + +@pytest.mark.anyio +async def test_async_create_connection_forwards_local_infile_gate(monkeypatch: pytest.MonkeyPatch) -> None: + """Async connection creation should forward mysql-connector's local infile gate.""" + calls: list[dict[str, Any]] = [] + connection = AsyncMock() + + async def connect(**kwargs: Any) -> AsyncMock: + calls.append(kwargs) + return connection + + monkeypatch.setattr("sqlspec.adapters.mysqlconnector.config.mysqlconnector_aio.connect", connect) + + config = MysqlConnectorAsyncConfig(connection_config={"allow_local_infile": True}) + await config.create_connection() + + assert calls[0]["allow_local_infile"] is True + assert "local_infile" not in calls[0] + + def test_sync_create_connection_passes_local_infile_gate(monkeypatch: pytest.MonkeyPatch) -> None: """Connection creation should pass the local infile gate explicitly to the driver.""" calls: list[dict[str, Any]] = [] @@ -208,3 +298,13 @@ def connect(**kwargs: Any) -> MagicMock: MysqlConnectorSyncConfig().create_connection() assert calls[0]["allow_local_infile"] is False + + +class _AsyncNullContext: + async def __aenter__(self) -> None: + return None + + async def __aexit__( + self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: object + ) -> None: + return None diff --git a/tests/unit/adapters/test_oracledb/test_arrow_batches.py b/tests/unit/adapters/test_oracledb/test_arrow_batches.py index 54cb78f7e..95ccb4ffd 100644 --- a/tests/unit/adapters/test_oracledb/test_arrow_batches.py +++ b/tests/unit/adapters/test_oracledb/test_arrow_batches.py @@ -26,6 +26,15 @@ def fetch_df_all(self, **_kwargs: object) -> object: raise AssertionError(msg) +class _OracleAllConnection: + def __init__(self) -> None: + self.fetch_all_calls: list[dict[str, object]] = [] + + def fetch_df_all(self, **kwargs: object) -> pa.Table: + self.fetch_all_calls.append(kwargs) + return pa.table({"ID": [1], "VALUE": [10]}) + + class _OracleAsyncBatchConnection: def __init__(self) -> None: self.fetch_batch_calls: list[dict[str, object]] = [] @@ -40,11 +49,20 @@ async def fetch_df_all(self, **_kwargs: object) -> object: raise AssertionError(msg) +class _OracleAsyncAllConnection: + def __init__(self) -> None: + self.fetch_all_calls: list[dict[str, object]] = [] + + async def fetch_df_all(self, **kwargs: object) -> pa.Table: + self.fetch_all_calls.append(kwargs) + return pa.table({"ID": [1], "VALUE": [10]}) + + def test_sync_select_to_arrow_batches_uses_fetch_df_batches() -> None: connection = _OracleBatchConnection() driver = OracleSyncDriver( cast("OracleSyncConnection", connection), - driver_features={"enable_lowercase_column_names": True, "fetch_decimals": True}, + driver_features={"enable_lowercase_column_names": True, "fetch_decimals": True, "fetch_lobs": False}, ) result = driver.select_to_arrow("SELECT id, value FROM example", return_format="batches", batch_size=2) @@ -55,6 +73,22 @@ def test_sync_select_to_arrow_batches_uses_fetch_df_batches() -> None: assert result.rows_affected == 3 assert connection.fetch_batch_calls[0]["size"] == 2 assert connection.fetch_batch_calls[0]["fetch_decimals"] is True + assert connection.fetch_batch_calls[0]["fetch_lobs"] is False + + +def test_sync_select_to_arrow_table_forwards_fetch_lobs_to_fetch_df_all() -> None: + connection = _OracleAllConnection() + driver = OracleSyncDriver( + cast("OracleSyncConnection", connection), + driver_features={"enable_lowercase_column_names": True, "fetch_decimals": True, "fetch_lobs": False}, + ) + + result = driver.select_to_arrow("SELECT id, value FROM example", batch_size=2) + + assert result.get_data().to_pydict() == {"id": [1], "value": [10]} + assert connection.fetch_all_calls[0]["arraysize"] == 2 + assert connection.fetch_all_calls[0]["fetch_decimals"] is True + assert connection.fetch_all_calls[0]["fetch_lobs"] is False @pytest.mark.anyio @@ -62,7 +96,7 @@ async def test_async_select_to_arrow_reader_uses_fetch_df_batches() -> None: connection = _OracleAsyncBatchConnection() driver = OracleAsyncDriver( cast("OracleAsyncConnection", connection), - driver_features={"enable_lowercase_column_names": True, "fetch_decimals": True}, + driver_features={"enable_lowercase_column_names": True, "fetch_decimals": True, "fetch_lobs": False}, ) result = await driver.select_to_arrow("SELECT id, value FROM example", return_format="reader", batch_size=2) @@ -71,3 +105,20 @@ async def test_async_select_to_arrow_reader_uses_fetch_df_batches() -> None: assert result.get_data().read_all().to_pydict() == {"id": [1, 2, 3], "value": [10, 20, 30]} assert connection.fetch_batch_calls[0]["size"] == 2 assert connection.fetch_batch_calls[0]["fetch_decimals"] is True + assert connection.fetch_batch_calls[0]["fetch_lobs"] is False + + +@pytest.mark.anyio +async def test_async_select_to_arrow_table_forwards_fetch_lobs_to_fetch_df_all() -> None: + connection = _OracleAsyncAllConnection() + driver = OracleAsyncDriver( + cast("OracleAsyncConnection", connection), + driver_features={"enable_lowercase_column_names": True, "fetch_decimals": True, "fetch_lobs": False}, + ) + + result = await driver.select_to_arrow("SELECT id, value FROM example", batch_size=2) + + assert result.get_data().to_pydict() == {"id": [1], "value": [10]} + assert connection.fetch_all_calls[0]["arraysize"] == 2 + assert connection.fetch_all_calls[0]["fetch_decimals"] is True + assert connection.fetch_all_calls[0]["fetch_lobs"] is False diff --git a/tests/unit/adapters/test_oracledb/test_core_driver_features.py b/tests/unit/adapters/test_oracledb/test_core_driver_features.py index 7989a7fab..d7b464e81 100644 --- a/tests/unit/adapters/test_oracledb/test_core_driver_features.py +++ b/tests/unit/adapters/test_oracledb/test_core_driver_features.py @@ -98,15 +98,21 @@ def test_apply_driver_features_preserves_user_raw_byte_limit() -> None: assert features["oracle_raw_byte_limit"] == 100 -def test_apply_driver_features_does_not_default_fetch_tuning_options() -> None: - """Absent fetch tuning keys should leave python-oracledb defaults untouched.""" +def test_apply_driver_features_defaults_fetch_lobs_false() -> None: + """``fetch_lobs`` defaults to ``False`` while unrelated fetch tuning stays unset.""" features = _apply_features({}) + assert features["fetch_lobs"] is False assert "arraysize" not in features assert "prefetchrows" not in features - assert "fetch_lobs" not in features assert "fetch_decimals" not in features +def test_apply_driver_features_preserves_user_fetch_lobs_true() -> None: + """User-supplied ``fetch_lobs=True`` remains an explicit opt-in escape hatch.""" + features = _apply_features({"fetch_lobs": True}) + assert features["fetch_lobs"] is True + + def test_apply_driver_features_preserves_user_fetch_tuning_options() -> None: """User-supplied fetch tuning options survive the defaults pass.""" features = _apply_features({"arraysize": 5000, "fetch_decimals": True}) diff --git a/tests/unit/adapters/test_oracledb/test_core_row_materialization.py b/tests/unit/adapters/test_oracledb/test_core_row_materialization.py index 7080f9ec7..a7b56a019 100644 --- a/tests/unit/adapters/test_oracledb/test_core_row_materialization.py +++ b/tests/unit/adapters/test_oracledb/test_core_row_materialization.py @@ -1,13 +1,16 @@ # pyright: reportArgumentType=false """Unit tests for Oracle row materialization helpers.""" +import inspect from typing import TYPE_CHECKING, Any, cast import pytest from sqlspec.adapters.oracledb.core import collect_async_rows, collect_sync_rows, resolve_row_metadata from sqlspec.adapters.oracledb.driver import OracleAsyncDriver, OracleSyncDriver +from sqlspec.core import OperationProfile, ParameterProfile, ProcessedState from sqlspec.driver import AsyncDriverAdapterBase, SyncDriverAdapterBase +from sqlspec.driver._query_cache import CachedQuery if TYPE_CHECKING: from sqlspec.adapters.oracledb._typing import OracleAsyncConnection, OracleSyncConnection @@ -30,6 +33,54 @@ def read(self) -> str: return self._value +class _AsyncReadableValue: + __slots__ = ("_value",) + + def __init__(self, value: str) -> None: + self._value = value + + async def read(self) -> str: + return self._value + + +class _AsyncCursor: + def __init__(self) -> None: + self.description = [("PAYLOAD", _TypeCode("DB_TYPE_CLOB"))] + self.rowcount = 1 + self.execute_calls: list[tuple[str, object, dict[str, object]]] = [] + + async def execute(self, sql: str, parameters: object, **kwargs: object) -> None: + self.execute_calls.append((sql, parameters, kwargs)) + + async def fetchall(self) -> list[tuple[object]]: + return [(_AsyncReadableValue("async text"),)] + + def close(self) -> None: + return None + + +class _AsyncConnection: + def __init__(self, cursor: _AsyncCursor) -> None: + self.cursor_obj = cursor + + def cursor(self) -> _AsyncCursor: + return self.cursor_obj + + +def _cached_select(compiled_sql: str = "SELECT payload FROM dual") -> CachedQuery: + return CachedQuery( + compiled_sql=compiled_sql, + parameter_profile=ParameterProfile(), + input_named_parameters=(), + applied_wrap_types=False, + parameter_casts={}, + operation_type="SELECT", + operation_profile=OperationProfile(returns_rows=True, modifies_rows=False), + param_count=0, + processed_state=ProcessedState(compiled_sql=compiled_sql, execution_parameters=[], operation_type="SELECT"), + ) + + def test_collect_sync_rows_reuses_original_rows_when_no_lob_columns() -> None: """collect_sync_rows should avoid row copy work for non-LOB result sets.""" rows = [(1, "alpha"), (2, "beta")] @@ -49,6 +100,15 @@ def test_collect_sync_rows_coerces_lob_values_when_lob_columns_present() -> None assert column_names == ["id", "payload"] +def test_collect_sync_rows_leaves_json_text_as_string_without_content_sniffing() -> None: + """Readable CLOB values should be read as strings, not auto-parsed as JSON.""" + rows = [(1, _ReadableValue('{"key": "value"}'))] + description = [("ID", _TypeCode("DB_TYPE_NUMBER")), ("PAYLOAD", _TypeCode("DB_TYPE_CLOB"))] + (data, _) = collect_sync_rows(rows, description, {"enable_lowercase_column_names": True}) + assert data == [(1, '{"key": "value"}')] + assert isinstance(data[0][1], str) + + async def test_collect_async_rows_reuses_original_rows_when_no_lob_columns() -> None: """collect_async_rows should avoid async coercion for non-LOB result sets.""" rows = [(1, "alpha"), (2, "beta")] @@ -67,6 +127,32 @@ async def test_collect_async_rows_coerces_lob_values_when_lob_columns_present() assert column_names == ["id", "payload"] +async def test_collect_async_rows_leaves_json_text_as_string_without_content_sniffing() -> None: + """Readable CLOB values should be read as strings, not auto-parsed as JSON.""" + rows = [(1, _ReadableValue('{"key": "value"}'))] + description = [("ID", _TypeCode("DB_TYPE_NUMBER")), ("PAYLOAD", _TypeCode("DB_TYPE_CLOB"))] + (data, _) = await collect_async_rows(rows, description, {"enable_lowercase_column_names": True}) + assert data == [(1, '{"key": "value"}')] + assert isinstance(data[0][1], str) + + +@pytest.mark.anyio +async def test_oracle_async_cached_collect_reads_async_lob_locators() -> None: + """Async direct collection should not leave coroutine objects in rows.""" + cursor = _AsyncCursor() + driver = OracleAsyncDriver( + connection=cast("OracleAsyncConnection", _AsyncConnection(cursor)), + driver_features={"enable_lowercase_column_names": True, "fetch_lobs": True}, + ) + + result = await driver._execute_cache_hit("SELECT payload FROM dual", {}, _cached_select()) + + value = result.get_data()[0]["payload"] + if inspect.iscoroutine(value): + value.close() + assert value == "async text" + + def test_resolve_row_metadata_reuses_cached_description() -> None: """resolve_row_metadata should cache normalized names and LOB flags by description identity.""" description = [("ID", _TypeCode("DB_TYPE_NUMBER")), ("NAME", _TypeCode("DB_TYPE_VARCHAR"))] diff --git a/tests/unit/adapters/test_oracledb/test_fetch_tuning.py b/tests/unit/adapters/test_oracledb/test_fetch_tuning.py index 72ca7dab2..3a656f3b8 100644 --- a/tests/unit/adapters/test_oracledb/test_fetch_tuning.py +++ b/tests/unit/adapters/test_oracledb/test_fetch_tuning.py @@ -1,11 +1,12 @@ """Unit tests for Oracle fetch tuning options.""" -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING, Any, cast import pytest from sqlspec.adapters.oracledb._typing import OracleAsyncCursor, OracleSyncCursor -from sqlspec.adapters.oracledb.core import build_fetch_kwargs +from sqlspec.adapters.oracledb.core import build_arrow_fetch_kwargs, build_fetch_kwargs +from sqlspec.adapters.oracledb.driver import OracleAsyncDriver, OracleSyncDriver if TYPE_CHECKING: from sqlspec.adapters.oracledb._typing import OracleAsyncConnection, OracleSyncConnection @@ -29,6 +30,31 @@ def cursor(self) -> _RawCursor: return self.cursor_obj +class _StreamCursor: + def __init__(self) -> None: + self.arraysize = 100 + self.prefetchrows = 2 + self.execute_kwargs: dict[str, object] = {} + self.closed = False + + def execute(self, _sql: str, _parameters: object, **kwargs: object) -> None: + self.execute_kwargs = kwargs + + def fetchmany(self, _size: int) -> list[tuple[object, ...]]: + return [] + + def close(self) -> None: + self.closed = True + + +class _AsyncStreamCursor(_StreamCursor): + async def execute(self, _sql: str, _parameters: object, **kwargs: object) -> None: # type: ignore[override] + self.execute_kwargs = kwargs + + async def fetchmany(self, _size: int) -> list[tuple[object, ...]]: # type: ignore[override] + return [] + + def test_build_fetch_kwargs_omits_absent_options() -> None: assert build_fetch_kwargs({}) == {} @@ -40,6 +66,13 @@ def test_build_fetch_kwargs_includes_explicit_fetch_options() -> None: } +def test_build_arrow_fetch_kwargs_matches_execute_fetch_options() -> None: + assert build_arrow_fetch_kwargs({"fetch_lobs": False, "fetch_decimals": True}) == { + "fetch_lobs": False, + "fetch_decimals": True, + } + + def test_oracle_sync_cursor_applies_arraysize_and_prefetchrows() -> None: raw_cursor = _RawCursor() cursor = OracleSyncCursor( @@ -68,6 +101,32 @@ def test_oracle_fetch_kwargs_omit_absent_fetch_tuning_keys() -> None: assert build_fetch_kwargs({"arraysize": 100, "prefetchrows": 200}) == {} +def test_oracle_sync_stream_forwards_driver_fetch_lobs_default() -> None: + raw_cursor = _StreamCursor() + driver = OracleSyncDriver( + cast("OracleSyncConnection", _SyncConnection(cast(_RawCursor, raw_cursor))), + driver_features={"fetch_lobs": False}, + ) + + stream = driver.select_stream("SELECT payload FROM example") + + assert list(stream) == [] + assert raw_cursor.execute_kwargs["fetch_lobs"] is False + + +def test_oracle_sync_stream_fetch_lobs_call_option_overrides_default() -> None: + raw_cursor = _StreamCursor() + driver = OracleSyncDriver( + cast("OracleSyncConnection", _SyncConnection(cast(_RawCursor, raw_cursor))), + driver_features={"fetch_lobs": False}, + ) + + stream = driver.select_stream("SELECT payload FROM example", fetch_lobs=True) + + assert list(stream) == [] + assert raw_cursor.execute_kwargs["fetch_lobs"] is True + + @pytest.mark.anyio async def test_oracle_async_cursor_applies_arraysize_and_prefetchrows() -> None: raw_cursor = _RawCursor() @@ -81,3 +140,33 @@ async def test_oracle_async_cursor_applies_arraysize_and_prefetchrows() -> None: assert raw_cursor.prefetchrows == 701 assert raw_cursor.closed is True + + +@pytest.mark.anyio +async def test_oracle_async_stream_forwards_driver_fetch_lobs_default() -> None: + raw_cursor = _AsyncStreamCursor() + driver = OracleAsyncDriver( + cast("OracleAsyncConnection", _SyncConnection(cast(_RawCursor, raw_cursor))), + driver_features={"fetch_lobs": False}, + ) + + stream = driver.select_stream("SELECT payload FROM example") + + rows: list[dict[str, Any]] = [row async for row in stream] + assert rows == [] + assert raw_cursor.execute_kwargs["fetch_lobs"] is False + + +@pytest.mark.anyio +async def test_oracle_async_stream_fetch_lobs_call_option_overrides_default() -> None: + raw_cursor = _AsyncStreamCursor() + driver = OracleAsyncDriver( + cast("OracleAsyncConnection", _SyncConnection(cast(_RawCursor, raw_cursor))), + driver_features={"fetch_lobs": False}, + ) + + stream = driver.select_stream("SELECT payload FROM example", fetch_lobs=True) + + rows: list[dict[str, Any]] = [row async for row in stream] + assert rows == [] + assert raw_cursor.execute_kwargs["fetch_lobs"] is True diff --git a/tests/unit/adapters/test_oracledb/test_json_handlers.py b/tests/unit/adapters/test_oracledb/test_json_handlers.py index b85a62511..8dd2572da 100644 --- a/tests/unit/adapters/test_oracledb/test_json_handlers.py +++ b/tests/unit/adapters/test_oracledb/test_json_handlers.py @@ -27,11 +27,13 @@ def _mock_cursor_with_major(major: int | None) -> Mock: return cursor -def _mock_metadata(type_code: object, type_name: str = "") -> Mock: +def _mock_metadata(type_code: object, type_name: str = "", *, is_json: bool = False, is_oson: bool = False) -> Mock: """Build a column-metadata mock matching python-oracledb's FetchInfo.""" metadata = Mock() metadata.type_code = type_code metadata.type_name = type_name + metadata.is_json = is_json + metadata.is_oson = is_oson return metadata @@ -208,35 +210,100 @@ def test_output_handler_short_circuits_db_type_json() -> None: def test_output_handler_claims_blob_is_json() -> None: - """BLOB columns whose type_name carries JSON should be parsed via outconverter.""" + """BLOB columns marked as JSON should be parsed via the blob outconverter.""" import oracledb cursor = Mock() cursor.arraysize = 100 cursor_var = Mock() cursor.var = Mock(return_value=cursor_var) - metadata = _mock_metadata(oracledb.DB_TYPE_BLOB, type_name="BLOB CHECK (… IS JSON)") + metadata = _mock_metadata(oracledb.DB_TYPE_BLOB, type_name="BLOB", is_json=True) result = json_output_type_handler(cursor, metadata) assert result is cursor_var - cursor.var.assert_called_once_with(oracledb.DB_TYPE_BLOB, arraysize=100, outconverter=json_converter_out_blob) + cursor.var.assert_called_once_with(oracledb.DB_TYPE_LONG_RAW, arraysize=100, outconverter=json_converter_out_blob) def test_output_handler_claims_clob_is_json() -> None: - """CLOB columns whose type_name carries JSON should be parsed via outconverter.""" + """CLOB columns marked as JSON should be parsed via the clob outconverter.""" import oracledb cursor = Mock() cursor.arraysize = 50 cursor_var = Mock() cursor.var = Mock(return_value=cursor_var) - metadata = _mock_metadata(oracledb.DB_TYPE_CLOB, type_name="CLOB CHECK (… IS JSON)") + metadata = _mock_metadata(oracledb.DB_TYPE_CLOB, type_name="CLOB", is_json=True) result = json_output_type_handler(cursor, metadata) assert result is cursor_var - cursor.var.assert_called_once_with(oracledb.DB_TYPE_CLOB, arraysize=50, outconverter=json_converter_out_clob) + cursor.var.assert_called_once_with(oracledb.DB_TYPE_LONG, arraysize=50, outconverter=json_converter_out_clob) + + +def test_output_handler_claims_varchar2_is_json() -> None: + """VARCHAR2 columns marked as JSON should be parsed via the clob outconverter.""" + import oracledb + + cursor = Mock() + cursor.arraysize = 7 + cursor_var = Mock() + cursor.var = Mock(return_value=cursor_var) + metadata = _mock_metadata(oracledb.DB_TYPE_VARCHAR, type_name="VARCHAR2", is_json=True) + + result = json_output_type_handler(cursor, metadata) + + assert result is cursor_var + cursor.var.assert_called_once_with(oracledb.DB_TYPE_VARCHAR, arraysize=7, outconverter=json_converter_out_clob) + + +def test_output_handler_claims_blob_is_oson() -> None: + """BLOB columns marked as OSON should decode through connection.decode_oson.""" + import oracledb + + cursor = Mock() + cursor.arraysize = 33 + connection = Mock() + connection.decode_oson = Mock(return_value={"decoded": True}) + cursor.connection = connection + cursor_var = Mock() + cursor.var = Mock(return_value=cursor_var) + metadata = _mock_metadata(oracledb.DB_TYPE_BLOB, type_name="BLOB", is_oson=True) + + result = json_output_type_handler(cursor, metadata) + + assert result is cursor_var + cursor.var.assert_called_once_with( + oracledb.DB_TYPE_LONG_RAW, arraysize=33, outconverter=cursor.var.call_args.kwargs["outconverter"] + ) + outconverter = cursor.var.call_args.kwargs["outconverter"] + assert len(inspect.signature(outconverter).parameters) == 1 + assert outconverter(b"\x01\x02") == {"decoded": True} + connection.decode_oson.assert_called_once_with(b"\x01\x02") + + +def test_output_handler_prefers_oson_when_blob_metadata_is_json_and_oson() -> None: + """OSON BLOB metadata can also report is_json; decode_oson must win.""" + import oracledb + + cursor = Mock() + cursor.arraysize = 34 + connection = Mock() + connection.decode_oson = Mock(return_value={"decoded": "oson"}) + cursor.connection = connection + cursor_var = Mock() + cursor.var = Mock(return_value=cursor_var) + metadata = _mock_metadata(oracledb.DB_TYPE_BLOB, type_name="BLOB", is_json=True, is_oson=True) + + result = json_output_type_handler(cursor, metadata) + + assert result is cursor_var + cursor.var.assert_called_once_with( + oracledb.DB_TYPE_LONG_RAW, arraysize=34, outconverter=cursor.var.call_args.kwargs["outconverter"] + ) + outconverter = cursor.var.call_args.kwargs["outconverter"] + assert outconverter(b"\x01\x02") == {"decoded": "oson"} + connection.decode_oson.assert_called_once_with(b"\x01\x02") def test_output_handler_ignores_plain_blob() -> None: @@ -260,7 +327,7 @@ def test_output_handler_ignores_plain_clob() -> None: def test_output_handler_ignores_varchar2() -> None: - """VARCHAR2 columns should never be claimed by JSON handler.""" + """VARCHAR2 columns without JSON metadata should not be claimed.""" import oracledb cursor = Mock() @@ -269,6 +336,16 @@ def test_output_handler_ignores_varchar2() -> None: assert json_output_type_handler(cursor, metadata) is None +def test_output_handler_ignores_non_json_metadata() -> None: + """Columns without JSON metadata should be ignored.""" + import oracledb + + cursor = Mock() + metadata = _mock_metadata(oracledb.DB_TYPE_BLOB, type_name="BLOB", is_json=False, is_oson=False) + + assert json_output_type_handler(cursor, metadata) is None + + def test_register_json_handlers_no_existing() -> None: """register_json_handlers should install both input and output handlers.""" connection = Mock() @@ -368,7 +445,7 @@ def test_output_handler_chain_prioritises_blob_is_json() -> None: cursor.arraysize = 10 cursor_var = Mock() cursor.var = Mock(return_value=cursor_var) - metadata = _mock_metadata(oracledb.DB_TYPE_BLOB, type_name="BLOB IS JSON") + metadata = _mock_metadata(oracledb.DB_TYPE_BLOB, type_name="BLOB", is_json=True) result = connection.outputtypehandler(cursor, metadata) @@ -387,6 +464,7 @@ def test_module_exports_required_symbols() -> None: assert "json_converter_in_clob" in _json_handlers.__all__ assert "json_converter_out_blob" in _json_handlers.__all__ assert "json_converter_out_clob" in _json_handlers.__all__ + assert "json_converter_out_oson" in _json_handlers.__all__ def test_roundtrip_clob() -> None: diff --git a/tests/unit/adapters/test_oracledb/test_type_converter_vectors.py b/tests/unit/adapters/test_oracledb/test_type_converter_vectors.py index b846e95dd..a465b5ef1 100644 --- a/tests/unit/adapters/test_oracledb/test_type_converter_vectors.py +++ b/tests/unit/adapters/test_oracledb/test_type_converter_vectors.py @@ -190,3 +190,11 @@ def test_convert_vector_to_numpy_uses_copy_true() -> None: oracle_array[0] = 999.0 assert result[0] == 1.0 + + +def test_oracle_output_converter_no_longer_exposes_json_storage_regex() -> None: + """Oracle type conversion should not publish the retired JSON storage regex.""" + import sqlspec.adapters.oracledb.type_converter as oracle_type_converter + + assert not hasattr(oracle_type_converter, "ORACLE_JSON_STORAGE_REGEX") + assert "ORACLE_JSON_STORAGE_REGEX" not in oracle_type_converter.__all__ diff --git a/tests/unit/adapters/test_psqlpy/test_config.py b/tests/unit/adapters/test_psqlpy/test_config.py index 90021af0d..af71eb21a 100644 --- a/tests/unit/adapters/test_psqlpy/test_config.py +++ b/tests/unit/adapters/test_psqlpy/test_config.py @@ -45,6 +45,18 @@ def serializer(_: object) -> str: assert parameter_config.json_serializer is serializer +def test_build_default_statement_config_custom_deserializer() -> None: + """Custom deserializer should propagate into the parameter configuration.""" + + def deserializer(_: str) -> object: + return {"deserialized": True} + + statement_config = build_statement_config(json_deserializer=deserializer) + + parameter_config = statement_config.parameter_config + assert parameter_config.json_deserializer is deserializer + + def test_psqlpy_config_applies_driver_feature_serializer() -> None: """Driver features should mutate the Psqlpy statement configuration.""" @@ -57,6 +69,18 @@ def serializer(_: object) -> str: assert parameter_config.json_serializer is serializer +def test_psqlpy_config_applies_driver_feature_deserializer() -> None: + """Driver features should mutate the Psqlpy JSON deserializer.""" + + def deserializer(_: str) -> object: + return {"feature": True} + + config = PsqlpyConfig(driver_features={"json_deserializer": deserializer}) + + parameter_config = config.statement_config.parameter_config + assert parameter_config.json_deserializer is deserializer + + def test_psqlpy_driver_features_declares_cast_detection() -> None: """Driver feature typing should expose every consumed feature flag.""" assert "enable_cast_detection" in PsqlpyDriverFeatures.__annotations__ diff --git a/tests/unit/adapters/test_psqlpy/test_core.py b/tests/unit/adapters/test_psqlpy/test_core.py index 66c688479..0b424a0bf 100644 --- a/tests/unit/adapters/test_psqlpy/test_core.py +++ b/tests/unit/adapters/test_psqlpy/test_core.py @@ -17,6 +17,7 @@ default_statement_config, encode_records_for_binary_copy, format_execute_many_parameters, + get_parameter_casts, prepare_parameters_with_casts, ) from sqlspec.adapters.psqlpy.driver import PsqlpyDriver @@ -104,6 +105,17 @@ def wrapped(*args: Any, **kwargs: Any) -> object: assert calls == 1 +def test_get_parameter_casts_reads_processed_state_from_cached_statement() -> None: + class _ProcessedState: + parameter_casts = {1: "JSONB"} + + class _Statement: + def get_processed_state(self) -> _ProcessedState: + return _ProcessedState() + + assert get_parameter_casts(_Statement()) == {1: "JSONB"} + + def test_format_execute_many_parameters_handles_scalar_input() -> None: """Scalar execute_many payloads should be normalized to a list containing one row.""" formatted = format_execute_many_parameters(5, coerce_numeric=False) @@ -218,6 +230,16 @@ class MyInt(int): assert prepared == [5] +def test_psqlpy_driver_no_longer_caches_output_converter() -> None: + """The driver should no longer construct the dead psqlpy output converter.""" + import sqlspec.adapters.psqlpy.driver as psqlpy_driver + import sqlspec.adapters.psqlpy.type_converter as psqlpy_type_converter + + assert not hasattr(psqlpy_driver, "_type_converter") + assert not hasattr(psqlpy_type_converter, "PostgreSQLOutputConverter") + assert "PostgreSQLOutputConverter" not in psqlpy_type_converter.__all__ + + def test_prepare_parameters_with_casts_supports_virtual_abc_dispatch() -> None: statement_config = build_statement_config() statement_config = statement_config.replace( diff --git a/tests/unit/adapters/test_psycopg/test_config.py b/tests/unit/adapters/test_psycopg/test_config.py index d8648047f..94816b295 100644 --- a/tests/unit/adapters/test_psycopg/test_config.py +++ b/tests/unit/adapters/test_psycopg/test_config.py @@ -90,6 +90,18 @@ def serializer(_: object) -> str: assert parameter_config.json_serializer is serializer +def test_build_default_statement_config_custom_deserializer() -> None: + """Custom deserializer should propagate into the parameter configuration.""" + + def deserializer(_: str) -> object: + return {"deserialized": True} + + statement_config = build_statement_config(json_deserializer=deserializer) + + parameter_config = statement_config.parameter_config + assert parameter_config.json_deserializer is deserializer + + def test_psycopg_sync_config_applies_driver_feature_serializer() -> None: """Driver features should mutate the sync Psycopg statement configuration.""" @@ -102,6 +114,18 @@ def serializer(_: object) -> str: assert parameter_config.json_serializer is serializer +def test_psycopg_sync_config_applies_driver_feature_deserializer() -> None: + """Driver features should mutate the sync Psycopg JSON deserializer.""" + + def deserializer(_: str) -> object: + return {"feature": True} + + config = PsycopgSyncConfig(driver_features={"json_deserializer": deserializer}) + + parameter_config = config.statement_config.parameter_config + assert parameter_config.json_deserializer is deserializer + + def test_psycopg_build_postgres_extension_probe_names_filters_disabled_features() -> None: """Only enabled extension probes should be returned.""" assert build_postgres_extension_probe_names({"enable_pgvector": True, "enable_paradedb": False}) == ["vector"] diff --git a/tests/unit/adapters/test_psycopg/test_core.py b/tests/unit/adapters/test_psycopg/test_core.py index 55549f769..78262079d 100644 --- a/tests/unit/adapters/test_psycopg/test_core.py +++ b/tests/unit/adapters/test_psycopg/test_core.py @@ -105,7 +105,7 @@ def returns_rows(self) -> bool: class _Cursor: def __init__(self) -> None: - self.description = [SimpleNamespace()] + self.description = [SimpleNamespace(name="id"), SimpleNamespace(name="name")] def fetchall(self) -> list[tuple[int, str]]: return [(1, "alice"), (2, "bob")] @@ -119,6 +119,37 @@ def fetchall(self) -> list[tuple[int, str]]: assert result.data_row_count == 2 +def test_build_pipeline_execution_result_detects_record_row_format() -> None: + + class _Statement: + def returns_rows(self) -> bool: + return True + + class _Row: + def __init__(self, data: dict[str, Any]) -> None: + self._data = data + + def keys(self) -> Any: + return self._data.keys() + + def __getitem__(self, key: str) -> Any: + return self._data[key] + + class _Cursor: + def __init__(self) -> None: + self.description = [SimpleNamespace(name="id"), SimpleNamespace(name="name")] + + def fetchall(self) -> list[_Row]: + return [_Row({"id": 1, "name": "alice"})] + + cursor = _Cursor() + result = build_pipeline_execution_result(cast("Any", _Statement()), cursor) + assert result.row_format == "record" + selected_data = result.selected_data + assert selected_data is not None + assert dict(selected_data[0]) == {"id": 1, "name": "alice"} + + def test_build_async_pipeline_execution_result_uses_column_resolver_fast_path() -> None: class _Statement: @@ -127,7 +158,7 @@ def returns_rows(self) -> bool: class _Cursor: def __init__(self) -> None: - self.description = [SimpleNamespace()] + self.description = [SimpleNamespace(name="id"), SimpleNamespace(name="name")] async def fetchall(self) -> list[tuple[int, str]]: return [(1, "alice"), (2, "bob")] @@ -144,6 +175,40 @@ async def _run() -> None: asyncio.run(_run()) +def test_build_async_pipeline_execution_result_detects_record_row_format() -> None: + + class _Statement: + def returns_rows(self) -> bool: + return True + + class _Row: + def __init__(self, data: dict[str, Any]) -> None: + self._data = data + + def keys(self) -> Any: + return self._data.keys() + + def __getitem__(self, key: str) -> Any: + return self._data[key] + + class _Cursor: + def __init__(self) -> None: + self.description = [SimpleNamespace(name="id"), SimpleNamespace(name="name")] + + async def fetchall(self) -> list[_Row]: + return [_Row({"id": 1, "name": "alice"})] + + async def _run() -> None: + cursor = _Cursor() + result = await build_async_pipeline_execution_result(cast("Any", _Statement()), cursor) + assert result.row_format == "record" + selected_data = result.selected_data + assert selected_data is not None + assert dict(selected_data[0]) == {"id": 1, "name": "alice"} + + asyncio.run(_run()) + + class _SyncCopyContext: def __init__(self, rows: list[bytes] | None = None) -> None: self.rows = rows or [] diff --git a/tests/unit/adapters/test_pymssql/conftest.py b/tests/unit/adapters/test_pymssql/conftest.py index e0fbaff81..077076e6e 100644 --- a/tests/unit/adapters/test_pymssql/conftest.py +++ b/tests/unit/adapters/test_pymssql/conftest.py @@ -26,6 +26,7 @@ def __init__( self.rowcount = rowcount self.closed = False self.calls: list[tuple[str, Any]] = [] + self.fetchmany_sizes: list[int] = [] self.many_calls: list[tuple[str, Any]] = [] def execute(self, sql: str, parameters: Any = None) -> None: @@ -37,6 +38,12 @@ def executemany(self, sql: str, parameters: Any = None) -> None: def fetchall(self) -> "list[Any]": return self.rows + def fetchmany(self, size: int) -> "list[Any]": + self.fetchmany_sizes.append(size) + chunk = self.rows[:size] + self.rows = self.rows[size:] + return chunk + def fetchone(self) -> Any: return self.rows[0] if self.rows else None diff --git a/tests/unit/adapters/test_pymssql/test_driver.py b/tests/unit/adapters/test_pymssql/test_driver.py index ea9f4248c..ba33f5cfa 100644 --- a/tests/unit/adapters/test_pymssql/test_driver.py +++ b/tests/unit/adapters/test_pymssql/test_driver.py @@ -113,6 +113,22 @@ def test_collect_rows_returns_column_names() -> None: assert row_count == 1 +def test_select_stream_uses_fetchmany_chunks() -> None: + """The pymssql driver should stream rows with cursor.fetchmany().""" + from sqlspec.adapters.pymssql.driver import PymssqlDriver + + cursor = FakeCursor(rows=[(1, "Ada"), (2, "Grace"), (3, "Linus")], description=[("id",), ("name",)]) + driver = PymssqlDriver(cast("PymssqlConnection", FakeConnection(cursor))) + + with driver.select_stream("SELECT id, name FROM dbo.users", native_only=True, chunk_size=2) as stream: + rows = list(stream) + + assert rows == [{"id": 1, "name": "Ada"}, {"id": 2, "name": "Grace"}, {"id": 3, "name": "Linus"}] + assert cursor.calls == [("SELECT id, name FROM dbo.users", ())] + assert cursor.fetchmany_sizes == [2, 2, 2] + assert cursor.closed is True + + def test_connection_in_transaction_is_false_without_supported_state() -> None: """pymssql does not expose a reliable transaction-state flag.""" from sqlspec.adapters.pymssql.driver import PymssqlDriver diff --git a/tests/unit/adapters/test_spanner/test_core.py b/tests/unit/adapters/test_spanner/test_core.py index 24a59931a..e177b5f9a 100644 --- a/tests/unit/adapters/test_spanner/test_core.py +++ b/tests/unit/adapters/test_spanner/test_core.py @@ -1,8 +1,25 @@ """Unit tests for Spanner core performance helpers.""" from types import SimpleNamespace +from typing import Any, cast -from sqlspec.adapters.spanner.core import build_param_type_signature, collect_rows, resolve_column_names +from google.cloud.spanner_v1.data_types import JsonObject +from google.cloud.spanner_v1.types.type import TypeCode + +from sqlspec.adapters.spanner.core import ( + build_param_type_signature, + collect_rows, + resolve_column_names, + resolve_row_plan, +) + + +def _field(name: str, code: int) -> SimpleNamespace: + return SimpleNamespace(name=name, type_=SimpleNamespace(code=code)) + + +def _json_object(value: str) -> object: + return cast("Any", JsonObject).from_str(value) def test_build_param_type_signature_empty_parameters() -> None: @@ -16,38 +33,77 @@ def test_build_param_type_signature_tracks_key_type_pairs() -> None: assert signature == (("id", int), ("name", str)) -def test_collect_rows_only_converts_string_values() -> None: - class Converter: - def __init__(self) -> None: - self.calls: list[str] = [] +def test_resolve_row_plan_returns_empty_plan_without_metadata() -> None: + cache: dict[int, tuple[object, list[str], tuple[tuple[int, object], ...] | None]] = {} + + column_names, column_plan = resolve_row_plan(None, cache, json_deserializer=lambda value: value) - def convert_if_detected(self, value: str) -> str: - self.calls.append(value) - return value.upper() + assert column_names == [] + assert column_plan is None + assert cache == {} - fields = [SimpleNamespace(name="id"), SimpleNamespace(name="name"), SimpleNamespace(name="score")] - rows = [(1, "alice", 98.5), (2, "bob", 91.0)] - converter = Converter() - data, column_names = collect_rows(rows, fields, converter) +def test_resolve_row_plan_returns_empty_plan_for_non_json_columns() -> None: + fields = [_field("id", TypeCode.INT64), _field("name", TypeCode.STRING)] + cache: dict[int, tuple[object, list[str], tuple[tuple[int, object], ...] | None]] = {} - assert column_names == ["id", "name", "score"] - assert data == [(1, "ALICE", 98.5), (2, "BOB", 91.0)] - assert converter.calls == ["alice", "bob"] + column_names, column_plan = resolve_row_plan(fields, cache, json_deserializer=lambda value: value) + assert column_names == ["id", "name"] + assert column_plan is None + assert len(cache) == 1 -def test_collect_rows_keeps_values_when_converter_makes_no_changes() -> None: - class Converter: - def convert_if_detected(self, value: str) -> str: - return value - fields = [SimpleNamespace(name="id"), SimpleNamespace(name="name")] - rows = [(1, "alice"), (2, "bob")] +def test_resolve_row_plan_caches_json_metadata_and_plan() -> None: + fields = [_field("id", TypeCode.INT64), _field("payload", TypeCode.JSON)] + cache: dict[int, tuple[object, list[str], tuple[tuple[int, object], ...] | None]] = {} - data, column_names = collect_rows(rows, fields, Converter()) + first_column_names, first_column_plan = resolve_row_plan(fields, cache, json_deserializer=lambda value: value) + second_column_names, second_column_plan = resolve_row_plan(fields, cache, json_deserializer=lambda value: value) - assert column_names == ["id", "name"] - assert data == rows + assert first_column_names == ["id", "payload"] + assert first_column_names is second_column_names + assert first_column_plan is not None + assert first_column_plan is second_column_plan + assert first_column_plan[0][0] == 1 + + +def test_collect_rows_returns_original_rows_without_a_plan() -> None: + fields = [_field("id", TypeCode.INT64), _field("payload", TypeCode.STRING)] + rows = [(1, '{"kind":"string"}')] + + data, column_names = collect_rows(rows, fields, column_names=["id", "payload"], column_plan=None) + + assert data is rows + assert column_names == ["id", "payload"] + + +def test_collect_rows_normalizes_list_rows_without_a_plan() -> None: + fields = [_field("id", TypeCode.INT64), _field("payload", TypeCode.STRING)] + rows: list[Any] = [[1, '{"kind":"string"}']] + + data, column_names = collect_rows(rows, fields, column_names=["id", "payload"], column_plan=None) + + assert data == [(1, '{"kind":"string"}')] + assert data is not rows + assert column_names == ["id", "payload"] + + +def test_collect_rows_applies_json_plan_to_metadata_only() -> None: + fields = [_field("id", TypeCode.INT64), _field("payload", TypeCode.JSON), _field("note", TypeCode.STRING)] + rows = [(1, _json_object('{"kind":"payload"}'), '{"kind":"string"}')] + json_calls: list[str] = [] + + def json_deserializer(value: str) -> dict[str, str]: + json_calls.append(value) + return {"decoded": value} + + column_names, column_plan = resolve_row_plan(fields, {}, json_deserializer=json_deserializer) + data, resolved_column_names = collect_rows(rows, fields, column_names=column_names, column_plan=column_plan) + + assert resolved_column_names == ["id", "payload", "note"] + assert data == [(1, {"decoded": '{"kind":"payload"}'}, '{"kind":"string"}')] + assert json_calls == ['{"kind":"payload"}'] def test_resolve_column_names_reuses_cached_fields() -> None: diff --git a/tests/unit/adapters/test_spanner/test_driver.py b/tests/unit/adapters/test_spanner/test_driver.py index c814af3ec..13d6edc32 100644 --- a/tests/unit/adapters/test_spanner/test_driver.py +++ b/tests/unit/adapters/test_spanner/test_driver.py @@ -1,10 +1,15 @@ import inspect -from unittest.mock import MagicMock, Mock +from collections.abc import Iterator, Sequence +from types import SimpleNamespace +from typing import Any, cast +from unittest.mock import MagicMock, Mock, call import pyarrow as pa import pytest from google.cloud.spanner_v1 import Transaction +from google.cloud.spanner_v1.data_types import JsonObject from google.cloud.spanner_v1.streamed import StreamedResultSet +from google.cloud.spanner_v1.types.type import TypeCode from sqlspec.adapters.spanner.driver import SpannerSyncDriver from sqlspec.exceptions import SQLConversionError @@ -34,6 +39,36 @@ def mock_transaction() -> MagicMock: return m +def _field(name: str, code: int) -> SimpleNamespace: + return SimpleNamespace(name=name, type_=SimpleNamespace(code=code)) + + +def _json_object(value: str) -> object: + return cast("Any", JsonObject).from_str(value) + + +def _result_set(fields: list[SimpleNamespace], rows: Sequence[tuple[object, ...]]) -> MagicMock: + result = MagicMock(spec=StreamedResultSet) + result.metadata = SimpleNamespace(row_type=SimpleNamespace(fields=fields)) + result.__iter__.return_value = iter(rows) + return result + + +class _LazyMetadataResultSet: + def __init__(self, fields: list[SimpleNamespace], rows: Sequence[tuple[object, ...]]) -> None: + self._fields = fields + self._rows = rows + self.close_called = False + self.metadata = SimpleNamespace(row_type=SimpleNamespace(fields=[])) + + def __iter__(self) -> Iterator[tuple[object, ...]]: + self.metadata = SimpleNamespace(row_type=SimpleNamespace(fields=self._fields)) + return iter(self._rows) + + def close(self) -> None: + self.close_called = True + + def test_driver_initialization(mock_connection: MagicMock) -> None: driver = SpannerSyncDriver(mock_connection) assert driver.connection == mock_connection @@ -102,6 +137,100 @@ def test_execute_statement_select_minimal_features_omit_retry_and_timeout(mock_c mock_connection.execute_sql.assert_called_once_with("SELECT id FROM users", params=None, param_types={}) +def test_execute_statement_select_uses_metadata_driven_json_conversion(mock_connection: MagicMock) -> None: + json_calls: list[str] = [] + + def json_deserializer(value: str) -> dict[str, str]: + json_calls.append(value) + return {"decoded": value} + + driver = SpannerSyncDriver(mock_connection, driver_features={"json_deserializer": json_deserializer}) + fields = [_field("id", TypeCode.INT64), _field("payload", TypeCode.JSON), _field("note", TypeCode.STRING)] + mock_connection.execute_sql.return_value = _result_set( + fields, [(1, _json_object('{"kind":"payload"}'), '{"kind":"string"}')] + ) + + statement = driver.prepare_statement( + "SELECT id, payload, note FROM users", statement_config=driver.statement_config + ) + result = driver.dispatch_execute(mock_connection, statement) # type: ignore[protected-access] + + assert result.selected_data == [(1, {"decoded": '{"kind":"payload"}'}, '{"kind":"string"}')] + assert result.column_names == ["id", "payload", "note"] + assert json_calls == ['{"kind":"payload"}'] + + +def test_select_stream_uses_native_plan_and_matches_buffered_rows(mock_connection: MagicMock) -> None: + retry = object() + json_calls: list[str] = [] + + def json_deserializer(value: str) -> dict[str, str]: + json_calls.append(value) + return {"decoded": value} + + driver = SpannerSyncDriver( + mock_connection, driver_features={"json_deserializer": json_deserializer, "retry": retry, "timeout": 12.5} + ) + fields = [_field("id", TypeCode.INT64), _field("payload", TypeCode.JSON), _field("note", TypeCode.STRING)] + buffered_rows: list[tuple[object, ...]] = [(1, _json_object('{"kind":"payload"}'), '{"kind":"string"}')] + stream_rows: list[tuple[object, ...]] = [(1, _json_object('{"kind":"payload"}'), '{"kind":"string"}')] + mock_connection.execute_sql.side_effect = [_result_set(fields, buffered_rows), _result_set(fields, stream_rows)] + + buffered = driver.execute("SELECT id, payload, note FROM users").get_data() + with driver.select_stream("SELECT id, payload, note FROM users", chunk_size=1) as stream: + assert stream._source.__class__.__name__ == "_SpannerSelectStreamSource" # pyright: ignore[reportPrivateUsage] + streamed = list(stream) + + assert buffered == [{"id": 1, "payload": {"decoded": '{"kind":"payload"}'}, "note": '{"kind":"string"}'}] + assert streamed == buffered + assert json_calls == ['{"kind":"payload"}', '{"kind":"payload"}'] + assert mock_connection.execute_sql.call_args_list == [ + call("SELECT id, payload, note FROM users", params=None, param_types={}, retry=retry, timeout=12.5), + call("SELECT id, payload, note FROM users", params=None, param_types={}, retry=retry, timeout=12.5), + ] + + +def test_select_stream_defers_metadata_until_rows_are_consumed(mock_connection: MagicMock) -> None: + driver = SpannerSyncDriver(mock_connection) + fields = [_field("id", TypeCode.INT64), _field("name", TypeCode.STRING)] + result_set = _LazyMetadataResultSet(fields, [(1, "Alice")]) + mock_connection.execute_sql.return_value = result_set + + with driver.select_stream("SELECT id, name FROM users", chunk_size=1) as stream: + rows = list(stream) + + assert rows == [{"id": 1, "name": "Alice"}] + assert result_set.close_called is True + + +def test_select_stream_per_call_options_survive_until_lazy_start(mock_connection: MagicMock) -> None: + retry = object() + request_options = {"priority": "high"} + directed_read_options = object() + driver = SpannerSyncDriver(mock_connection) + mock_connection.execute_sql.return_value = _result_set([_field("id", TypeCode.INT64)], [(1,)]) + + stream = driver.select_stream( + "SELECT id FROM users", + chunk_size=1, + request_options=request_options, + directed_read_options=directed_read_options, + retry=retry, + timeout=3.5, + ) + + assert list(stream) == [{"id": 1}] + mock_connection.execute_sql.assert_called_once_with( + "SELECT id FROM users", + params=None, + param_types={}, + request_options=request_options, + directed_read_options=directed_read_options, + retry=retry, + timeout=3.5, + ) + + def test_execute_statement_dml_in_transaction(mock_transaction: MagicMock) -> None: driver = SpannerSyncDriver(mock_transaction) mock_transaction.execute_update.return_value = 10 diff --git a/tests/unit/adapters/test_spanner/test_type_converter.py b/tests/unit/adapters/test_spanner/test_type_converter.py index c607030f6..fed7e791a 100644 --- a/tests/unit/adapters/test_spanner/test_type_converter.py +++ b/tests/unit/adapters/test_spanner/test_type_converter.py @@ -1,44 +1,20 @@ from datetime import datetime, timezone from uuid import UUID -from sqlspec.adapters.spanner.type_converter import SpannerOutputConverter, coerce_params_for_spanner +from sqlspec.adapters.spanner.type_converter import coerce_params_for_spanner, spanner_json, spanner_to_uuid from sqlspec.core import TypedParameter -def test_uuid_conversion() -> None: - """Test UUID string auto-conversion.""" - converter = SpannerOutputConverter(enable_uuid_conversion=True) - uuid_str = "550e8400-e29b-41d4-a716-446655440000" - result = converter.convert(uuid_str) - assert isinstance(result, UUID) - assert str(result) == uuid_str +def test_spanner_to_uuid_converts_bytes() -> None: + uuid_value = UUID("550e8400-e29b-41d4-a716-446655440000") + assert spanner_to_uuid(uuid_value.bytes) == uuid_value -def test_json_detection() -> None: - """Test JSON string auto-detection.""" - converter = SpannerOutputConverter() - json_str = '{"key": "value"}' - result = converter.convert(json_str) - assert isinstance(result, dict) - assert result == {"key": "value"} - - -def test_disabled_uuid_conversion() -> None: - """Test UUID conversion when disabled.""" - converter = SpannerOutputConverter(enable_uuid_conversion=False) - uuid_str = "550e8400-e29b-41d4-a716-446655440000" - result = converter.convert(uuid_str) - assert isinstance(result, str) - assert result == uuid_str +def test_spanner_json_uses_native_json_object() -> None: + result = spanner_json({"key": "value"}) -def test_spanner_output_converter_is_final() -> None: - assert getattr(SpannerOutputConverter, "__final__", False) is True - - -def test_spanner_output_converter_instantiates() -> None: - converter = SpannerOutputConverter() - assert isinstance(converter, SpannerOutputConverter) + assert result == {"key": "value"} def test_coerce_params_unwraps_typed_datetime_parameter() -> None: diff --git a/tests/unit/adapters/test_sqlite/test_config.py b/tests/unit/adapters/test_sqlite/test_config.py index aacb9a5fa..06dc860e2 100644 --- a/tests/unit/adapters/test_sqlite/test_config.py +++ b/tests/unit/adapters/test_sqlite/test_config.py @@ -219,4 +219,4 @@ def test_runtime_keys_removed_from_driver_features() -> None: assert "row_factory" not in config.driver_features assert "pragmas" not in config.driver_features assert "extensions" not in config.driver_features - assert config.driver_features["enable_custom_adapters"] is True + assert config.driver_features["enable_custom_adapters"] is False diff --git a/tests/unit/adapters/test_sqlite/test_driver.py b/tests/unit/adapters/test_sqlite/test_driver.py index 534e75899..30f44dfc5 100644 --- a/tests/unit/adapters/test_sqlite/test_driver.py +++ b/tests/unit/adapters/test_sqlite/test_driver.py @@ -1,6 +1,10 @@ import inspect +from typing import Any, cast + +import pytest from sqlspec.adapters.sqlite.driver import SqliteDriver +from sqlspec.core import SQL def test_driver_cache_execute_cache_hit_has_no_unreachable_returns_rows_guard() -> None: @@ -48,3 +52,39 @@ def test_pool_no_duplicate_typedef_pool_creates_connection_after_cleanup() -> No pool.close() assert row is not None assert row[0] == 1 + + +def test_dispatch_execute_detects_record_row_format(monkeypatch: pytest.MonkeyPatch) -> None: + class _Row: + def __init__(self, data: dict[str, object]) -> None: + self._data = data + + def keys(self) -> object: + return self._data.keys() + + def __getitem__(self, key: str) -> object: + return self._data[key] + + class _Cursor: + def __init__(self) -> None: + self.description = [("id",), ("name",)] + self.executed: list[tuple[str, object]] = [] + + def execute(self, sql: str, parameters: object) -> None: + self.executed.append((sql, parameters)) + + def fetchall(self) -> list[_Row]: + return [_Row({"id": 1, "name": "alice"})] + + monkeypatch.setattr(SqliteDriver, "_compiled_sql", lambda *_args, **_kwargs: ("SELECT id, name FROM users", [])) + + driver = SqliteDriver(connection=cast("Any", object())) + statement = SQL("SELECT id, name FROM users") + cursor = _Cursor() + + result = driver.dispatch_execute(cursor, statement) + + assert result.row_format == "record" + selected_data = result.selected_data + assert selected_data is not None + assert dict(selected_data[0]) == {"id": 1, "name": "alice"} diff --git a/tests/unit/adapters/test_sqlite/test_type_handlers.py b/tests/unit/adapters/test_sqlite/test_type_handlers.py index 4304c15ff..6e6fbe78b 100644 --- a/tests/unit/adapters/test_sqlite/test_type_handlers.py +++ b/tests/unit/adapters/test_sqlite/test_type_handlers.py @@ -42,3 +42,14 @@ def test_unregister_type_handlers_is_noop() -> None: from sqlspec.adapters.sqlite.type_converter import unregister_type_handlers unregister_type_handlers() + + +def test_sqlite_config_does_not_register_type_handlers_by_default() -> None: + import sqlspec.adapters.sqlite.config as config_mod + + register_mock = patch.object(config_mod, "register_type_handlers") + with register_mock as mock_register: + config = config_mod.SqliteConfig(connection_config={"database": ":memory:"}) + config._create_pool() # pyright: ignore[reportPrivateUsage] + + mock_register.assert_not_called() diff --git a/tests/unit/core/test_result.py b/tests/unit/core/test_result.py index 1532ea2f8..d37db7bc3 100644 --- a/tests/unit/core/test_result.py +++ b/tests/unit/core/test_result.py @@ -278,11 +278,50 @@ class UserDict(TypedDict): users = result.get_data(schema_type=UserDict) assert isinstance(users, list) assert len(users) == 2 + assert users is test_data assert isinstance(users[0], dict) assert users[0]["name"] == "Alice" assert users[1]["name"] == "Bob" +def test_sql_result_get_data_with_schema_type_releases_materialized_dict_cache() -> None: + """Schema conversion should release the cached dict materialization.""" + + @dataclass + class User: + id: int + name: str + + sql_stmt = SQL("SELECT id, name FROM users") + raw_tuples: list[Any] = [(1, "Alice"), (2, "Bob")] + result = SQLResult( + statement=sql_stmt, data=raw_tuples, rows_affected=2, column_names=["id", "name"], row_format="tuple" + ) + users = result.get_data(schema_type=User) + assert len(users) == 2 + assert result._materialized_dicts is None + assert result._schema_rows_cache is not None + assert result._schema_rows_cache[User] is users + + +@pytest.mark.skipif(not __debug__, reason="row_format assertions are disabled under optimization") +def test_sql_result_get_data_with_mismatched_row_format_raises() -> None: + """The row_format hook should reject mismatched first-row shapes in debug mode.""" + + @dataclass + class User: + id: int + name: str + + sql_stmt = SQL("SELECT id, name FROM users") + raw_rows: list[Any] = [(1, "Alice")] + result = SQLResult( + statement=sql_stmt, data=raw_rows, rows_affected=1, column_names=["id", "name"], row_format="dict" + ) + with pytest.raises(AssertionError, match="row_format"): + result.get_data(schema_type=User) + + def test_stack_result_from_sql_result() -> None: sql_stmt = SQL("SELECT * FROM users") sql_result = SQLResult(statement=sql_stmt, data=[{"id": 1}], rows_affected=1, metadata={"warning": "slow"}) diff --git a/tests/unit/core/test_type_conversion.py b/tests/unit/core/test_type_conversion.py index 4046939e2..688f808d6 100644 --- a/tests/unit/core/test_type_conversion.py +++ b/tests/unit/core/test_type_conversion.py @@ -1,20 +1,13 @@ -"""Tests for centralized type conversion functionality. +"""Tests for centralized scalar conversion helpers.""" -Tests for the BaseTypeConverter class and related conversion utilities, -ensuring consistent type handling across all database adapters. -""" - -import inspect from datetime import date, datetime, time, timezone from decimal import Decimal -from typing import cast -from uuid import UUID, uuid4 +from uuid import UUID import pytest import sqlspec.utils.serializers._json as json_serialization from sqlspec.core import ( - BaseTypeConverter, convert_decimal, convert_iso_date, convert_iso_datetime, @@ -24,192 +17,6 @@ format_datetime_rfc3339, parse_datetime_rfc3339, ) -from tests.conftest import requires_interpreted - - -@pytest.fixture -def detector() -> "BaseTypeConverter": - """Provide a fresh type converter for each test.""" - return BaseTypeConverter() - - -def test_uuid_detection(detector: "BaseTypeConverter") -> None: - """Test UUID string detection.""" - test_uuid = str(uuid4()) - detected = detector.detect_type(test_uuid) - assert detected == "uuid" - - -def test_uuid_conversion(detector: "BaseTypeConverter") -> None: - """Test UUID string conversion.""" - test_uuid_str = "123e4567-e89b-12d3-a456-426614174000" - detected = detector.detect_type(test_uuid_str) - assert detected == "uuid" - - if detected is not None: - converted = detector.convert_value(test_uuid_str, detected) - assert isinstance(converted, UUID) - assert str(converted) == test_uuid_str - - -def test_datetime_detection(detector: "BaseTypeConverter") -> None: - """Test ISO datetime detection.""" - test_cases = [ - "2023-12-25T10:30:00Z", - "2023-12-25T10:30:00+00:00", - "2023-12-25T10:30:00.123456Z", - "2023-12-25 10:30:00", - ] - - for dt_str in test_cases: - detected = detector.detect_type(dt_str) - assert detected == "iso_datetime", f"Failed for: {dt_str}" - - -def test_datetime_conversion(detector: "BaseTypeConverter") -> None: - """Test ISO datetime conversion.""" - dt_str = "2023-12-25T10:30:00Z" - detected = detector.detect_type(dt_str) - assert detected == "iso_datetime" - - if detected is not None: - converted = detector.convert_value(dt_str, detected) - assert isinstance(converted, datetime) - assert converted.year == 2023 - assert converted.month == 12 - assert converted.day == 25 - - -def test_date_detection(detector: "BaseTypeConverter") -> None: - """Test ISO date detection.""" - date_str = "2023-12-25" - detected = detector.detect_type(date_str) - assert detected == "iso_date" - - -def test_date_conversion(detector: "BaseTypeConverter") -> None: - """Test ISO date conversion.""" - date_str = "2023-12-25" - detected = detector.detect_type(date_str) - assert detected == "iso_date" - - if detected is not None: - converted = detector.convert_value(date_str, detected) - assert isinstance(converted, date) - assert converted.year == 2023 - assert converted.month == 12 - assert converted.day == 25 - - -def test_time_detection(detector: "BaseTypeConverter") -> None: - """Test ISO time detection.""" - time_str = "10:30:00" - detected = detector.detect_type(time_str) - assert detected == "iso_time" - - -def test_time_conversion(detector: "BaseTypeConverter") -> None: - """Test ISO time conversion.""" - time_str = "10:30:45" - detected = detector.detect_type(time_str) - assert detected == "iso_time" - - if detected is not None: - converted = detector.convert_value(time_str, detected) - assert isinstance(converted, time) - assert converted.hour == 10 - assert converted.minute == 30 - assert converted.second == 45 - - -def test_json_detection(detector: "BaseTypeConverter") -> None: - """Test JSON string detection.""" - test_cases = ['{"key": "value"}', "[1, 2, 3]", '{"nested": {"data": true}}'] - - for json_str in test_cases: - detected = detector.detect_type(json_str) - assert detected == "json", f"Failed for: {json_str}" - - -def test_json_conversion(detector: "BaseTypeConverter") -> None: - """Test JSON string conversion.""" - json_str = '{"key": "value", "number": 42}' - detected = detector.detect_type(json_str) - assert detected == "json" - - if detected is not None: - converted = detector.convert_value(json_str, detected) - assert isinstance(converted, dict) - assert converted["key"] == "value" - assert converted["number"] == 42 - - -def test_ipv4_detection(detector: "BaseTypeConverter") -> None: - """Test IPv4 address detection.""" - ip_str = "192.168.1.1" - detected = detector.detect_type(ip_str) - assert detected == "ipv4" - - -def test_ipv6_detection(detector: "BaseTypeConverter") -> None: - """Test IPv6 address detection.""" - ip_str = "2001:0db8:85a3:0000:0000:8a2e:0370:7334" - detected = detector.detect_type(ip_str) - assert detected == "ipv6" - - -def test_mac_address_detection(detector: "BaseTypeConverter") -> None: - """Test MAC address detection.""" - mac_str = "00:1B:44:11:3A:B7" - detected = detector.detect_type(mac_str) - assert detected == "mac" - - -def test_detect_type_uses_lastgroup_fast_path(detector: "BaseTypeConverter") -> None: - """Detection should preserve outputs without allocating a groupdict per match.""" - samples = { - "uuid": "123e4567-e89b-12d3-a456-426614174000", - "iso_datetime": "2023-12-25T10:30:00Z", - "iso_date": "2023-12-25", - "iso_time": "10:30:00", - "json": '{"key": "value"}', - "ipv4": "192.168.1.1", - "ipv6": "2001:0db8:85a3:0000:0000:8a2e:0370:7334", - "mac": "00:1B:44:11:3A:B7", - } - - for expected_type, value in samples.items(): - assert detector.detect_type(value) == expected_type - - source = inspect.getsource(BaseTypeConverter.detect_type) - assert ".lastgroup" in source - assert "groupdict()" not in source - - -def test_non_special_type(detector: "BaseTypeConverter") -> None: - """Test that regular strings are not detected as special types.""" - regular_strings = ["hello world", "123", "not-a-uuid", "invalid-date"] - - for string in regular_strings: - detected = detector.detect_type(string) - assert detected is None, f"Incorrectly detected: {string}" - - -@requires_interpreted -def test_none_input(detector: "BaseTypeConverter") -> None: - """Test that None input is handled correctly.""" - detected = detector.detect_type(cast("str", None)) - assert detected is None - - -@requires_interpreted -def test_non_string_input(detector: "BaseTypeConverter") -> None: - """Test that non-string input is handled correctly.""" - non_strings = [123, [], {}, True] - - for value in non_strings: - detected = detector.detect_type(cast("str", value)) - assert detected is None def test_convert_uuid() -> None: @@ -323,7 +130,7 @@ def fake_convert(value: str) -> datetime: assert type_converter.parse_datetime_rfc3339("2023-12-25T10:30:00Z") is expected -def test_datetime_round_trip(detector: "BaseTypeConverter") -> None: +def test_datetime_round_trip() -> None: """Test datetime formatting and parsing round trip.""" original = datetime(2023, 12, 25, 10, 30, 0, tzinfo=timezone.utc) formatted = format_datetime_rfc3339(original) @@ -351,24 +158,3 @@ def test_invalid_json() -> None: invalid_json = "not valid json" with pytest.raises((ValueError, Exception)): convert_json(invalid_json) - - -def test_empty_string(detector: "BaseTypeConverter") -> None: - """Test empty string handling.""" - detected = detector.detect_type("") - assert detected is None - - -def test_whitespace_string(detector: "BaseTypeConverter") -> None: - """Test whitespace-only string.""" - detected = detector.detect_type(" ") - assert detected is None - - -def test_case_insensitive_patterns(detector: "BaseTypeConverter") -> None: - """Test that patterns work case-insensitively where appropriate.""" - uuid_lower = "123e4567-e89b-12d3-a456-426614174000" - uuid_upper = "123E4567-E89B-12D3-A456-426614174000" - - assert detector.detect_type(uuid_lower) == "uuid" - assert detector.detect_type(uuid_upper) == "uuid" diff --git a/tests/unit/core/test_type_converter_dead_code.py b/tests/unit/core/test_type_converter_dead_code.py new file mode 100644 index 000000000..55af31403 --- /dev/null +++ b/tests/unit/core/test_type_converter_dead_code.py @@ -0,0 +1,86 @@ +"""Dead-code regression tests for the adapter type converter sweep.""" + +import importlib + +import pytest + +import sqlspec.core as core +import sqlspec.core.type_converter as type_converter +from sqlspec.adapters.oracledb.type_converter import OracleOutputConverter + + +def _name(*parts: str) -> str: + return "".join(parts) + + +def test_core_type_converter_module_no_longer_exposes_cached_path() -> None: + """The shared converter module should no longer export the retired cache path.""" + cached_output = _name("Cached", "Output", "Converter") + cached_class = _name("_", "Cached", "Converter") + cached_factory = _name("_make_", "cached", "_converter") + + assert not hasattr(core, cached_output) + assert not hasattr(core, "DEFAULT_CACHE_SIZE") + assert not hasattr(core, "DEFAULT_SPECIAL_CHARS") + + assert not hasattr(type_converter, cached_output) + assert not hasattr(type_converter, cached_class) + assert not hasattr(type_converter, cached_factory) + assert not hasattr(type_converter, "DEFAULT_CACHE_SIZE") + assert not hasattr(type_converter, "DEFAULT_SPECIAL_CHARS") + assert cached_output not in type_converter.__all__ + assert "DEFAULT_CACHE_SIZE" not in type_converter.__all__ + assert "DEFAULT_SPECIAL_CHARS" not in type_converter.__all__ + + +def test_core_type_converter_module_no_longer_exposes_content_sniffing_surface() -> None: + """The shared regex-driven content-detection surface is retired in C7.""" + for name in ( + _name("Base", "Type", "Converter"), + "DEFAULT_DETECTION_CHARS", + _name("SPECIAL", "_TYPE", "_REGEX"), + "_TYPE_CONVERTERS", + ): + assert not hasattr(core, name) + assert not hasattr(type_converter, name) + assert name not in type_converter.__all__ + + +def test_adapter_detection_methods_are_removed() -> None: + """Adapter converters must not retain content-sniffing methods after C7.""" + from sqlspec.adapters.adbc.type_converter import ADBCOutputConverter + + detect_method = _name("detect", "_type") + convert_method = _name("convert", "_value") + convert_if_method = _name("convert", "_if", "_detected") + dialect_method = _name("get", "_dialect", "_specific", "_converter") + + for converter_cls in (ADBCOutputConverter, OracleOutputConverter): + assert not hasattr(converter_cls, convert_if_method) + assert not hasattr(converter_cls, detect_method) + assert not hasattr(converter_cls, convert_method) + + assert not hasattr(ADBCOutputConverter, dialect_method) + + +def test_bigquery_type_converter_module_is_removed() -> None: + """The dead BigQuery output converter module should no longer be importable.""" + module_name = _name("sqlspec.adapters.bigquery.", "type", "_converter") + with pytest.raises(ModuleNotFoundError): + importlib.import_module(module_name) + + +def test_duckdb_type_converter_module_is_removed() -> None: + """DuckDB no longer keeps a string-content-sniffing converter module.""" + module_name = _name("sqlspec.adapters.duckdb.", "type", "_converter") + with pytest.raises(ModuleNotFoundError): + importlib.import_module(module_name) + + +def test_spanner_output_converter_is_removed() -> None: + """Spanner fetch conversion is column-plan driven, not content-sniffing class driven.""" + import sqlspec.adapters.spanner.type_converter as spanner_type_converter + + converter_name = _name("Spanner", "Output", "Converter") + assert not hasattr(spanner_type_converter, converter_name) + assert converter_name not in spanner_type_converter.__all__ diff --git a/tests/unit/driver/test_result_tools.py b/tests/unit/driver/test_result_tools.py index c7b4c3639..b3513f545 100644 --- a/tests/unit/driver/test_result_tools.py +++ b/tests/unit/driver/test_result_tools.py @@ -320,7 +320,7 @@ def test_to_schema_mixin_with_typeddict_single_record() -> None: def test_to_schema_mixin_with_typeddict_multiple_records() -> None: - """Test CommonDriverAttributesMixin.to_schema with TypedDict for multiple records.""" + """TypedDict list conversion should preserve the original list object.""" test_data = [ {"name": "user1", "age": 25, "optional_field": "value1"}, {"name": "user2", "age": 30, "optional_field": "value2"}, @@ -330,25 +330,26 @@ def test_to_schema_mixin_with_typeddict_multiple_records() -> None: assert isinstance(result, list) assert len(result) == 2 + assert result is test_data for item in result: assert isinstance(item, dict) assert result == test_data def test_to_schema_mixin_with_typeddict_mixed_data() -> None: - """Test CommonDriverAttributesMixin.to_schema with TypedDict filters non-dict items.""" + """TypedDict list conversion should preserve non-dict values when present.""" test_data = [ {"name": "user1", "age": 25, "optional_field": "value1"}, - "not_a_dict", # This should be filtered out + "not_a_dict", {"name": "user2", "age": 30, "optional_field": "value2"}, ] result = CommonDriverAttributesMixin.to_schema(test_data, schema_type=SampleTypedDict) assert isinstance(result, list) - assert len(result) == 2 # Only dict items should be included - for item in result: - assert isinstance(item, dict) + assert len(result) == 3 + assert result is test_data + assert result[1] == "not_a_dict" def test_to_schema_mixin_with_typeddict_non_dict_data() -> None: diff --git a/tests/unit/utils/test_bench.py b/tests/unit/utils/test_bench.py index 33c8c9d2d..e81f8ba22 100644 --- a/tests/unit/utils/test_bench.py +++ b/tests/unit/utils/test_bench.py @@ -35,6 +35,19 @@ def test_bench_summary_reports_stddev_iqr_and_noise() -> None: assert summary["noisy"] is False +def test_schema_type_numpy_benchmark_scenario_is_registered_and_runnable() -> None: + module = _load_script_module("bench.py", "bench_for_tests") + + assert module.SCENARIO_REGISTRY[("raw", "sqlite", "schema_type_numpy")] is module.raw_sqlite_schema_type_numpy + assert module.SCENARIO_REGISTRY[("sqlspec", "sqlite", "schema_type_numpy")] is ( + module.sqlspec_sqlite_schema_type_numpy + ) + + module.raw_sqlite_schema_type_numpy() + module.sqlspec_sqlite_schema_type_numpy() + module.assert_schema_type_numpy_vector_fallback() + + def test_async_benchmark_iterations_reuse_one_event_loop() -> None: module = _load_script_module("bench.py", "bench_for_tests") loop_ids: list[int] = [] diff --git a/tests/unit/utils/test_serializers.py b/tests/unit/utils/test_serializers.py index 0efa9a3e0..32cc8e490 100644 --- a/tests/unit/utils/test_serializers.py +++ b/tests/unit/utils/test_serializers.py @@ -759,6 +759,35 @@ def count_fields(schema_type: type[Any]) -> Any: assert call_count == 1 +def test_schema_dump_dataclass_fields_are_precomputed_per_serializer(monkeypatch: pytest.MonkeyPatch) -> None: + """Serializer registration should precompute dataclass field metadata once.""" + + from sqlspec.utils.serializers import _schema as schema_module + + @dataclass + class _User: + user_id: str + display_name: str + + original_fields = schema_module.dataclasses_fields + call_count = 0 + + def count_fields(schema_type: type[Any]) -> Any: + nonlocal call_count + call_count += 1 + return original_fields(schema_type) + + reset_serializer_cache() + monkeypatch.setattr(schema_module, "dataclasses_fields", count_fields) + + first = schema_dump(_User(user_id="first", display_name="Alice"), exclude_unset=False) + second = schema_dump(_User(user_id="second", display_name="Bob"), exclude_unset=False) + + assert first == {"user_id": "first", "display_name": "Alice"} + assert second == {"user_id": "second", "display_name": "Bob"} + assert call_count == 1 + + def test_schema_dump_wire_format_opt_out_msgspec_dump_modes_use_one_parameterized_function() -> None: """All msgspec serializer modes should share one dump implementation.""" import msgspec diff --git a/tests/unit/utils/test_to_value_type.py b/tests/unit/utils/test_to_value_type.py index 5739a6cdc..a5b71dcb5 100644 --- a/tests/unit/utils/test_to_value_type.py +++ b/tests/unit/utils/test_to_value_type.py @@ -205,6 +205,47 @@ def test_convert_numpy_recursive_preserves_tuple_shape() -> None: assert converted == {"items": ([1, 2], {"values": [3, 4]})} +def test_msgspec_conversion_does_not_walk_numpy_for_plain_payload(monkeypatch: pytest.MonkeyPatch) -> None: + """Plain msgspec conversion should not recurse through numpy helpers.""" + if not schema_utils.NUMPY_INSTALLED: + pytest.skip("numpy is not installed") + + class User(msgspec.Struct): + id: int + name: str + + def fail_walk(_obj: object) -> object: + msg = "numpy walk should not run for plain msgspec payloads" + raise AssertionError(msg) + + monkeypatch.setattr(schema_utils, "_convert_numpy_recursive", fail_walk) + result = schema_utils._convert_msgspec([{"id": 1, "name": "Alice"}], User) + assert result == [User(id=1, name="Alice")] + + +@pytest.mark.skipif(not schema_utils.NUMPY_INSTALLED, reason="numpy is not installed") +def test_msgspec_conversion_falls_back_to_numpy_walk_for_ndarray_payload(monkeypatch: pytest.MonkeyPatch) -> None: + """Ndarray payloads should still convert through the numpy fallback path.""" + import numpy as np + + class Measurement(msgspec.Struct): + values: list[float] + + original_walk = schema_utils._convert_numpy_recursive + call_count = 0 + + def count_walk(obj: object) -> object: + nonlocal call_count + if isinstance(obj, list): + call_count += 1 + return original_walk(obj) + + monkeypatch.setattr(schema_utils, "_convert_numpy_recursive", count_walk) + result = schema_utils._convert_msgspec([{"values": np.array([1.0, 2.0])}], Measurement) + assert result == [Measurement(values=[1.0, 2.0])] + assert call_count == 1 + + def test_int_conversion_float_to_int() -> None: """Float truncates to int.""" assert to_value_type(3.7, int) == 3 diff --git a/tests/unit/utils/test_type_guards.py b/tests/unit/utils/test_type_guards.py index 48a8b2ff3..db21a9c7d 100644 --- a/tests/unit/utils/test_type_guards.py +++ b/tests/unit/utils/test_type_guards.py @@ -970,6 +970,25 @@ def test_get_msgspec_rename_config_with_pascal_rename() -> None: assert result == "pascal" +def test_get_msgspec_rename_config_caches_per_type(monkeypatch: pytest.MonkeyPatch) -> None: + """Repeated rename-config lookups for one type should hit the cache.""" + original_fields = msgspec.structs.fields + call_count = 0 + + class CachedMsgspecStruct(msgspec.Struct, rename="camel"): + test_name: str = "test" + + def count_fields(schema_type: Any) -> Any: + nonlocal call_count + call_count += 1 + return original_fields(schema_type) + + monkeypatch.setattr(msgspec.structs, "fields", count_fields) + assert get_msgspec_rename_config(CachedMsgspecStruct) == "camel" + assert get_msgspec_rename_config(CachedMsgspecStruct) == "camel" + assert call_count == 1 + + def test_is_typed_dict_with_typeddict_class() -> None: """Test is_typed_dict returns True for TypedDict classes.""" assert is_typed_dict(SampleTypedDict) is True diff --git a/tools/scripts/bench.py b/tools/scripts/bench.py index 617100c4f..08fe554b9 100644 --- a/tools/scripts/bench.py +++ b/tools/scripts/bench.py @@ -8,12 +8,13 @@ import gc import inspect import json +import os import pstats import sqlite3 import statistics import tempfile import time -from collections.abc import Sequence # noqa: TC003 +from collections.abc import Sequence from contextlib import suppress from pathlib import Path from typing import Any, TypedDict @@ -26,12 +27,13 @@ from sqlspec import SQLSpec from sqlspec.adapters.duckdb import DuckDBConfig from sqlspec.adapters.sqlite import SqliteConfig -from sqlspec.utils.schema import transform_dict_keys +from sqlspec.utils.schema import _convert_numpy_recursive, to_schema, transform_dict_keys from sqlspec.utils.text import camelize __all__ = ( "main", "print_benchmark_table", + "raw_aiosqlite_worker_hops", "raw_asyncpg_initialization", "raw_asyncpg_read_heavy", "raw_asyncpg_write_heavy", @@ -40,6 +42,8 @@ "raw_duckdb_read_heavy", "raw_duckdb_repeated_queries", "raw_duckdb_write_heavy", + "raw_oracle_lob_fetch_1k", + "raw_oracle_lob_fetch_100k", "raw_sqlite_complex_parameters", "raw_sqlite_dict_key_transform", "raw_sqlite_initialization", @@ -47,6 +51,7 @@ "raw_sqlite_read_heavy", "raw_sqlite_repeated_queries", "raw_sqlite_schema_mapping", + "raw_sqlite_schema_type_numpy", "raw_sqlite_thin_path_stress", "raw_sqlite_write_heavy", "run_benchmark", @@ -64,6 +69,7 @@ "sqlalchemy_sqlite_read_heavy", "sqlalchemy_sqlite_repeated_queries", "sqlalchemy_sqlite_write_heavy", + "sqlspec_aiosqlite_worker_hops", "sqlspec_asyncpg_initialization", "sqlspec_asyncpg_read_heavy", "sqlspec_asyncpg_write_heavy", @@ -72,6 +78,14 @@ "sqlspec_duckdb_read_heavy", "sqlspec_duckdb_repeated_queries", "sqlspec_duckdb_write_heavy", + "sqlspec_oracle_lob_fetch_1k", + "sqlspec_oracle_lob_fetch_100k", + "sqlspec_oracle_lob_fetch_async_1k", + "sqlspec_oracle_lob_fetch_async_100k", + "sqlspec_oracle_lob_fetch_async_fetch_lobs_true_1k", + "sqlspec_oracle_lob_fetch_async_fetch_lobs_true_100k", + "sqlspec_oracle_lob_fetch_fetch_lobs_true_1k", + "sqlspec_oracle_lob_fetch_fetch_lobs_true_100k", "sqlspec_sqlite_complex_parameters", "sqlspec_sqlite_dict_key_transform", "sqlspec_sqlite_initialization", @@ -79,6 +93,7 @@ "sqlspec_sqlite_read_heavy", "sqlspec_sqlite_repeated_queries", "sqlspec_sqlite_schema_mapping", + "sqlspec_sqlite_schema_type_numpy", "sqlspec_sqlite_thin_path_stress", "sqlspec_sqlite_write_heavy", ) @@ -105,6 +120,47 @@ def _is_compiled() -> bool: DEFAULT_BENCH_ITERATIONS = 7 DEFAULT_BENCH_WARMUP = 3 NOISY_STDDEV_RATIO = 0.10 +CORE_LIBRARIES = ("raw", "sqlspec", "sqlalchemy") +CORE_SCENARIOS = ("initialization", "write_heavy", "read_heavy", "iterative_inserts", "repeated_queries") +ORACLE_LOB_ROWS = 100 +ORACLE_LOB_PAYLOAD_SIZES = {"1k": 1024, "100k": 100 * 1024} +ORACLE_LOB_ENV_VARS = ( + "SQLSPEC_BENCH_ORACLE_HOST", + "SQLSPEC_BENCH_ORACLE_PORT", + "SQLSPEC_BENCH_ORACLE_SERVICE_NAME", + "SQLSPEC_BENCH_ORACLE_USER", + "SQLSPEC_BENCH_ORACLE_PASSWORD", +) +SQLITE_EXTENDED_SCENARIOS = ( + ("raw", "dict_key_transform"), + ("sqlspec", "dict_key_transform"), + ("raw", "schema_mapping"), + ("sqlspec", "schema_mapping"), + ("raw", "complex_parameters"), + ("sqlspec", "complex_parameters"), + ("raw", "thin_path_stress"), + ("sqlspec", "thin_path_stress"), + ("raw", "schema_type_numpy"), + ("sqlspec", "schema_type_numpy"), +) +ORACLE_EXTENDED_SCENARIOS = ( + ("raw", "lob_fetch_1k"), + ("sqlspec", "lob_fetch_1k"), + ("sqlspec_fetch_lobs_true", "lob_fetch_1k"), + ("sqlspec_async", "lob_fetch_1k"), + ("sqlspec_async_fetch_lobs_true", "lob_fetch_1k"), + ("raw", "lob_fetch_100k"), + ("sqlspec", "lob_fetch_100k"), + ("sqlspec_fetch_lobs_true", "lob_fetch_100k"), + ("sqlspec_async", "lob_fetch_100k"), + ("sqlspec_async_fetch_lobs_true", "lob_fetch_100k"), +) +AIOSQLITE_EXTENDED_SCENARIOS = (("raw", "worker_hops"), ("sqlspec", "worker_hops")) +EXTENDED_SCENARIOS_BY_DRIVER = { + "sqlite": SQLITE_EXTENDED_SCENARIOS, + "aiosqlite": AIOSQLITE_EXTENDED_SCENARIOS, + "oracle": ORACLE_EXTENDED_SCENARIOS, +} @click.command() @@ -167,13 +223,14 @@ def main( f"Running benchmark for driver: {drv} " f"(rows={rows}, pool_size={pool_size}, iterations={iterations}, warmup={warmup})" ) - if profile: + has_core_scenarios = _driver_has_core_scenarios(drv) + if profile and has_core_scenarios: results.extend( run_benchmark_profiled( drv, errors, iterations=iterations, warmup=warmup, profile_scenario=profile_scenario ) ) - else: + elif has_core_scenarios or not extended: results.extend(run_benchmark(drv, errors, iterations=iterations, warmup=warmup)) if extended: click.echo(f"Running extended benchmarks for driver: {drv}") @@ -265,6 +322,26 @@ def _summarize_times(times: list[float]) -> dict[str, float | bool]: } +def _benchmark_library_label(library: str) -> str: + """Return a display label for a benchmark library key.""" + if library == "sqlspec": + return SQLSPEC_LABEL + if library == "sqlspec_fetch_lobs_true": + return f"{SQLSPEC_LABEL} fetch_lobs=True" + if library == "sqlspec_async": + return f"{SQLSPEC_LABEL} async" + if library == "sqlspec_async_fetch_lobs_true": + return f"{SQLSPEC_LABEL} async fetch_lobs=True" + return library + + +def _driver_has_core_scenarios(driver: str) -> bool: + """Return True when a driver has any core benchmark scenario registered.""" + return any( + (library, driver, scenario) in SCENARIO_REGISTRY for scenario in CORE_SCENARIOS for library in CORE_LIBRARIES + ) + + def run_benchmark( driver: str, errors: list[str], *, iterations: int = DEFAULT_BENCH_ITERATIONS, warmup: int = DEFAULT_BENCH_WARMUP ) -> list[dict[str, Any]]: @@ -279,12 +356,10 @@ def run_benchmark( Returns: List of benchmark result dictionaries """ - libraries = ["raw", "sqlspec", "sqlalchemy"] - scenarios = ["initialization", "write_heavy", "read_heavy", "iterative_inserts", "repeated_queries"] results: list[dict[str, Any]] = [] - for scenario in scenarios: - for lib in libraries: + for scenario in CORE_SCENARIOS: + for lib in CORE_LIBRARIES: func = SCENARIO_REGISTRY.get((lib, driver, scenario)) if func is None: errors.append(f"No implementation for library={lib}, driver={driver}, scenario={scenario}") @@ -295,7 +370,7 @@ def run_benchmark( try: times = _run_benchmark_iterations(func, is_async=is_async, iterations=iterations, warmup=warmup) stats = _summarize_times(times) - label = SQLSPEC_LABEL if lib == "sqlspec" else lib + label = _benchmark_library_label(lib) results.append({"driver": driver, "library": label, "scenario": scenario, "times": times, **stats}) except Exception as exc: errors.append(f"{lib}/{driver}/{scenario}: {exc}") @@ -329,13 +404,11 @@ def run_benchmark_profiled( profiles_dir = Path(__file__).parent / "profiles" profiles_dir.mkdir(exist_ok=True) - libraries = ["raw", "sqlspec", "sqlalchemy"] - scenarios = ["initialization", "write_heavy", "read_heavy", "iterative_inserts", "repeated_queries"] results: list[dict[str, Any]] = [] console = Console() - for scenario in scenarios: - for lib in libraries: + for scenario in CORE_SCENARIOS: + for lib in CORE_LIBRARIES: # If profiling a specific scenario, skip others if profile_scenario and scenario != profile_scenario: continue @@ -346,7 +419,6 @@ def run_benchmark_profiled( continue is_async = inspect.iscoroutinefunction(func) - label = SQLSPEC_LABEL if lib == "sqlspec" else lib prof_name = f"{driver}_{lib}_{scenario}" try: @@ -355,7 +427,13 @@ def run_benchmark_profiled( func, is_async=is_async, iterations=iterations, warmup=warmup, profiler=profiler ) summary = _summarize_times(times) - results.append({"driver": driver, "library": label, "scenario": scenario, "times": times, **summary}) + results.append({ + "driver": driver, + "library": _benchmark_library_label(lib), + "scenario": scenario, + "times": times, + **summary, + }) # Save profile data prof_path = profiles_dir / f"{prof_name}.prof" @@ -393,26 +471,24 @@ def run_extended_benchmark( Returns: List of benchmark result dictionaries """ - libraries = ["raw", "sqlspec"] - scenarios = ["dict_key_transform", "schema_mapping", "complex_parameters", "thin_path_stress"] + scenario_entries = EXTENDED_SCENARIOS_BY_DRIVER.get(driver, SQLITE_EXTENDED_SCENARIOS) results: list[dict[str, Any]] = [] - for scenario in scenarios: - for lib in libraries: - func = SCENARIO_REGISTRY.get((lib, driver, scenario)) - if func is None: - errors.append(f"No extended implementation for library={lib}, driver={driver}, scenario={scenario}") - continue + for lib, scenario in scenario_entries: + func = SCENARIO_REGISTRY.get((lib, driver, scenario)) + if func is None: + errors.append(f"No extended implementation for library={lib}, driver={driver}, scenario={scenario}") + continue - is_async = inspect.iscoroutinefunction(func) + is_async = inspect.iscoroutinefunction(func) - try: - times = _run_benchmark_iterations(func, is_async=is_async, iterations=iterations, warmup=warmup) - stats = _summarize_times(times) - label = SQLSPEC_LABEL if lib == "sqlspec" else lib - results.append({"driver": driver, "library": label, "scenario": scenario, "times": times, **stats}) - except Exception as exc: - errors.append(f"{lib}/{driver}/{scenario}: {exc}") + try: + times = _run_benchmark_iterations(func, is_async=is_async, iterations=iterations, warmup=warmup) + stats = _summarize_times(times) + label = _benchmark_library_label(lib) + results.append({"driver": driver, "library": label, "scenario": scenario, "times": times, **stats}) + except Exception as exc: + errors.append(f"{lib}/{driver}/{scenario}: {exc}") return results @@ -1259,6 +1335,35 @@ async def sqlspec_aiosqlite_repeated_queries() -> None: await anyio.Path(tmp_path).unlink() +async def raw_aiosqlite_worker_hops() -> None: + aiosqlite = _get_aiosqlite() + if aiosqlite is None: + return + async with aiosqlite.connect(":memory:") as conn: + for i in range(ROWS_TO_INSERT): + cursor = await conn.execute("SELECT ?", (i,)) + row = await cursor.fetchone() + assert row[0] == i + + +async def sqlspec_aiosqlite_worker_hops() -> None: + AiosqliteConfig = _get_aiosqlite_config() # noqa: N806 + if AiosqliteConfig is None: + return + spec = SQLSpec() + config = AiosqliteConfig(database=":memory:", pool_size=1) + try: + async with spec.provide_session(config) as session: + for i in range(ROWS_TO_INSERT): + value = await session.select_value("SELECT ?", i) + assert value == i + _check_pool_leak(config.connection_instance, "aiosqlite/worker_hops") + await config.close_pool() + finally: + if config.connection_instance is not None: + await config.close_pool() + + async def sqlalchemy_aiosqlite_initialization() -> None: create_async_engine, text = _get_async_sqlalchemy() if create_async_engine is None: @@ -1562,6 +1667,64 @@ class WideRowDict(TypedDict): is_active: str +SCHEMA_TYPE_NUMPY_ROW_COUNT = 5_000 +SCHEMA_TYPE_NUMPY_FIRST_FIELD_4 = 4 +_SCHEMA_TYPE_NUMPY_PAYLOAD: list[dict[str, int]] | None = None +_SCHEMA_TYPE_NUMPY_ROW_TYPE: type[Any] | None = None +_SCHEMA_TYPE_NUMPY_VECTOR_TYPE: type[Any] | None = None + + +def _schema_type_numpy_payload() -> list[dict[str, int]]: + """Return the stable 5000x5 payload used by the schema_type numpy benchmark.""" + global _SCHEMA_TYPE_NUMPY_PAYLOAD + if _SCHEMA_TYPE_NUMPY_PAYLOAD is None: + _SCHEMA_TYPE_NUMPY_PAYLOAD = [ + {"field_0": index, "field_1": index + 1, "field_2": index + 2, "field_3": index + 3, "field_4": index + 4} + for index in range(SCHEMA_TYPE_NUMPY_ROW_COUNT) + ] + return _SCHEMA_TYPE_NUMPY_PAYLOAD + + +def _schema_type_numpy_row_type() -> type[Any]: + """Return the cached msgspec row type used by the schema_type numpy benchmark.""" + global _SCHEMA_TYPE_NUMPY_ROW_TYPE + if _SCHEMA_TYPE_NUMPY_ROW_TYPE is None: + import msgspec + + class SchemaTypeNumpyRow(msgspec.Struct): + field_0: int + field_1: int + field_2: int + field_3: int + field_4: int + + _SCHEMA_TYPE_NUMPY_ROW_TYPE = SchemaTypeNumpyRow + return _SCHEMA_TYPE_NUMPY_ROW_TYPE + + +def _schema_type_numpy_vector_type() -> type[Any]: + """Return the cached msgspec vector row type used to assert ndarray fallback behavior.""" + global _SCHEMA_TYPE_NUMPY_VECTOR_TYPE + if _SCHEMA_TYPE_NUMPY_VECTOR_TYPE is None: + import msgspec + + class SchemaTypeNumpyVectorRow(msgspec.Struct): + values: list[float] + + _SCHEMA_TYPE_NUMPY_VECTOR_TYPE = SchemaTypeNumpyVectorRow + return _SCHEMA_TYPE_NUMPY_VECTOR_TYPE + + +def assert_schema_type_numpy_vector_fallback() -> None: + """Assert ndarray -> list[float] fallback still works for msgspec schema conversion.""" + import numpy as np + + vector_type = _schema_type_numpy_vector_type() + converted = to_schema([{"values": np.array([1.0, 2.0])}], schema_type=vector_type) + assert len(converted) == 1 + assert converted[0].values == [1.0, 2.0] + + def _generate_wide_row(i: int) -> tuple[str, ...]: """Generate a single row tuple for the wide_test table.""" return ( @@ -1662,6 +1825,25 @@ def sqlspec_sqlite_schema_mapping() -> None: assert "first_name" in first +# --- schema_type_numpy scenarios --- + + +def raw_sqlite_schema_type_numpy() -> None: + """Legacy baseline: pre-walk a 5000x5 payload before msgspec schema conversion.""" + row_type = _schema_type_numpy_row_type() + rows = to_schema(_convert_numpy_recursive(_schema_type_numpy_payload()), schema_type=row_type) + assert len(rows) == SCHEMA_TYPE_NUMPY_ROW_COUNT + assert rows[0].field_4 == SCHEMA_TYPE_NUMPY_FIRST_FIELD_4 + + +def sqlspec_sqlite_schema_type_numpy() -> None: + """Convert a 5000x5 payload through the current msgspec schema_type path.""" + row_type = _schema_type_numpy_row_type() + rows = to_schema(_schema_type_numpy_payload(), schema_type=row_type) + assert len(rows) == SCHEMA_TYPE_NUMPY_ROW_COUNT + assert rows[0].field_4 == SCHEMA_TYPE_NUMPY_FIRST_FIELD_4 + + # --- complex_parameters scenarios --- @@ -1723,6 +1905,226 @@ def sqlspec_sqlite_thin_path_stress() -> None: session.execute(INSERT_TEST_VALUE, (f"value_{i}",)) +# --- Oracle LOB fetch scenarios --- + +ORACLE_LOB_TABLES = {"1k": "SQLSPEC_LOB_1K", "100k": "SQLSPEC_LOB_100K"} + + +def _get_oracledb() -> Any: + """Import python-oracledb lazily.""" + try: + import oracledb + except ImportError: + return None + else: + return oracledb + + +def _oracle_connection_config_from_env() -> dict[str, Any]: + """Build Oracle connection config from service-safe benchmark env vars.""" + missing = [name for name in ORACLE_LOB_ENV_VARS if not os.environ.get(name)] + if missing: + joined = ", ".join(missing) + msg = f"Oracle benchmarks require service-derived environment variables: {joined}" + raise RuntimeError(msg) + + return { + "host": os.environ["SQLSPEC_BENCH_ORACLE_HOST"], + "port": int(os.environ["SQLSPEC_BENCH_ORACLE_PORT"]), + "service_name": os.environ["SQLSPEC_BENCH_ORACLE_SERVICE_NAME"], + "user": os.environ["SQLSPEC_BENCH_ORACLE_USER"], + "password": os.environ["SQLSPEC_BENCH_ORACLE_PASSWORD"], + "min": 1, + "max": max(1, POOL_SIZE), + "increment": 1, + } + + +def _oracle_connect_config_from_env() -> dict[str, Any]: + config = _oracle_connection_config_from_env() + for key in ("min", "max", "increment"): + config.pop(key, None) + return config + + +def _oracle_lob_payload(size_key: str) -> str: + return "x" * ORACLE_LOB_PAYLOAD_SIZES[size_key] + + +def _oracle_lob_table(size_key: str) -> str: + return ORACLE_LOB_TABLES[size_key] + + +def _oracle_drop_table_sql(table_name: str) -> str: + return ( + f"BEGIN EXECUTE IMMEDIATE 'DROP TABLE {table_name}'; " + "EXCEPTION WHEN OTHERS THEN IF SQLCODE != -942 THEN RAISE; END IF; END;" + ) + + +def _oracle_create_lob_table_sql(table_name: str) -> str: + return f"CREATE TABLE {table_name} (id NUMBER PRIMARY KEY, payload CLOB)" + + +def _oracle_insert_lob_sql(table_name: str) -> str: + return f"INSERT INTO {table_name} (id, payload) VALUES (:1, :2)" + + +def _oracle_select_lob_sql(table_name: str) -> str: + return f"SELECT id, payload FROM {table_name} ORDER BY id" + + +def _oracle_lob_rows(size_key: str) -> list[tuple[int, str]]: + payload = _oracle_lob_payload(size_key) + return [(index, payload) for index in range(1, ORACLE_LOB_ROWS + 1)] + + +def _read_oracle_lob_value(value: Any) -> Any: + read = getattr(value, "read", None) + if callable(read): + return read() + return value + + +async def _read_oracle_lob_value_async(value: Any) -> Any: + read = getattr(value, "read", None) + if not callable(read): + return value + result = read() + if inspect.isawaitable(result): + return await result + return result + + +def _assert_oracle_lob_rows(rows: Sequence[Any], size_key: str) -> None: + expected = _oracle_lob_payload(size_key) + assert len(rows) == ORACLE_LOB_ROWS + first_row = rows[0] + payload = first_row["payload"] if isinstance(first_row, dict) else first_row[1] + assert _read_oracle_lob_value(payload) == expected + + +async def _assert_oracle_lob_rows_async(rows: Sequence[Any], size_key: str) -> None: + expected = _oracle_lob_payload(size_key) + assert len(rows) == ORACLE_LOB_ROWS + first_row = rows[0] + payload = first_row["payload"] if isinstance(first_row, dict) else first_row[1] + assert await _read_oracle_lob_value_async(payload) == expected + + +def _run_raw_oracle_lob_fetch(size_key: str) -> None: + oracledb = _get_oracledb() + if oracledb is None: + return + table_name = _oracle_lob_table(size_key) + conn = oracledb.connect(**_oracle_connect_config_from_env()) + try: + with conn.cursor() as cursor: + cursor.execute(_oracle_drop_table_sql(table_name)) + cursor.execute(_oracle_create_lob_table_sql(table_name)) + cursor.executemany(_oracle_insert_lob_sql(table_name), _oracle_lob_rows(size_key)) + conn.commit() + cursor.execute(_oracle_select_lob_sql(table_name)) + rows = cursor.fetchall() + _assert_oracle_lob_rows(rows, size_key) + finally: + conn.close() + + +def _run_sqlspec_oracle_lob_fetch(size_key: str, *, fetch_lobs: bool) -> None: + from sqlspec.adapters.oracledb import OracleSyncConfig + + table_name = _oracle_lob_table(size_key) + spec = SQLSpec() + config = OracleSyncConfig( + connection_config=_oracle_connection_config_from_env(), driver_features={"fetch_lobs": fetch_lobs} + ) + try: + with spec.provide_session(config) as session: + session.execute_script(_oracle_drop_table_sql(table_name)) + session.execute(_oracle_create_lob_table_sql(table_name)) + session.execute_many(_oracle_insert_lob_sql(table_name), _oracle_lob_rows(size_key)) + rows = session.fetch(_oracle_select_lob_sql(table_name)) + _assert_oracle_lob_rows(rows, size_key) + _check_pool_leak(config.connection_instance, f"oracle/lob_fetch/{size_key}/fetch_lobs={fetch_lobs}") + config.close_pool() + finally: + if config.connection_instance is not None: + config.close_pool() + + +async def _run_sqlspec_oracle_lob_fetch_async(size_key: str, *, fetch_lobs: bool) -> None: + from sqlspec.adapters.oracledb import OracleAsyncConfig + + table_name = _oracle_lob_table(size_key) + spec = SQLSpec() + config = OracleAsyncConfig( + connection_config=_oracle_connection_config_from_env(), driver_features={"fetch_lobs": fetch_lobs} + ) + try: + async with spec.provide_session(config) as session: + await session.execute_script(_oracle_drop_table_sql(table_name)) + await session.execute(_oracle_create_lob_table_sql(table_name)) + await session.execute_many(_oracle_insert_lob_sql(table_name), _oracle_lob_rows(size_key)) + rows = await session.fetch(_oracle_select_lob_sql(table_name)) + await _assert_oracle_lob_rows_async(rows, size_key) + _check_pool_leak(config.connection_instance, f"oracle/lob_fetch_async/{size_key}/fetch_lobs={fetch_lobs}") + await config.close_pool() + finally: + if config.connection_instance is not None: + await config.close_pool() + + +def raw_oracle_lob_fetch_1k() -> None: + """Fetch 100 1 KiB Oracle CLOB rows through raw python-oracledb.""" + _run_raw_oracle_lob_fetch("1k") + + +def raw_oracle_lob_fetch_100k() -> None: + """Fetch 100 100 KiB Oracle CLOB rows through raw python-oracledb.""" + _run_raw_oracle_lob_fetch("100k") + + +def sqlspec_oracle_lob_fetch_1k() -> None: + """Fetch 100 1 KiB Oracle CLOB rows with sqlspec direct LOB fetching.""" + _run_sqlspec_oracle_lob_fetch("1k", fetch_lobs=False) + + +def sqlspec_oracle_lob_fetch_100k() -> None: + """Fetch 100 100 KiB Oracle CLOB rows with sqlspec direct LOB fetching.""" + _run_sqlspec_oracle_lob_fetch("100k", fetch_lobs=False) + + +def sqlspec_oracle_lob_fetch_fetch_lobs_true_1k() -> None: + """Fetch 100 1 KiB Oracle CLOB rows with sqlspec LOB locators enabled.""" + _run_sqlspec_oracle_lob_fetch("1k", fetch_lobs=True) + + +def sqlspec_oracle_lob_fetch_fetch_lobs_true_100k() -> None: + """Fetch 100 100 KiB Oracle CLOB rows with sqlspec LOB locators enabled.""" + _run_sqlspec_oracle_lob_fetch("100k", fetch_lobs=True) + + +async def sqlspec_oracle_lob_fetch_async_1k() -> None: + """Fetch 100 1 KiB Oracle CLOB rows with async sqlspec direct LOB fetching.""" + await _run_sqlspec_oracle_lob_fetch_async("1k", fetch_lobs=False) + + +async def sqlspec_oracle_lob_fetch_async_100k() -> None: + """Fetch 100 100 KiB Oracle CLOB rows with async sqlspec direct LOB fetching.""" + await _run_sqlspec_oracle_lob_fetch_async("100k", fetch_lobs=False) + + +async def sqlspec_oracle_lob_fetch_async_fetch_lobs_true_1k() -> None: + """Fetch 100 1 KiB Oracle CLOB rows with async sqlspec LOB locators enabled.""" + await _run_sqlspec_oracle_lob_fetch_async("1k", fetch_lobs=True) + + +async def sqlspec_oracle_lob_fetch_async_fetch_lobs_true_100k() -> None: + """Fetch 100 100 KiB Oracle CLOB rows with async sqlspec LOB locators enabled.""" + await _run_sqlspec_oracle_lob_fetch_async("100k", fetch_lobs=True) + + SCENARIO_REGISTRY: dict[tuple[str, str, str], Any] = { # SQLite scenarios ("raw", "sqlite", "initialization"): raw_sqlite_initialization, @@ -1772,6 +2174,8 @@ def sqlspec_sqlite_thin_path_stress() -> None: ("sqlalchemy", "aiosqlite", "read_heavy"): sqlalchemy_aiosqlite_read_heavy, ("sqlalchemy", "aiosqlite", "iterative_inserts"): sqlalchemy_aiosqlite_iterative_inserts, ("sqlalchemy", "aiosqlite", "repeated_queries"): sqlalchemy_aiosqlite_repeated_queries, + ("raw", "aiosqlite", "worker_hops"): raw_aiosqlite_worker_hops, + ("sqlspec", "aiosqlite", "worker_hops"): sqlspec_aiosqlite_worker_hops, # Asyncpg scenarios ("raw", "asyncpg", "initialization"): raw_asyncpg_initialization, ("raw", "asyncpg", "write_heavy"): raw_asyncpg_write_heavy, @@ -1787,10 +2191,23 @@ def sqlspec_sqlite_thin_path_stress() -> None: ("sqlspec", "sqlite", "dict_key_transform"): sqlspec_sqlite_dict_key_transform, ("raw", "sqlite", "schema_mapping"): raw_sqlite_schema_mapping, ("sqlspec", "sqlite", "schema_mapping"): sqlspec_sqlite_schema_mapping, + ("raw", "sqlite", "schema_type_numpy"): raw_sqlite_schema_type_numpy, + ("sqlspec", "sqlite", "schema_type_numpy"): sqlspec_sqlite_schema_type_numpy, ("raw", "sqlite", "complex_parameters"): raw_sqlite_complex_parameters, ("sqlspec", "sqlite", "complex_parameters"): sqlspec_sqlite_complex_parameters, ("raw", "sqlite", "thin_path_stress"): raw_sqlite_thin_path_stress, ("sqlspec", "sqlite", "thin_path_stress"): sqlspec_sqlite_thin_path_stress, + # Extended Oracle LOB fetch scenarios + ("raw", "oracle", "lob_fetch_1k"): raw_oracle_lob_fetch_1k, + ("raw", "oracle", "lob_fetch_100k"): raw_oracle_lob_fetch_100k, + ("sqlspec", "oracle", "lob_fetch_1k"): sqlspec_oracle_lob_fetch_1k, + ("sqlspec", "oracle", "lob_fetch_100k"): sqlspec_oracle_lob_fetch_100k, + ("sqlspec_fetch_lobs_true", "oracle", "lob_fetch_1k"): sqlspec_oracle_lob_fetch_fetch_lobs_true_1k, + ("sqlspec_fetch_lobs_true", "oracle", "lob_fetch_100k"): sqlspec_oracle_lob_fetch_fetch_lobs_true_100k, + ("sqlspec_async", "oracle", "lob_fetch_1k"): sqlspec_oracle_lob_fetch_async_1k, + ("sqlspec_async", "oracle", "lob_fetch_100k"): sqlspec_oracle_lob_fetch_async_100k, + ("sqlspec_async_fetch_lobs_true", "oracle", "lob_fetch_1k"): sqlspec_oracle_lob_fetch_async_fetch_lobs_true_1k, + ("sqlspec_async_fetch_lobs_true", "oracle", "lob_fetch_100k"): sqlspec_oracle_lob_fetch_async_fetch_lobs_true_100k, }