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

sqlspec/exceptions.py

@@ -0,0 +1,338 @@
+from collections.abc import Generator
+from contextlib import contextmanager
+from typing import Any, Final
+
+STACK_SQL_PREVIEW_LIMIT: Final[int] = 120
+
+__all__ = (
+    "CheckViolationError",
+    "ConfigResolverError",
+    "DataError",
+    "DatabaseConnectionError",
+    "DialectNotSupportedError",
+    "EventChannelError",
+    "FileNotFoundInStorageError",
+    "ForeignKeyViolationError",
+    "ImproperConfigurationError",
+    "IntegrityError",
+    "InvalidVersionFormatError",
+    "MigrationError",
+    "MissingDependencyError",
+    "MultipleResultsFoundError",
+    "NotFoundError",
+    "NotNullViolationError",
+    "OperationalError",
+    "OutOfOrderMigrationError",
+    "RepositoryError",
+    "SQLBuilderError",
+    "SQLConversionError",
+    "SQLFileNotFoundError",
+    "SQLFileParseError",
+    "SQLParsingError",
+    "SQLSpecError",
+    "SerializationConflictError",
+    "SerializationError",
+    "StackExecutionError",
+    "StorageCapabilityError",
+    "StorageOperationFailedError",
+    "TransactionError",
+    "TransactionRetryError",
+    "UniqueViolationError",
+)
+
+
+class SQLSpecError(Exception):
+    """Base exception class for SQLSpec exceptions."""
+
+    detail: str = ""
+
+    def __init__(self, *args: Any, detail: str = "") -> None:
+        """Initialize SQLSpecError.
+
+        Args:
+            *args: args are converted to :class:`str` before passing to :class:`Exception`
+            detail: detail of the exception.
+        """
+        str_args = [str(arg) for arg in args if arg]
+        if not detail:
+            detail = str_args[0] if str_args else self.detail
+        self.detail = detail
+        if detail and detail not in str_args:
+            str_args = [detail, *str_args]
+        super().__init__(*str_args)
+
+    def __repr__(self) -> str:
+        if self.detail:
+            return f"{self.__class__.__name__} - {self.detail}"
+        return self.__class__.__name__
+
+    def __str__(self) -> str:
+        parts = list(self.args)
+        if self.detail and self.detail not in self.args:
+            parts.append(self.detail)
+        return " ".join(parts).strip()
+
+
+class MissingDependencyError(SQLSpecError, ImportError):
+    """Raised when a required dependency is not installed."""
+
+    def __init__(self, package: str, install_package: str | None = None) -> None:
+        super().__init__(
+            f"Package {package!r} is not installed but required. You can install it by running "
+            f"'pip install sqlspec[{install_package or package}]' to install sqlspec with the required extra "
+            f"or 'pip install {install_package or package}' to install the package separately"
+        )
+
+
+class BackendNotRegisteredError(SQLSpecError):
+    """Raised when a requested storage backend key is not registered."""
+
+    def __init__(self, backend_key: str) -> None:
+        super().__init__(f"Storage backend '{backend_key}' is not registered. Please register it before use.")
+
+
+class EventChannelError(SQLSpecError):
+    """Raised when event channel operations fail."""
+
+
+class ConfigResolverError(SQLSpecError):
+    """Exception raised when config resolution fails."""
+
+
+class SQLParsingError(SQLSpecError):
+    """Issues parsing SQL statements."""
+
+    def __init__(self, message: str | None = None) -> None:
+        if message is None:
+            message = "Issues parsing SQL statement."
+        super().__init__(message)
+
+
+class SQLBuilderError(SQLSpecError):
+    """Issues Building or Generating SQL statements."""
+
+    def __init__(self, message: str | None = None) -> None:
+        if message is None:
+            message = "Issues building SQL statement."
+        super().__init__(message)
+
+
+class SQLConversionError(SQLSpecError):
+    """Issues converting SQL statements."""
+
+    def __init__(self, message: str | None = None) -> None:
+        if message is None:
+            message = "Issues converting SQL statement."
+        super().__init__(message)
+
+
+class ImproperConfigurationError(SQLSpecError):
+    """Raised when configuration is invalid or incomplete."""
+
+
+class DialectNotSupportedError(SQLBuilderError):
+    """Raised when a SQL dialect does not support a specific feature."""
+
+
+class SerializationError(SQLSpecError):
+    """Encoding or decoding of an object failed."""
+
+
+class RepositoryError(SQLSpecError):
+    """Base repository exception type."""
+
+
+class IntegrityError(RepositoryError):
+    """Data integrity error."""
+
+
+class NotFoundError(RepositoryError):
+    """An identity does not exist."""
+
+
+class MultipleResultsFoundError(RepositoryError):
+    """A single database result was required but more than one were found."""
+
+
+class UniqueViolationError(IntegrityError):
+    """A unique constraint was violated."""
+
+
+class ForeignKeyViolationError(IntegrityError):
+    """A foreign key constraint was violated."""
+
+
+class CheckViolationError(IntegrityError):
+    """A check constraint was violated."""
+
+
+class NotNullViolationError(IntegrityError):
+    """A not-null constraint was violated."""
+
+
+class DatabaseConnectionError(SQLSpecError):
+    """Database connection error (invalid credentials, network failure, etc.)."""
+
+
+class TransactionError(SQLSpecError):
+    """Transaction error (rollback, deadlock, serialization failure)."""
+
+
+class SerializationConflictError(TransactionError):
+    """Serialization conflict (SQLSTATE 40001) requiring retry."""
+
+
+class TransactionRetryError(TransactionError):
+    """Transaction failed after retries were exhausted."""
+
+
+class DataError(SQLSpecError):
+    """Invalid data type or format for database operation."""
+
+
+class StackExecutionError(SQLSpecError):
+    """Raised when a statement stack operation fails."""
+
+    def __init__(
+        self,
+        operation_index: int,
+        sql: str,
+        original_error: Exception,
+        *,
+        adapter: str | None = None,
+        mode: str = "fail-fast",
+        native_pipeline: bool | None = None,
+        downgrade_reason: str | None = None,
+    ) -> None:
+        pipeline_state = "enabled" if native_pipeline else "disabled"
+        adapter_label = adapter or "unknown-adapter"
+        preview = " ".join(sql.strip().split())
+        if len(preview) > STACK_SQL_PREVIEW_LIMIT:
+            preview = f"{preview[: STACK_SQL_PREVIEW_LIMIT - 3]}..."
+        detail = (
+            f"Stack operation {operation_index} failed on {adapter_label} "
+            f"(mode={mode}, pipeline={pipeline_state}) sql={preview}"
+        )
+        super().__init__(detail)
+        self.operation_index = operation_index
+        self.sql = sql
+        self.original_error = original_error
+        self.adapter = adapter
+        self.mode = mode
+        self.native_pipeline = native_pipeline
+        self.downgrade_reason = downgrade_reason
+
+    def __str__(self) -> str:
+        base = super().__str__()
+        return f"{base}: {self.original_error}" if self.original_error else base
+
+
+class OperationalError(SQLSpecError):
+    """Operational database error (timeout, disk full, resource limit)."""
+
+
+class StorageOperationFailedError(SQLSpecError):
+    """Raised when a storage backend operation fails (e.g., network, permission, API error)."""
+
+
+class StorageCapabilityError(SQLSpecError):
+    """Raised when a requested storage bridge capability is unavailable."""
+
+    def __init__(self, message: str, *, capability: str | None = None, remediation: str | None = None) -> None:
+        parts = [message]
+        if capability:
+            parts.append(f"(capability: {capability})")
+        if remediation:
+            parts.append(remediation)
+        detail = " ".join(parts)
+        super().__init__(detail)
+        self.capability = capability
+        self.remediation = remediation
+
+
+class FileNotFoundInStorageError(StorageOperationFailedError):
+    """Raised when a file or object is not found in the storage backend."""
+
+
+class SQLFileNotFoundError(SQLSpecError):
+    """Raised when a SQL file cannot be found."""
+
+    def __init__(self, name: str, path: "str | None" = None) -> None:
+        """Initialize the error.
+
+        Args:
+            name: Name of the SQL file.
+            path: Optional path where the file was expected.
+        """
+        message = f"SQL file '{name}' not found at path: {path}" if path else f"SQL file '{name}' not found"
+        super().__init__(message)
+        self.name = name
+        self.path = path
+
+
+class SQLFileParseError(SQLSpecError):
+    """Raised when a SQL file cannot be parsed."""
+
+    def __init__(self, name: str, path: str, original_error: "Exception") -> None:
+        """Initialize the error.
+
+        Args:
+            name: Name of the SQL file.
+            path: Path to the SQL file.
+            original_error: The underlying parsing error.
+        """
+        message = f"Failed to parse SQL file '{name}' at {path}: {original_error}"
+        super().__init__(message)
+        self.name = name
+        self.path = path
+        self.original_error = original_error
+
+
+class MigrationError(SQLSpecError):
+    """Base exception for migration-related errors."""
+
+
+class InvalidVersionFormatError(MigrationError):
+    """Raised when a migration version format is invalid.
+
+    Invalid formats include versions that don't match sequential (0001)
+    or timestamp (YYYYMMDDHHmmss) patterns, or timestamps with invalid dates.
+    """
+
+
+class OutOfOrderMigrationError(MigrationError):
+    """Raised when an out-of-order migration is detected in strict mode.
+
+    Out-of-order migrations occur when a pending migration has a timestamp
+    earlier than already-applied migrations, typically from late-merging branches.
+    """
+
+
+@contextmanager
+def wrap_exceptions(
+    wrap_exceptions: bool = True, suppress: "type[Exception] | tuple[type[Exception], ...] | None" = None
+) -> Generator[None, None, None]:
+    """Context manager for exception handling with optional suppression.
+
+    Args:
+        wrap_exceptions: If True, wrap exceptions in RepositoryError. If False, let them pass through.
+        suppress: Exception type(s) to suppress completely (like contextlib.suppress).
+                 If provided, these exceptions are caught and ignored.
+    """
+    try:
+        yield
+
+    except Exception as exc:
+        if suppress is not None and (
+            (isinstance(suppress, type) and isinstance(exc, suppress))
+            or (isinstance(suppress, tuple) and isinstance(exc, suppress))
+        ):
+            return
+
+        if isinstance(exc, SQLSpecError):
+            raise
+
+        if wrap_exceptions is False:
+            raise
+        msg = "An error occurred during the operation."
+        raise RepositoryError(detail=msg) from exc

