sqlspec 0.26.0__py3-none-any.whl → 0.28.0__py3-none-any.whl

sqlspec/adapters/psycopg/_type_handlers.py

@@ -0,0 +1,80 @@
+"""Psycopg pgvector type handlers for vector data type support.
+
+Provides automatic conversion between NumPy arrays and PostgreSQL vector types
+via pgvector-python library. Supports both sync and async connections.
+"""
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+from sqlspec.typing import NUMPY_INSTALLED, PGVECTOR_INSTALLED
+
+if TYPE_CHECKING:
+    from psycopg import AsyncConnection, Connection
+
+__all__ = ("register_pgvector_async", "register_pgvector_sync")
+
+
+logger = logging.getLogger(__name__)
+
+
+def register_pgvector_sync(connection: "Connection[Any]") -> None:
+    """Register pgvector type handlers on psycopg sync connection.
+
+    Enables automatic conversion between NumPy arrays and PostgreSQL vector types
+    using the pgvector-python library.
+
+    Args:
+        connection: Psycopg sync connection.
+    """
+    if not PGVECTOR_INSTALLED:
+        logger.debug("pgvector not installed - skipping vector type handlers")
+        return
+
+    if not NUMPY_INSTALLED:
+        logger.debug("NumPy not installed - registering pgvector without NumPy support")
+
+    try:
+        import pgvector.psycopg
+
+        pgvector.psycopg.register_vector(connection)
+        logger.debug("Registered pgvector type handlers on psycopg sync connection")
+    except ValueError as error:
+        message = str(error).lower()
+        if "vector type not found" in message:
+            logger.debug("Skipping pgvector registration - extension not enabled in database")
+            return
+        logger.warning("Unexpected error during pgvector registration: %s", error)
+    except Exception:
+        logger.exception("Failed to register pgvector for psycopg sync")
+
+
+async def register_pgvector_async(connection: "AsyncConnection[Any]") -> None:
+    """Register pgvector type handlers on psycopg async connection.
+
+    Enables automatic conversion between NumPy arrays and PostgreSQL vector types
+    using the pgvector-python library.
+
+    Args:
+        connection: Psycopg async connection.
+    """
+    if not PGVECTOR_INSTALLED:
+        logger.debug("pgvector not installed - skipping vector type handlers")
+        return
+
+    if not NUMPY_INSTALLED:
+        logger.debug("NumPy not installed - registering pgvector without NumPy support")
+
+    try:
+        from pgvector.psycopg import register_vector_async
+
+        await register_vector_async(connection)
+        logger.debug("Registered pgvector type handlers on psycopg async connection")
+    except ValueError as error:
+        message = str(error).lower()
+        if "vector type not found" in message:
+            logger.debug("Skipping pgvector registration - extension not enabled in database")
+            return
+        logger.warning("Unexpected error during pgvector registration: %s", error)
+    except Exception:
+        logger.exception("Failed to register pgvector for psycopg async")

sqlspec/adapters/psycopg/_types.py

@@ -3,8 +3,9 @@ from typing import TYPE_CHECKING
 from psycopg.rows import DictRow as PsycopgDictRow
 
 if TYPE_CHECKING:
+    from typing import TypeAlias
+
     from psycopg import AsyncConnection, Connection
-    from typing_extensions import TypeAlias
 
     PsycopgSyncConnection: TypeAlias = Connection[PsycopgDictRow]
     PsycopgAsyncConnection: TypeAlias = AsyncConnection[PsycopgDictRow]

sqlspec/adapters/psycopg/adk/__init__.py

@@ -0,0 +1,5 @@
+"""Psycopg ADK store module."""
+
+from sqlspec.adapters.psycopg.adk.store import PsycopgAsyncADKStore, PsycopgSyncADKStore
+
+__all__ = ("PsycopgAsyncADKStore", "PsycopgSyncADKStore")