pg-partsmith 0.1.0__py3-none-any.whl

pg_partsmith/__init__.py

@@ -0,0 +1,61 @@
+"""PostgreSQL partition lifecycle management with extensible hooks."""
+
+from .__version__ import __version__
+from .entities import (
+    MaintenanceIssueStep,
+    MaintenanceResult,
+    PartitionGranularity,
+    PartitionInfo,
+    PartitionStrategy,
+    PartitionType,
+    Period,
+    TablePartitionConfig,
+)
+from .exceptions import (
+    DropRetryExhaustedError,
+    InvalidPartitionConfigError,
+    LockAcquisitionError,
+    PartitionAlreadyExistsError,
+    PartitionAttachedError,
+    PartitionDetachInProgressError,
+    PartitionError,
+    PartitionNotFoundError,
+    UnmanagedPartitionDropError,
+)
+from .protocols import PeriodCalculator
+from .strategies import (
+    BasePeriodCalculator,
+    DayPeriodCalculator,
+    MonthPeriodCalculator,
+    WeekPeriodCalculator,
+    YearPeriodCalculator,
+)
+from .strategies.selector import get_period_calculator
+
+__all__ = [
+    "BasePeriodCalculator",
+    "DayPeriodCalculator",
+    "DropRetryExhaustedError",
+    "InvalidPartitionConfigError",
+    "LockAcquisitionError",
+    "MaintenanceIssueStep",
+    "MaintenanceResult",
+    "MonthPeriodCalculator",
+    "PartitionAlreadyExistsError",
+    "PartitionAttachedError",
+    "PartitionDetachInProgressError",
+    "PartitionError",
+    "PartitionGranularity",
+    "PartitionInfo",
+    "PartitionNotFoundError",
+    "PartitionStrategy",
+    "PartitionType",
+    "Period",
+    "PeriodCalculator",
+    "TablePartitionConfig",
+    "UnmanagedPartitionDropError",
+    "WeekPeriodCalculator",
+    "YearPeriodCalculator",
+    "__version__",
+    "get_period_calculator",
+]

pg_partsmith/__version__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

pg_partsmith/aio/__init__.py

@@ -0,0 +1,24 @@
+"""Async implementations for partition management."""
+
+from .hooks import BasePartitionLifecycleHooks, PartitionLifecycleHooks
+from .lock import PostgresAdvisoryLockManager, RedisDistributedLockManager
+from .maintainer import PartitionMaintainer, maintain_partitions
+from .metadata import PostgresMetadataProvider
+from .protocols import LockManager, PartitionMetadataProvider, PartitionRepository
+from .repositories import PostgresPartitionRepository
+from .service import PartitionLifecycleService
+
+__all__ = [
+    "BasePartitionLifecycleHooks",
+    "LockManager",
+    "PartitionLifecycleHooks",
+    "PartitionLifecycleService",
+    "PartitionMaintainer",
+    "PartitionMetadataProvider",
+    "PartitionRepository",
+    "PostgresAdvisoryLockManager",
+    "PostgresMetadataProvider",
+    "PostgresPartitionRepository",
+    "RedisDistributedLockManager",
+    "maintain_partitions",
+]

pg_partsmith/aio/hooks.py

