"""mssql-python sync and async drivers."""
import contextlib
from typing import TYPE_CHECKING, Any, TypedDict, cast
from typing_extensions import NotRequired
from sqlspec.adapters.mssql_python._typing import (
MSSQL_PYTHON_MODULE,
MssqlPythonConnection,
MssqlPythonCursor,
MssqlPythonRawCursor,
MssqlPythonSessionContext,
)
from sqlspec.adapters.mssql_python.core import (
create_mapped_exception,
default_statement_config,
driver_profile,
materialize_tuple_rows,
)
from sqlspec.adapters.mssql_python.data_dictionary import MssqlPythonSyncDataDictionary
from sqlspec.core import (
build_arrow_result_from_reader,
build_arrow_result_from_table,
get_cache_config,
register_driver_profile,
)
from sqlspec.driver import BaseSyncExceptionHandler, SyncDriverAdapterBase, SyncRowStream, rows_to_dicts
from sqlspec.driver._common import validate_savepoint_name
from sqlspec.exceptions import SQLSpecError
from sqlspec.utils.arrow_helpers import arrow_reader_with_deferred_close
from sqlspec.utils.logging import get_logger
from sqlspec.utils.module_loader import ensure_pyarrow
from sqlspec.utils.text import quote_identifier, split_qualified_identifier
if TYPE_CHECKING:
from collections.abc import Iterable
from sqlspec.builder import QueryBuilder
from sqlspec.core import SQL, ArrowResult, Statement, StatementConfig, StatementFilter
from sqlspec.driver import ExecutionResult
from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry
from sqlspec.typing import ArrowRecordBatchReader, ArrowReturnFormat, StatementParameters
__all__ = (
"MssqlPythonBulkCopyResult",
"MssqlPythonCursor",
"MssqlPythonDriver",
"MssqlPythonExceptionHandler",
"MssqlPythonSessionContext",
)
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):
"""BulkCopy statistics returned by mssql-python."""
rows_copied: int
batch_count: NotRequired[int]
elapsed_time: NotRequired[float]
class MssqlPythonExceptionHandler(BaseSyncExceptionHandler):
"""Sync context manager handling mssql-python exceptions."""
__slots__ = ()
def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "BaseException") -> bool:
if exc_type is None:
return False
if isinstance(exc_val, _MSSQL_ERROR):
self.pending_exception = create_mapped_exception(cast("Exception", exc_val), logger=logger)
return True
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, error: bool = False) -> 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 MssqlPythonDriver(SyncDriverAdapterBase):
"""mssql-python sync driver."""
__slots__ = ("_column_name_cache", "_data_dictionary", "_restore_autocommit", "_transaction_active")
dialect = "tsql"
def __init__(
self,
connection: "MssqlPythonConnection",
statement_config: "StatementConfig | None" = None,
driver_features: "dict[str, Any] | None" = None,
) -> None:
if statement_config is None:
statement_config = default_statement_config.replace(
enable_caching=get_cache_config().compiled_cache_enabled
)
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]]] = {}
self._restore_autocommit = False
self._transaction_active = False