omnibase_infra 0.2.1__py3-none-any.whl

omnibase_infra/handlers/models/vault/registry_payload_vault.py

@@ -0,0 +1,213 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2025 OmniNode Team
+"""Registry for Vault payload types with decorator-based registration.
+
+This module provides:
+- RegistryPayloadVault: Decorator-based registry for payload type discovery
+
+Design Pattern:
+    Instead of maintaining explicit union types like:
+        VaultPayload = ModelVaultSecretPayload | ModelVaultWritePayload | ...
+
+    Payload models self-register via decorator:
+        @RegistryPayloadVault.register("read_secret")
+        class ModelVaultSecretPayload(ModelPayloadVault): ...
+
+    The registry resolves types dynamically during Pydantic validation,
+    enabling new payload types to be added without modifying existing code.
+
+Related:
+    - ModelPayloadVault: Base model for Vault payloads
+    - EnumVaultOperationType: Operation type enum for discriminator
+    - OMN-1007: Union reduction refactoring
+"""
+
+from __future__ import annotations
+
+import threading
+import warnings
+from collections.abc import Callable
+from typing import TYPE_CHECKING, ClassVar, TypeVar
+
+from omnibase_infra.enums import EnumInfraTransportType
+from omnibase_infra.errors import ModelInfraErrorContext, ProtocolConfigurationError
+
+if TYPE_CHECKING:
+    from omnibase_infra.handlers.models.vault.model_payload_vault import (
+        ModelPayloadVault,
+    )
+
+# TypeVar bound to ModelPayloadVault to preserve subclass type through decorator
+_VaultPayloadT = TypeVar("_VaultPayloadT", bound="ModelPayloadVault")
+
+
+class RegistryPayloadVault:
+    """Registry for Vault payload model types with decorator-based registration.
+
+    This registry enables dynamic type resolution for Pydantic validation
+    without requiring explicit union type definitions.
+
+    Example:
+        @RegistryPayloadVault.register("read_secret")
+        class ModelVaultSecretPayload(ModelPayloadVault):
+            operation_type: Literal[EnumVaultOperationType.READ_SECRET]
+            ...
+
+        # Later, resolve type from operation_type
+        payload_cls = RegistryPayloadVault.get_type("read_secret")
+        payload = payload_cls.model_validate(data)
+
+    Thread Safety:
+        All registration and mutation operations are protected by a class-level
+        threading.Lock to ensure thread-safe access. Read operations (get_type,
+        get_all_types, is_registered) are atomic dictionary operations and safe
+        for concurrent access. The lock ensures that check-and-set operations
+        in register() are atomic, preventing race conditions.
+    """
+
+    _types: ClassVar[dict[str, type[ModelPayloadVault]]] = {}
+    _lock: ClassVar[threading.Lock] = threading.Lock()
+
+    @classmethod
+    def register(
+        cls, operation_type: str
+    ) -> Callable[[type[_VaultPayloadT]], type[_VaultPayloadT]]:
+        """Decorator to register a Vault payload model type.
+
+        Args:
+            operation_type: The Vault operation type identifier (e.g., "read_secret").
+                           Must match the model's `operation_type` field value.
+
+        Returns:
+            Decorator function that registers the class and returns it unchanged.
+
+        Raises:
+            ValueError: If operation_type is already registered (prevents duplicates).
+            TypeError: If payload_cls is not a subclass of ModelPayloadVault.
+
+        Example:
+            @RegistryPayloadVault.register("read_secret")
+            class ModelVaultSecretPayload(ModelPayloadVault):
+                operation_type: Literal[EnumVaultOperationType.READ_SECRET]
+                path: str
+                ...
+        """
+
+        def decorator(
+            payload_cls: type[_VaultPayloadT],
+        ) -> type[_VaultPayloadT]:
+            # Runtime type validation (outside lock - class inspection is read-only)
+            from omnibase_infra.handlers.models.vault.model_payload_vault import (
+                ModelPayloadVault,
+            )
+
+            if not issubclass(payload_cls, ModelPayloadVault):
+                context = ModelInfraErrorContext(
+                    transport_type=EnumInfraTransportType.VAULT,
+                    operation="register_payload_type",
+                )
+                raise ProtocolConfigurationError(
+                    f"Registered class {payload_cls.__name__!r} must be a subclass "
+                    f"of ModelPayloadVault, got {payload_cls.__mro__}",
+                    context=context,
+                )
+
+            # Thread-safe registration with atomic check-and-set
+            with cls._lock:
+                if operation_type in cls._types:
+                    context = ModelInfraErrorContext(
+                        transport_type=EnumInfraTransportType.VAULT,
+                        operation="register_payload_type",
+                    )
+                    raise ProtocolConfigurationError(
+                        f"Vault payload operation_type '{operation_type}' already registered "
+                        f"to {cls._types[operation_type].__name__}. "
+                        f"Cannot register {payload_cls.__name__}.",
+                        context=context,
+                    )
+                cls._types[operation_type] = payload_cls
+            return payload_cls
+
+        return decorator
+
+    @classmethod
+    def get_type(cls, operation_type: str) -> type[ModelPayloadVault]:
+        """Get the payload model class for a given operation type.
+
+        Args:
+            operation_type: The Vault operation type identifier.
+
+        Returns:
+            The registered payload model class.
+
+        Raises:
+            KeyError: If operation_type is not registered.
+
+        Example:
+            payload_cls = RegistryPayloadVault.get_type("read_secret")
+            payload = payload_cls.model_validate({"operation_type": "read_secret", ...})
+        """
+        if operation_type not in cls._types:
+            registered = ", ".join(sorted(cls._types.keys())) or "(none)"
+            raise KeyError(
+                f"Unknown Vault operation_type '{operation_type}'. "
+                f"Registered types: {registered}"
+            )
+        return cls._types[operation_type]
+
+    @classmethod
+    def get_all_types(cls) -> dict[str, type[ModelPayloadVault]]:
+        """Get all registered Vault payload types.
+
+        Returns:
+            Dict mapping operation_type strings to payload model classes.
+
+        Example:
+            all_types = RegistryPayloadVault.get_all_types()
+            for op_type, payload_cls in all_types.items():
+                print(f"{op_type}: {payload_cls.__name__}")
+        """
+        return dict(cls._types)
+
+    @classmethod
+    def is_registered(cls, operation_type: str) -> bool:
+        """Check if an operation type is registered.
+
+        Args:
+            operation_type: The Vault operation type identifier.
+
+        Returns:
+            True if registered, False otherwise.
+
+        Example:
+            if RegistryPayloadVault.is_registered("read_secret"):
+                payload_cls = RegistryPayloadVault.get_type("read_secret")
+        """
+        return operation_type in cls._types
+
+    @classmethod
+    def clear(cls) -> None:
+        """Clear all registered types.
+
+        Warning:
+            This method is intended for **testing purposes only**.
+            Calling it in production code will emit a warning.
+            It breaks the immutability guarantee after startup.
+
+        Thread Safety:
+            This method is protected by the class-level lock to ensure
+            thread-safe clearing of the registry.
+        """
+        warnings.warn(
+            "RegistryPayloadVault.clear() is intended for testing only. "
+            "Do not use in production code.",
+            UserWarning,
+            stacklevel=2,
+        )
+        with cls._lock:
+            cls._types.clear()
+
+
+__all__ = [
+    "RegistryPayloadVault",  # Vault payload registry
+]