@@ -0,0 +1,168 @@
+"""Lifecycle hooks (middleware) and protocol for partition operations.
+
+Hooks let you inject custom logic at each step of the partition lifecycle -
+for example, exporting data before a partition is dropped, publishing events
+to Kafka after a partition is created, or archiving rows before detachment.
+
+Usage example::
+
+    class KafkaExportHooks(BasePartitionLifecycleHooks):
+        def __init__(self, producer: KafkaProducer) -> None:
+            self._producer = producer
+
+        async def before_drop(self, table_name: str, partition_name: str) -> None:
+            await self._producer.send("partition.before_drop", {
+                "table": table_name,
+                "partition": partition_name,
+            })
+
+    service = PartitionLifecycleService(
+        repo=repo,
+        metadata=metadata,
+        locks=locks,
+        period_calculator=calculator,
+        hooks=[KafkaExportHooks(producer)],
+    )
+"""
+
+from typing import Protocol, runtime_checkable
+
+from pg_partsmith.entities import PartitionInfo, TablePartitionConfig
+
+
+@runtime_checkable
+class PartitionLifecycleHooks(Protocol):
+    """Protocol for partition lifecycle hooks.
+
+    Implement this protocol to inject custom logic at each step of the partition
+    lifecycle without inheriting from a specific base class.
+    """
+
+    async def before_create(
+        self,
+        config: TablePartitionConfig,
+        partition_name: str,
+        from_value: str,
+        to_value: str,
+    ) -> None: ...
+
+    async def after_create(
+        self,
+        config: TablePartitionConfig,
+        partition: PartitionInfo,
+    ) -> None: ...
+
+    async def before_detach(
+        self,
+        table_name: str,
+        partition: PartitionInfo,
+    ) -> None: ...
+
+    async def after_detach(
+        self,
+        table_name: str,
+        partition_name: str,
+    ) -> None: ...
+
+    async def before_drop(
+        self,
+        table_name: str,
+        partition_name: str,
+    ) -> None: ...
+
+    async def after_drop(
+        self,
+        table_name: str,
+        partition_name: str,
+    ) -> None: ...
+
+
+class BasePartitionLifecycleHooks:
+    """No-op base implementation of partition lifecycle hooks.
+
+    Subclass and override only the methods you need.
+    All methods are no-ops by default so you can selectively add behaviour
+    without implementing every step.
+    """
+
+    async def before_create(
+        self,
+        config: TablePartitionConfig,
+        partition_name: str,
+        from_value: str,
+        to_value: str,
+    ) -> None:
+        """Called before a partition is created.
+
+        Args:
+            config: Table partition configuration.
+            partition_name: Name the new partition will be given.
+            from_value: Start boundary value.
+            to_value: End boundary value.
+        """
+
+    async def after_create(
+        self,
+        config: TablePartitionConfig,
+        partition: PartitionInfo,
+    ) -> None:
+        """Called after a partition has been created (and optionally attached).
+
+        Args:
+            config: Table partition configuration.
+            partition: Info about the newly created partition.
+        """
+
+    async def before_detach(
+        self,
+        table_name: str,
+        partition: PartitionInfo,
+    ) -> None:
+        """Called before a partition is detached from its parent table.
+
+        This is a good place to export or archive data while the partition
+        is still accessible via the parent table's indexes and constraints.
+
+        Args:
+            table_name: Parent table name.
+            partition: Info about the partition being detached.
+        """
+
+    async def after_detach(
+        self,
+        table_name: str,
+        partition_name: str,
+    ) -> None:
+        """Called after a partition has been detached.
+
+        Args:
+            table_name: Parent table name.
+            partition_name: Name of the detached partition.
+        """
+
+    async def before_drop(
+        self,
+        table_name: str,
+        partition_name: str,
+    ) -> None:
+        """Called before a partition table is dropped.
+
+        This is the last chance to read or export data from the partition
+        before it is permanently destroyed.
+
+        Args:
+            table_name: Parent table name.
+            partition_name: Name of the partition about to be dropped.
+        """
+
+    async def after_drop(
+        self,
+        table_name: str,
+        partition_name: str,
+    ) -> None:
+        """Called after a partition table has been dropped.
+
+        Args:
+            table_name: Parent table name.
+            partition_name: Name of the dropped partition.
+        """

pg_partsmith/aio/lock/__init__.py

@@ -0,0 +1,9 @@
+"""Lock manager implementations."""
+
+from .postgres import PostgresAdvisoryLockManager
+from .redis import RedisDistributedLockManager
+
+__all__ = [
+    "PostgresAdvisoryLockManager",
+    "RedisDistributedLockManager",
+]