sqlspec/extensions/adk/__init__.py

@@ -0,0 +1,70 @@
+"""Google ADK session backend extension for SQLSpec.
+
+Provides session, event, and memory storage for Google Agent Development Kit using
+SQLSpec database adapters.
+
+Public API exports:
+    - ADKConfig: TypedDict for extension config (type-safe configuration)
+    - SQLSpecSessionService: Main service class implementing BaseSessionService
+    - SQLSpecMemoryService: Main async service class implementing BaseMemoryService
+    - SQLSpecSyncMemoryService: Sync memory service for sync adapters
+    - BaseAsyncADKStore: Base class for async database store implementations
+    - BaseSyncADKStore: Base class for sync database store implementations
+    - BaseAsyncADKMemoryStore: Base class for async memory store implementations
+    - BaseSyncADKMemoryStore: Base class for sync memory store implementations
+    - SessionRecord: TypedDict for session database records
+    - EventRecord: TypedDict for event database records
+    - MemoryRecord: TypedDict for memory database records
+
+Example (with extension_config):
+    from sqlspec.adapters.asyncpg import AsyncpgConfig
+    from sqlspec.adapters.asyncpg.adk.store import AsyncpgADKStore
+    from sqlspec.extensions.adk import SQLSpecSessionService
+
+    config = AsyncpgConfig(
+        connection_config={"dsn": "postgresql://..."},
+        extension_config={
+            "adk": {
+                "session_table": "my_sessions",
+                "events_table": "my_events",
+                "owner_id_column": "tenant_id INTEGER REFERENCES tenants(id)"
+            }
+        }
+    )
+
+    store = AsyncpgADKStore(config)
+    await store.ensure_tables()
+
+    service = SQLSpecSessionService(store)
+    session = await service.create_session(
+        app_name="my_app",
+        user_id="user123",
+        state={"key": "value"}
+    )
+"""
+
+from sqlspec.config import ADKConfig
+from sqlspec.extensions.adk._types import EventRecord, SessionRecord
+from sqlspec.extensions.adk.memory import (
+    BaseAsyncADKMemoryStore,
+    BaseSyncADKMemoryStore,
+    MemoryRecord,
+    SQLSpecMemoryService,
+    SQLSpecSyncMemoryService,
+)
+from sqlspec.extensions.adk.service import SQLSpecSessionService
+from sqlspec.extensions.adk.store import BaseAsyncADKStore, BaseSyncADKStore
+
+__all__ = (
+    "ADKConfig",
+    "BaseAsyncADKMemoryStore",
+    "BaseAsyncADKStore",
+    "BaseSyncADKMemoryStore",
+    "BaseSyncADKStore",
+    "EventRecord",
+    "MemoryRecord",
+    "SQLSpecMemoryService",
+    "SQLSpecSessionService",
+    "SQLSpecSyncMemoryService",
+    "SessionRecord",
+)

