plain.postgres 0.84.0__py3-none-any.whl

plain/postgres/otel.py

@@ -0,0 +1,231 @@
+from __future__ import annotations
+
+import re
+import traceback
+from collections.abc import Generator
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any
+
+from opentelemetry import context as otel_context
+from opentelemetry import trace
+
+if TYPE_CHECKING:
+    from opentelemetry.trace import Span
+
+    from plain.postgres.connection import DatabaseConnection
+
+from opentelemetry.semconv._incubating.attributes.db_attributes import (
+    DB_QUERY_PARAMETER_TEMPLATE,
+    DB_USER,
+)
+from opentelemetry.semconv.attributes.code_attributes import (
+    CODE_COLUMN_NUMBER,
+    CODE_FILE_PATH,
+    CODE_FUNCTION_NAME,
+    CODE_LINE_NUMBER,
+    CODE_STACKTRACE,
+)
+from opentelemetry.semconv.attributes.db_attributes import (
+    DB_COLLECTION_NAME,
+    DB_NAMESPACE,
+    DB_OPERATION_NAME,
+    DB_QUERY_SUMMARY,
+    DB_QUERY_TEXT,
+    DB_SYSTEM_NAME,
+)
+from opentelemetry.semconv.attributes.network_attributes import (
+    NETWORK_PEER_ADDRESS,
+    NETWORK_PEER_PORT,
+)
+from opentelemetry.semconv.trace import DbSystemValues
+from opentelemetry.trace import SpanKind
+
+from plain.runtime import settings
+
+# Use a stable string key so OpenTelemetry context APIs receive the expected type.
+_SUPPRESS_KEY = "plain.postgres.suppress_db_tracing"
+
+tracer = trace.get_tracer("plain.postgres")
+
+
+DB_SYSTEM = DbSystemValues.POSTGRESQL.value
+
+
+def extract_operation_and_target(sql: str) -> tuple[str, str | None, str | None]:
+    """Extract operation, table name, and collection from SQL.
+
+    Returns: (operation, summary, collection_name)
+    """
+    sql_upper = sql.upper().strip()
+
+    # Strip leading parentheses (e.g. UNION queries: "(SELECT ... UNION ...)")
+    operation = sql_upper.lstrip("(").split()[0] if sql_upper else "UNKNOWN"
+
+    # Pattern to match quoted and unquoted identifiers
+    # Matches: "quoted" (PostgreSQL), unquoted.name
+    identifier_pattern = r'("([^"]+)"|([\w.]+))'
+
+    # Map operations to the SQL keyword that precedes the table name.
+    keyword_by_operation = {
+        "SELECT": "FROM",
+        "DELETE": "FROM",
+        "INSERT": "INTO",
+        "UPDATE": "UPDATE",
+    }
+
+    # Extract table/collection name based on operation
+    collection_name = None
+    summary = operation
+
+    keyword = keyword_by_operation.get(operation)
+    if keyword:
+        match = re.search(rf"{keyword}\s+{identifier_pattern}", sql, re.IGNORECASE)
+        if match:
+            collection_name = _clean_identifier(match.group(1))
+            summary = f"{operation} {collection_name}"
+
+    # Detect UNION queries
+    if " UNION " in sql_upper and summary:
+        summary = f"{summary} UNION"
+
+    return operation, summary, collection_name
+
+
+def _clean_identifier(identifier: str) -> str:
+    """Remove quotes from SQL identifiers."""
+    if identifier.startswith('"') and identifier.endswith('"'):
+        return identifier[1:-1]
+    return identifier
+
+
+@contextmanager
+def db_span(
+    db: DatabaseConnection, sql: Any, *, many: bool = False, params: Any = None
+) -> Generator[Span | None]:
+    """Open an OpenTelemetry CLIENT span for a database query.
+
+    All common attributes (`db.*`, `network.*`, etc.) are set automatically.
+    Follows OpenTelemetry semantic conventions for database instrumentation.
+    """
+
+    # Fast-exit if instrumentation suppression flag set in context.
+    if otel_context.get_value(_SUPPRESS_KEY):
+        yield None
+        return
+
+    sql = str(sql)  # Ensure SQL is a string for span attributes.
+
+    # Extract operation and target information
+    operation, summary, collection_name = extract_operation_and_target(sql)
+
+    if many:
+        summary = f"{summary} many"
+
+    # Span name follows semantic conventions: {target} or {db.operation.name} {target}
+    if summary:
+        span_name = summary[:255]
+    else:
+        span_name = operation
+
+    # Build attribute set following semantic conventions
+    attrs: dict[str, Any] = {
+        DB_SYSTEM_NAME: DB_SYSTEM,
+        DB_NAMESPACE: db.settings_dict.get("DATABASE"),
+        DB_QUERY_TEXT: sql,  # Already parameterized from Django/Plain
+        DB_QUERY_SUMMARY: summary,
+        DB_OPERATION_NAME: operation,
+    }
+
+    attrs.update(_get_code_attributes())
+
+    # Add collection name if detected
+    if collection_name:
+        attrs[DB_COLLECTION_NAME] = collection_name
+
+    # Add user attribute
+    if user := db.settings_dict.get("USER"):
+        attrs[DB_USER] = user
+
+    # Network attributes
+    if host := db.settings_dict.get("HOST"):
+        attrs[NETWORK_PEER_ADDRESS] = host
+
+    if port := db.settings_dict.get("PORT"):
+        try:
+            attrs[NETWORK_PEER_PORT] = int(port)
+        except (TypeError, ValueError):
+            pass
+
+    # Add query parameters as attributes when DEBUG is True
+    if settings.DEBUG and params is not None:
+        # Convert params to appropriate format based on type
+        if isinstance(params, dict):
+            # Dictionary params (e.g., for named placeholders)
+            for key, value in params.items():
+                attrs[f"{DB_QUERY_PARAMETER_TEMPLATE}.{key}"] = str(value)
+        elif isinstance(params, list | tuple):
+            # Sequential params (e.g., for %s or ? placeholders)
+            for i, value in enumerate(params):
+                attrs[f"{DB_QUERY_PARAMETER_TEMPLATE}.{i + 1}"] = str(value)
+        else:
+            # Single param (rare but possible)
+            attrs[f"{DB_QUERY_PARAMETER_TEMPLATE}.1"] = str(params)
+
+    with tracer.start_as_current_span(
+        span_name, kind=SpanKind.CLIENT, attributes=attrs
+    ) as span:
+        yield span
+        span.set_status(trace.StatusCode.OK)
+
+
+@contextmanager
+def suppress_db_tracing() -> Generator[None]:
+    token = otel_context.attach(otel_context.set_value(_SUPPRESS_KEY, True))
+    try:
+        yield
+    finally:
+        otel_context.detach(token)
+
+
+def _is_internal_frame(frame: traceback.FrameSummary) -> bool:
+    """Return True if the frame is internal to plain.postgres or contextlib."""
+    filepath = frame.filename
+    if not filepath:
+        return True
+    if "/plain/postgres/" in filepath:
+        return True
+    if filepath.endswith("contextlib.py"):
+        return True
+    return False
+
+
+def _get_code_attributes() -> dict[str, Any]:
+    """Extract code context attributes for the current database query.
+
+    Returns a dict of OpenTelemetry code attributes.
+    """
+    stack = traceback.extract_stack()
+
+    # Find the first user code frame (outermost non-internal frame from the top of the call stack)
+    for frame in reversed(stack):
+        if _is_internal_frame(frame):
+            continue
+
+        attrs: dict[str, Any] = {
+            CODE_FILE_PATH: frame.filename,
+        }
+        if frame.lineno:
+            attrs[CODE_LINE_NUMBER] = frame.lineno
+        if frame.name:
+            attrs[CODE_FUNCTION_NAME] = frame.name
+        if frame.colno:
+            attrs[CODE_COLUMN_NUMBER] = frame.colno
+
+        # Add full stack trace only in DEBUG mode (expensive)
+        if settings.DEBUG:
+            filtered_stack = [f for f in stack if not _is_internal_frame(f)]
+            attrs[CODE_STACKTRACE] = "".join(traceback.format_list(filtered_stack))
+
+        return attrs
+
+    return {}

