sqlspec 0.36.0__cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl

sqlspec/adapters/mysqlconnector/litestar/store.py

@@ -0,0 +1,426 @@
+"""MysqlConnector session store for Litestar integration."""
+
+from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING, Any, Final, cast
+
+import mysql.connector
+
+from sqlspec.extensions.litestar.store import BaseSQLSpecStore
+from sqlspec.utils.logging import get_logger
+from sqlspec.utils.sync_tools import async_
+
+if TYPE_CHECKING:
+    from sqlspec.adapters.mysqlconnector.config import MysqlConnectorAsyncConfig, MysqlConnectorSyncConfig
+
+logger = get_logger("sqlspec.adapters.mysqlconnector.litestar.store")
+
+__all__ = ("MysqlConnectorAsyncStore", "MysqlConnectorSyncStore")
+
+MYSQL_TABLE_NOT_FOUND_ERROR: Final = 1146
+
+
+class MysqlConnectorAsyncStore(BaseSQLSpecStore["MysqlConnectorAsyncConfig"]):
+    """MySQL/MariaDB session store using mysql-connector async driver."""
+
+    __slots__ = ()
+
+    def __init__(self, config: "MysqlConnectorAsyncConfig") -> None:
+        super().__init__(config)
+
+    def _get_create_table_sql(self) -> str:
+        return f"""
+        CREATE TABLE IF NOT EXISTS {self._table_name} (
+            session_id VARCHAR(255) PRIMARY KEY,
+            data LONGBLOB NOT NULL,
+            expires_at DATETIME(6),
+            created_at DATETIME(6) DEFAULT CURRENT_TIMESTAMP(6),
+            updated_at DATETIME(6) DEFAULT CURRENT_TIMESTAMP(6) ON UPDATE CURRENT_TIMESTAMP(6),
+            INDEX idx_{self._table_name}_expires_at (expires_at)
+        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
+        """
+
+    def _get_drop_table_sql(self) -> "list[str]":
+        return [
+            f"DROP INDEX idx_{self._table_name}_expires_at ON {self._table_name}",
+            f"DROP TABLE IF EXISTS {self._table_name}",
+        ]
+
+    async def create_table(self) -> None:
+        sql = self._get_create_table_sql()
+        async with self._config.provide_session() as driver:
+            await driver.execute_script(sql)
+        self._log_table_created()
+
+    async def get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | None":
+        sql = f"""
+        SELECT data, expires_at FROM {self._table_name}
+        WHERE session_id = %s
+        AND (expires_at IS NULL OR expires_at > UTC_TIMESTAMP(6))
+        """
+
+        try:
+            async with self._config.provide_connection() as conn:
+                cursor = await conn.cursor()
+                try:
+                    await cursor.execute(sql, (key,))
+                    row = await cursor.fetchone()
+                finally:
+                    await cursor.close()
+
+                if row is None:
+                    return None
+
+                data_value: Any = row[0]
+                expires_at: Any = row[1]
+
+                if renew_for is not None and expires_at is not None:
+                    new_expires_at = self._calculate_expires_at(renew_for)
+                    if new_expires_at is not None:
+                        naive_expires_at = new_expires_at.replace(tzinfo=None)
+                        update_sql = f"""
+                            UPDATE {self._table_name}
+                            SET expires_at = %s, updated_at = UTC_TIMESTAMP(6)
+                            WHERE session_id = %s
+                            """
+                        update_cursor = await conn.cursor()
+                        try:
+                            await update_cursor.execute(update_sql, (naive_expires_at, key))
+                        finally:
+                            await update_cursor.close()
+                        await conn.commit()
+
+                return bytes(cast("bytes", data_value))
+        except mysql.connector.Error as exc:
+            if "doesn't exist" in str(exc) or getattr(exc, "errno", None) == MYSQL_TABLE_NOT_FOUND_ERROR:
+                return None
+            raise
+
+    async def set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | None" = None) -> None:
+        data = self._value_to_bytes(value)
+        expires_at = self._calculate_expires_at(expires_in)
+        naive_expires_at = expires_at.replace(tzinfo=None) if expires_at else None
+
+        sql = f"""
+        INSERT INTO {self._table_name} (session_id, data, expires_at)
+        VALUES (%s, %s, %s) AS new
+        ON DUPLICATE KEY UPDATE
+            data = new.data,
+            expires_at = new.expires_at,
+            updated_at = UTC_TIMESTAMP(6)
+        """
+
+        async with self._config.provide_connection() as conn:
+            cursor = await conn.cursor()
+            try:
+                await cursor.execute(sql, (key, data, naive_expires_at))
+            finally:
+                await cursor.close()
+            await conn.commit()
+
+    async def delete(self, key: str) -> None:
+        sql = f"DELETE FROM {self._table_name} WHERE session_id = %s"
+
+        async with self._config.provide_connection() as conn:
+            cursor = await conn.cursor()
+            try:
+                await cursor.execute(sql, (key,))
+            finally:
+                await cursor.close()
+            await conn.commit()
+
+    async def delete_all(self) -> None:
+        sql = f"DELETE FROM {self._table_name}"
+
+        try:
+            async with self._config.provide_connection() as conn:
+                cursor = await conn.cursor()
+                try:
+                    await cursor.execute(sql)
+                finally:
+                    await cursor.close()
+                await conn.commit()
+            self._log_delete_all()
+        except mysql.connector.Error as exc:
+            if "doesn't exist" in str(exc) or getattr(exc, "errno", None) == MYSQL_TABLE_NOT_FOUND_ERROR:
+                logger.debug("Table %s does not exist, skipping delete_all", self._table_name)
+                return
+            raise
+
+    async def exists(self, key: str) -> bool:
+        sql = f"""
+        SELECT 1 FROM {self._table_name}
+        WHERE session_id = %s
+        AND (expires_at IS NULL OR expires_at > UTC_TIMESTAMP(6))
+        """
+
+        try:
+            async with self._config.provide_connection() as conn:
+                cursor = await conn.cursor()
+                try:
+                    await cursor.execute(sql, (key,))
+                    result = await cursor.fetchone()
+                finally:
+                    await cursor.close()
+                return result is not None
+        except mysql.connector.Error as exc:
+            if "doesn't exist" in str(exc) or getattr(exc, "errno", None) == MYSQL_TABLE_NOT_FOUND_ERROR:
+                return False
+            raise
+
+    async def expires_in(self, key: str) -> "int | None":
+        sql = f"""
+        SELECT expires_at FROM {self._table_name}
+        WHERE session_id = %s
+        """
+
+        async with self._config.provide_connection() as conn:
+            cursor = await conn.cursor()
+            try:
+                await cursor.execute(sql, (key,))
+                row = await cursor.fetchone()
+            finally:
+                await cursor.close()
+
+            if row is None or row[0] is None:
+                return None
+
+            expires_at_naive: datetime = cast("datetime", row[0])
+            expires_at_utc = expires_at_naive.replace(tzinfo=timezone.utc)
+            now = datetime.now(timezone.utc)
+
+            if expires_at_utc <= now:
+                return 0
+
+            delta = expires_at_utc - now
+            return int(delta.total_seconds())
+
+    async def delete_expired(self) -> int:
+        sql = f"DELETE FROM {self._table_name} WHERE expires_at <= UTC_TIMESTAMP(6)"
+
+        async with self._config.provide_connection() as conn:
+            cursor = await conn.cursor()
+            try:
+                await cursor.execute(sql)
+                await conn.commit()
+                count: int = cursor.rowcount
+            finally:
+                await cursor.close()
+            if count > 0:
+                self._log_delete_expired(count)
+            return count
+
+
+class MysqlConnectorSyncStore(BaseSQLSpecStore["MysqlConnectorSyncConfig"]):
+    """MySQL/MariaDB session store using mysql-connector sync driver."""
+
+    __slots__ = ()
+
+    def __init__(self, config: "MysqlConnectorSyncConfig") -> None:
+        super().__init__(config)
+
+    def _get_create_table_sql(self) -> str:
+        return f"""
+        CREATE TABLE IF NOT EXISTS {self._table_name} (
+            session_id VARCHAR(255) PRIMARY KEY,
+            data LONGBLOB NOT NULL,
+            expires_at DATETIME(6),
+            created_at DATETIME(6) DEFAULT CURRENT_TIMESTAMP(6),
+            updated_at DATETIME(6) DEFAULT CURRENT_TIMESTAMP(6) ON UPDATE CURRENT_TIMESTAMP(6),
+            INDEX idx_{self._table_name}_expires_at (expires_at)
+        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
+        """
+
+    def _get_drop_table_sql(self) -> "list[str]":
+        return [
+            f"DROP INDEX idx_{self._table_name}_expires_at ON {self._table_name}",
+            f"DROP TABLE IF EXISTS {self._table_name}",
+        ]
+
+    def _create_table(self) -> None:
+        sql = self._get_create_table_sql()
+        with self._config.provide_session() as driver:
+            driver.execute_script(sql)
+            driver.commit()
+        self._log_table_created()
+
+    async def create_table(self) -> None:
+        await async_(self._create_table)()
+
+    def _get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | None":
+        sql = f"""
+        SELECT data, expires_at FROM {self._table_name}
+        WHERE session_id = %s
+        AND (expires_at IS NULL OR expires_at > UTC_TIMESTAMP(6))
+        """
+
+        try:
+            with self._config.provide_connection() as conn:
+                cursor = conn.cursor()
+                try:
+                    cursor.execute(sql, (key,))
+                    row = cursor.fetchone()
+                finally:
+                    cursor.close()
+
+                if row is None:
+                    return None
+
+                data_value: Any = row[0]
+                expires_at: Any = row[1]
+
+                if renew_for is not None and expires_at is not None:
+                    new_expires_at = self._calculate_expires_at(renew_for)
+                    if new_expires_at is not None:
+                        naive_expires_at = new_expires_at.replace(tzinfo=None)
+                        update_sql = f"""
+                            UPDATE {self._table_name}
+                            SET expires_at = %s, updated_at = UTC_TIMESTAMP(6)
+                            WHERE session_id = %s
+                            """
+                        update_cursor = conn.cursor()
+                        try:
+                            update_cursor.execute(update_sql, (naive_expires_at, key))
+                        finally:
+                            update_cursor.close()
+                        conn.commit()
+
+                return bytes(cast("bytes", data_value))
+        except mysql.connector.Error as exc:
+            if "doesn't exist" in str(exc) or getattr(exc, "errno", None) == MYSQL_TABLE_NOT_FOUND_ERROR:
+                return None
+            raise
+
+    async def get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | None":
+        return await async_(self._get)(key, renew_for)
+
+    def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | None" = None) -> None:
+        data = self._value_to_bytes(value)
+        expires_at = self._calculate_expires_at(expires_in)
+        naive_expires_at = expires_at.replace(tzinfo=None) if expires_at else None
+
+        sql = f"""
+        INSERT INTO {self._table_name} (session_id, data, expires_at)
+        VALUES (%s, %s, %s) AS new
+        ON DUPLICATE KEY UPDATE
+            data = new.data,
+            expires_at = new.expires_at,
+            updated_at = UTC_TIMESTAMP(6)
+        """
+
+        with self._config.provide_connection() as conn:
+            cursor = conn.cursor()
+            try:
+                cursor.execute(sql, (key, data, naive_expires_at))
+            finally:
+                cursor.close()
+            conn.commit()
+
+    async def set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | None" = None) -> None:
+        await async_(self._set)(key, value, expires_in)
+
+    def _delete(self, key: str) -> None:
+        sql = f"DELETE FROM {self._table_name} WHERE session_id = %s"
+
+        with self._config.provide_connection() as conn:
+            cursor = conn.cursor()
+            try:
+                cursor.execute(sql, (key,))
+            finally:
+                cursor.close()
+            conn.commit()
+
+    async def delete(self, key: str) -> None:
+        await async_(self._delete)(key)
+
+    def _delete_all(self) -> None:
+        sql = f"DELETE FROM {self._table_name}"
+
+        try:
+            with self._config.provide_connection() as conn:
+                cursor = conn.cursor()
+                try:
+                    cursor.execute(sql)
+                finally:
+                    cursor.close()
+                conn.commit()
+            self._log_delete_all()
+        except mysql.connector.Error as exc:
+            if "doesn't exist" in str(exc) or getattr(exc, "errno", None) == MYSQL_TABLE_NOT_FOUND_ERROR:
+                logger.debug("Table %s does not exist, skipping delete_all", self._table_name)
+                return
+            raise
+
+    async def delete_all(self) -> None:
+        await async_(self._delete_all)()
+
+    def _exists(self, key: str) -> bool:
+        sql = f"""
+        SELECT 1 FROM {self._table_name}
+        WHERE session_id = %s
+        AND (expires_at IS NULL OR expires_at > UTC_TIMESTAMP(6))
+        """
+
+        try:
+            with self._config.provide_connection() as conn:
+                cursor = conn.cursor()
+                try:
+                    cursor.execute(sql, (key,))
+                    result = cursor.fetchone()
+                finally:
+                    cursor.close()
+                return result is not None
+        except mysql.connector.Error as exc:
+            if "doesn't exist" in str(exc) or getattr(exc, "errno", None) == MYSQL_TABLE_NOT_FOUND_ERROR:
+                return False
+            raise
+
+    async def exists(self, key: str) -> bool:
+        return await async_(self._exists)(key)
+
+    def _expires_in(self, key: str) -> "int | None":
+        sql = f"""
+        SELECT expires_at FROM {self._table_name}
+        WHERE session_id = %s
+        """
+
+        with self._config.provide_connection() as conn:
+            cursor = conn.cursor()
+            try:
+                cursor.execute(sql, (key,))
+                row = cursor.fetchone()
+            finally:
+                cursor.close()
+
+            if row is None or row[0] is None:
+                return None
+
+            expires_at_naive: datetime = cast("datetime", row[0])
+            expires_at_utc = expires_at_naive.replace(tzinfo=timezone.utc)
+            now = datetime.now(timezone.utc)
+
+            if expires_at_utc <= now:
+                return 0
+
+            delta = expires_at_utc - now
+            return int(delta.total_seconds())
+
+    async def expires_in(self, key: str) -> "int | None":
+        return await async_(self._expires_in)(key)
+
+    def _delete_expired(self) -> int:
+        sql = f"DELETE FROM {self._table_name} WHERE expires_at <= UTC_TIMESTAMP(6)"
+
+        with self._config.provide_connection() as conn:
+            cursor = conn.cursor()
+            try:
+                cursor.execute(sql)
+                conn.commit()
+                count: int = cursor.rowcount
+            finally:
+                cursor.close()
+            if count > 0:
+                self._log_delete_expired(count)
+            return count
+
+    async def delete_expired(self) -> int:
+        return await async_(self._delete_expired)()