sqlspec/extensions/adk/_types.py

@@ -0,0 +1,51 @@
+"""Type definitions for ADK extension.
+
+These types define the database record structures for storing sessions and events.
+They are separate from the Pydantic models to keep mypyc compilation working.
+"""
+
+from datetime import datetime
+from typing import Any, TypedDict
+
+__all__ = ("EventRecord", "SessionRecord")
+
+
+class SessionRecord(TypedDict):
+    """Database record for a session.
+
+    Represents the schema for sessions stored in the database.
+    """
+
+    id: str
+    app_name: str
+    user_id: str
+    state: "dict[str, Any]"
+    create_time: datetime
+    update_time: datetime
+
+
+class EventRecord(TypedDict):
+    """Database record for an event.
+
+    Represents the schema for events stored in the database.
+    Follows the ADK Event model plus session metadata.
+    """
+
+    id: str
+    app_name: str
+    user_id: str
+    session_id: str
+    invocation_id: str
+    author: str
+    branch: "str | None"
+    actions: bytes
+    long_running_tool_ids_json: "str | None"
+    timestamp: datetime
+    content: "dict[str, Any] | None"
+    grounding_metadata: "dict[str, Any] | None"
+    custom_metadata: "dict[str, Any] | None"
+    partial: "bool | None"
+    turn_complete: "bool | None"
+    interrupted: "bool | None"
+    error_code: "str | None"
+    error_message: "str | None"

