altcodepro-polydb-python 2.3.20__py3-none-any.whl → 2.3.22__py3-none-any.whl

polydb/errors.py

@@ -80,3 +80,9 @@ class OperationNotSupportedError(PolyDBError):
     """Raised when an adapter cannot perform a requested operation."""
 
     pass
+
+
+class InsufficientBalanceError(PolyDBError):
+    """Raised when an atomic decrement is attempted but the balance is insufficient."""
+
+    pass

polydb/observability/__init__.py

@@ -0,0 +1,3 @@
+from .logging import configure_logging, set_polydb_log_context
+
+__all__ = ["configure_logging", "set_polydb_log_context"]

polydb/observability/logging.py

@@ -0,0 +1,124 @@
+"""
+Structured JSON logging for PolyDB.
+
+Usage::
+
+    from polydb.observability.logging import configure_logging, set_polydb_log_context
+
+    configure_logging()
+    set_polydb_log_context(tenant_id="acme", model="User", operation="select")
+"""
+
+from __future__ import annotations
+
+import contextvars
+import json
+import logging
+import os
+import time
+from typing import Any, Optional
+
+# ---------------------------------------------------------------------------
+# Context variables
+# ---------------------------------------------------------------------------
+
+_ctx_tenant_id: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
+    "polydb_tenant_id", default=None
+)
+_ctx_model: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
+    "polydb_model", default=None
+)
+_ctx_operation: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
+    "polydb_operation", default=None
+)
+_ctx_duration_ms: contextvars.ContextVar[Optional[float]] = contextvars.ContextVar(
+    "polydb_duration_ms", default=None
+)
+
+
+def set_polydb_log_context(
+    *,
+    tenant_id: Optional[str] = None,
+    model: Optional[str] = None,
+    operation: Optional[str] = None,
+    duration_ms: Optional[float] = None,
+) -> None:
+    """Set per-request context fields that are injected into every log record."""
+    if tenant_id is not None:
+        _ctx_tenant_id.set(tenant_id)
+    if model is not None:
+        _ctx_model.set(model)
+    if operation is not None:
+        _ctx_operation.set(operation)
+    if duration_ms is not None:
+        _ctx_duration_ms.set(duration_ms)
+
+
+# ---------------------------------------------------------------------------
+# JSON Formatter
+# ---------------------------------------------------------------------------
+
+
+class _PolyDBJsonFormatter(logging.Formatter):
+    """Emit each log record as a single-line JSON object."""
+
+    def format(self, record: logging.LogRecord) -> str:  # type: ignore[override]
+        payload: dict[str, Any] = {
+            "timestamp": self.formatTime(record, self.datefmt),
+            "level": record.levelname,
+            "logger": record.name,
+            "message": record.getMessage(),
+        }
+
+        # Inject context vars
+        tenant_id = _ctx_tenant_id.get()
+        if tenant_id is not None:
+            payload["tenant_id"] = tenant_id
+
+        model = _ctx_model.get()
+        if model is not None:
+            payload["model"] = model
+
+        operation = _ctx_operation.get()
+        if operation is not None:
+            payload["operation"] = operation
+
+        duration_ms = _ctx_duration_ms.get()
+        if duration_ms is not None:
+            payload["duration_ms"] = duration_ms
+
+        # Attach any extra fields the caller passed directly to the record
+        for key in ("operation", "table", "duration_ms", "tenant_id", "model"):
+            val = getattr(record, key, None)
+            if val is not None:
+                payload[key] = val
+
+        if record.exc_info:
+            payload["exc_info"] = self.formatException(record.exc_info)
+
+        return json.dumps(payload, default=str)
+
+
+# ---------------------------------------------------------------------------
+# configure_logging
+# ---------------------------------------------------------------------------
+
+
+def configure_logging(level: Optional[str] = None) -> None:
+    """
+    Configure root / polydb logger to emit structured JSON.
+
+    Reads ``POLYDB_LOG_LEVEL`` env var (default ``INFO``).
+    """
+    resolved_level = level or os.getenv("POLYDB_LOG_LEVEL", "INFO").upper()
+    numeric_level = getattr(logging, resolved_level, logging.INFO)
+
+    handler = logging.StreamHandler()
+    handler.setFormatter(_PolyDBJsonFormatter())
+
+    polydb_logger = logging.getLogger("polydb")
+    polydb_logger.setLevel(numeric_level)
+
+    # Replace existing handlers to avoid duplicates
+    polydb_logger.handlers = [handler]
+    polydb_logger.propagate = False

polydb/services/__init__.py

@@ -0,0 +1,9 @@
+from .compliance_service import ComplianceService, AuditEvent
+from .security_service import SecurityService, EncryptionConfigError
+
+__all__ = [
+    "ComplianceService",
+    "AuditEvent",
+    "SecurityService",
+    "EncryptionConfigError",
+]