omnibase_infra/handlers/registration_storage/__init__.py

@@ -0,0 +1,43 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2025 OmniNode Team
+"""Registration Storage Handlers Module.
+
+This module provides pluggable handler implementations for registration
+storage operations, supporting the capability-oriented node architecture.
+
+Handlers:
+    - HandlerRegistrationStoragePostgres: PostgreSQL-backed storage
+    - HandlerRegistrationStorageMock: In-memory mock for testing
+
+Models (from omnibase_infra.nodes.node_registration_storage_effect.models):
+    - ModelRegistrationRecord: Registration record model
+    - ModelUpsertResult: Upsert operation result
+    - ModelStorageResult: Storage query result
+
+Protocols:
+    - ProtocolRegistrationPersistence: Registration persistence protocol definition
+"""
+
+from omnibase_infra.handlers.registration_storage.handler_registration_storage_mock import (
+    HandlerRegistrationStorageMock,
+)
+from omnibase_infra.handlers.registration_storage.handler_registration_storage_postgres import (
+    HandlerRegistrationStoragePostgres,
+)
+from omnibase_infra.handlers.registration_storage.protocol_registration_persistence import (
+    ProtocolRegistrationPersistence,
+)
+from omnibase_infra.nodes.node_registration_storage_effect.models import (
+    ModelRegistrationRecord,
+    ModelStorageResult,
+    ModelUpsertResult,
+)
+
+__all__: list[str] = [
+    "HandlerRegistrationStorageMock",
+    "HandlerRegistrationStoragePostgres",
+    "ModelRegistrationRecord",
+    "ModelStorageResult",
+    "ModelUpsertResult",
+    "ProtocolRegistrationPersistence",
+]

