oxyde-postgres 0.1.0__py3-none-any.whl

oxyde/db/__init__.py

@@ -0,0 +1,230 @@
+"""Database connection management and transaction support.
+
+This module provides the primary interface for database connectivity in Oxyde.
+It wraps the Rust connection pool and provides Python-friendly APIs.
+
+High-Level Functions:
+    init(**databases, settings=None):
+        Initialize one or more database connections at startup.
+        Example: await db.init(default="postgresql://...")
+
+    close():
+        Disconnect all database connections.
+        Example: await db.close()
+
+    connect(url, name="default"):
+        Context manager for single database connection.
+        Example: async with db.connect("postgresql://localhost/test"): ...
+
+    lifespan(**databases):
+        FastAPI/Starlette lifespan context manager.
+        Example: app = FastAPI(lifespan=db.lifespan(default="..."))
+
+Classes:
+    AsyncDatabase: Connection pool wrapper with execute() method.
+    AsyncTransaction: Transaction wrapper with execute() method.
+    PoolSettings: Configuration for connection pool and PostgreSQL settings.
+
+Transaction API (Django-style):
+    from oxyde.db import transaction
+
+    async with transaction.atomic():
+        await User.objects.create(name="Alice")
+        # Auto-commits on success, rolls back on exception
+
+    Nested transactions use savepoints:
+    async with transaction.atomic():
+        await User.objects.create(name="Alice")
+        async with transaction.atomic():
+            await User.objects.create(name="Bob")
+            # Inner failure rolls back to savepoint, not entire transaction
+
+Registry Functions:
+    register_connection(): Add connection to global registry.
+    get_connection(): Get connection by name (default: "default").
+    disconnect_all(): Close all connections.
+
+Connection Resolution:
+    Queries automatically use the active transaction if inside atomic(),
+    otherwise they use the connection from the registry.
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Callable
+from contextlib import AbstractAsyncContextManager, asynccontextmanager
+from typing import Any
+
+from . import transaction
+from .pool import AsyncDatabase, PoolSettings
+from .registry import (
+    disconnect_all,
+    get_connection,
+    register_connection,
+    unregister_connection,
+)
+from .schema import create_tables, drop_tables
+from .transaction import (
+    AsyncTransaction,
+    TransactionTimeoutError,
+    atomic,
+    get_active_transaction,
+    get_transaction_depth,
+)
+
+
+async def init(
+    *,
+    settings: PoolSettings | None = None,
+    **databases: str,
+) -> None:
+    """
+    Initialize database connections.
+
+    Args:
+        settings: Default pool settings for all connections.
+        **databases: Mapping of name -> URL (e.g., default="postgresql://...")
+
+    Raises:
+        ValueError: If no databases provided.
+        Exception: If any database fails to connect (all connections are rolled back).
+
+    Example:
+        await db.init(
+            default="postgresql://user:pass@localhost/mydb",
+            analytics="postgresql://user:pass@localhost/analytics",
+            settings=PoolSettings(max_connections=20),
+        )
+    """
+    if not databases:
+        raise ValueError(
+            "At least one database URL required (e.g., default='postgresql://...')"
+        )
+
+    connected: list[AsyncDatabase] = []
+    try:
+        for name, url in databases.items():
+            database = AsyncDatabase(
+                url,
+                name=name,
+                settings=settings,
+                auto_register=True,
+                overwrite=True,
+            )
+            await database.connect()
+            connected.append(database)
+    except Exception:
+        # Rollback all successfully connected databases
+        for db in connected:
+            try:
+                await db.disconnect()
+            except Exception:
+                pass  # Ignore errors during cleanup
+        raise
+
+
+async def close() -> None:
+    """
+    Close all database connections.
+
+    Example:
+        await db.close()
+    """
+    await disconnect_all()
+
+
+@asynccontextmanager
+async def connect(
+    url: str,
+    *,
+    name: str = "default",
+    settings: PoolSettings | None = None,
+) -> AsyncIterator[AsyncDatabase]:
+    """
+    Context manager for a single database connection.
+
+    Useful for scripts and tests where you want automatic cleanup.
+
+    Args:
+        url: Database connection URL.
+        name: Connection name (default: "default").
+        settings: Optional pool settings.
+
+    Example:
+        async with db.connect("postgresql://localhost/test") as conn:
+            users = await User.objects.all()
+    """
+    database = AsyncDatabase(
+        url,
+        name=name,
+        settings=settings,
+        auto_register=True,
+        overwrite=True,
+    )
+    try:
+        await database.connect()
+        yield database
+    finally:
+        await database.disconnect()
+        unregister_connection(name)
+
+
+def lifespan(
+    *,
+    settings: PoolSettings | None = None,
+    **databases: str,
+) -> Callable[[Any], AbstractAsyncContextManager[None]]:
+    """
+    Create a lifespan context manager for FastAPI/Starlette.
+
+    Args:
+        settings: Default pool settings for all connections.
+        **databases: Mapping of name -> URL.
+
+    Returns:
+        Async context manager compatible with FastAPI's lifespan parameter.
+
+    Example:
+        from fastapi import FastAPI
+        from oxyde import db
+
+        app = FastAPI(lifespan=db.lifespan(
+            default="postgresql://user:pass@localhost/mydb",
+        ))
+    """
+
+    @asynccontextmanager
+    async def _lifespan(app: Any) -> AsyncIterator[None]:
+        await init(settings=settings, **databases)
+        try:
+            yield
+        finally:
+            await close()
+
+    return _lifespan
+
+
+__all__ = [
+    # High-level API
+    "init",
+    "close",
+    "connect",
+    "lifespan",
+    "get_connection",
+    # Classes
+    "AsyncDatabase",
+    "AsyncTransaction",
+    "PoolSettings",
+    "TransactionTimeoutError",
+    # Transactions (Django-style: transaction.atomic())
+    "transaction",
+    "atomic",
+    "get_active_transaction",
+    "get_transaction_depth",
+    # Schema management
+    "create_tables",
+    "drop_tables",
+    # Low-level (for advanced use)
+    "register_connection",
+    "disconnect_all",
+]

oxyde/db/pool.py

@@ -0,0 +1,305 @@
+"""Connection pool wrapper and configuration.
+
+This module provides AsyncDatabase - the Python wrapper around the Rust
+connection pool managed by sqlx. It implements the SupportsExecute protocol.
+
+AsyncDatabase:
+    Wraps a Rust connection pool identified by name. The actual pool lives
+    in Rust; this class provides Python async interface.
+
+    Methods:
+        connect(): Initialize Rust pool with URL and settings.
+        disconnect(): Close Rust pool.
+        execute(ir): Send query IR to Rust, get MessagePack bytes back.
+
+    Usage:
+        db = AsyncDatabase("postgresql://...", name="default")
+        await db.connect()
+        result_bytes = await db.execute({"type": "select", ...})
+        await db.disconnect()
+
+PoolSettings:
+    Configuration dataclass for pool tuning.
+
+    PostgreSQL pool settings:
+        max_connections: Maximum pool size (default: auto)
+        min_connections: Minimum idle connections
+        acquire_timeout: Max wait time for connection
+        idle_timeout: Close idle connections after
+        max_lifetime: Max connection age
+        test_before_acquire: Ping before using connection
+
+    Transaction settings:
+        transaction_timeout: Max transaction duration (default: 5 min)
+        transaction_cleanup_interval: Cleanup check interval (default: 1 min)
+
+Rust Integration:
+    _init_pool(): Create pool in Rust registry
+    _execute(): Send IR bytes to Rust, execute via sqlx
+    close_pool(): Close specific pool
+    close_all_pools(): Close all pools (used by disconnect_all())
+
+URL Schemes:
+    postgres://, postgresql://  → PostgreSQL
+"""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+from datetime import timedelta
+from typing import Any
+
+from oxyde._msgpack import msgpack
+from oxyde.core.types import serialize_value
+from oxyde.db.registry import register_connection
+
+try:
+    from oxyde.core import close_all_pools, close_pool
+    from oxyde.core import execute as _execute
+    from oxyde.core import init_pool as _init_pool
+    from oxyde.core import init_pool_overwrite as _init_pool_overwrite
+    from oxyde.core import pool_backend as _pool_backend
+except ImportError:
+    # Stub for when the Rust module is not built
+    async def _execute(pool_name: str, ir_bytes: bytes) -> bytes:
+        raise RuntimeError(
+            "Rust core module not available. Please install oxyde-postgres-core."
+        )
+
+    async def _init_pool(name: str, url: str, settings: dict[str, Any] | None) -> None:
+        raise RuntimeError(
+            "Rust core module not available. Please install oxyde-postgres-core."
+        )
+
+    async def _init_pool_overwrite(
+        name: str, url: str, settings: dict[str, Any] | None
+    ) -> None:
+        raise RuntimeError(
+            "Rust core module not available. Please install oxyde-postgres-core."
+        )
+
+    async def close_pool(name: str) -> None:
+        raise RuntimeError(
+            "Rust core module not available. Please install oxyde-postgres-core."
+        )
+
+    async def close_all_pools() -> None:
+        raise RuntimeError(
+            "Rust core module not available. Please install oxyde-postgres-core."
+        )
+
+    async def _pool_backend(pool_name: str) -> str:
+        raise RuntimeError(
+            "Rust core module not available. Please install oxyde-postgres-core."
+        )
+
+
+def _msgpack_encoder(obj: Any) -> Any:
+    """Encode non-native types for msgpack via TYPE_REGISTRY."""
+    return serialize_value(obj)
+
+
+def _normalize_duration(value: float | int | timedelta | None) -> float | None:
+    """Convert duration to float seconds."""
+    if value is None:
+        return None
+    if isinstance(value, timedelta):
+        result = value.total_seconds()
+    elif isinstance(value, (int, float)):
+        result = float(value)
+    else:
+        raise TypeError(
+            f"Duration value must be int, float, or timedelta, got {type(value).__name__}"
+        )
+    if result < 0:
+        raise ValueError(f"Duration must be non-negative, got {result}")
+    return result
+
+
+def _validate_url_scheme(url: str) -> None:
+    """Accept PostgreSQL connection URL schemes only."""
+    if not url.startswith(("postgres://", "postgresql://")):
+        raise ValueError(
+            "PostgreSQL 18 or newer is required; use postgres:// or postgresql://."
+        )
+
+
+@dataclass
+class PoolSettings:
+    """Convenience container for pool configuration."""
+
+    max_connections: int | None = None
+    min_connections: int | None = None
+    acquire_timeout: float | int | timedelta | None = None
+    idle_timeout: float | int | timedelta | None = None
+    max_lifetime: float | int | timedelta | None = None
+    test_before_acquire: bool | None = None
+
+    # Transaction cleanup settings
+    transaction_timeout: float | int | timedelta | None = (
+        300  # 5 minutes (max age before cleanup)
+    )
+    transaction_cleanup_interval: float | int | timedelta | None = (
+        60  # 1 minute (cleanup interval)
+    )
+
+    # TLS settings (PostgreSQL)
+    # ssl_mode controls both TLS requirement and certificate verification:
+    #   PG:    "disable", "allow", "prefer", "require", "verify-ca", "verify-full"
+    ssl_mode: str | None = None
+    ssl_root_cert: str | None = None  # Path to CA certificate
+    ssl_client_cert: str | None = None  # Path to client certificate (mTLS)
+    ssl_client_key: str | None = None  # Path to client private key (mTLS)
+
+    # PostgreSQL-specific
+    pg_application_name: str | None = None  # Visible in pg_stat_activity
+    pg_statement_cache_capacity: int | None = None  # Prepared statement cache size
+
+    def to_payload(self) -> dict[str, Any] | None:
+        payload: dict[str, Any] = {}
+
+        if self.max_connections is not None:
+            payload["max_connections"] = int(self.max_connections)
+        if self.min_connections is not None:
+            payload["min_connections"] = int(self.min_connections)
+        if (value := _normalize_duration(self.idle_timeout)) is not None:
+            payload["idle_timeout"] = value
+        if (value := _normalize_duration(self.acquire_timeout)) is not None:
+            payload["acquire_timeout"] = value
+        if (value := _normalize_duration(self.max_lifetime)) is not None:
+            payload["max_lifetime"] = value
+        if self.test_before_acquire is not None:
+            payload["test_before_acquire"] = bool(self.test_before_acquire)
+
+        # Add transaction cleanup settings
+        if (value := _normalize_duration(self.transaction_timeout)) is not None:
+            payload["transaction_timeout"] = value
+        if (
+            value := _normalize_duration(self.transaction_cleanup_interval)
+        ) is not None:
+            payload["transaction_cleanup_interval"] = value
+
+        # TLS settings
+        if self.ssl_mode is not None:
+            payload["ssl_mode"] = str(self.ssl_mode)
+        if self.ssl_root_cert is not None:
+            payload["ssl_root_cert"] = str(self.ssl_root_cert)
+        if self.ssl_client_cert is not None:
+            payload["ssl_client_cert"] = str(self.ssl_client_cert)
+        if self.ssl_client_key is not None:
+            payload["ssl_client_key"] = str(self.ssl_client_key)
+
+        # PostgreSQL-specific
+        if self.pg_application_name is not None:
+            payload["pg_application_name"] = str(self.pg_application_name)
+        if self.pg_statement_cache_capacity is not None:
+            payload["pg_statement_cache_capacity"] = int(
+                self.pg_statement_cache_capacity
+            )
+
+        return payload or None
+
+
+class AsyncDatabase:
+    """Async database connection manager."""
+
+    def __init__(
+        self,
+        url: str,
+        *,
+        name: str = "default",
+        settings: PoolSettings | None = None,
+        auto_register: bool = True,
+        overwrite: bool = False,
+    ):
+        """
+        Initialize database connection wrapper.
+
+        Args:
+            url: Database connection URL (e.g., "postgresql://user:pass@host/db").
+            name: Identifier of the pool (used to look it up in Rust registry).
+            settings: Optional pool configuration.
+            auto_register: Automatically store this instance for later retrieval.
+            overwrite: Replace previously registered connection with the same name.
+        """
+        _validate_url_scheme(url)
+
+        self.url = url
+        self.name = name
+        self.settings = settings or PoolSettings()
+        self.backend: str | None = None
+        self._connected = False
+        self._connect_lock = asyncio.Lock()
+        self._overwrite = overwrite
+
+        if auto_register:
+            register_connection(self, overwrite=overwrite)
+
+    @property
+    def connected(self) -> bool:
+        """Return True if the connection pool has been initialised."""
+        return self._connected
+
+    async def connect(self) -> None:
+        """Establish database connection pool."""
+        async with self._connect_lock:
+            if self._connected:
+                return
+
+            payload = self.settings.to_payload()
+            if self._overwrite:
+                await _init_pool_overwrite(self.name, self.url, payload)
+            else:
+                await _init_pool(self.name, self.url, payload)
+            self.backend = await _pool_backend(self.name)
+            self._connected = True
+
+    async def disconnect(self) -> None:
+        """Close database connection pool."""
+        async with self._connect_lock:
+            if not self._connected:
+                return
+
+            await close_pool(self.name)
+            self._connected = False
+
+    async def ensure_connected(self) -> None:
+        """Connect on demand if not connected yet."""
+        if not self._connected:
+            await self.connect()
+
+    async def execute(self, ir: dict[str, Any]) -> bytes:
+        """
+        Execute a query using the IR format against this database.
+
+        Args:
+            ir: Query intermediate representation.
+
+        Returns:
+            MessagePack bytes containing query results.
+        """
+        if not self._connected:
+            raise RuntimeError(
+                f"Database '{self.name}' not connected. Call connect() first."
+            )
+
+        ir_bytes = msgpack.packb(ir, default=_msgpack_encoder)
+        result_bytes: bytes = await _execute(self.name, ir_bytes)
+        return result_bytes
+
+    async def __aenter__(self) -> AsyncDatabase:
+        """Context manager entry."""
+        await self.connect()
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """Context manager exit."""
+        await self.disconnect()
+
+
+__all__ = [
+    "AsyncDatabase",
+    "PoolSettings",
+    "close_all_pools",
+]

oxyde/db/registry.py

@@ -0,0 +1,113 @@
+"""Global registry for named database connections.
+
+This module maintains a dict mapping connection names to AsyncDatabase instances.
+The "default" connection is used when no explicit name is specified.
+
+Why a Registry?
+    Queries need database connections, but models don't store connection refs.
+    The registry provides a global lookup: Model.objects uses get_connection()
+    to find the right pool.
+
+    # At startup
+    await db.init(default="postgresql://...", analytics="postgresql://...")
+
+    # In queries (implicit lookup via registry)
+    users = await User.objects.all()  # Uses "default"
+    stats = await Metric.objects.all(using="analytics")  # Uses "analytics"
+
+Functions:
+    register_connection(database, overwrite=False):
+        Add AsyncDatabase to registry. Raises if name exists and not overwrite.
+
+    unregister_connection(name):
+        Remove from registry (used during cleanup).
+
+    get_connection(name="default", ensure_connected=True):
+        Look up by name. Auto-connects if ensure_connected=True.
+        Raises KeyError if not found.
+
+    disconnect_all():
+        Close all pools (both Python registry and Rust pools).
+        Clears the registry. Used at shutdown.
+
+Thread Safety:
+    The registry is a simple dict. In async context, this is safe because
+    Python's GIL prevents concurrent dict mutations. Registration typically
+    happens at startup before concurrent access.
+
+Transaction Integration:
+    QueryManager._resolve_client() checks get_active_transaction() before
+    falling back to get_connection(). This ensures queries inside atomic()
+    use the transaction, not a new pool connection.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from oxyde.core import close_all_pools
+
+if TYPE_CHECKING:
+    from oxyde.db.pool import AsyncDatabase
+
+# Global connection registry
+_CONNECTIONS: dict[str, AsyncDatabase] = {}
+
+
+def register_connection(
+    database: AsyncDatabase, *, overwrite: bool = False
+) -> AsyncDatabase:
+    """Register a connection instance so it can be retrieved later."""
+    if not overwrite and database.name in _CONNECTIONS:
+        raise ValueError(f"Connection '{database.name}' already registered")
+    _CONNECTIONS[database.name] = database
+    return database
+
+
+def unregister_connection(name: str) -> None:
+    """Remove a connection from the registry."""
+    _CONNECTIONS.pop(name, None)
+
+
+async def get_connection(
+    name: str = "default",
+    *,
+    ensure_connected: bool = True,
+) -> AsyncDatabase:
+    """
+    Retrieve a registered connection by name.
+
+    Args:
+        name: Connection identifier.
+        ensure_connected: Automatically call connect() if the pool is not ready.
+    """
+    try:
+        database = _CONNECTIONS[name]
+    except KeyError as exc:
+        raise KeyError(f"Connection '{name}' is not registered") from exc
+
+    if ensure_connected and not database.connected:
+        await database.connect()
+
+    return database
+
+
+async def disconnect_all() -> None:
+    """Disconnect all registered pools and clear the registry."""
+    # Mark all Python-side connections as disconnected
+    for db in _CONNECTIONS.values():
+        db._connected = False
+
+    # Clear Python registry
+    _CONNECTIONS.clear()
+
+    # Close all Rust pools (also rolls back active transactions)
+    await close_all_pools()
+
+
+__all__ = [
+    "register_connection",
+    "unregister_connection",
+    "get_connection",
+    "disconnect_all",
+]