polydb/services/compliance_service.py

@@ -0,0 +1,141 @@
+"""
+polydb.services.compliance_service
+====================================
+Compliance audit log writer with retry queue and CRITICAL-level alerting.
+
+Fix M14: audit failures were silently swallowed (except Exception: logger.warning).
+This implementation:
+  - Emits a compliance.audit_failure metric counter on first failure
+  - Enqueues a retry in an async queue
+  - After MAX_RETRIES consecutive failures logs at CRITICAL and calls the alert fn
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import time
+from dataclasses import dataclass, field
+from typing import Any, Awaitable, Callable, Dict, Optional
+
+logger = logging.getLogger("polydb.services.compliance")
+
+
+@dataclass
+class AuditEvent:
+    event_type: str
+    tenant_id: str
+    actor_id: str
+    resource_type: str
+    resource_id: str
+    action: str
+    metadata: Dict[str, Any] = field(default_factory=dict)
+    attempt: int = 0
+    created_at: float = field(default_factory=time.time)
+
+
+WriterFn = Callable[[AuditEvent], Awaitable[None]]
+MetricsFn = Callable[[str, Dict[str, str]], None]
+AlertFn = Callable[[AuditEvent], Awaitable[None]]
+
+
+class ComplianceService:
+    """Audit log writer with async retry queue and escalation."""
+
+    MAX_RETRIES = 3
+    RETRY_DELAY_S = 2.0
+
+    def __init__(
+        self,
+        *,
+        metrics_fn: Optional[MetricsFn] = None,
+        alert_fn: Optional[AlertFn] = None,
+    ) -> None:
+        self._metrics = metrics_fn
+        self._alert = alert_fn
+        self._retry_queue: asyncio.Queue[AuditEvent] = asyncio.Queue(maxsize=10_000)
+        self._running = False
+        self._worker_task: Optional[asyncio.Task] = None
+
+    async def log(self, event: AuditEvent, *, writer: WriterFn) -> None:
+        """Write audit event. On failure, enqueue for retry."""
+        try:
+            await writer(event)
+        except Exception as exc:
+            self._record_failure(event, exc, step="initial")
+            event.attempt = 1
+            try:
+                self._retry_queue.put_nowait(event)
+            except asyncio.QueueFull:
+                logger.critical(
+                    "compliance.audit_failure retry queue full — event dropped: "
+                    "event_type=%s tenant_id=%s",
+                    event.event_type, event.tenant_id,
+                )
+
+    async def start(self, *, writer: WriterFn) -> None:
+        """Start the background retry worker. Call once at service startup."""
+        self._running = True
+        self._worker_task = asyncio.create_task(
+            self._drain_retry_queue(writer=writer),
+            name="compliance_retry_worker",
+        )
+
+    def stop(self) -> None:
+        self._running = False
+        if self._worker_task:
+            self._worker_task.cancel()
+
+    async def _drain_retry_queue(self, *, writer: WriterFn) -> None:
+        while self._running:
+            try:
+                event = await asyncio.wait_for(self._retry_queue.get(), timeout=5.0)
+            except asyncio.TimeoutError:
+                continue
+
+            if event.attempt >= self.MAX_RETRIES:
+                logger.critical(
+                    "compliance.audit_failure UNRECOVERABLE after %d retries: "
+                    "event_type=%s tenant_id=%s resource=%s/%s",
+                    event.attempt,
+                    event.event_type,
+                    event.tenant_id,
+                    event.resource_type,
+                    event.resource_id,
+                )
+                if self._alert:
+                    try:
+                        await self._alert(event)
+                    except Exception as alert_exc:
+                        logger.error("Alert function failed: %s", alert_exc)
+                continue
+
+            await asyncio.sleep(self.RETRY_DELAY_S)
+            try:
+                await writer(event)
+            except Exception as exc:
+                self._record_failure(event, exc, step=f"retry_{event.attempt}")
+                event.attempt += 1
+                try:
+                    self._retry_queue.put_nowait(event)
+                except asyncio.QueueFull:
+                    logger.critical(
+                        "compliance.audit_failure retry queue full on re-enqueue — event lost: "
+                        "event_type=%s tenant_id=%s",
+                        event.event_type, event.tenant_id,
+                    )
+
+    def _record_failure(self, event: AuditEvent, exc: Exception, step: str) -> None:
+        logger.error(
+            "compliance.audit_failure step=%s event_type=%s tenant_id=%s resource=%s/%s error=%s",
+            step, event.event_type, event.tenant_id,
+            event.resource_type, event.resource_id, exc,
+        )
+        if self._metrics:
+            try:
+                self._metrics(
+                    "compliance.audit_failure",
+                    {"event_type": event.event_type, "tenant_id": event.tenant_id, "step": step},
+                )
+            except Exception:
+                pass

polydb/services/security_service.py

@@ -0,0 +1,133 @@
+"""
+polydb.services.security_service
+==================================
+Field-level encryption for sensitive model fields.
+
+Fix M15: Missing POLYDB_ENCRYPTION_KEY previously caused encrypted fields to
+be stored in plaintext with only a warning. Now raises EncryptionConfigError
+at startup if any registered model has encrypted=true fields and the key is
+not set.
+"""
+
+from __future__ import annotations
+
+import base64
+import logging
+import os
+from typing import Any, Dict, List, Optional
+
+logger = logging.getLogger("polydb.services.security")
+
+
+class EncryptionConfigError(RuntimeError):
+    """Raised on startup when encryption is required but not configured."""
+
+
+class SecurityService:
+    """Field-level encryption.
+
+    Raises EncryptionConfigError at startup if any registered model has
+    encrypted=true fields and POLYDB_ENCRYPTION_KEY is not set.
+    """
+
+    ENV_KEY = "POLYDB_ENCRYPTION_KEY"
+
+    def __init__(
+        self,
+        *,
+        encryption_key: Optional[bytes] = None,
+    ) -> None:
+        self._key: Optional[bytes] = encryption_key or self._read_env_key()
+
+    # ------------------------------------------------------------------
+    # Startup validation
+    # ------------------------------------------------------------------
+
+    def validate_startup(self, registered_models: List[Dict[str, Any]]) -> None:
+        """Call during application startup after models are registered.
+
+        Raises EncryptionConfigError if any model declares encrypted fields
+        and the encryption key is absent.  Fail-fast is intentional: running
+        without encryption when it is configured is a silent data-security
+        violation.
+        """
+        if self._key:
+            return
+
+        encrypted_models: List[str] = []
+        for m in registered_models:
+            fields = m.get("fields") or {}
+            if isinstance(fields, dict):
+                for _fname, fdef in fields.items():
+                    if isinstance(fdef, dict) and fdef.get("encrypted"):
+                        encrypted_models.append(m.get("name", "<unknown>"))
+                        break
+
+        if encrypted_models:
+            raise EncryptionConfigError(
+                f"{self.ENV_KEY} is not set but the following models have "
+                f"encrypted=true fields: {', '.join(encrypted_models)}. "
+                f"Set {self.ENV_KEY} to a base64url-encoded 32-byte Fernet key "
+                f"or remove the encrypted=true flag from those model fields."
+            )
+
+    # ------------------------------------------------------------------
+    # Encrypt / Decrypt
+    # ------------------------------------------------------------------
+
+    def encrypt(self, value: str) -> str:
+        """Encrypt *value*. Returns Fernet token (URL-safe base64 string)."""
+        if not self._key:
+            raise EncryptionConfigError(
+                f"Cannot encrypt: {self.ENV_KEY} is not set."
+            )
+        try:
+            from cryptography.fernet import Fernet
+            return Fernet(self._key).encrypt(value.encode("utf-8")).decode("ascii")
+        except ImportError:
+            raise EncryptionConfigError(
+                "cryptography package is required for field encryption. "
+                "Install it: pip install cryptography"
+            )
+
+    def decrypt(self, value: str) -> str:
+        """Decrypt a Fernet token produced by :meth:`encrypt`."""
+        if not self._key:
+            raise EncryptionConfigError(
+                f"Cannot decrypt: {self.ENV_KEY} is not set."
+            )
+        try:
+            from cryptography.fernet import Fernet
+            return Fernet(self._key).decrypt(value.encode("ascii")).decode("utf-8")
+        except ImportError:
+            raise EncryptionConfigError(
+                "cryptography package is required for field encryption."
+            )
+        except Exception as exc:
+            raise ValueError(f"Decryption failed: {exc}") from exc
+
+    # ------------------------------------------------------------------
+    # Internal
+    # ------------------------------------------------------------------
+
+    @classmethod
+    def _read_env_key(cls) -> Optional[bytes]:
+        raw = os.environ.get(cls.ENV_KEY, "").strip()
+        if not raw:
+            return None
+        try:
+            # Fernet requires a 32-byte URL-safe base64 key (44 chars with padding)
+            key_bytes = raw.encode("ascii")
+            decoded = base64.urlsafe_b64decode(key_bytes + b"=" * (-len(raw) % 4))
+            if len(decoded) != 32:
+                logger.error(
+                    "%s decoded to %d bytes; Fernet requires exactly 32 bytes. "
+                    "Generate a valid key with: python -c \""
+                    "from cryptography.fernet import Fernet; print(Fernet.generate_key().decode())\"",
+                    cls.ENV_KEY, len(decoded),
+                )
+                return None
+            return key_bytes
+        except Exception as exc:
+            logger.error("Failed to decode %s: %s", cls.ENV_KEY, exc)
+            return None