sqlspec/adapters/oracledb/__init__.py

@@ -0,0 +1,60 @@
+import sqlspec.adapters.oracledb._numpy_handlers as numpy_handlers
+from sqlspec.adapters.oracledb._numpy_handlers import (
+    DTYPE_TO_ARRAY_CODE,
+    numpy_converter_in,
+    numpy_converter_out,
+    numpy_input_type_handler,
+    numpy_output_type_handler,
+    register_numpy_handlers,
+)
+from sqlspec.adapters.oracledb._typing import OracleAsyncConnection, OracleSyncConnection
+from sqlspec.adapters.oracledb._uuid_handlers import (
+    register_uuid_handlers,
+    uuid_converter_in,
+    uuid_converter_out,
+    uuid_input_type_handler,
+    uuid_output_type_handler,
+)
+from sqlspec.adapters.oracledb.config import (
+    OracleAsyncConfig,
+    OracleConnectionParams,
+    OraclePoolParams,
+    OracleSyncConfig,
+)
+from sqlspec.adapters.oracledb.core import default_statement_config
+from sqlspec.adapters.oracledb.driver import (
+    OracleAsyncCursor,
+    OracleAsyncDriver,
+    OracleAsyncExceptionHandler,
+    OracleSyncCursor,
+    OracleSyncDriver,
+    OracleSyncExceptionHandler,
+)
+
+__all__ = (
+    "DTYPE_TO_ARRAY_CODE",
+    "OracleAsyncConfig",
+    "OracleAsyncConnection",
+    "OracleAsyncCursor",
+    "OracleAsyncDriver",
+    "OracleAsyncExceptionHandler",
+    "OracleConnectionParams",
+    "OraclePoolParams",
+    "OracleSyncConfig",
+    "OracleSyncConnection",
+    "OracleSyncCursor",
+    "OracleSyncDriver",
+    "OracleSyncExceptionHandler",
+    "default_statement_config",
+    "numpy_converter_in",
+    "numpy_converter_out",
+    "numpy_handlers",
+    "numpy_input_type_handler",
+    "numpy_output_type_handler",
+    "register_numpy_handlers",
+    "register_uuid_handlers",
+    "uuid_converter_in",
+    "uuid_converter_out",
+    "uuid_input_type_handler",
+    "uuid_output_type_handler",
+)

