diff --git a/docs/changelog.rst b/docs/changelog.rst index ababb3879..6b4b5df51 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -14,6 +14,11 @@ Unreleased **Fixed:** +* MySQL adapters accept either ``local_infile=True`` or + ``allow_local_infile=True`` to enable eligible native bulk loads. A separate + bulk-load opt-in is no longer required; set + ``enable_local_infile_bulk_load=False`` to retain ``executemany``. + * ``sqlspec.extensions.litestar.LitestarConfig`` now exposes the complete plugin configuration, including ``session_table=True``, through the same type as ``sqlspec.config.LitestarConfig``. diff --git a/docs/usage/bulk_ingest.rst b/docs/usage/bulk_ingest.rst index 80242b02f..e31f4f453 100644 --- a/docs/usage/bulk_ingest.rst +++ b/docs/usage/bulk_ingest.rst @@ -77,8 +77,8 @@ Capability matrix * - MySQL family (pymysql, asyncmy, aiomysql, mysql-connector) - ``executemany`` (default); ``LOAD DATA LOCAL INFILE`` (opt-in) - Server-managed - - ``enable_local_infile_bulk_load`` + connection ``local_infile`` / - ``allow_local_infile`` + - Connection ``local_infile=True`` or ``allow_local_infile=True``; + ``enable_local_infile_bulk_load=False`` forces fallback * - bigquery - Parquet load job (default); Arrow Storage Write API (opt-in) - All-or-nothing load job / PENDING write stream @@ -101,14 +101,29 @@ Security and opt-in paths Some fast paths are opt-in because they read local files or change semantics: -- **MySQL ``LOAD DATA LOCAL INFILE``** requires both the adapter feature - ``enable_local_infile_bulk_load`` and the connection's local-infile setting - (``local_infile=True`` for pymysql/aiomysql/asyncmy, ``allow_local_infile=True`` - for mysql-connector). Enabling the feature without the connection gate raises - :class:`~sqlspec.exceptions.ImproperConfigurationError` at config construction. +- **MySQL ``LOAD DATA LOCAL INFILE``** is enabled by either + ``local_infile=True`` or ``allow_local_infile=True`` in the connection + configuration. Both names work for each MySQL adapter. If both are set, + either true value enables loading. If neither is true, loading stays off. + SQLSpec sends only the driver's native flag. + + This one flag also turns on bulk loads for supported data. Set + ``driver_features={"enable_local_infile_bulk_load": False}`` to keep using + ``executemany``. Setting that feature to true with no connection opt-in raises + :class:`~sqlspec.exceptions.ImproperConfigurationError` when you create the config. The MySQL server must also have ``local_infile`` enabled. mysql-connector additionally honors ``allow_local_infile_in_path`` -- the staged temp file must live under that directory when it is set. + Connection opt-in trusts the configured MySQL server to request client files. + + For asyncmy bulk loads, the requested filename must match that operation's + payload; this is not a connection-wide file restriction for other queries. + SQLSpec uses asyncmy's native sender + (version 0.2.13 or newer), removes its private UTF-8 payload after each attempt, + and closes the connection if native loading fails or is cancelled. + When bulk loading is disabled, and for nested, binary or duration values, + asyncmy uses ``executemany``. As with the other MySQL adapters, ``overwrite=True`` first + truncates the table; a later load failure does not restore those rows. - **Oracle direct path load** is the default bulk-ingest transport in Thin mode. Set ``enable_direct_path_load=False`` to force ``executemany``. Connections that do not expose the Direct Path Load API, including Thick-mode connections, @@ -134,11 +149,25 @@ MySQL ``LOAD DATA LOCAL INFILE``: config = PyMysqlConfig( connection_config={"host": "localhost", "local_infile": True}, - driver_features={"enable_local_infile_bulk_load": True}, ) with config.provide_session() as driver: driver.load_from_arrow("orders", arrow_table) +Asyncmy with explicit LOCAL INFILE consent: + +.. code-block:: python + + from sqlspec.adapters.asyncmy import AsyncmyConfig + + config = AsyncmyConfig( + connection_config={ + "host": "localhost", + "allow_local_infile": True, + }, + ) + async with config.provide_session() as driver: + await driver.load_from_arrow("orders", arrow_table) + Oracle per-call batch error and array-DML row-count reporting: .. code-block:: python diff --git a/pyproject.toml b/pyproject.toml index 063b2f456..ad1fe44f5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,7 +38,7 @@ aiomysql = ["aiomysql"] aiosqlite = ["aiosqlite"] alloydb = ["google-cloud-alloydb-connector"] arrow-odbc = ["arrow-odbc>=10.4", "pyarrow"] -asyncmy = ["asyncmy"] +asyncmy = ["asyncmy>=0.2.13"] asyncpg = ["asyncpg"] attrs = ["attrs", "cattrs"] bigquery = ["google-cloud-bigquery", "google-cloud-bigquery-storage", "google-cloud-storage"] diff --git a/sqlspec/adapters/aiomysql/config.py b/sqlspec/adapters/aiomysql/config.py index 34123ba62..e4ebdafc1 100644 --- a/sqlspec/adapters/aiomysql/config.py +++ b/sqlspec/adapters/aiomysql/config.py @@ -35,7 +35,6 @@ __all__ = ("AiomysqlConfig", "AiomysqlConnectionParams", "AiomysqlDriverFeatures", "AiomysqlPoolParams") _POOL_ONLY_CONFIG_KEYS = frozenset({"maxsize", "minsize", "pool_recycle"}) -_AIOMYSQL_LOCAL_INFILE_GATE = "allow_local_infile" aiomysql: "AiomysqlModule" = cast("AiomysqlModule", AiomysqlModule) @@ -84,28 +83,19 @@ class AiomysqlPoolParams(AiomysqlConnectionParams): pool_recycle: NotRequired[int] -def _normalize_local_infile(connection_config: "Mapping[str, Any]", *, strip_consent_gate: bool) -> "dict[str, Any]": - """Normalize aiomysql local-infile settings and SQLSpec's consent gate.""" +def _normalize_local_infile(connection_config: "Mapping[str, Any]") -> "dict[str, Any]": + """Normalize aiomysql local-infile aliases to the native connection flag.""" config = dict(connection_config) 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)) - if local_infile and not allow_local_infile: - msg = ( - "Aiomysql local_infile=True requires allow_local_infile=True because " - "LOAD DATA LOCAL INFILE can read client files." - ) - raise ImproperConfigurationError(msg) - config["local_infile"] = bool(local_infile and allow_local_infile) - if strip_consent_gate: - config.pop(_AIOMYSQL_LOCAL_INFILE_GATE, None) + allow_local_infile = bool(config.pop("allow_local_infile", False)) + config["local_infile"] = bool(config.get("local_infile", False) or allow_local_infile) return config def _normalize_connection_kwargs(connection_config: "Mapping[str, Any]") -> "dict[str, Any]": """Build aiomysql.connect-compatible kwargs from SQLSpec connection config.""" - config = _normalize_local_infile(connection_config, strip_consent_gate=True) + config = _normalize_local_infile(connection_config) for key in _POOL_ONLY_CONFIG_KEYS: config.pop(key, None) @@ -253,9 +243,7 @@ def __init__( observability_config: Adapter-level observability overrides for lifecycle hooks and observers **kwargs: Additional keyword arguments """ - connection_config = _normalize_local_infile( - normalize_connection_config(connection_config), strip_consent_gate=False - ) + connection_config = _normalize_local_infile(normalize_connection_config(connection_config)) connection_config.setdefault("host", "localhost") connection_config.setdefault("port", 3306) @@ -271,11 +259,9 @@ def __init__( # Track initialized connections to ensure callback runs exactly once per physical connection self._initialized_connections: WeakSet[Any] = WeakSet() + features_dict.setdefault("enable_local_infile_bulk_load", connection_config["local_infile"]) if features_dict.get("enable_local_infile_bulk_load") and not connection_config.get("local_infile"): - msg = ( - "enable_local_infile_bulk_load requires local_infile=True and " - "allow_local_infile=True in connection_config." - ) + msg = "enable_local_infile_bulk_load requires local_infile=True or allow_local_infile=True in connection_config." raise ImproperConfigurationError(msg) super().__init__( diff --git a/sqlspec/adapters/asyncmy/_typing.py b/sqlspec/adapters/asyncmy/_typing.py index 93dbbdc2c..c97bec4c0 100644 --- a/sqlspec/adapters/asyncmy/_typing.py +++ b/sqlspec/adapters/asyncmy/_typing.py @@ -5,18 +5,24 @@ """ import contextlib +import os from typing import TYPE_CHECKING, Any import asyncmy as _asyncmy # pyright: ignore from asyncmy import Connection # pyright: ignore from asyncmy import errors as _asyncmy_errors # pyright: ignore +from asyncmy.connection import LoadLocalFile as _LoadLocalFile # pyright: ignore +from asyncmy.connection import MySQLResult as _AsyncmyResult # pyright: ignore from asyncmy.constants import FIELD_TYPE as _ASYNCMY_FIELD_TYPE # pyright: ignore from asyncmy.cursors import Cursor as _AsyncmyCursor # pyright: ignore from asyncmy.cursors import DictCursor as _AsyncmyDictCursor # pyright: ignore from asyncmy.pool import Pool as _AsyncmyPool # pyright: ignore +from asyncmy.protocol import LoadLocalPacketWrapper as _LoadLocalPacketWrapper # pyright: ignore + +from sqlspec.exceptions import SQLSpecError if TYPE_CHECKING: - from collections.abc import Awaitable, Callable + from collections.abc import Awaitable, Callable, Iterator from types import TracebackType from typing import Protocol, TypeAlias @@ -30,7 +36,7 @@ async def commit(self) -> object: ... async def rollback(self) -> object: ... - async def close(self) -> object: ... + def close(self) -> None: ... class AsyncmyModuleProtocol(Protocol): def connect(self, *args: Any, **kwargs: Any) -> "AsyncmyConnection": ... @@ -70,6 +76,7 @@ class AsyncmyFieldTypeProtocol(Protocol): "AsyncmyPool", "AsyncmyRawCursor", "AsyncmySessionContext", + "asyncmy_local_infile", ) @@ -148,3 +155,72 @@ async def __aexit__( await self._release_connection(self._connection) self._connection = None return None + + +class _AsyncmyLocalInfileResult(_AsyncmyResult): + """Normalize the upstream filename handoff while retaining its native sender.""" + + __slots__ = ("_filename",) + + def __init__(self, connection: Any, filename: str) -> None: + super().__init__(connection) + self._filename = filename + + async def _read_load_local_packet(self, first_packet: Any) -> None: + request = _LoadLocalPacketWrapper(first_packet).filename + if not self.connection._local_infile or os.fsdecode(request) != self._filename: + msg = "MySQL requested an unexpected LOCAL INFILE payload." + raise SQLSpecError(msg) + await _LoadLocalFile(self._filename, self.connection).send_data() # type: ignore[no-untyped-call] + packet = await self.connection.read_packet() + if not packet.is_ok_packet(): + msg = "MySQL did not acknowledge the LOCAL INFILE payload." + raise SQLSpecError(msg) + self._read_ok_packet(packet) # type: ignore[attr-defined] + + +@contextlib.contextmanager +def asyncmy_local_infile(connection: "AsyncmyConnection", filename: str) -> "Iterator[None]": + """Scope the asyncmy 0.2.13/0.2.14 filename handoff fix to one native load. + + Args: + connection: Physical connection exclusively held by this operation. + filename: Owned payload path expected in the server's file request. + + Yields: + Control while native result reading uses the filename adapter. + """ + raw: Any = connection + missing = object() + previous = raw.__dict__.get("_read_query_result", missing) + + async def read_result(unbuffered: bool = False) -> None: + raw._result = None + result = _AsyncmyLocalInfileResult(raw, filename) + if unbuffered: + try: + await result.init_unbuffered_query() # type: ignore[no-untyped-call] + except BaseException: + result.unbuffered_active = False + result.connection = None + raise + else: + await result.read() # type: ignore[no-untyped-call] + raw._result = result + raw._affected_rows = result.affected_rows + if result.server_status: + raw.server_status = result.server_status + + raw._read_query_result = read_result + try: + yield + except BaseException: + with contextlib.suppress(Exception): + raw.close() + raw._connected = False + raise + finally: + if previous is missing: + del raw._read_query_result + else: + raw._read_query_result = previous diff --git a/sqlspec/adapters/asyncmy/config.py b/sqlspec/adapters/asyncmy/config.py index 7204ac731..429c4fdf9 100644 --- a/sqlspec/adapters/asyncmy/config.py +++ b/sqlspec/adapters/asyncmy/config.py @@ -38,7 +38,6 @@ _ASYNCMY_POOL_ONLY_KEYS = frozenset(("minsize", "maxsize", "pool_recycle")) _ASYNCMY_POOL_KEYS = _ASYNCMY_POOL_ONLY_KEYS | {"echo"} -_ASYNCMY_LOCAL_INFILE_GATE = "allow_local_infile" asyncmy: "AsyncmyModule" = cast("AsyncmyModule", AsyncmyModule) @@ -120,12 +119,8 @@ def _normalize_connection_config(connection_config: "Mapping[str, Any] | None") raise ImproperConfigurationError(msg) config["cursor_cls"] = cursor_class - allow_local_infile = bool(config.pop(_ASYNCMY_LOCAL_INFILE_GATE, False)) - local_infile = bool(config.get("local_infile", False)) - if local_infile and not allow_local_infile: - msg = "Asyncmy local_infile=True requires allow_local_infile=True because LOAD DATA LOCAL INFILE can read client files." - raise ImproperConfigurationError(msg) - config["local_infile"] = bool(local_infile and allow_local_infile) + allow_local_infile = bool(config.pop("allow_local_infile", False)) + config["local_infile"] = bool(config.get("local_infile", False) or allow_local_infile) return config @@ -159,6 +154,9 @@ class AsyncmyDriverFeatures(TypedDict): MySQL/MariaDB handle JSON natively, but custom serializers can be provided for specialized use cases. + enable_local_infile_bulk_load: Use native LOCAL INFILE for eligible Arrow rows. + Defaults to the connection's local_infile or allow_local_infile opt-in. + Set False to force executemany on an opted-in connection. json_serializer: Custom JSON serializer function. Defaults to sqlspec.utils.serializers.to_json. Use for performance (orjson) or custom encoding. @@ -178,6 +176,7 @@ class AsyncmyDriverFeatures(TypedDict): Defaults to "poll_queue". """ + enable_local_infile_bulk_load: NotRequired[bool] json_serializer: NotRequired["Callable[[Any], str]"] json_deserializer: NotRequired["Callable[[str], Any]"] on_connection_create: "NotRequired[Callable[[AsyncmyConnection], Awaitable[None]]]" @@ -296,12 +295,9 @@ def __init__( # Track initialized connections to ensure callback runs exactly once per physical connection self._initialized_connections: WeakSet[Any] = WeakSet() - if features_dict.get("enable_local_infile_bulk_load"): - msg = ( - "asyncmy does not currently support SQLSpec's LOAD DATA LOCAL INFILE bulk path reliably. " - "Use aiomysql, mysql-connector, or pymysql for LOCAL INFILE bulk loads, or omit " - "enable_local_infile_bulk_load to use asyncmy batched executemany." - ) + features_dict.setdefault("enable_local_infile_bulk_load", connection_config["local_infile"]) + if features_dict.get("enable_local_infile_bulk_load") and not connection_config.get("local_infile"): + msg = "enable_local_infile_bulk_load requires local_infile=True or allow_local_infile=True in connection_config." raise ImproperConfigurationError(msg) super().__init__( diff --git a/sqlspec/adapters/asyncmy/core.py b/sqlspec/adapters/asyncmy/core.py index a325b06c7..53af4bbb2 100644 --- a/sqlspec/adapters/asyncmy/core.py +++ b/sqlspec/adapters/asyncmy/core.py @@ -34,6 +34,7 @@ "AsyncmyStreamSource", "apply_driver_features", "build_insert_statement", + "build_load_data_statement", "build_profile", "build_statement_config", "collect_rows", @@ -43,6 +44,7 @@ "detect_json_columns", "detect_json_columns_from_description", "driver_profile", + "encode_records_for_local_infile", "format_identifier", "normalize_execute_many_parameters", "normalize_execute_parameters", @@ -135,6 +137,58 @@ def build_insert_statement(table: str, columns: "list[str]") -> str: return f"INSERT INTO {format_identifier(table)} ({column_clause}) VALUES ({placeholders})" +def encode_records_for_local_infile(records: "list[tuple[Any, ...]]") -> bytes: + """Encode rows as escaped UTF-8 TSV for MySQL LOCAL INFILE. + + Args: + records: Scalar rows without nested values requiring preparation. + + Returns: + Encoded payload with MySQL NULL and field escaping. + """ + lines: list[str] = [] + for record in records: + fields: list[str] = [] + for value in record: + if value is None: + fields.append("\\N") + continue + if isinstance(value, bool): + value = int(value) + text = str(value) + text = ( + text + .replace("\\", "\\\\") + .replace("\x00", "\\0") + .replace("\t", "\\t") + .replace("\n", "\\n") + .replace("\r", "\\r") + .replace("\x1a", "\\Z") + ) + fields.append(text) + lines.append("\t".join(fields)) + return ("\n".join(lines) + "\n").encode("utf-8") + + +def build_load_data_statement(table: str, columns: "list[str]") -> str: + """Build native LOAD DATA SQL with a bound filename. + + Args: + table: Destination table identifier. + columns: Destination column names. + + Returns: + SQL with one positional filename placeholder. + """ + table_sql = format_identifier(table).replace("%", "%%") + column_sql = ", ".join(quote_backtick_identifier(column).replace("%", "%%") for column in columns) + return ( + f"LOAD DATA LOCAL INFILE %s INTO TABLE {table_sql} " + "CHARACTER SET utf8mb4 FIELDS TERMINATED BY '\\t' ESCAPED BY '\\\\' " + f"LINES TERMINATED BY '\\n' ({column_sql})" + ) + + def normalize_execute_parameters(parameters: Any) -> Any: """Normalize parameters for AsyncMy execute calls. diff --git a/sqlspec/adapters/asyncmy/driver.py b/sqlspec/adapters/asyncmy/driver.py index 3a9be414d..7991385d3 100644 --- a/sqlspec/adapters/asyncmy/driver.py +++ b/sqlspec/adapters/asyncmy/driver.py @@ -4,7 +4,9 @@ type coercion, error handling, and transaction management. """ +import tempfile from collections.abc import Sized +from datetime import timedelta from typing import TYPE_CHECKING, Any, Final, cast from sqlspec.adapters.asyncmy._typing import ( @@ -13,14 +15,17 @@ AsyncmyFieldType, AsyncmyMySQLError, AsyncmySessionContext, + asyncmy_local_infile, ) from sqlspec.adapters.asyncmy.core import ( AsyncmyStreamSource, build_insert_statement, + build_load_data_statement, collect_rows, create_mapped_exception, default_statement_config, driver_profile, + encode_records_for_local_infile, format_identifier, normalize_execute_many_parameters, normalize_execute_parameters, @@ -336,7 +341,7 @@ async def load_from_arrow( overwrite: bool = False, telemetry: "StorageTelemetry | None" = None, ) -> "StorageBridgeJob": - """Load Arrow data into MySQL using batched inserts.""" + """Load Arrow data using batched inserts or opt-in native LOCAL INFILE.""" self._require_capability("arrow_import_enabled") arrow_table = self._coerce_arrow_table(source) @@ -351,23 +356,48 @@ async def load_from_arrow( columns, records = self._arrow_table_to_rows(arrow_table) if records: needs_preparation = self._arrow_rows_need_preparation(arrow_table) - insert_sql = build_insert_statement(table, columns) - prepared_records = ( - self.prepare_driver_parameters(records, self.statement_config, is_many=True) - if needs_preparation - else records + use_infile = ( + self.driver_features.get("enable_local_infile_bulk_load") + and not needs_preparation + and not any( + isinstance(value, (bytes, bytearray, memoryview, timedelta)) for row in records for value in row + ) ) - exc_handler = self.handle_database_exceptions() - async with exc_handler, self.with_cursor(self.connection) as cursor: - await cursor.executemany(insert_sql, prepared_records) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + if use_infile: + await self._load_from_arrow_via_local_infile(table, columns, records) + else: + insert_sql = build_insert_statement(table, columns) + prepared_records = ( + self.prepare_driver_parameters(records, self.statement_config, is_many=True) + if needs_preparation + else records + ) + exc_handler = self.handle_database_exceptions() + async with exc_handler, self.with_cursor(self.connection) as cursor: + await cursor.executemany(insert_sql, prepared_records) + if exc_handler.pending_exception is not None: + raise exc_handler.pending_exception from None telemetry_payload = self._ingest_telemetry(arrow_table) telemetry_payload["destination"] = table self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) + async def _load_from_arrow_via_local_infile( + self, table: str, columns: "list[str]", records: "list[tuple[Any, ...]]" + ) -> None: + with tempfile.TemporaryDirectory(prefix="sqlspec-asyncmy-") as directory: + with tempfile.NamedTemporaryFile(dir=directory, suffix=".tsv", delete=False) as payload: + payload.write(encode_records_for_local_infile(records)) + filename = payload.name + statement = build_load_data_statement(table, columns) + exc_handler = self.handle_database_exceptions() + async with exc_handler, self.with_cursor(self.connection) as cursor: + with asyncmy_local_infile(self.connection, filename): + await cursor.execute(statement, (filename,)) + if exc_handler.pending_exception is not None: + raise exc_handler.pending_exception from exc_handler.pending_exception.__cause__ + async def load_from_storage( self, table: str, diff --git a/sqlspec/adapters/mysqlconnector/config.py b/sqlspec/adapters/mysqlconnector/config.py index 84bbe9c02..82310d7e4 100644 --- a/sqlspec/adapters/mysqlconnector/config.py +++ b/sqlspec/adapters/mysqlconnector/config.py @@ -140,6 +140,7 @@ class _MysqlConnectorBaseConnectionParams(TypedDict): option_files: NotRequired[MysqlConnectorPathSequence] option_groups: NotRequired[MysqlConnectorStringSequence] allow_local_infile: NotRequired[bool] + local_infile: NotRequired[bool] allow_local_infile_in_path: NotRequired[str] use_pure: NotRequired[bool] dsn: NotRequired[str] @@ -197,8 +198,8 @@ class MysqlConnectorDriverFeatures(TypedDict): 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)) + local_infile = bool(config.pop("local_infile", False)) + config["allow_local_infile"] = bool(config.get("allow_local_infile", False) or local_infile) return config @@ -324,8 +325,9 @@ def __init__( # Track initialized connections to ensure callback runs exactly once per physical connection self._initialized_connections: WeakSet[Any] = WeakSet() + features_dict.setdefault("enable_local_infile_bulk_load", connection_config["allow_local_infile"]) if features_dict.get("enable_local_infile_bulk_load") and not connection_config.get("allow_local_infile"): - msg = "enable_local_infile_bulk_load requires allow_local_infile=True in connection_config." + msg = "enable_local_infile_bulk_load requires local_infile=True or allow_local_infile=True in connection_config." raise ImproperConfigurationError(msg) super().__init__( @@ -440,8 +442,9 @@ def __init__( features_dict.pop("on_connection_create", None) ) + features_dict.setdefault("enable_local_infile_bulk_load", self.connection_config["allow_local_infile"]) if features_dict.get("enable_local_infile_bulk_load") and not self.connection_config.get("allow_local_infile"): - msg = "enable_local_infile_bulk_load requires allow_local_infile=True in connection_config." + msg = "enable_local_infile_bulk_load requires local_infile=True or allow_local_infile=True in connection_config." raise ImproperConfigurationError(msg) super().__init__( diff --git a/sqlspec/adapters/psycopg/_typing.py b/sqlspec/adapters/psycopg/_typing.py index 9af670d66..9fdeb96cc 100644 --- a/sqlspec/adapters/psycopg/_typing.py +++ b/sqlspec/adapters/psycopg/_typing.py @@ -140,7 +140,7 @@ class PsycopgSyncSessionContext: def __init__( self, acquire_connection: "Callable[[], Any]", - release_connection: "Callable[[Any], Any]", + release_connection: "Callable[..., Any]", statement_config: "StatementConfig | Callable[[], StatementConfig]", driver_features: "dict[str, Any]", prepare_driver: "Callable[[PsycopgSyncDriver], PsycopgSyncDriver]", @@ -167,7 +167,7 @@ def __exit__( self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None" ) -> "bool | None": if self._connection is not None: - self._release_connection(self._connection) + self._release_connection(self._connection, exc_type=exc_type, exc_val=exc_val, exc_tb=exc_tb) self._connection = None return None @@ -196,7 +196,7 @@ class PsycopgAsyncSessionContext: def __init__( self, acquire_connection: "Callable[[], Any]", - release_connection: "Callable[[Any], Any]", + release_connection: "Callable[..., Any]", statement_config: "StatementConfig | Callable[[], StatementConfig]", driver_features: "dict[str, Any]", prepare_driver: "Callable[[PsycopgAsyncDriver], PsycopgAsyncDriver]", @@ -223,6 +223,6 @@ async def __aexit__( self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None" ) -> "bool | None": if self._connection is not None: - await self._release_connection(self._connection) + await self._release_connection(self._connection, exc_type=exc_type, exc_val=exc_val, exc_tb=exc_tb) self._connection = None return None diff --git a/sqlspec/adapters/psycopg/config.py b/sqlspec/adapters/psycopg/config.py index 95d957e19..b045beefc 100644 --- a/sqlspec/adapters/psycopg/config.py +++ b/sqlspec/adapters/psycopg/config.py @@ -245,7 +245,7 @@ def acquire_connection(self) -> "PsycopgSyncConnection": def release_connection(self, _conn: "PsycopgSyncConnection", **kwargs: Any) -> None: if self._ctx is not None: - self._ctx.__exit__(None, None, None) + self._ctx.__exit__(kwargs.get("exc_type"), kwargs.get("exc_val"), kwargs.get("exc_tb")) self._ctx = None return if self._conn is not None: @@ -597,7 +597,7 @@ async def acquire_connection(self) -> "PsycopgAsyncConnection": async def release_connection(self, _conn: "PsycopgAsyncConnection", **kwargs: Any) -> None: if self._ctx is None: return - await self._ctx.__aexit__(None, None, None) + await self._ctx.__aexit__(kwargs.get("exc_type"), kwargs.get("exc_val"), kwargs.get("exc_tb")) self._ctx = None diff --git a/sqlspec/adapters/pymssql/core.py b/sqlspec/adapters/pymssql/core.py index 578878209..d382e5372 100644 --- a/sqlspec/adapters/pymssql/core.py +++ b/sqlspec/adapters/pymssql/core.py @@ -197,10 +197,12 @@ def collect_rows( description: "Sequence[Any] | None", column_name_cache: "dict[int, tuple[Any, list[str]]] | None" = None, ) -> "tuple[list[Any], list[str], Literal['dict', 'tuple', 'record']]": - """Collect pymssql rows, preserving tuple row shape.""" + """Collect pymssql rows, preserving dictionary or tuple row shape.""" column_names = resolve_column_names(description, column_name_cache) if not fetched_data: return [], column_names, "tuple" + if isinstance(fetched_data[0], dict): + return list(fetched_data), column_names, "dict" return list(fetched_data), column_names, "tuple" diff --git a/sqlspec/adapters/pymysql/config.py b/sqlspec/adapters/pymysql/config.py index 2235582ad..226e9f4e6 100644 --- a/sqlspec/adapters/pymysql/config.py +++ b/sqlspec/adapters/pymysql/config.py @@ -120,7 +120,8 @@ class PyMysqlDriverFeatures(TypedDict): enable_events: Enable database event channel support. events_backend: Event channel backend selection. enable_local_infile_bulk_load: Route load_from_arrow through LOAD DATA LOCAL INFILE. - Requires local_infile=True in connection_config. + Defaults to the connection's local_infile or allow_local_infile opt-in. + Set False to force executemany on an opted-in connection. enable_cloud_sql: Enable Google Cloud SQL connector integration. Requires cloud-sql-python-connector package. Defaults to False (explicit opt-in required). @@ -159,17 +160,10 @@ class PyMysqlDriverFeatures(TypedDict): def _normalize_local_infile(connection_config: Mapping[str, Any]) -> dict[str, Any]: - """Normalize PyMySQL local-infile configuration and SQLSpec's consent gate.""" + """Normalize PyMySQL local-infile aliases to the native connection flag.""" config = dict(connection_config) allow_local_infile = bool(config.pop("allow_local_infile", False)) - local_infile = bool(config.get("local_infile", False)) - if local_infile and not allow_local_infile: - msg = ( - "PyMySQL local_infile=True requires allow_local_infile=True because " - "LOAD DATA LOCAL INFILE can read client files." - ) - raise ImproperConfigurationError(msg) - config["local_infile"] = bool(local_infile and allow_local_infile) + config["local_infile"] = bool(config.get("local_infile", False) or allow_local_infile) return config @@ -266,8 +260,9 @@ def __init__( "on_connection_create", None ) + features_dict.setdefault("enable_local_infile_bulk_load", connection_config["local_infile"]) if features_dict.get("enable_local_infile_bulk_load") and not connection_config.get("local_infile"): - msg = "enable_local_infile_bulk_load requires local_infile=True in connection_config." + msg = "enable_local_infile_bulk_load requires local_infile=True or allow_local_infile=True in connection_config." raise ImproperConfigurationError(msg) super().__init__( diff --git a/sqlspec/adapters/spanner/type_converter.py b/sqlspec/adapters/spanner/type_converter.py index 161f90e7b..5151b8f80 100644 --- a/sqlspec/adapters/spanner/type_converter.py +++ b/sqlspec/adapters/spanner/type_converter.py @@ -9,7 +9,7 @@ - JSON detection and deserialization Input conversion handles: - - UUID → base64-encoded bytes + - UUID → 36-character strings (when automatic conversion is enabled) - bytes → base64-encoded bytes - datetime timezone awareness - dict/list → JsonObject wrapping diff --git a/tests/integration/adapters/_shared/_driver_type_system.py b/tests/integration/adapters/_shared/_driver_type_system.py index d7b977d4b..351d01992 100644 --- a/tests/integration/adapters/_shared/_driver_type_system.py +++ b/tests/integration/adapters/_shared/_driver_type_system.py @@ -127,7 +127,14 @@ class SourceEquivalenceCase: "enable_events", "on_connection_create", ), - "asyncmy": ("json_serializer", "json_deserializer", "on_connection_create", "enable_events", "events_backend"), + "asyncmy": ( + "json_serializer", + "json_deserializer", + "on_connection_create", + "enable_events", + "events_backend", + "enable_local_infile_bulk_load", + ), "bigquery": ( "connection_instance", "on_connection_create", diff --git a/tests/integration/adapters/mysql/asyncmy/test_local_infile_bulk_load.py b/tests/integration/adapters/mysql/asyncmy/test_local_infile_bulk_load.py new file mode 100644 index 000000000..80462e200 --- /dev/null +++ b/tests/integration/adapters/mysql/asyncmy/test_local_infile_bulk_load.py @@ -0,0 +1,137 @@ +"""Native asyncmy LOCAL INFILE behavior against MySQL.""" + +import asyncio +from collections.abc import AsyncGenerator +from pathlib import Path + +import pyarrow as pa +import pytest +from asyncmy.cursors import Cursor, SSCursor, SSDictCursor +from pytest_databases.docker.mysql import MySQLService + +from sqlspec.adapters.asyncmy import AsyncmyConfig +from sqlspec.exceptions import SQLSpecError + +pytestmark = [pytest.mark.xdist_group("mysql"), pytest.mark.mysql, pytest.mark.asyncmy] + + +@pytest.fixture(params=[Cursor, SSCursor, SSDictCursor], ids=["buffered", "unbuffered", "unbuffered_dict"]) +async def asyncmy_infile_config( + mysql_service: MySQLService, request: pytest.FixtureRequest +) -> AsyncGenerator[AsyncmyConfig, None]: + config = AsyncmyConfig( + connection_config={ + "host": mysql_service.host, + "port": mysql_service.port, + "user": mysql_service.user, + "password": mysql_service.password, + "db": mysql_service.db, + "autocommit": True, + "cursor_cls": request.param, + "allow_local_infile": True, + } + ) + original = 0 + try: + async with config.provide_session() as driver: + original = await driver.select_value("SELECT @@GLOBAL.local_infile") + await driver.execute("SET GLOBAL local_infile = 1") + await driver.execute("DROP TABLE IF EXISTS asyncmy_native_infile") + await driver.execute( + "CREATE TABLE asyncmy_native_infile (id INT PRIMARY KEY, text_value LONGTEXT, flag BOOLEAN) " + "CHARACTER SET utf8mb4" + ) + yield config + finally: + try: + async with config.provide_session() as driver: + await driver.execute("DROP TABLE IF EXISTS asyncmy_native_infile") + await driver.execute("SET GLOBAL local_infile = 1" if original else "SET GLOBAL local_infile = 0") + finally: + await config.close_pool() + + +async def test_native_bulk_roundtrip_multichunk_and_overwrite( + asyncmy_infile_config: AsyncmyConfig, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + import tempfile + + tmp_path = tmp_path / "quoted'percent%雪" + tmp_path.mkdir() + monkeypatch.setattr(tempfile, "tempdir", str(tmp_path)) + texts = ["é\tline\nreturn\rback\\slash\x00\x1a", None, "", "雪" * 400_000] + arrow = pa.table({"id": [1, 2, 3, 4], "text_value": texts, "flag": [True, False, True, False]}) + async with asyncmy_infile_config.provide_session() as driver: + before = await driver.select_one("SHOW SESSION STATUS LIKE 'Com_load'") + job = await driver.load_from_arrow("asyncmy_native_infile", arrow) + after = await driver.select_one("SHOW SESSION STATUS LIKE 'Com_load'") + assert int(after["Value"]) == int(before["Value"]) + 1 + assert job.telemetry["rows_processed"] == 4 + rows = await driver.select("SELECT id, text_value, flag FROM asyncmy_native_infile ORDER BY id") + assert [row["text_value"] for row in rows] == texts + assert [row["flag"] for row in rows] == [1, 0, 1, 0] + assert list(tmp_path.iterdir()) == [] + await driver.load_from_arrow("asyncmy_native_infile", arrow.slice(0, 1), overwrite=True) + assert await driver.select_value("SELECT COUNT(*) FROM asyncmy_native_infile") == 1 + await driver.load_from_arrow("asyncmy_native_infile", arrow.slice(0, 0)) + assert await driver.select_value("SELECT COUNT(*) FROM asyncmy_native_infile") == 1 + await driver.execute( + "INSERT INTO asyncmy_native_infile (id, text_value) VALUES (:id, :value)", {"id": 5, "value": "after"} + ) + assert await driver.select_value("SELECT COUNT(*) FROM asyncmy_native_infile") == 2 + + +async def test_native_server_error_cleans_payload_and_pool_recovers( + asyncmy_infile_config: AsyncmyConfig, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + import tempfile + + monkeypatch.setattr(tempfile, "tempdir", str(tmp_path)) + async with asyncmy_infile_config.provide_session() as driver: + connection = driver.connection + with pytest.raises(SQLSpecError): + await driver.load_from_arrow("asyncmy_native_infile_missing", pa.table({"id": [1]})) + assert not connection.connected + assert list(tmp_path.iterdir()) == [] # noqa: ASYNC240 + async with asyncmy_infile_config.provide_session() as driver: + assert await driver.select_value("SELECT 1") == 1 + + +async def test_native_cancellation_discards_connection_and_pool_recovers( + asyncmy_infile_config: AsyncmyConfig, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + import tempfile + + monkeypatch.setattr(tempfile, "tempdir", str(tmp_path)) + async with asyncmy_infile_config.provide_session() as blocker: + await blocker.execute("LOCK TABLES asyncmy_native_infile WRITE") + try: + async with asyncmy_infile_config.provide_session() as driver: + connection = driver.connection + connection_id = await driver.select_value("SELECT CONNECTION_ID()") + task = asyncio.create_task(driver.load_from_arrow("asyncmy_native_infile", pa.table({"id": [1]}))) + try: + async with asyncmy_infile_config.provide_session() as observer: + for _ in range(200): + process = await observer.select_one_or_none( + "SELECT INFO FROM information_schema.PROCESSLIST WHERE ID = :id", {"id": connection_id} + ) + if process and str(process["INFO"]).startswith("LOAD DATA LOCAL INFILE"): + break + await asyncio.sleep(0.01) + else: + pytest.fail("Native LOAD DATA query did not reach MySQL") + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert not connection.connected + finally: + if not task.done(): + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + finally: + await blocker.execute("UNLOCK TABLES") + assert list(tmp_path.iterdir()) == [] # noqa: ASYNC240 + async with asyncmy_infile_config.provide_session() as driver: + assert await driver.select_value("SELECT 1") == 1 diff --git a/tests/unit/adapters/test_aiomysql/test_config.py b/tests/unit/adapters/test_aiomysql/test_config.py index 68d089049..8a0397543 100644 --- a/tests/unit/adapters/test_aiomysql/test_config.py +++ b/tests/unit/adapters/test_aiomysql/test_config.py @@ -9,7 +9,6 @@ from sqlspec.adapters.aiomysql._typing import AiomysqlCursor, AiomysqlDictCursor, AiomysqlRawCursor from sqlspec.adapters.aiomysql.config import AiomysqlConfig from sqlspec.adapters.aiomysql.core import build_statement_config -from sqlspec.exceptions import ImproperConfigurationError def test_build_default_statement_config_custom_serializers() -> None: @@ -79,16 +78,35 @@ def test_aiomysql_connection_kwargs_normalize_cursor_alias_and_omit_pool_only_ke assert "allow_local_infile" not in connect_kwargs -def test_aiomysql_local_infile_requires_explicit_security_gate() -> None: - """LOAD DATA LOCAL INFILE should require a separate consent gate.""" - with pytest.raises(ImproperConfigurationError, match="allow_local_infile=True"): - AiomysqlConfig(connection_config={"local_infile": True}) - - config = AiomysqlConfig(connection_config={"allow_local_infile": True, "local_infile": True}) - connect_kwargs = config._connection_kwargs() # pyright: ignore[reportPrivateUsage] - - assert connect_kwargs["local_infile"] is True - assert "allow_local_infile" not in connect_kwargs +@pytest.mark.parametrize( + ("connection_config", "enabled"), + [ + ({}, False), + ({"local_infile": False}, False), + ({"allow_local_infile": False}, False), + ({"local_infile": False, "allow_local_infile": False}, False), + ({"local_infile": True}, True), + ({"allow_local_infile": True}, True), + ({"local_infile": True, "allow_local_infile": False}, True), + ({"local_infile": False, "allow_local_infile": True}, True), + ({"local_infile": True, "allow_local_infile": True}, True), + ], +) +def test_local_infile_aliases_enable_native_bulk(connection_config: dict[str, bool], enabled: bool) -> None: + config = AiomysqlConfig(connection_config=connection_config) + assert config.connection_config["local_infile"] is enabled + assert "allow_local_infile" not in config.connection_config + assert config.driver_features["enable_local_infile_bulk_load"] is enabled + kwargs = config._connection_kwargs() # pyright: ignore[reportPrivateUsage] + assert kwargs["local_infile"] is enabled + assert "allow_local_infile" not in kwargs + + +@pytest.mark.parametrize("flag", ["local_infile", "allow_local_infile"]) +def test_local_infile_explicit_bulk_disable(flag: str) -> None: + config = AiomysqlConfig(connection_config={flag: True}, driver_features={"enable_local_infile_bulk_load": False}) + assert config.connection_config["local_infile"] is True + assert config.driver_features["enable_local_infile_bulk_load"] is False def test_aiomysql_connection_kwargs_default_local_infile_disabled() -> None: diff --git a/tests/unit/adapters/test_aiomysql/test_load_from_arrow.py b/tests/unit/adapters/test_aiomysql/test_load_from_arrow.py index 6b94fce26..ece33bdcb 100644 --- a/tests/unit/adapters/test_aiomysql/test_load_from_arrow.py +++ b/tests/unit/adapters/test_aiomysql/test_load_from_arrow.py @@ -7,6 +7,7 @@ import pyarrow as pa import pyarrow.parquet as pq +from sqlspec.adapters.aiomysql.config import AiomysqlConfig from sqlspec.adapters.aiomysql.driver import AiomysqlDriver _CAPS: dict[str, Any] = { @@ -47,9 +48,12 @@ async def cursor(self, *_args: Any, **_kwargs: Any) -> _FakeCursor: def _make_driver(connection: _FakeConnection, *, enable_local_infile: bool) -> AiomysqlDriver: + config = AiomysqlConfig( + connection_config={"local_infile": True}, + driver_features={} if enable_local_infile else {"enable_local_infile_bulk_load": False}, + ) return AiomysqlDriver( - connection=cast("Any", connection), - driver_features={"storage_capabilities": _CAPS, "enable_local_infile_bulk_load": enable_local_infile}, + connection=cast("Any", connection), driver_features={**config.driver_features, "storage_capabilities": _CAPS} ) @@ -65,7 +69,7 @@ async def test_load_from_arrow_local_infile_writes_tsv_and_loads() -> None: assert conn._cursor.executemany_calls == [] -async def test_load_from_arrow_without_feature_uses_executemany() -> None: +async def test_load_from_arrow_explicit_bulk_disable_uses_executemany() -> None: conn = _FakeConnection() driver = _make_driver(conn, enable_local_infile=False) 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 1195c8d0a..14e4d13ea 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 @@ -32,8 +32,5 @@ def test_config_gate_raises_when_local_infile_disabled() -> None: def test_config_gate_allows_when_local_infile_set() -> None: - config = AiomysqlConfig( - connection_config={"local_infile": True, "allow_local_infile": True}, - driver_features={"enable_local_infile_bulk_load": True}, - ) + config = AiomysqlConfig(connection_config={"local_infile": True}) assert config.driver_features["enable_local_infile_bulk_load"] is True diff --git a/tests/unit/adapters/test_asyncmy/test_config.py b/tests/unit/adapters/test_asyncmy/test_config.py index 50789ef92..6f1c8f4aa 100644 --- a/tests/unit/adapters/test_asyncmy/test_config.py +++ b/tests/unit/adapters/test_asyncmy/test_config.py @@ -79,24 +79,38 @@ def test_asyncmy_cursor_cls_and_cursor_class_conflict_raises() -> None: AsyncmyConfig(connection_config={"cursor_cls": object, "cursor_class": AsyncmyDictCursor}) -def test_asyncmy_local_infile_requires_explicit_security_gate() -> None: - """LOAD DATA LOCAL INFILE should stay disabled unless separately gated.""" - with pytest.raises(ImproperConfigurationError, match="allow_local_infile=True"): - AsyncmyConfig(connection_config={"local_infile": True}) +@pytest.mark.parametrize( + ("connection_config", "enabled"), + [ + ({}, False), + ({"local_infile": False}, False), + ({"allow_local_infile": False}, False), + ({"local_infile": False, "allow_local_infile": False}, False), + ({"local_infile": True}, True), + ({"allow_local_infile": True}, True), + ({"local_infile": True, "allow_local_infile": False}, True), + ({"local_infile": False, "allow_local_infile": True}, True), + ({"local_infile": True, "allow_local_infile": True}, True), + ], +) +def test_local_infile_aliases_enable_native_bulk(connection_config: dict[str, bool], enabled: bool) -> None: + config = AsyncmyConfig(connection_config=connection_config) + assert config.connection_config["local_infile"] is enabled + assert "allow_local_infile" not in config.connection_config + assert config.driver_features["enable_local_infile_bulk_load"] is enabled - config = AsyncmyConfig(connection_config={"allow_local_infile": True, "local_infile": True}) +@pytest.mark.parametrize("flag", ["local_infile", "allow_local_infile"]) +def test_local_infile_explicit_bulk_disable(flag: str) -> None: + config = AsyncmyConfig(connection_config={flag: True}, driver_features={"enable_local_infile_bulk_load": False}) assert config.connection_config["local_infile"] is True - assert "allow_local_infile" not in config.connection_config + assert config.driver_features["enable_local_infile_bulk_load"] is False -def test_asyncmy_rejects_local_infile_bulk_load_feature() -> None: - """Asyncmy exposes local_infile, but its LOAD DATA LOCAL INFILE protocol path is not usable.""" - with pytest.raises(ImproperConfigurationError, match="asyncmy does not currently support"): - AsyncmyConfig( - connection_config={"local_infile": True, "allow_local_infile": True}, - driver_features={"enable_local_infile_bulk_load": True}, - ) +@pytest.mark.parametrize("connection_config", [{}, {"local_infile": False}, {"allow_local_infile": False}]) +def test_bulk_load_requires_connection_opt_in(connection_config: dict[str, bool]) -> None: + with pytest.raises(ImproperConfigurationError, match="local_infile=True"): + AsyncmyConfig(connection_config=connection_config, driver_features={"enable_local_infile_bulk_load": True}) async def test_asyncmy_create_pool_normalizes_connection_and_pool_kwargs(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/tests/unit/adapters/test_asyncmy/test_load_from_arrow.py b/tests/unit/adapters/test_asyncmy/test_load_from_arrow.py index f6ae77385..f85b5bb83 100644 --- a/tests/unit/adapters/test_asyncmy/test_load_from_arrow.py +++ b/tests/unit/adapters/test_asyncmy/test_load_from_arrow.py @@ -1,11 +1,15 @@ """Asyncmy load_from_arrow ingest paths.""" +import asyncio +from datetime import timedelta from pathlib import Path from typing import Any, cast import pyarrow as pa import pyarrow.parquet as pq +import pytest +from sqlspec.adapters.asyncmy.config import AsyncmyConfig from sqlspec.adapters.asyncmy.driver import AsyncmyDriver _CAPS: dict[str, Any] = { @@ -22,9 +26,17 @@ def __init__(self) -> None: self.execute_calls: list[str] = [] self.executemany_calls: list[tuple[str, list[Any]]] = [] self.rowcount = 0 + self.payload: bytes | None = None + self.payload_path: Path | None = None + self.failure: BaseException | None = None async def execute(self, sql: str, *_args: Any) -> None: self.execute_calls.append(sql) + if sql.startswith("LOAD DATA"): + self.payload_path = Path(_args[0][0]) + self.payload = self.payload_path.read_bytes() + if self.failure is not None: + raise self.failure async def executemany(self, sql: str, params: Any) -> None: self.executemany_calls.append((sql, [tuple(row) for row in params])) @@ -36,6 +48,13 @@ async def close(self) -> None: class _FakeConnection: def __init__(self) -> None: self._cursor = _FakeCursor() + self.closed = False + + def close(self) -> None: + self.closed = True + + async def _read_query_result(self, unbuffered: bool = False) -> None: + pass def cursor(self, *_args: Any, **_kwargs: Any) -> _FakeCursor: return self._cursor @@ -45,9 +64,15 @@ def _make_driver(connection: _FakeConnection) -> AsyncmyDriver: return AsyncmyDriver(connection=cast("Any", connection), driver_features={"storage_capabilities": _CAPS}) -async def test_load_from_arrow_uses_executemany() -> None: +@pytest.mark.parametrize("connection_opt_in", [False, True]) +async def test_load_from_arrow_explicit_bulk_disable_uses_executemany(connection_opt_in: bool) -> None: conn = _FakeConnection() - driver = _make_driver(conn) + config = AsyncmyConfig( + connection_config={"local_infile": connection_opt_in}, driver_features={"enable_local_infile_bulk_load": False} + ) + driver = AsyncmyDriver( + connection=cast("Any", conn), driver_features={**config.driver_features, "storage_capabilities": _CAPS} + ) job = await driver.load_from_arrow("orders", pa.table({"id": [1, 2], "name": ["a", "b"]})) @@ -80,3 +105,167 @@ async def test_load_from_storage_reads_parquet_and_delegates(tmp_path: Path) -> insert_sql, rows = conn._cursor.executemany_calls[0] assert insert_sql.startswith("INSERT INTO") assert rows == [(1, "a"), (2, "b")] + + +async def test_local_infile_payload_roundtrip_and_cleanup() -> None: + conn = _FakeConnection() + config = AsyncmyConfig(connection_config={"allow_local_infile": True}) + driver = AsyncmyDriver( + connection=cast("Any", conn), driver_features={**config.driver_features, "storage_capabilities": _CAPS} + ) + job = await driver.load_from_arrow( + "order%`table", pa.table({"id": [1, 2], "text.with%tick`": ["é\t\n\r\\\x00\x1a", None], "flag": [True, False]}) + ) + assert job.telemetry["rows_processed"] == 2 + assert conn._cursor.executemany_calls == [] + assert conn._cursor.payload == "1\té\\t\\n\\r\\\\\\0\\Z\t1\n2\t\\N\t0\n".encode() + assert "LOCAL INFILE %s" in conn._cursor.execute_calls[0] + assert "`text.with%%tick```" in conn._cursor.execute_calls[0] + assert conn._cursor.payload_path is not None + assert not conn._cursor.payload_path.parent.exists() + assert not conn.closed + + +@pytest.mark.parametrize("failure", [RuntimeError("transfer failed"), asyncio.CancelledError()]) +async def test_local_infile_failure_discards_connection_and_payload(failure: BaseException) -> None: + conn = _FakeConnection() + conn._cursor.failure = failure + driver = AsyncmyDriver( + connection=cast("Any", conn), + driver_features={"storage_capabilities": _CAPS, "enable_local_infile_bulk_load": True}, + ) + with pytest.raises(type(failure)) as caught: + await driver.load_from_arrow("orders", pa.table({"id": [1]})) + assert caught.value is failure + assert conn.closed + assert conn._cursor.payload_path is not None + assert not conn._cursor.payload_path.parent.exists() + + +async def test_local_infile_empty_table_does_not_send_payload() -> None: + conn = _FakeConnection() + driver = AsyncmyDriver( + connection=cast("Any", conn), + driver_features={"storage_capabilities": _CAPS, "enable_local_infile_bulk_load": True}, + ) + job = await driver.load_from_arrow("orders", pa.table({"id": pa.array([], type=pa.int64())}), overwrite=True) + assert job.telemetry["rows_processed"] == 0 + assert conn._cursor.execute_calls == ["TRUNCATE TABLE `orders`"] + assert conn._cursor.executemany_calls == [] + assert conn._cursor.payload_path is None + + +@pytest.mark.parametrize( + ("values", "expected"), + [([[1, 2]], "[1,2]"), ([b"\x00\xff"], b"\x00\xff"), ([timedelta(days=1)], timedelta(days=1))], +) +async def test_local_infile_unsupported_values_fall_back_to_executemany(values: Any, expected: Any) -> None: + conn = _FakeConnection() + driver = AsyncmyDriver( + connection=cast("Any", conn), + driver_features={"storage_capabilities": _CAPS, "enable_local_infile_bulk_load": True}, + ) + await driver.load_from_arrow("orders", pa.table({"data": values})) + assert conn._cursor.execute_calls == [] + assert len(conn._cursor.executemany_calls) == 1 + assert conn._cursor.executemany_calls[0][1] == [(expected,)] + assert conn._cursor.payload_path is None + + +@pytest.mark.parametrize("stage", ["encode", "create", "write"]) +async def test_local_infile_preparation_failure_removes_private_directory( + stage: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + import tempfile + from unittest.mock import Mock + + import sqlspec.adapters.asyncmy.driver as driver_module + + monkeypatch.setattr(tempfile, "tempdir", str(tmp_path)) + failure = OSError("payload preparation failed") + if stage == "encode": + monkeypatch.setattr(driver_module, "encode_records_for_local_infile", Mock(side_effect=failure)) + elif stage == "create": + monkeypatch.setattr(tempfile, "NamedTemporaryFile", Mock(side_effect=failure)) + else: + original = tempfile.NamedTemporaryFile + + def failing_writer(*args: Any, **kwargs: Any) -> Any: + file = original(*args, **kwargs) + file.write = Mock(side_effect=failure) + return file + + monkeypatch.setattr(tempfile, "NamedTemporaryFile", failing_writer) + conn = _FakeConnection() + driver = AsyncmyDriver( + connection=cast("Any", conn), + driver_features={"storage_capabilities": _CAPS, "enable_local_infile_bulk_load": True}, + ) + with pytest.raises(OSError, match="payload preparation failed"): + await driver.load_from_arrow("orders", pa.table({"id": [1]})) + assert list(tmp_path.iterdir()) == [] # noqa: ASYNC240 + assert conn._cursor.execute_calls == [] + assert not conn.closed + + +@pytest.mark.parametrize("unbuffered", [False, True]) +@pytest.mark.parametrize("existing_hook", [False, True]) +@pytest.mark.parametrize("outcome", ["success", "unexpected_filename", "invalid_ack", "sender_error", "cancelled"]) +async def test_local_infile_native_handoff_and_hook_restoration( + outcome: str, existing_hook: bool, unbuffered: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + from types import SimpleNamespace + from unittest.mock import AsyncMock, Mock + + from asyncmy import Connection + from asyncmy.protocol import MysqlPacket + + import sqlspec.adapters.asyncmy._typing as native + from sqlspec.exceptions import SQLSpecError + + connection: Any = Connection(local_infile=True) # type: ignore[no-untyped-call] + connection._connected = True + original_reader = AsyncMock() + if existing_hook: + connection._read_query_result = original_reader + request = b"\xfb/tmp/unexpected" if outcome == "unexpected_filename" else b"\xfb/tmp/payload" + ack = b"\xfe\x00\x00\x00\x00" if outcome == "invalid_ack" else b"\x00\x01\x00\x02\x00\x00\x00" + connection.read_packet = AsyncMock(side_effect=[MysqlPacket(request, "utf8"), MysqlPacket(ack, "utf8")]) + failure = asyncio.CancelledError() if outcome == "cancelled" else OSError("native sender failed") + send = AsyncMock(side_effect=failure if outcome in {"sender_error", "cancelled"} else None) + sender = Mock(return_value=SimpleNamespace(send_data=send)) + monkeypatch.setattr(native, "_LoadLocalFile", sender) + result_type = native._AsyncmyLocalInfileResult + results: list[Any] = [] + + def capture_result(raw: Any, filename: str) -> Any: + result = result_type(raw, filename) + results.append(result) + return result + + monkeypatch.setattr(native, "_AsyncmyLocalInfileResult", capture_result) + if outcome == "success": + with native.asyncmy_local_infile(connection, "/tmp/payload"): + await connection._read_query_result(unbuffered=unbuffered) + assert connection._affected_rows == 1 + assert connection.server_status == 2 + assert connection.connected + sender.assert_called_once_with("/tmp/payload", connection) + send.assert_awaited_once() + else: + error_type = type(failure) if outcome in {"sender_error", "cancelled"} else SQLSpecError + with pytest.raises(error_type): + with native.asyncmy_local_infile(connection, "/tmp/payload"): + await connection._read_query_result(unbuffered=unbuffered) + assert not connection.connected + if unbuffered: + assert len(results) == 1 + assert not results[0].unbuffered_active + assert results[0].connection is None + if outcome == "unexpected_filename": + sender.assert_not_called() + if existing_hook: + assert connection._read_query_result is original_reader + else: + assert "_read_query_result" not in connection.__dict__ + original_reader.assert_not_called() diff --git a/tests/unit/adapters/test_cockroach_retry.py b/tests/unit/adapters/test_cockroach_retry.py new file mode 100644 index 000000000..e4ef29802 --- /dev/null +++ b/tests/unit/adapters/test_cockroach_retry.py @@ -0,0 +1,76 @@ +"""A secondary rollback failure must not replace a transaction's outcome.""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from sqlspec.adapters.cockroach_asyncpg import CockroachAsyncpgDriver +from sqlspec.adapters.cockroach_psycopg import CockroachPsycopgAsyncDriver, CockroachPsycopgSyncDriver + + +class _RetryableError(Exception): + sqlstate = "40001" + + +@pytest.mark.parametrize("outcome", ["retry", "exhausted", "nonretryable"]) +def test_sync_retry_preserves_outcome_when_rollback_fails(outcome: str, monkeypatch: pytest.MonkeyPatch) -> None: + driver_type = CockroachPsycopgSyncDriver + monkeypatch.setattr(driver_type, "_connection_in_transaction", lambda _self: False) + begin = MagicMock() + commit = MagicMock() + rollback = MagicMock(side_effect=RuntimeError("secondary rollback failure")) + monkeypatch.setattr(driver_type, "begin", begin) + monkeypatch.setattr(driver_type, "commit", commit) + monkeypatch.setattr(driver_type, "rollback", rollback) + driver = driver_type( + connection=MagicMock(), + driver_features={"max_retries": 1, "retry_delay_base_ms": 0, "enable_retry_logging": False}, + ) + original = ValueError("operation failed") if outcome == "nonretryable" else _RetryableError("restart transaction") + operation = MagicMock(side_effect=[original, "ok"] if outcome == "retry" else original) + + if outcome == "retry": + assert driver.run_transaction_with_retry(operation) == "ok" + else: + with pytest.raises(type(original)) as caught: + driver.run_transaction_with_retry(operation) + assert caught.value is original + + attempts = 1 if outcome == "nonretryable" else 2 + assert operation.call_count == begin.call_count == attempts + assert rollback.call_count == (1 if outcome == "retry" else attempts) + assert commit.call_count == (1 if outcome == "retry" else 0) + + +@pytest.mark.parametrize("driver_type", [CockroachAsyncpgDriver, CockroachPsycopgAsyncDriver]) +@pytest.mark.parametrize("outcome", ["retry", "exhausted", "nonretryable"]) +async def test_async_retry_preserves_outcome_when_rollback_fails( + driver_type: type[CockroachAsyncpgDriver | CockroachPsycopgAsyncDriver], + outcome: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(driver_type, "_connection_in_transaction", lambda _self: False) + begin = AsyncMock() + commit = AsyncMock() + rollback = AsyncMock(side_effect=RuntimeError("secondary rollback failure")) + monkeypatch.setattr(driver_type, "begin", begin) + monkeypatch.setattr(driver_type, "commit", commit) + monkeypatch.setattr(driver_type, "rollback", rollback) + driver = driver_type( + connection=MagicMock(), + driver_features={"max_retries": 1, "retry_delay_base_ms": 0, "enable_retry_logging": False}, + ) + original = ValueError("operation failed") if outcome == "nonretryable" else _RetryableError("restart transaction") + operation = AsyncMock(side_effect=[original, "ok"] if outcome == "retry" else original) + + if outcome == "retry": + assert await driver.run_transaction_with_retry(operation) == "ok" + else: + with pytest.raises(type(original)) as caught: + await driver.run_transaction_with_retry(operation) + assert caught.value is original + + attempts = 1 if outcome == "nonretryable" else 2 + assert operation.await_count == begin.await_count == attempts + assert rollback.await_count == (1 if outcome == "retry" else attempts) + assert commit.await_count == (1 if outcome == "retry" else 0) diff --git a/tests/unit/adapters/test_mysqlconnector/test_config.py b/tests/unit/adapters/test_mysqlconnector/test_config.py index bf4af0aab..d270cd141 100644 --- a/tests/unit/adapters/test_mysqlconnector/test_config.py +++ b/tests/unit/adapters/test_mysqlconnector/test_config.py @@ -80,12 +80,35 @@ def _fake_connect(**kwargs: Any) -> Any: @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}) +@pytest.mark.parametrize( + ("connection_config", "enabled"), + [ + ({}, False), + ({"local_infile": False}, False), + ({"allow_local_infile": False}, False), + ({"local_infile": False, "allow_local_infile": False}, False), + ({"local_infile": True}, True), + ({"allow_local_infile": True}, True), + ({"local_infile": True, "allow_local_infile": False}, True), + ({"local_infile": False, "allow_local_infile": True}, True), + ({"local_infile": True, "allow_local_infile": True}, True), + ], +) +def test_local_infile_aliases_enable_native_bulk( + config_cls: type[Any], connection_config: dict[str, bool], enabled: bool +) -> None: + config = config_cls(connection_config=connection_config) + assert config.connection_config["allow_local_infile"] is enabled + assert "local_infile" not in config.connection_config + assert config.driver_features["enable_local_infile_bulk_load"] is enabled + +@pytest.mark.parametrize("config_cls", [MysqlConnectorSyncConfig, MysqlConnectorAsyncConfig]) +@pytest.mark.parametrize("flag", ["local_infile", "allow_local_infile"]) +def test_local_infile_explicit_bulk_disable(config_cls: type[Any], flag: str) -> None: + config = config_cls(connection_config={flag: True}, driver_features={"enable_local_infile_bulk_load": False}) assert config.connection_config["allow_local_infile"] is True - assert "local_infile" not in config.connection_config + assert config.driver_features["enable_local_infile_bulk_load"] is False def test_sync_connection_params_type_accepts_modern_connector_options() -> None: diff --git a/tests/unit/adapters/test_mysqlconnector/test_load_from_arrow.py b/tests/unit/adapters/test_mysqlconnector/test_load_from_arrow.py index cb2c2b07c..a50a530f5 100644 --- a/tests/unit/adapters/test_mysqlconnector/test_load_from_arrow.py +++ b/tests/unit/adapters/test_mysqlconnector/test_load_from_arrow.py @@ -7,6 +7,7 @@ import pyarrow as pa import pyarrow.parquet as pq +from sqlspec.adapters.mysqlconnector.config import MysqlConnectorAsyncConfig, MysqlConnectorSyncConfig from sqlspec.adapters.mysqlconnector.driver import MysqlConnectorAsyncDriver, MysqlConnectorSyncDriver _CAPS: dict[str, Any] = { @@ -74,9 +75,9 @@ async def cursor(self, *_args: Any, **_kwargs: Any) -> _FakeAsyncCursor: def test_sync_load_from_arrow_local_infile_writes_tsv_and_loads() -> None: conn = _FakeSyncConnection() + config = MysqlConnectorSyncConfig(connection_config={"local_infile": True}) driver = MysqlConnectorSyncDriver( - connection=cast("Any", conn), - driver_features={"storage_capabilities": _CAPS, "enable_local_infile_bulk_load": True}, + connection=cast("Any", conn), driver_features={**config.driver_features, "storage_capabilities": _CAPS} ) job = driver.load_from_arrow("orders", pa.table({"id": [1, 2], "name": ["a", "b"]})) @@ -86,9 +87,14 @@ def test_sync_load_from_arrow_local_infile_writes_tsv_and_loads() -> None: assert conn._cursor.execute_calls[0].startswith("LOAD DATA LOCAL INFILE") -def test_sync_load_from_arrow_without_feature_uses_executemany() -> None: +def test_sync_load_from_arrow_explicit_bulk_disable_uses_executemany() -> None: conn = _FakeSyncConnection() - driver = MysqlConnectorSyncDriver(connection=cast("Any", conn), driver_features={"storage_capabilities": _CAPS}) + config = MysqlConnectorSyncConfig( + connection_config={"allow_local_infile": True}, driver_features={"enable_local_infile_bulk_load": False} + ) + driver = MysqlConnectorSyncDriver( + connection=cast("Any", conn), driver_features={**config.driver_features, "storage_capabilities": _CAPS} + ) driver.load_from_arrow("orders", pa.table({"id": [1, 2], "name": ["a", "b"]})) @@ -111,9 +117,9 @@ def test_sync_load_from_arrow_overwrite_truncates_first() -> None: async def test_async_load_from_arrow_local_infile_writes_tsv_and_loads() -> None: conn = _FakeAsyncConnection() + config = MysqlConnectorAsyncConfig(connection_config={"local_infile": True}) driver = MysqlConnectorAsyncDriver( - connection=cast("Any", conn), - driver_features={"storage_capabilities": _CAPS, "enable_local_infile_bulk_load": True}, + connection=cast("Any", conn), driver_features={**config.driver_features, "storage_capabilities": _CAPS} ) job = await driver.load_from_arrow("orders", pa.table({"id": [1, 2], "name": ["a", "b"]})) @@ -123,9 +129,14 @@ async def test_async_load_from_arrow_local_infile_writes_tsv_and_loads() -> None assert conn._cursor.execute_calls[0].startswith("LOAD DATA LOCAL INFILE") -async def test_async_load_from_arrow_without_feature_uses_executemany() -> None: +async def test_async_load_from_arrow_explicit_bulk_disable_uses_executemany() -> None: conn = _FakeAsyncConnection() - driver = MysqlConnectorAsyncDriver(connection=cast("Any", conn), driver_features={"storage_capabilities": _CAPS}) + config = MysqlConnectorAsyncConfig( + connection_config={"allow_local_infile": True}, driver_features={"enable_local_infile_bulk_load": False} + ) + driver = MysqlConnectorAsyncDriver( + connection=cast("Any", conn), driver_features={**config.driver_features, "storage_capabilities": _CAPS} + ) await driver.load_from_arrow("orders", pa.table({"id": [1, 2], "name": ["a", "b"]})) diff --git a/tests/unit/adapters/test_mysqlconnector/test_local_infile_bulk_load.py b/tests/unit/adapters/test_mysqlconnector/test_local_infile_bulk_load.py index 680a2ddf5..805a1dc85 100644 --- a/tests/unit/adapters/test_mysqlconnector/test_local_infile_bulk_load.py +++ b/tests/unit/adapters/test_mysqlconnector/test_local_infile_bulk_load.py @@ -28,9 +28,7 @@ def test_sync_config_gate_raises_when_allow_local_infile_disabled() -> None: def test_sync_config_gate_allows_when_allow_local_infile_enabled() -> None: - config = MysqlConnectorSyncConfig( - connection_config={"allow_local_infile": True}, driver_features={"enable_local_infile_bulk_load": True} - ) + config = MysqlConnectorSyncConfig(connection_config={"allow_local_infile": True}) assert config.driver_features["enable_local_infile_bulk_load"] is True @@ -40,7 +38,5 @@ def test_async_config_gate_raises_when_allow_local_infile_disabled() -> None: def test_async_config_gate_allows_when_allow_local_infile_enabled() -> None: - config = MysqlConnectorAsyncConfig( - connection_config={"allow_local_infile": True}, driver_features={"enable_local_infile_bulk_load": True} - ) + config = MysqlConnectorAsyncConfig(connection_config={"allow_local_infile": True}) assert config.driver_features["enable_local_infile_bulk_load"] is True diff --git a/tests/unit/adapters/test_psycopg/test_config.py b/tests/unit/adapters/test_psycopg/test_config.py index aa3c07788..b340ac64f 100644 --- a/tests/unit/adapters/test_psycopg/test_config.py +++ b/tests/unit/adapters/test_psycopg/test_config.py @@ -367,7 +367,7 @@ def test_psycopg_sync_session_context_resolves_callable_statement_config() -> No expected_config = StatementConfig(dialect="pgvector") context = PsycopgSyncSessionContext( acquire_connection=lambda: object(), - release_connection=lambda _conn: None, + release_connection=lambda _conn, **_kwargs: None, statement_config=lambda: expected_config, driver_features={}, prepare_driver=lambda driver: driver, @@ -398,7 +398,7 @@ def test_psycopg_sync_session_context_preserves_explicit_statement_config() -> N explicit_config = StatementConfig(dialect="postgres") context = PsycopgSyncSessionContext( acquire_connection=lambda: object(), - release_connection=lambda _conn: None, + release_connection=lambda _conn, **_kwargs: None, statement_config=explicit_config, driver_features={}, prepare_driver=lambda driver: driver, diff --git a/tests/unit/adapters/test_psycopg/test_driver.py b/tests/unit/adapters/test_psycopg/test_driver.py index 09ecb6a6f..74b2e9536 100644 --- a/tests/unit/adapters/test_psycopg/test_driver.py +++ b/tests/unit/adapters/test_psycopg/test_driver.py @@ -1,11 +1,15 @@ """Unit tests for psycopg driver transaction behavior.""" +from asyncio import CancelledError +from contextlib import nullcontext from types import SimpleNamespace from typing import TYPE_CHECKING, cast +from unittest.mock import AsyncMock, MagicMock import psycopg import pytest +from sqlspec.adapters.psycopg.config import PsycopgAsyncConfig, PsycopgSyncConfig from sqlspec.adapters.psycopg.driver import PsycopgAsyncDriver, PsycopgSyncDriver from sqlspec.exceptions import SQLSpecError @@ -13,6 +17,67 @@ from sqlspec.adapters.psycopg._typing import PsycopgAsyncConnection, PsycopgSyncConnection +@pytest.mark.parametrize("error", [None, ValueError("session failed"), CancelledError("session cancelled")]) +def test_psycopg_sync_session_forwards_pool_exit_exception(error: BaseException | None) -> None: + pool = MagicMock() + pool_context = pool.connection.return_value + pool_context.__exit__.return_value = False + config = PsycopgSyncConfig(connection_instance=pool) + + with pytest.raises(type(error)) if error is not None else nullcontext(): + with config.provide_session(): + if error is not None: + raise error + + pool_context.__exit__.assert_called_once_with( + type(error) if error is not None else None, error, error.__traceback__ if error is not None else None + ) + + +@pytest.mark.parametrize("error", [None, ValueError("session failed"), CancelledError("session cancelled")]) +async def test_psycopg_async_session_forwards_pool_exit_exception(error: BaseException | None) -> None: + pool = MagicMock() + pool_context = pool.connection.return_value + pool_context.__aenter__ = AsyncMock(return_value=MagicMock()) + pool_context.__aexit__ = AsyncMock(return_value=False) + config = PsycopgAsyncConfig(connection_instance=pool) + + with pytest.raises(type(error)) if error is not None else nullcontext(): + async with config.provide_session(): + if error is not None: + raise error + + pool_context.__aexit__.assert_awaited_once_with( + type(error) if error is not None else None, error, error.__traceback__ if error is not None else None + ) + + +def test_psycopg_sync_session_in_except_handler_exits_successfully() -> None: + pool = MagicMock() + config = PsycopgSyncConfig(connection_instance=pool) + + try: + raise ValueError("previous operation failed") + except ValueError: + with config.provide_session(): + pass + + pool.connection.return_value.__exit__.assert_called_once_with(None, None, None) + + +async def test_psycopg_async_session_in_except_handler_exits_successfully() -> None: + pool = MagicMock() + config = PsycopgAsyncConfig(connection_instance=pool) + + try: + raise ValueError("previous operation failed") + except ValueError: + async with config.provide_session(): + pass + + pool.connection.return_value.__aexit__.assert_awaited_once_with(None, None, None) + + class _SyncTransactionConnection: def __init__( self, diff --git a/tests/unit/adapters/test_pymssql/test_driver.py b/tests/unit/adapters/test_pymssql/test_driver.py index fd8886572..6b8189de3 100644 --- a/tests/unit/adapters/test_pymssql/test_driver.py +++ b/tests/unit/adapters/test_pymssql/test_driver.py @@ -17,6 +17,33 @@ UNSAFE_SAVEPOINT_NAMES = ["1; DROP TABLE users", "sp-1", "sp 1", "", '"sp"'] +@pytest.mark.parametrize( + ("rows", "expected"), + [ + ( + [{"name": "Ada", "id": 1}, {"name": "Grace", "id": 2}], + [{"id": 1, "name": "Ada"}, {"id": 2, "name": "Grace"}], + ), + ([(1, "Ada"), (2, "Grace")], [{"id": 1, "name": "Ada"}, {"id": 2, "name": "Grace"}]), + ([], []), + ], + ids=["dict", "tuple", "empty"], +) +def test_execute_maps_pymssql_row_formats( + rows: list[tuple[int, str] | dict[str, int | str]], expected: list[dict[str, int | str]] +) -> None: + from sqlspec.adapters.pymssql.driver import PymssqlDriver + + cursor = FakeCursor(rows=rows, description=[("id",), ("name",)]) + driver = PymssqlDriver(cast("PymssqlConnection", FakeConnection(cursor))) + + result = driver.execute("SELECT id, name FROM dbo.users") + + assert result.get_data() == expected + assert result.column_names == ["id", "name"] + assert len(cursor.calls) == 1 + + @pytest.mark.parametrize("bad_name", UNSAFE_SAVEPOINT_NAMES) def test_pymssql_savepoint_overrides_reject_unsafe_names(bad_name: str) -> None: """The T-SQL savepoint overrides must reject unsafe identifiers before interpolation.""" diff --git a/tests/unit/adapters/test_pymysql/test_config.py b/tests/unit/adapters/test_pymysql/test_config.py index 416413d96..3ccd622f5 100644 --- a/tests/unit/adapters/test_pymysql/test_config.py +++ b/tests/unit/adapters/test_pymysql/test_config.py @@ -8,7 +8,6 @@ from typing_extensions import NotRequired from sqlspec.adapters.pymysql.config import PyMysqlConfig, PyMysqlConnectionParams -from sqlspec.exceptions import ImproperConfigurationError def _unwrap_not_required(annotation: object) -> object: @@ -80,28 +79,35 @@ def test_create_pool_defaults_local_infile_off() -> None: assert pool._connection_parameters["local_infile"] is False -def test_create_pool_preserves_local_infile_opt_in_with_security_gate() -> None: - """Explicit local infile opt-in should still pass through to PyMySQL after consent.""" - config = PyMysqlConfig(connection_config={"allow_local_infile": True, "local_infile": True}) +@pytest.mark.parametrize( + ("connection_config", "enabled"), + [ + ({}, False), + ({"local_infile": False}, False), + ({"allow_local_infile": False}, False), + ({"local_infile": False, "allow_local_infile": False}, False), + ({"local_infile": True}, True), + ({"allow_local_infile": True}, True), + ({"local_infile": True, "allow_local_infile": False}, True), + ({"local_infile": False, "allow_local_infile": True}, True), + ({"local_infile": True, "allow_local_infile": True}, True), + ], +) +def test_local_infile_aliases_enable_native_bulk(connection_config: dict[str, bool], enabled: bool) -> None: + config = PyMysqlConfig(connection_config=connection_config) + assert config.connection_config["local_infile"] is enabled + assert "allow_local_infile" not in config.connection_config + assert config.driver_features["enable_local_infile_bulk_load"] is enabled pool = config._create_pool() - - assert pool._connection_parameters["local_infile"] is True + assert pool._connection_parameters["local_infile"] is enabled assert "allow_local_infile" not in pool._connection_parameters -def test_create_pool_rejects_local_infile_without_security_gate() -> None: - """PyMySQL should match asyncmy's separate local-infile consent gate.""" - with pytest.raises(ImproperConfigurationError, match="allow_local_infile=True"): - PyMysqlConfig(connection_config={"local_infile": True}) - - -def test_create_pool_does_not_enable_local_infile_for_gate_only() -> None: - """The consent gate alone should not enable client-file reads.""" - config = PyMysqlConfig(connection_config={"allow_local_infile": True}) - pool = config._create_pool() - - assert pool._connection_parameters["local_infile"] is False - assert "allow_local_infile" not in pool._connection_parameters +@pytest.mark.parametrize("flag", ["local_infile", "allow_local_infile"]) +def test_local_infile_explicit_bulk_disable(flag: str) -> None: + config = PyMysqlConfig(connection_config={flag: True}, driver_features={"enable_local_infile_bulk_load": False}) + assert config.connection_config["local_infile"] is True + assert config.driver_features["enable_local_infile_bulk_load"] is False def test_create_pool_preserves_ssl_context_and_flat_tls_options() -> None: diff --git a/tests/unit/adapters/test_pymysql/test_load_from_arrow.py b/tests/unit/adapters/test_pymysql/test_load_from_arrow.py index a5b80014c..b84ccc117 100644 --- a/tests/unit/adapters/test_pymysql/test_load_from_arrow.py +++ b/tests/unit/adapters/test_pymysql/test_load_from_arrow.py @@ -6,6 +6,7 @@ import pyarrow as pa import pyarrow.parquet as pq +from sqlspec.adapters.pymysql.config import PyMysqlConfig from sqlspec.adapters.pymysql.driver import PyMysqlDriver _CAPS: dict[str, Any] = { @@ -46,9 +47,12 @@ def cursor(self, *_args: Any, **_kwargs: Any) -> _FakeCursor: def _make_driver(connection: _FakeConnection, *, enable_local_infile: bool) -> PyMysqlDriver: + config = PyMysqlConfig( + connection_config={"local_infile": True}, + driver_features={} if enable_local_infile else {"enable_local_infile_bulk_load": False}, + ) return PyMysqlDriver( - connection=cast("Any", connection), - driver_features={"storage_capabilities": _CAPS, "enable_local_infile_bulk_load": enable_local_infile}, + connection=cast("Any", connection), driver_features={**config.driver_features, "storage_capabilities": _CAPS} ) @@ -66,7 +70,7 @@ def test_load_from_arrow_local_infile_writes_tsv_and_loads() -> None: assert conn._cursor.executemany_calls == [] -def test_load_from_arrow_without_feature_uses_executemany() -> None: +def test_load_from_arrow_explicit_bulk_disable_uses_executemany() -> None: conn = _FakeConnection() driver = _make_driver(conn, enable_local_infile=False) diff --git a/tests/unit/adapters/test_pymysql/test_local_infile_bulk_load.py b/tests/unit/adapters/test_pymysql/test_local_infile_bulk_load.py index e2e30bb79..e0084f9ec 100644 --- a/tests/unit/adapters/test_pymysql/test_local_infile_bulk_load.py +++ b/tests/unit/adapters/test_pymysql/test_local_infile_bulk_load.py @@ -41,8 +41,5 @@ def test_config_gate_raises_when_local_infile_disabled() -> None: def test_config_gate_allows_when_local_infile_enabled() -> None: - config = PyMysqlConfig( - connection_config={"allow_local_infile": True, "local_infile": True}, - driver_features={"enable_local_infile_bulk_load": True}, - ) + config = PyMysqlConfig(connection_config={"allow_local_infile": True}) assert config.driver_features["enable_local_infile_bulk_load"] is True diff --git a/uv.lock b/uv.lock index 84d1fee70..242151b5a 100644 --- a/uv.lock +++ b/uv.lock @@ -7195,7 +7195,7 @@ requires-dist = [ { name = "aiomysql", marker = "extra == 'aiomysql'" }, { name = "aiosqlite", marker = "extra == 'aiosqlite'" }, { name = "arrow-odbc", marker = "extra == 'arrow-odbc'", specifier = ">=10.4" }, - { name = "asyncmy", marker = "extra == 'asyncmy'" }, + { name = "asyncmy", marker = "extra == 'asyncmy'", specifier = ">=0.2.13" }, { name = "asyncpg", marker = "extra == 'asyncpg'" }, { name = "asyncpg", marker = "extra == 'cockroachdb'" }, { name = "attrs", marker = "extra == 'attrs'" },