sqlspec/extensions/adk/converters.py

@@ -0,0 +1,172 @@
+"""Conversion functions between ADK models and database records."""
+
+import json
+import pickle
+from datetime import datetime, timezone
+from typing import Any
+
+from google.adk.events.event import Event
+from google.adk.sessions import Session
+from google.genai import types
+
+from sqlspec.extensions.adk._types import EventRecord, SessionRecord
+from sqlspec.utils.logging import get_logger
+
+logger = get_logger("sqlspec.extensions.adk.converters")
+
+__all__ = ("event_to_record", "record_to_event", "record_to_session", "session_to_record")
+
+
+def session_to_record(session: "Session") -> SessionRecord:
+    """Convert ADK Session to database record.
+
+    Args:
+        session: ADK Session object.
+
+    Returns:
+        SessionRecord for database storage.
+    """
+    return SessionRecord(
+        id=session.id,
+        app_name=session.app_name,
+        user_id=session.user_id,
+        state=session.state,
+        create_time=datetime.now(timezone.utc),
+        update_time=datetime.fromtimestamp(session.last_update_time, tz=timezone.utc),
+    )
+
+
+def record_to_session(record: SessionRecord, events: "list[EventRecord]") -> "Session":
+    """Convert database record to ADK Session.
+
+    Args:
+        record: Session database record.
+        events: List of event records for this session.
+
+    Returns:
+        ADK Session object.
+    """
+    event_objects = [record_to_event(event_record) for event_record in events]
+
+    return Session(
+        id=record["id"],
+        app_name=record["app_name"],
+        user_id=record["user_id"],
+        state=record["state"],
+        events=event_objects,
+        last_update_time=record["update_time"].timestamp(),
+    )
+
+
+def event_to_record(event: "Event", session_id: str, app_name: str, user_id: str) -> EventRecord:
+    """Convert ADK Event to database record.
+
+    Args:
+        event: ADK Event object.
+        session_id: ID of the parent session.
+        app_name: Name of the application.
+        user_id: ID of the user.
+
+    Returns:
+        EventRecord for database storage.
+    """
+    actions_bytes = pickle.dumps(event.actions)
+
+    long_running_tool_ids_json = None
+    if event.long_running_tool_ids:
+        long_running_tool_ids_json = json.dumps(list(event.long_running_tool_ids))
+
+    content_dict = None
+    if event.content:
+        content_dict = event.content.model_dump(exclude_none=True, mode="json")
+
+    grounding_metadata_dict = None
+    if event.grounding_metadata:
+        grounding_metadata_dict = event.grounding_metadata.model_dump(exclude_none=True, mode="json")
+
+    custom_metadata_dict = event.custom_metadata
+
+    return EventRecord(
+        id=event.id,
+        app_name=app_name,
+        user_id=user_id,
+        session_id=session_id,
+        invocation_id=event.invocation_id,
+        author=event.author,
+        branch=event.branch,
+        actions=actions_bytes,
+        long_running_tool_ids_json=long_running_tool_ids_json,
+        timestamp=datetime.fromtimestamp(event.timestamp, tz=timezone.utc),
+        content=content_dict,
+        grounding_metadata=grounding_metadata_dict,
+        custom_metadata=custom_metadata_dict,
+        partial=event.partial,
+        turn_complete=event.turn_complete,
+        interrupted=event.interrupted,
+        error_code=event.error_code,
+        error_message=event.error_message,
+    )
+
+
+def record_to_event(record: "EventRecord") -> "Event":
+    """Convert database record to ADK Event.
+
+    Args:
+        record: Event database record.
+
+    Returns:
+        ADK Event object.
+    """
+    actions = pickle.loads(record["actions"])  # noqa: S301
+
+    long_running_tool_ids = None
+    if record["long_running_tool_ids_json"]:
+        long_running_tool_ids = set(json.loads(record["long_running_tool_ids_json"]))
+
+    return Event(
+        id=record["id"],
+        invocation_id=record["invocation_id"],
+        author=record["author"],
+        branch=record["branch"],
+        actions=actions,
+        timestamp=record["timestamp"].timestamp(),
+        content=_decode_content(record["content"]),
+        long_running_tool_ids=long_running_tool_ids,
+        partial=record["partial"],
+        turn_complete=record["turn_complete"],
+        error_code=record["error_code"],
+        error_message=record["error_message"],
+        interrupted=record["interrupted"],
+        grounding_metadata=_decode_grounding_metadata(record["grounding_metadata"]),
+        custom_metadata=record["custom_metadata"],
+    )
+
+
+def _decode_content(content_dict: "dict[str, Any] | None") -> Any:
+    """Decode content dictionary from database to ADK Content object.
+
+    Args:
+        content_dict: Content dictionary from database.
+
+    Returns:
+        ADK Content object or None.
+    """
+    if not content_dict:
+        return None
+
+    return types.Content.model_validate(content_dict)
+
+
+def _decode_grounding_metadata(grounding_dict: "dict[str, Any] | None") -> Any:
+    """Decode grounding metadata dictionary from database to ADK object.
+
+    Args:
+        grounding_dict: Grounding metadata dictionary from database.
+
+    Returns:
+        ADK GroundingMetadata object or None.
+    """
+    if not grounding_dict:
+        return None
+
+    return types.GroundingMetadata.model_validate(grounding_dict)

