oxyde-postgres 0.1.0__py3-none-any.whl

oxyde/db/transaction.py

@@ -0,0 +1,392 @@
+"""Transaction management with Django-style atomic() context manager.
+
+This module provides transaction support with automatic commit/rollback
+and nested transaction handling via savepoints.
+
+Usage:
+    from oxyde.db import transaction
+
+    # Basic transaction
+    async with transaction.atomic():
+        await User.objects.create(name="Alice")
+        await Order.objects.create(user_id=1)
+        # Commits on exit, rolls back on exception
+
+    # Explicit rollback
+    async with transaction.atomic() as tx:
+        await User.objects.create(name="Alice")
+        tx.set_rollback()  # Force rollback even without exception
+
+    # Nested transactions (savepoints)
+    async with transaction.atomic():
+        await User.objects.create(name="Alice")
+        try:
+            async with transaction.atomic():
+                await User.objects.create(name="Bob")
+                raise ValueError("oops")
+        except ValueError:
+            pass  # Bob rolled back to savepoint, Alice still pending
+        # Alice commits
+
+Classes:
+    AsyncTransaction:
+        Low-level transaction wrapper. Holds tx_id from Rust.
+        execute() sends queries through the transaction.
+
+    AtomicTransactionContext:
+        High-level context manager for atomic(). Manages nesting
+        via savepoints. Tracks depth per connection alias.
+
+    TransactionTimeoutError:
+        Raised when transaction exceeds configured timeout.
+
+Functions:
+    atomic(using="default", timeout=None):
+        Create atomic context manager. Django-compatible API.
+
+    get_active_transaction(using="default"):
+        Get current transaction for connection (or None).
+        Used by QueryManager to route queries through active tx.
+
+    get_transaction_depth(using="default"):
+        Get atomic() depth for the connection alias and current task.
+        Returns 0 outside atomic(), 1 in the outermost block, and 2+ when nested.
+
+ContextVar Storage:
+    _ACTIVE_TRANSACTIONS stores {alias: {transaction, depth, force_rollback, task}}
+    per async context. This allows concurrent requests to have separate
+    transaction state. The 'task' field tracks ownership via asyncio.current_task()
+    to prevent child tasks from inheriting parent transactions through ContextVar.
+
+Savepoint Handling:
+    Nested atomic() calls create savepoints (sp_1, sp_2, ...).
+    Inner exception rolls back to savepoint, not entire transaction.
+    This matches Django/PostgreSQL behavior.
+
+Rust Integration:
+    _begin_transaction(pool_name) → tx_id
+    _execute_in_transaction(pool_name, tx_id, ir_bytes) → bytes
+    _commit_transaction(tx_id)
+    _rollback_transaction(tx_id)
+    _create_savepoint(tx_id, name)
+    _rollback_to_savepoint(tx_id, name)
+    _release_savepoint(tx_id, name)
+"""
+
+from __future__ import annotations
+
+import asyncio
+from contextvars import ContextVar
+from datetime import timedelta
+from typing import TYPE_CHECKING, Any
+
+from oxyde._msgpack import msgpack
+from oxyde.db.pool import _msgpack_encoder, _normalize_duration
+from oxyde.db.registry import get_connection
+
+if TYPE_CHECKING:
+    from oxyde.db.pool import AsyncDatabase
+
+try:
+    from oxyde.core import begin_transaction as _begin_transaction
+    from oxyde.core import commit_transaction as _commit_transaction
+    from oxyde.core import create_savepoint as _create_savepoint
+    from oxyde.core import execute_in_transaction as _execute_in_transaction
+    from oxyde.core import release_savepoint as _release_savepoint
+    from oxyde.core import rollback_to_savepoint as _rollback_to_savepoint
+    from oxyde.core import rollback_transaction as _rollback_transaction
+except ImportError:
+    # Stub for when the Rust module is not built
+    async def _execute_in_transaction(
+        pool_name: str, tx_id: int, ir_bytes: bytes
+    ) -> bytes:
+        raise RuntimeError(
+            "Rust core module not available. Please install oxyde-postgres-core."
+        )
+
+    async def _begin_transaction(pool_name: str) -> int:
+        raise RuntimeError(
+            "Rust core module not available. Please install oxyde-postgres-core."
+        )
+
+    async def _commit_transaction(tx_id: int) -> None:
+        raise RuntimeError(
+            "Rust core module not available. Please install oxyde-postgres-core."
+        )
+
+    async def _rollback_transaction(tx_id: int) -> None:
+        raise RuntimeError(
+            "Rust core module not available. Please install oxyde-postgres-core."
+        )
+
+    async def _create_savepoint(tx_id: int, savepoint_name: str) -> None:
+        raise RuntimeError(
+            "Rust core module not available. Please install oxyde-postgres-core."
+        )
+
+    async def _rollback_to_savepoint(tx_id: int, savepoint_name: str) -> None:
+        raise RuntimeError(
+            "Rust core module not available. Please install oxyde-postgres-core."
+        )
+
+    async def _release_savepoint(tx_id: int, savepoint_name: str) -> None:
+        raise RuntimeError(
+            "Rust core module not available. Please install oxyde-postgres-core."
+        )
+
+
+# Active transactions per connection (ContextVar for async safety)
+_ACTIVE_TRANSACTIONS: ContextVar[dict[str, dict[str, Any]]] = ContextVar(
+    "oxyde_active_transactions",
+    default={},
+)
+
+
+def _get_owned_entry(using: str) -> dict[str, Any] | None:
+    """Return the transaction entry only if it belongs to the current task."""
+    entry = _ACTIVE_TRANSACTIONS.get().get(using)
+    if not entry:
+        return None
+    try:
+        task = asyncio.current_task()
+    except RuntimeError:
+        # No running event loop (e.g. a worker thread with a copied context)
+        return None
+    if entry.get("task") is task:
+        return entry
+    return None
+
+
+class TransactionTimeoutError(TimeoutError):
+    """Raised when a transaction exceeds the configured timeout."""
+
+
+class AsyncTransaction:
+    """Wrapper for executing queries within a database transaction."""
+
+    def __init__(self, database: AsyncDatabase, timeout: float | None = None):
+        self._database = database
+        self._tx_id: int | None = None
+        self._timeout = timeout
+        self._deadline: float | None = None
+        self._timed_out = False
+
+    @property
+    def id(self) -> int:
+        if self._tx_id is None:
+            raise RuntimeError("Transaction not started")
+        return self._tx_id
+
+    def _remaining_timeout(self) -> float | None:
+        if self._deadline is None:
+            return None
+        loop = asyncio.get_running_loop()
+        return self._deadline - loop.time()
+
+    async def execute(self, ir: dict[str, Any]) -> bytes:
+        if self._tx_id is None:
+            raise RuntimeError("Transaction not started")
+        if self._timed_out:
+            raise TransactionTimeoutError(
+                "Transaction can no longer be used because its timeout elapsed"
+            )
+
+        ir_bytes = msgpack.packb(ir, default=_msgpack_encoder)
+        coro = _execute_in_transaction(self._database.name, self._tx_id, ir_bytes)
+
+        remaining = self._remaining_timeout()
+        if remaining is not None:
+            if remaining <= 0:
+                self._timed_out = True
+                msg = f"Transaction {self._tx_id} exceeded timeout"
+                raise TransactionTimeoutError(msg)
+            try:
+                return await asyncio.wait_for(coro, remaining)
+            except asyncio.TimeoutError as exc:
+                self._timed_out = True
+                msg = f"Transaction {self._tx_id} exceeded timeout"
+                raise TransactionTimeoutError(msg) from exc
+
+        result: bytes = await coro
+        return result
+
+    async def __aenter__(self) -> AsyncTransaction:
+        await self._database.ensure_connected()
+        self._tx_id = await _begin_transaction(self._database.name)
+        if self._timeout is not None:
+            loop = asyncio.get_running_loop()
+            self._deadline = loop.time() + self._timeout
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        tx_id = self._tx_id
+        self._tx_id = None
+        self._deadline = None
+
+        if tx_id is None:
+            return
+
+        try:
+            if self._timed_out or exc_type is not None:
+                await _rollback_transaction(tx_id)
+            else:
+                await _commit_transaction(tx_id)
+        finally:
+            self._timed_out = False
+
+
+class AtomicTransactionContext:
+    """Context manager that reuses or creates transactions per connection alias."""
+
+    def __init__(
+        self,
+        using: str = "default",
+        timeout: float | None = None,
+        database: AsyncDatabase | None = None,
+    ):
+        self.using = using
+        self.timeout = timeout
+        self._database = database
+        self._transaction: AsyncTransaction | None = None
+        self._savepoint_name: str | None = None
+
+    async def __aenter__(self) -> AtomicTransactionContext:
+        state = dict(_ACTIVE_TRANSACTIONS.get())
+        entry = _get_owned_entry(self.using)
+        if entry:
+            # Nested transaction - create savepoint
+            depth = entry["depth"]
+            self._transaction = entry["transaction"]
+
+            # Create savepoint BEFORE incrementing depth —
+            # if this fails, __aexit__ won't be called and depth must stay unchanged
+            self._savepoint_name = f"sp_{depth + 1}"
+            await _create_savepoint(self._transaction.id, self._savepoint_name)
+
+            entry["depth"] += 1
+            state[self.using] = entry
+            _ACTIVE_TRANSACTIONS.set(state)
+
+            return self
+
+        # Use provided database instance or get from registry
+        database = (
+            self._database
+            if self._database is not None
+            else await get_connection(self.using)
+        )
+        tx = AsyncTransaction(database, timeout=self.timeout)
+        await tx.__aenter__()
+        state[self.using] = {
+            "transaction": tx,
+            "depth": 1,
+            "force_rollback": False,
+            "task": asyncio.current_task(),
+        }
+        _ACTIVE_TRANSACTIONS.set(state)
+        self._transaction = tx
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        state = dict(_ACTIVE_TRANSACTIONS.get())
+        entry = _get_owned_entry(self.using)
+        if entry is None:
+            raise RuntimeError("transaction.atomic context mismatch")
+
+        entry["depth"] -= 1
+        force_rollback = entry.get("force_rollback", False)
+
+        if entry["depth"] == 0:
+            # Top-level transaction - commit or rollback entire transaction
+            state.pop(self.using, None)
+            _ACTIVE_TRANSACTIONS.set(state)
+            tx = entry["transaction"]
+            if exc_type is None and force_rollback:
+                exc_val = RuntimeError("transaction marked for rollback")
+                exc_type = RuntimeError
+            await tx.__aexit__(exc_type, exc_val, exc_tb)
+        else:
+            # Nested transaction - handle savepoint
+            if self._savepoint_name:
+                if exc_type:
+                    # Rollback to savepoint - this is the end of this nested transaction
+                    # Don't mark outer transaction for rollback, savepoint handles it
+                    await _rollback_to_savepoint(
+                        entry["transaction"].id, self._savepoint_name
+                    )
+                else:
+                    # Release savepoint - nested transaction succeeded
+                    await _release_savepoint(
+                        entry["transaction"].id, self._savepoint_name
+                    )
+                # Don't modify force_rollback - savepoint already handled the rollback
+            else:
+                # No savepoint (shouldn't happen with correct nesting)
+                if exc_type:
+                    entry["force_rollback"] = True
+
+            state[self.using] = entry
+            _ACTIVE_TRANSACTIONS.set(state)
+
+    def set_rollback(self, rollback: bool = True) -> None:
+        entry = _get_owned_entry(self.using)
+        if entry is None:
+            raise RuntimeError("No active transaction to mark for rollback")
+        state = dict(_ACTIVE_TRANSACTIONS.get())
+        entry["force_rollback"] = bool(rollback)
+        state[self.using] = entry
+        _ACTIVE_TRANSACTIONS.set(state)
+
+    @property
+    def transaction(self) -> AsyncTransaction:
+        if self._transaction is None:
+            raise RuntimeError("transaction.atomic context not entered")
+        return self._transaction
+
+    async def execute(self, ir: dict[str, Any]) -> bytes:
+        return await self.transaction.execute(ir)
+
+    def __getattr__(self, item: str) -> Any:
+        return getattr(self.transaction, item)
+
+
+def atomic(
+    *,
+    using: str = "default",
+    timeout: float | int | timedelta | None = None,
+) -> AtomicTransactionContext:
+    """Create an atomic transaction context."""
+    normalized = _normalize_duration(timeout)
+    return AtomicTransactionContext(using=using, timeout=normalized)
+
+
+def get_active_transaction(using: str = "default") -> AsyncTransaction | None:
+    """Get the currently active transaction for a connection alias."""
+    entry = _get_owned_entry(using)
+    if entry:
+        tx: AsyncTransaction = entry["transaction"]
+        return tx
+    return None
+
+
+def get_transaction_depth(using: str = "default") -> int:
+    """Return the atomic() depth for a connection alias in the calling task.
+
+    Returns 0 when no atomic() block is active, 1 in the outermost block,
+    and 2+ in nested blocks. Transactions owned by other tasks are excluded.
+    """
+    entry = _get_owned_entry(using)
+    if entry:
+        depth: int = entry["depth"]
+        return depth
+    return 0
+
+
+__all__ = [
+    "AsyncTransaction",
+    "TransactionTimeoutError",
+    "AtomicTransactionContext",
+    "atomic",
+    "get_active_transaction",
+    "get_transaction_depth",
+]