oxyde/db/schema.py

@@ -0,0 +1,51 @@
+"""Schema management. create/drop tables from registered models."""
+
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING, Any
+
+from oxyde.core import migration_compute_diff, migration_to_sql
+from oxyde.migrations.extract import extract_current_schema
+from oxyde.migrations.utils import detect_dialect
+from oxyde.queries.raw import execute_raw
+
+if TYPE_CHECKING:
+    from oxyde.db.pool import AsyncDatabase
+
+
+async def create_tables(database: AsyncDatabase) -> None:
+    """Create all tables for registered models.
+
+    Generates dialect-aware DDL via sea-query and executes it.
+    No migration files or tracking table needed.
+
+    Args:
+        database: Connected AsyncDatabase instance.
+    """
+    dialect = detect_dialect(database.url)
+    current = extract_current_schema(dialect=dialect)
+    empty: dict[str, Any] = {"version": 1, "tables": {}}
+    ops_json = migration_compute_diff(json.dumps(empty), json.dumps(current))
+    statements = migration_to_sql(ops_json, dialect)
+    for sql in statements:
+        await execute_raw(sql, client=database)
+
+
+async def drop_tables(database: AsyncDatabase) -> None:
+    """Drop all tables for registered models.
+
+    Handles FK constraints per dialect.
+
+    Args:
+        database: Connected AsyncDatabase instance.
+    """
+    dialect = detect_dialect(database.url)
+    current = extract_current_schema(dialect=dialect)
+    table_names = list(current.get("tables", {}).keys())
+
+    if not table_names:
+        return
+
+    tables = ", ".join(table_names)
+    await execute_raw(f"DROP TABLE IF EXISTS {tables} CASCADE", client=database)