sqlspec/extensions/adk/memory/__init__.py

@@ -0,0 +1,69 @@
+"""Google ADK memory backend extension for SQLSpec.
+
+Provides memory entry storage for Google Agent Development Kit using
+SQLSpec database adapters. Memory stores are used to persist conversational
+context across agent sessions for long-term recall.
+
+Public API exports:
+    - SQLSpecMemoryService: Main async service class implementing BaseMemoryService
+    - SQLSpecSyncMemoryService: Sync service for sync adapters
+    - BaseAsyncADKMemoryStore: Base class for async database store implementations
+    - BaseSyncADKMemoryStore: Base class for sync database store implementations
+    - MemoryRecord: TypedDict for memory database records
+    - extract_content_text: Helper to extract searchable text from Content
+    - session_to_memory_records: Convert Session to memory records
+    - record_to_memory_entry: Convert database record to MemoryEntry
+
+Example (async):
+    from sqlspec.adapters.asyncpg import AsyncpgConfig
+    from sqlspec.adapters.asyncpg.adk.store import AsyncpgADKMemoryStore
+    from sqlspec.extensions.adk.memory import SQLSpecMemoryService
+
+    config = AsyncpgConfig(
+        connection_config={"dsn": "postgresql://..."},
+        extension_config={
+            "adk": {
+                "memory_table": "adk_memory_entries",
+                "memory_use_fts": True,
+                "memory_max_results": 50,
+            }
+        }
+    )
+
+    store = AsyncpgADKMemoryStore(config)
+    await store.ensure_tables()
+
+    service = SQLSpecMemoryService(store)
+
+    # Store completed session as memories
+    await service.add_session_to_memory(completed_session)
+
+    # Search memories
+    response = await service.search_memory(
+        app_name="my_app",
+        user_id="user123",
+        query="previous discussion about Python"
+    )
+    for entry in response.memories:
+        print(entry.content)
+"""
+
+from sqlspec.extensions.adk.memory._types import MemoryRecord
+from sqlspec.extensions.adk.memory.converters import (
+    extract_content_text,
+    record_to_memory_entry,
+    session_to_memory_records,
+)
+from sqlspec.extensions.adk.memory.service import SQLSpecMemoryService, SQLSpecSyncMemoryService
+from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore, BaseSyncADKMemoryStore
+
+__all__ = (
+    "BaseAsyncADKMemoryStore",
+    "BaseSyncADKMemoryStore",
+    "MemoryRecord",
+    "SQLSpecMemoryService",
+    "SQLSpecSyncMemoryService",
+    "extract_content_text",
+    "record_to_memory_entry",
+    "session_to_memory_records",
+)