oxyde/exceptions.py

@@ -0,0 +1,89 @@
+"""Exception hierarchy for Oxyde ORM.
+
+All Oxyde exceptions inherit from OxydeError, allowing catch-all handling:
+
+    try:
+        user = await User.objects.get(id=999)
+    except OxydeError as e:
+        print(f"ORM error: {e}")
+
+Exception hierarchy:
+    OxydeError (base)
+    ├── FieldError                - Invalid field definition or access
+    ├── FieldLookupError          - Unknown lookup operator (e.g., __xyz)
+    │   └── FieldLookupValueError - Invalid value for lookup (e.g., __in=None)
+    └── ManagerError              - Query execution errors
+        ├── NotFoundError         - get() returned no rows
+        ├── MultipleObjectsReturned - get() returned multiple rows
+        └── IntegrityError        - Constraint violation
+            ├── UniqueViolationError     - UNIQUE / PK constraint
+            ├── ForeignKeyViolationError - FOREIGN KEY constraint
+            ├── NotNullViolationError    - NOT NULL constraint
+            └── CheckViolationError      - CHECK constraint
+"""
+
+from __future__ import annotations
+
+
+class OxydeError(Exception):
+    """Base exception for all Oxyde-related errors."""
+
+
+class FieldError(OxydeError):
+    """Raised when a model field is invalid or missing."""
+
+
+class FieldLookupError(OxydeError):
+    """Raised when an unsupported field lookup is requested."""
+
+
+class FieldLookupValueError(FieldLookupError):
+    """Raised when the lookup value is not compatible with the operator."""
+
+
+class ManagerError(OxydeError):
+    """Raised for issues inside the ORM manager layer."""
+
+
+class NotFoundError(ManagerError):
+    """Raised when a query expecting a single row finds none."""
+
+
+class MultipleObjectsReturned(ManagerError):
+    """Raised when a query expecting a single row finds more than one."""
+
+
+class IntegrityError(ManagerError):
+    """Raised when database integrity constraints are violated."""
+
+
+class UniqueViolationError(IntegrityError):
+    """Raised when a UNIQUE (or PRIMARY KEY) constraint is violated."""
+
+
+class ForeignKeyViolationError(IntegrityError):
+    """Raised when a FOREIGN KEY constraint is violated."""
+
+
+class NotNullViolationError(IntegrityError):
+    """Raised when a NOT NULL constraint is violated."""
+
+
+class CheckViolationError(IntegrityError):
+    """Raised when a CHECK constraint is violated."""
+
+
+__all__ = [
+    "OxydeError",
+    "FieldError",
+    "FieldLookupError",
+    "FieldLookupValueError",
+    "ManagerError",
+    "NotFoundError",
+    "MultipleObjectsReturned",
+    "IntegrityError",
+    "UniqueViolationError",
+    "ForeignKeyViolationError",
+    "NotNullViolationError",
+    "CheckViolationError",
+]