pg_partsmith/aio/lock/postgres.py

@@ -0,0 +1,211 @@
+"""PostgreSQL advisory lock manager."""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import logging
+import time
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING
+
+from sqlalchemy import text
+from sqlalchemy.exc import SQLAlchemyError
+
+from pg_partsmith.exceptions import LockAcquisitionError
+from pg_partsmith.utils import calculate_lock_id
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+    from contextlib import AbstractAsyncContextManager
+
+    from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
+
+logger = logging.getLogger(__name__)
+
+
+class PostgresAdvisoryLockManager:
+    """Lock manager using PostgreSQL advisory locks.
+
+    Holds the advisory lock on a dedicated AUTOCOMMIT connection from the
+    engine pool. This guarantees the lock survives any number of commits or
+    rollbacks on the caller's session, which is required when the caller
+    needs to commit DDL (e.g. ATTACH PARTITION) before running
+    DETACH PARTITION CONCURRENTLY.
+
+    Override `_compute_lock_id` to customise the lock ID derivation.
+    """
+
+    def __init__(
+        self,
+        engine: AsyncEngine,
+        prefix: str = "partitioner",
+        acquire_min_interval_seconds: float = 0.0,
+    ) -> None:
+        """Initialize lock manager.
+
+        Args:
+            engine: SQLAlchemy async engine used to open a dedicated connection
+                for the advisory lock.
+            prefix: Prefix for lock key generation.
+            acquire_min_interval_seconds: Minimum seconds between acquire attempts
+                per table (rate limiting). 0 disables.
+        """
+        self._engine = engine
+        self._prefix = prefix
+        self._acquire_min_interval = max(0.0, acquire_min_interval_seconds)
+        self._last_acquire_time: dict[str, float] = {}
+        self._rate_limit_lock = asyncio.Lock()
+
+    def _compute_lock_id(self, table_name: str) -> int:
+        """Compute the advisory lock ID for a table name.
+
+        Override this method to customise the ID derivation strategy.
+
+        Args:
+            table_name: Table name to lock.
+
+        Returns:
+            Advisory lock ID.
+        """
+        return calculate_lock_id(table_name, prefix=self._prefix)
+
+    def acquire_lock(self, table_name: str) -> AbstractAsyncContextManager[None]:
+        """Acquire advisory lock for a table.
+
+        Opens a dedicated AUTOCOMMIT connection from the engine pool and
+        acquires a session-level advisory lock on it. The lock is released
+        when the context manager exits, with cancellation-safe cleanup.
+
+        Args:
+            table_name: Table name to lock.
+
+        Returns:
+            Async context manager for the lock.
+
+        Raises:
+            LockAcquisitionError: If the lock cannot be acquired.
+        """
+        return self._lock_scope(table_name)
+
+    @asynccontextmanager
+    async def _lock_scope(self, table_name: str) -> AsyncIterator[None]:
+        """Internal acquire/release flow for a single advisory lock."""
+        await self._respect_rate_limit(table_name)
+        lock_id = self._compute_lock_id(table_name)
+
+        async with self._engine.connect() as base_conn:
+            conn = await base_conn.execution_options(isolation_level="AUTOCOMMIT")
+            await self._try_acquire(conn, lock_id, table_name)
+
+            body_exc: BaseException | None = None
+            try:
+                yield
+            except BaseException as exc:
+                body_exc = exc
+                raise
+            finally:
+                await self._release_safely(conn, lock_id, table_name, body_exc)
+
+    async def _respect_rate_limit(self, table_name: str) -> None:
+        """Sleep enough to enforce the configured min-interval between acquires."""
+        if self._acquire_min_interval <= 0:
+            return
+        async with self._rate_limit_lock:
+            now = time.monotonic()
+            last = self._last_acquire_time.get(table_name, 0.0)
+            delay = last + self._acquire_min_interval - now
+            if delay > 0:
+                await asyncio.sleep(delay)
+            self._last_acquire_time[table_name] = time.monotonic()
+
+    async def _try_acquire(self, conn: AsyncConnection, lock_id: int, table_name: str) -> None:
+        """Run ``pg_try_advisory_lock`` and raise if not granted."""
+        result = await conn.execute(text("SELECT pg_try_advisory_lock(:lock_id)"), {"lock_id": lock_id})
+        if not result.scalar():
+            raise LockAcquisitionError(table_name, "advisory lock unavailable")
+
+    async def _release_safely(
+        self,
+        conn: AsyncConnection,
+        lock_id: int,
+        table_name: str,
+        body_exc: BaseException | None,
+    ) -> None:
+        """Release the lock with shielding so cancellation cannot leak a held lock."""
+        try:
+            await asyncio.shield(self._unlock(conn, lock_id, table_name))
+        except (asyncio.CancelledError, KeyboardInterrupt, SystemExit):
+            # Defensively invalidate so the connection is not returned to the pool with a dangling lock.
+            with contextlib.suppress(Exception):
+                await asyncio.shield(conn.invalidate())
+            raise
+        except Exception:
+            # Body exception takes precedence; otherwise propagate the unlock failure.
+            if body_exc is None:
+                raise
+            logger.warning(
+                "Failed to release advisory lock",
+                extra={"table_name": table_name, "lock_id": lock_id},
+            )
+
+    async def _unlock(self, conn: AsyncConnection, lock_id: int, table_name: str) -> None:
+        """Run ``pg_advisory_unlock``; invalidate the connection on any failure."""
+        try:
+            await conn.execute(text("SELECT pg_advisory_unlock(:lock_id)"), {"lock_id": lock_id})
+        except (asyncio.CancelledError, KeyboardInterrupt, SystemExit):
+            raise
+        except (SQLAlchemyError, OSError) as e:
+            logger.warning(
+                "Failed to release advisory lock (recoverable)",
+                extra={"table_name": table_name, "lock_id": lock_id, "error": str(e)},
+            )
+            with contextlib.suppress(Exception):
+                await conn.invalidate()
+            raise
+        except Exception:
+            logger.exception(
+                "Unexpected error while releasing advisory lock",
+                extra={"table_name": table_name, "lock_id": lock_id},
+            )
+            with contextlib.suppress(Exception):
+                await conn.invalidate()
+            raise
+
+    async def is_locked(self, table_name: str) -> bool:
+        """Check if lock is held by any session.
+
+        Args:
+            table_name: Table name.
+
+        Returns:
+            True if the advisory lock for the given table is currently held.
+        """
+        lock_id = self._compute_lock_id(table_name)
+        # Split 64-bit lock_id into classid and objid as stored in pg_locks for int8 advisory locks (objsubid=1).
+        class_id = (lock_id >> 32) & 0xFFFFFFFF
+        if class_id > 0x7FFFFFFF:
+            class_id -= 0x100000000
+        obj_id = lock_id & 0xFFFFFFFF
+        if obj_id > 0x7FFFFFFF:
+            obj_id -= 0x100000000
+
+        async with self._engine.connect() as base_conn:
+            conn = await base_conn.execution_options(isolation_level="AUTOCOMMIT")
+            result = await conn.execute(
+                text(
+                    """
+                    SELECT count(*)
+                    FROM pg_locks
+                    WHERE locktype = 'advisory'
+                      AND granted = true
+                      AND database = (SELECT oid FROM pg_database WHERE datname = current_database())
+                      AND classid = CAST(:class_id AS int4)
+                      AND objid = CAST(:obj_id AS int4)
+                      AND objsubid = 1
+                    """
+                ),
+                {"class_id": class_id, "obj_id": obj_id},
+            )
+            count = result.scalar()
+        return bool(count is not None and count > 0)