omnibase_infra/handlers/registration_storage/handler_registration_storage_mock.py

@@ -0,0 +1,392 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2025 OmniNode Team
+"""Mock Registration Storage Handler.
+
+This module provides an in-memory mock implementation of the registration
+storage handler protocol for testing purposes.
+
+Thread Safety:
+    This handler uses asyncio.Lock for coroutine-safe access to the
+    in-memory record store.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import time
+from datetime import UTC, datetime
+from uuid import UUID, uuid4
+
+from omnibase_infra.handlers.registration_storage.models import (
+    ModelDeleteRegistrationRequest,
+    ModelUpdateRegistrationRequest,
+)
+from omnibase_infra.nodes.node_registration_storage_effect.models import (
+    ModelDeleteResult,
+    ModelRegistrationRecord,
+    ModelStorageHealthCheckDetails,
+    ModelStorageHealthCheckResult,
+    ModelStorageQuery,
+    ModelStorageResult,
+    ModelUpsertResult,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class HandlerRegistrationStorageMock:
+    """In-memory mock for testing registration storage.
+
+    Provides a simple in-memory implementation of the registration storage
+    protocol for unit and integration testing.
+
+    Thread Safety:
+        This handler is coroutine-safe. All operations on the internal
+        record store are protected by asyncio.Lock.
+
+    Attributes:
+        handler_type: Returns "mock" identifier.
+
+    Example:
+        >>> handler = HandlerRegistrationStorageMock()
+        >>> result = await handler.store_registration(record)
+        >>> assert result.success
+    """
+
+    def __init__(self) -> None:
+        """Initialize HandlerRegistrationStorageMock with empty record store."""
+        self._records: dict[UUID, ModelRegistrationRecord] = {}
+        self._lock = asyncio.Lock()
+
+        logger.debug("HandlerRegistrationStorageMock initialized")
+
+    @property
+    def handler_type(self) -> str:
+        """Return the handler type identifier.
+
+        Returns:
+            "mock" identifier string.
+        """
+        return "mock"
+
+    async def store_registration(
+        self,
+        record: ModelRegistrationRecord,
+        correlation_id: UUID | None = None,
+    ) -> ModelUpsertResult:
+        """Store a registration record in the mock store.
+
+        Args:
+            record: Registration record to store.
+            correlation_id: Optional correlation ID for tracing.
+
+        Returns:
+            ModelUpsertResult with upsert outcome.
+        """
+        correlation_id = correlation_id or uuid4()
+        start_time = time.monotonic()
+
+        async with self._lock:
+            was_insert = record.node_id not in self._records
+            now = datetime.now(UTC)
+
+            # Update timestamps
+            updated_record = ModelRegistrationRecord(
+                node_id=record.node_id,
+                node_type=record.node_type,
+                node_version=record.node_version,
+                capabilities=record.capabilities,
+                endpoints=record.endpoints,
+                metadata=record.metadata,
+                created_at=record.created_at or now,
+                updated_at=now,
+                correlation_id=correlation_id,
+            )
+
+            self._records[record.node_id] = updated_record
+
+        duration_ms = (time.monotonic() - start_time) * 1000
+
+        logger.debug(
+            "Mock registration stored",
+            extra={
+                "node_id": str(record.node_id),
+                "was_insert": was_insert,
+                "correlation_id": str(correlation_id),
+            },
+        )
+
+        return ModelUpsertResult(
+            success=True,
+            node_id=record.node_id,
+            operation="insert" if was_insert else "update",
+            duration_ms=duration_ms,
+            backend_type=self.handler_type,
+            correlation_id=correlation_id,
+        )
+
+    async def query_registrations(
+        self,
+        query: ModelStorageQuery,
+        correlation_id: UUID | None = None,
+    ) -> ModelStorageResult:
+        """Query registration records from the mock store.
+
+        Args:
+            query: Query parameters including filters and pagination.
+            correlation_id: Optional correlation ID for tracing.
+
+        Returns:
+            ModelStorageResult with list of matching records.
+        """
+        correlation_id = correlation_id or uuid4()
+        start_time = time.monotonic()
+
+        async with self._lock:
+            matching_records: list[ModelRegistrationRecord] = []
+
+            for record in self._records.values():
+                # Filter by node_id (exact match)
+                if query.node_id is not None and record.node_id != query.node_id:
+                    continue
+
+                # Filter by node type
+                if query.node_type is not None and record.node_type != query.node_type:
+                    continue
+
+                # Filter by capability (contains match)
+                if query.capability_filter is not None:
+                    if query.capability_filter not in record.capabilities:
+                        continue
+
+                matching_records.append(record)
+
+            # Sort by updated_at descending
+            matching_records.sort(
+                key=lambda r: r.updated_at or datetime.min.replace(tzinfo=UTC),
+                reverse=True,
+            )
+
+            total_count = len(matching_records)
+
+            # Apply pagination
+            paginated = matching_records[query.offset : query.offset + query.limit]
+
+        duration_ms = (time.monotonic() - start_time) * 1000
+
+        logger.debug(
+            "Mock registration query completed",
+            extra={
+                "record_count": len(paginated),
+                "total_count": total_count,
+                "correlation_id": str(correlation_id),
+            },
+        )
+
+        return ModelStorageResult(
+            success=True,
+            records=tuple(paginated),
+            total_count=total_count,
+            duration_ms=duration_ms,
+            backend_type=self.handler_type,
+            correlation_id=correlation_id,
+        )
+
+    async def update_registration(
+        self,
+        request: ModelUpdateRegistrationRequest,
+    ) -> ModelUpsertResult:
+        """Update an existing registration record in the mock store.
+
+        Args:
+            request: ModelUpdateRegistrationRequest containing:
+                - node_id: ID of the node to update
+                - updates: ModelRegistrationUpdate with fields to update
+                - correlation_id: Optional correlation ID for tracing
+
+        Returns:
+            ModelUpsertResult with update outcome.
+        """
+        # Extract fields from request model
+        node_id = request.node_id
+        updates = request.updates
+        correlation_id = request.correlation_id or uuid4()
+        start_time = time.monotonic()
+
+        async with self._lock:
+            if node_id not in self._records:
+                duration_ms = (time.monotonic() - start_time) * 1000
+                return ModelUpsertResult(
+                    success=False,
+                    node_id=node_id,
+                    operation="update",
+                    error="Record not found",
+                    duration_ms=duration_ms,
+                    backend_type=self.handler_type,
+                    correlation_id=correlation_id,
+                )
+
+            existing = self._records[node_id]
+            now = datetime.now(UTC)
+
+            # Update with new values from ModelRegistrationUpdate
+            updated_record = ModelRegistrationRecord(
+                node_id=existing.node_id,
+                node_type=existing.node_type,
+                node_version=(
+                    updates.node_version
+                    if updates.node_version is not None
+                    else existing.node_version
+                ),
+                capabilities=(
+                    updates.capabilities
+                    if updates.capabilities is not None
+                    else existing.capabilities
+                ),
+                endpoints=(
+                    updates.endpoints
+                    if updates.endpoints is not None
+                    else existing.endpoints
+                ),
+                metadata=(
+                    updates.metadata
+                    if updates.metadata is not None
+                    else existing.metadata
+                ),
+                created_at=existing.created_at,
+                updated_at=now,
+                correlation_id=correlation_id,
+            )
+
+            self._records[node_id] = updated_record
+
+        duration_ms = (time.monotonic() - start_time) * 1000
+
+        logger.debug(
+            "Mock registration updated",
+            extra={
+                "node_id": str(node_id),
+                "correlation_id": str(correlation_id),
+            },
+        )
+
+        return ModelUpsertResult(
+            success=True,
+            node_id=node_id,
+            operation="update",
+            duration_ms=duration_ms,
+            backend_type=self.handler_type,
+            correlation_id=correlation_id,
+        )
+
+    async def delete_registration(
+        self,
+        request: ModelDeleteRegistrationRequest,
+    ) -> ModelDeleteResult:
+        """Delete a registration record from the mock store.
+
+        Args:
+            request: ModelDeleteRegistrationRequest containing:
+                - node_id: ID of the node to delete
+                - correlation_id: Optional correlation ID for tracing
+
+        Returns:
+            ModelDeleteResult with deletion outcome.
+        """
+        # Extract fields from request model
+        node_id = request.node_id
+        correlation_id = request.correlation_id or uuid4()
+        start_time = time.monotonic()
+
+        async with self._lock:
+            if node_id in self._records:
+                del self._records[node_id]
+                deleted = True
+            else:
+                deleted = False
+
+        duration_ms = (time.monotonic() - start_time) * 1000
+
+        logger.debug(
+            "Mock registration deletion completed",
+            extra={
+                "node_id": str(node_id),
+                "deleted": deleted,
+                "correlation_id": str(correlation_id),
+            },
+        )
+
+        return ModelDeleteResult(
+            success=True,
+            node_id=node_id,
+            deleted=deleted,
+            duration_ms=duration_ms,
+            backend_type=self.handler_type,
+            correlation_id=correlation_id,
+        )
+
+    async def health_check(
+        self,
+        correlation_id: UUID | None = None,
+    ) -> ModelStorageHealthCheckResult:
+        """Perform a health check on the mock handler.
+
+        Args:
+            correlation_id: Optional correlation ID for tracing.
+
+        Returns:
+            ModelStorageHealthCheckResult with health status information.
+        """
+        correlation_id = correlation_id or uuid4()
+        start_time = time.monotonic()
+
+        async with self._lock:
+            pass  # Lock is acquired here; duration measures acquisition latency
+
+        latency_ms = (time.monotonic() - start_time) * 1000
+
+        return ModelStorageHealthCheckResult(
+            healthy=True,
+            backend_type=self.handler_type,
+            latency_ms=latency_ms,
+            reason="ok",
+            details=ModelStorageHealthCheckDetails(
+                server_version="mock-1.0.0",
+            ),
+            correlation_id=correlation_id,
+        )
+
+    async def clear(self) -> None:
+        """Clear all records from the mock store.
+
+        Utility method for test cleanup.
+        """
+        async with self._lock:
+            self._records.clear()
+
+        logger.debug("Mock registration store cleared")
+
+    async def get_record_count(self) -> int:
+        """Get the number of stored records.
+
+        Returns:
+            Number of records in the store.
+        """
+        async with self._lock:
+            return len(self._records)
+
+    async def get_record(self, node_id: UUID) -> ModelRegistrationRecord | None:
+        """Get a specific record by node ID.
+
+        Args:
+            node_id: ID of the node to retrieve.
+
+        Returns:
+            The record if found, None otherwise.
+        """
+        async with self._lock:
+            return self._records.get(node_id)
+
+
+__all__ = ["HandlerRegistrationStorageMock"]