oxyde/migrations/__init__.py

@@ -0,0 +1,49 @@
+"""Oxyde migrations system.
+
+This module provides database migration functionality:
+- Schema extraction from models
+- Migration replay in memory
+- MigrationContext for upgrade/downgrade operations
+- Migration file generation
+- Migration tracking and execution
+"""
+
+from __future__ import annotations
+
+from oxyde.migrations.context import MigrationContext
+from oxyde.migrations.executor import (
+    apply_migrations,
+    rollback_migration,
+    rollback_migrations,
+)
+from oxyde.migrations.extract import extract_current_schema
+from oxyde.migrations.generator import generate_migration_file
+from oxyde.migrations.replay import SchemaState, replay_migrations
+from oxyde.migrations.squash import SquashResult, squash_migrations
+from oxyde.migrations.tracker import (
+    ensure_migrations_table,
+    get_applied_migrations,
+    get_migration_files,
+    get_pending_migrations,
+    record_migration,
+    remove_migration,
+)
+
+__all__ = [
+    "SquashResult",
+    "squash_migrations",
+    "extract_current_schema",
+    "SchemaState",
+    "replay_migrations",
+    "MigrationContext",
+    "generate_migration_file",
+    "apply_migrations",
+    "rollback_migration",
+    "rollback_migrations",
+    "ensure_migrations_table",
+    "get_applied_migrations",
+    "get_migration_files",
+    "get_pending_migrations",
+    "record_migration",
+    "remove_migration",
+]