sqlspec/adapters/oracledb/_numpy_handlers.py

@@ -0,0 +1,141 @@
+"""Oracle NumPy vector type handlers for VECTOR data type support.
+
+Provides automatic conversion between NumPy arrays and Oracle VECTOR types
+via connection type handlers. Requires Oracle Database 23ai or higher.
+"""
+
+import array
+from typing import TYPE_CHECKING, Any
+
+from sqlspec.typing import NUMPY_INSTALLED
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from oracledb import AsyncConnection, AsyncCursor, Connection, Cursor
+
+__all__ = (
+    "numpy_converter_in",
+    "numpy_converter_out",
+    "numpy_input_type_handler",
+    "numpy_output_type_handler",
+    "register_numpy_handlers",
+)
+
+
+logger = get_logger(__name__)
+
+
+DTYPE_TO_ARRAY_CODE: "dict[str, str]" = {"float64": "d", "float32": "f", "uint8": "B", "int8": "b"}
+
+
+def numpy_converter_in(value: Any) -> "array.array[Any]":
+    """Convert NumPy array to Oracle array for VECTOR insertion.
+
+    Args:
+        value: NumPy ndarray to convert.
+
+    Returns:
+        Python array.array compatible with Oracle VECTOR type.
+
+    Raises:
+        ImportError: If NumPy is not installed.
+        TypeError: If NumPy dtype is not supported for Oracle VECTOR.
+    """
+    if not NUMPY_INSTALLED:
+        msg = "NumPy is not installed - cannot convert vectors"
+        raise ImportError(msg)
+
+    dtype_name = value.dtype.name
+    array_code = DTYPE_TO_ARRAY_CODE.get(dtype_name)
+
+    if not array_code:
+        supported = ", ".join(DTYPE_TO_ARRAY_CODE.keys())
+        msg = f"Unsupported NumPy dtype for Oracle VECTOR: {dtype_name}. Supported: {supported}"
+        raise TypeError(msg)
+
+    return array.array(array_code, value)
+
+
+def numpy_converter_out(value: "array.array[Any]") -> Any:
+    """Convert Oracle array to NumPy array for VECTOR retrieval.
+
+    Args:
+        value: Oracle array.array from VECTOR column.
+
+    Returns:
+        NumPy ndarray with appropriate dtype, or original value if NumPy not installed.
+    """
+    if not NUMPY_INSTALLED:
+        return value
+
+    import numpy as np
+
+    return np.array(value, copy=True, dtype=value.typecode)
+
+
+def _input_type_handler(cursor: "Cursor | AsyncCursor", value: Any, arraysize: int) -> Any:
+    """Oracle input type handler for NumPy arrays.
+
+    Args:
+        cursor: Oracle cursor (sync or async).
+        value: Value being inserted.
+        arraysize: Array size for the cursor variable.
+
+    Returns:
+        Cursor variable with NumPy converter if value is ndarray, None otherwise.
+    """
+    if not NUMPY_INSTALLED:
+        return None
+
+    import numpy as np
+    import oracledb
+
+    if isinstance(value, np.ndarray):
+        return cursor.var(oracledb.DB_TYPE_VECTOR, arraysize=arraysize, inconverter=numpy_converter_in)
+    return None
+
+
+def _output_type_handler(cursor: "Cursor | AsyncCursor", metadata: Any) -> Any:
+    """Oracle output type handler for VECTOR columns.
+
+    Args:
+        cursor: Oracle cursor (sync or async).
+        metadata: Column metadata from Oracle.
+
+    Returns:
+        Cursor variable with NumPy converter if column is VECTOR, None otherwise.
+    """
+    if not NUMPY_INSTALLED:
+        return None
+
+    import oracledb
+
+    if metadata.type_code is oracledb.DB_TYPE_VECTOR:
+        return cursor.var(metadata.type_code, arraysize=cursor.arraysize, outconverter=numpy_converter_out)
+    return None
+
+
+def numpy_input_type_handler(cursor: "Cursor | AsyncCursor", value: Any, arraysize: int) -> Any:
+    """Public input type handler for NumPy arrays."""
+    return _input_type_handler(cursor, value, arraysize)
+
+
+def numpy_output_type_handler(cursor: "Cursor | AsyncCursor", metadata: Any) -> Any:
+    """Public output type handler for NumPy VECTOR columns."""
+    return _output_type_handler(cursor, metadata)
+
+
+def register_numpy_handlers(connection: "Connection | AsyncConnection") -> None:
+    """Register NumPy type handlers on Oracle connection.
+
+    Enables automatic conversion between NumPy arrays and Oracle VECTOR types.
+    Works for both sync and async connections.
+
+    Args:
+        connection: Oracle connection (sync or async).
+    """
+    if not NUMPY_INSTALLED:
+        return
+
+    connection.inputtypehandler = numpy_input_type_handler
+    connection.outputtypehandler = numpy_output_type_handler