plain/postgres/preflight.py

@@ -0,0 +1,336 @@
+from __future__ import annotations
+
+import inspect
+from collections import defaultdict
+from collections.abc import Callable
+from typing import Any
+
+from plain.packages import packages_registry
+from plain.postgres.db import get_connection
+from plain.postgres.migrations.recorder import MIGRATION_TABLE_NAME
+from plain.postgres.registry import ModelsRegistry, models_registry
+from plain.preflight import PreflightCheck, PreflightResult, register_check
+
+
+@register_check("postgres.all_models")
+class CheckAllModels(PreflightCheck):
+    """Validates all model definitions for common issues."""
+
+    def run(self) -> list[PreflightResult]:
+        db_table_models = defaultdict(list)
+        indexes = defaultdict(list)
+        constraints = defaultdict(list)
+        errors = []
+        models = models_registry.get_models()
+        for model in models:
+            db_table_models[model.model_options.db_table].append(
+                model.model_options.label
+            )
+            if not inspect.ismethod(model.preflight):
+                errors.append(
+                    PreflightResult(
+                        fix=f"The '{model.__name__}.preflight()' class method is currently overridden by {model.preflight!r}.",
+                        obj=model,
+                        id="postgres.preflight_method_overridden",
+                    )
+                )
+            else:
+                errors.extend(model.preflight())
+            for model_index in model.model_options.indexes:
+                indexes[model_index.name].append(model.model_options.label)
+            for model_constraint in model.model_options.constraints:
+                constraints[model_constraint.name].append(model.model_options.label)
+        for db_table, model_labels in db_table_models.items():
+            if len(model_labels) != 1:
+                model_labels_str = ", ".join(model_labels)
+                errors.append(
+                    PreflightResult(
+                        fix=f"db_table '{db_table}' is used by multiple models: {model_labels_str}.",
+                        obj=db_table,
+                        id="postgres.duplicate_db_table",
+                    )
+                )
+        for index_name, model_labels in indexes.items():
+            if len(model_labels) > 1:
+                model_labels = set(model_labels)
+                errors.append(
+                    PreflightResult(
+                        fix="index name '{}' is not unique {} {}.".format(
+                            index_name,
+                            "for model" if len(model_labels) == 1 else "among models:",
+                            ", ".join(sorted(model_labels)),
+                        ),
+                        id="postgres.index_name_not_unique_single"
+                        if len(model_labels) == 1
+                        else "postgres.index_name_not_unique_multiple",
+                    ),
+                )
+        for constraint_name, model_labels in constraints.items():
+            if len(model_labels) > 1:
+                model_labels = set(model_labels)
+                errors.append(
+                    PreflightResult(
+                        fix="constraint name '{}' is not unique {} {}.".format(
+                            constraint_name,
+                            "for model" if len(model_labels) == 1 else "among models:",
+                            ", ".join(sorted(model_labels)),
+                        ),
+                        id="postgres.constraint_name_not_unique_single"
+                        if len(model_labels) == 1
+                        else "postgres.constraint_name_not_unique_multiple",
+                    ),
+                )
+        return errors
+
+
+def _check_lazy_references(
+    models_registry: ModelsRegistry, packages_registry: Any
+) -> list[PreflightResult]:
+    """
+    Ensure all lazy (i.e. string) model references have been resolved.
+
+    Lazy references are used in various places throughout Plain, primarily in
+    related fields and model signals. Identify those common cases and provide
+    more helpful error messages for them.
+    """
+    pending_models = set(models_registry._pending_operations)
+
+    # Short circuit if there aren't any errors.
+    if not pending_models:
+        return []
+
+    def extract_operation(
+        obj: Any,
+    ) -> tuple[Callable[..., Any], list[Any], dict[str, Any]]:
+        """
+        Take a callable found in Packages._pending_operations and identify the
+        original callable passed to Packages.lazy_model_operation(). If that
+        callable was a partial, return the inner, non-partial function and
+        any arguments and keyword arguments that were supplied with it.
+
+        obj is a callback defined locally in Packages.lazy_model_operation() and
+        annotated there with a `func` attribute so as to imitate a partial.
+        """
+        operation, args, keywords = obj, [], {}
+        while hasattr(operation, "func"):
+            args.extend(getattr(operation, "args", []))
+            keywords.update(getattr(operation, "keywords", {}))
+            operation = operation.func
+        return operation, args, keywords
+
+    def app_model_error(model_key: tuple[str, str]) -> str:
+        try:
+            packages_registry.get_package_config(model_key[0])
+            model_error = "app '{}' doesn't provide model '{}'".format(*model_key)
+        except LookupError:
+            model_error = f"app '{model_key[0]}' isn't installed"
+        return model_error
+
+    # Here are several functions which return CheckMessage instances for the
+    # most common usages of lazy operations throughout Plain. These functions
+    # take the model that was being waited on as an (package_label, modelname)
+    # pair, the original lazy function, and its positional and keyword args as
+    # determined by extract_operation().
+
+    def field_error(
+        model_key: tuple[str, str],
+        func: Callable[..., Any],
+        args: list[Any],
+        keywords: dict[str, Any],
+    ) -> PreflightResult:
+        error_msg = (
+            "The field %(field)s was declared with a lazy reference "
+            "to '%(model)s', but %(model_error)s."
+        )
+        params = {
+            "model": ".".join(model_key),
+            "field": keywords["field"],
+            "model_error": app_model_error(model_key),
+        }
+        return PreflightResult(
+            fix=error_msg % params,
+            obj=keywords["field"],
+            id="fields.lazy_reference_not_resolvable",
+        )
+
+    def default_error(
+        model_key: tuple[str, str],
+        func: Callable[..., Any],
+        args: list[Any],
+        keywords: dict[str, Any],
+    ) -> PreflightResult:
+        error_msg = (
+            "%(op)s contains a lazy reference to %(model)s, but %(model_error)s."
+        )
+        params = {
+            "op": func,
+            "model": ".".join(model_key),
+            "model_error": app_model_error(model_key),
+        }
+        return PreflightResult(
+            fix=error_msg % params,
+            obj=func,
+            id="postgres.lazy_reference_resolution_failed",
+        )
+
+    # Maps common uses of lazy operations to corresponding error functions
+    # defined above. If a key maps to None, no error will be produced.
+    # default_error() will be used for usages that don't appear in this dict.
+    known_lazy = {
+        ("plain.postgres.fields.related", "resolve_related_class"): field_error,
+    }
+
+    def build_error(
+        model_key: tuple[str, str],
+        func: Callable[..., Any],
+        args: list[Any],
+        keywords: dict[str, Any],
+    ) -> PreflightResult | None:
+        key = (func.__module__, func.__name__)  # type: ignore[attr-defined]
+        error_fn = known_lazy.get(key, default_error)
+        return error_fn(model_key, func, args, keywords) if error_fn else None
+
+    return sorted(
+        filter(
+            None,
+            (
+                build_error(model_key, *extract_operation(func))
+                for model_key in pending_models
+                for func in models_registry._pending_operations[model_key]
+            ),
+        ),
+        key=lambda error: error.fix,
+    )
+
+
+@register_check("postgres.lazy_references")
+class CheckLazyReferences(PreflightCheck):
+    """Ensures all lazy (string) model references have been resolved."""
+
+    def run(self) -> list[PreflightResult]:
+        return _check_lazy_references(models_registry, packages_registry)
+
+
+@register_check("postgres.postgres_version")
+class CheckPostgresVersion(PreflightCheck):
+    """Checks that the PostgreSQL server meets the minimum version requirement."""
+
+    MINIMUM_VERSION = 16
+
+    def run(self) -> list[PreflightResult]:
+        conn = get_connection()
+        major, minor = divmod(conn.pg_version, 10000)
+        if major < self.MINIMUM_VERSION:
+            return [
+                PreflightResult(
+                    fix=f"PostgreSQL {self.MINIMUM_VERSION} or later is required (found {major}.{minor}).",
+                    id="postgres.postgres_version_too_old",
+                )
+            ]
+        return []
+
+
+@register_check("postgres.database_tables")
+class CheckDatabaseTables(PreflightCheck):
+    """Checks for unknown tables in the database when plain.postgres is available."""
+
+    def run(self) -> list[PreflightResult]:
+        conn = get_connection()
+        unknown_tables = (
+            set(conn.table_names())
+            - set(conn.plain_table_names())
+            - {MIGRATION_TABLE_NAME}
+        )
+
+        if not unknown_tables:
+            return []
+
+        table_names = ", ".join(sorted(unknown_tables))
+        return [
+            PreflightResult(
+                fix=f"Unknown tables in default database: {table_names}. "
+                "Tables may be from packages/models that have been uninstalled. "
+                "Make sure you have a backup, then run `plain db drop-unknown-tables` to remove them.",
+                id="postgres.unknown_database_tables",
+                warning=True,
+            )
+        ]
+
+
+@register_check("postgres.prunable_migrations")
+class CheckPrunableMigrations(PreflightCheck):
+    """Warns about stale migration records in the database."""
+
+    def run(self) -> list[PreflightResult]:
+        # Import here to avoid circular import issues
+        from plain.postgres.migrations.loader import MigrationLoader
+        from plain.postgres.migrations.recorder import MigrationRecorder
+
+        errors = []
+
+        # Load migrations from disk and database
+        conn = get_connection()
+        loader = MigrationLoader(conn, ignore_no_migrations=True)
+        recorder = MigrationRecorder(conn)
+        recorded_migrations = recorder.applied_migrations()
+
+        # disk_migrations should not be None after MigrationLoader initialization,
+        # but check to satisfy type checker
+        if loader.disk_migrations is None:
+            return errors
+
+        # Find all prunable migrations (recorded but not on disk)
+        all_prunable = [
+            migration
+            for migration in recorded_migrations
+            if migration not in loader.disk_migrations
+        ]
+
+        if not all_prunable:
+            return errors
+
+        # Separate into existing packages vs orphaned packages
+        existing_packages = set(loader.migrated_packages)
+        prunable_existing: list[tuple[str, str]] = []
+        prunable_orphaned: list[tuple[str, str]] = []
+
+        for migration in all_prunable:
+            package, name = migration
+            if package in existing_packages:
+                prunable_existing.append(migration)
+            else:
+                prunable_orphaned.append(migration)
+
+        # Build the warning message
+        total_count = len(all_prunable)
+        message_parts = [
+            f"Found {total_count} stale migration record{'s' if total_count != 1 else ''} in the database."
+        ]
+
+        if prunable_existing:
+            existing_list = ", ".join(
+                f"{pkg}.{name}" for pkg, name in prunable_existing[:3]
+            )
+            if len(prunable_existing) > 3:
+                existing_list += f" (and {len(prunable_existing) - 3} more)"
+            message_parts.append(f"From existing packages: {existing_list}.")
+
+        if prunable_orphaned:
+            orphaned_list = ", ".join(
+                f"{pkg}.{name}" for pkg, name in prunable_orphaned[:3]
+            )
+            if len(prunable_orphaned) > 3:
+                orphaned_list += f" (and {len(prunable_orphaned) - 3} more)"
+            message_parts.append(f"From removed packages: {orphaned_list}.")
+
+        message_parts.append("Run 'plain migrations prune' to review and remove them.")
+
+        errors.append(
+            PreflightResult(
+                fix=" ".join(message_parts),
+                id="postgres.prunable_migrations",
+                warning=True,
+            )
+        )
+
+        return errors