async-orm-kit 1.0.0__py3-none-any.whl

async_orm_kit/query/planning/planner.py

@@ -0,0 +1,160 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from async_orm_kit._logging import get_logger
+from async_orm_kit.exceptions import RelationshipError
+
+from .directives import LoadDirective, LoadStrategy, QueryContext
+from .passes.base import PassResult, PlanningPass
+from .protocols import JoinGraphReadable
+
+logger = get_logger(__name__)
+
+
+class QueryPlanner:
+    """
+    Orchestrates planning passes in a defined order.
+
+    The planner itself holds NO planning logic — it only:
+    1. Manages the pass pipeline
+    2. Threads state between passes
+    3. Collects audit information
+
+    This is the Strategy + Chain of Responsibility pattern:
+    passes can be added, removed, or reordered without
+    touching any other code.
+    """
+
+    def __init__(self, passes: Sequence[PlanningPass] | None = None):
+        self._passes: list[PlanningPass] = list(passes or [])
+        self._last_results: list[PassResult] = []
+
+    def add_pass(self, planning_pass: PlanningPass) -> QueryPlanner:
+        """Fluent API for building the pipeline."""
+        self._passes.append(planning_pass)
+        return self
+
+    def plan(
+        self,
+        graph: JoinGraphReadable,
+        context: QueryContext,
+    ) -> list[LoadDirective]:
+        """
+        Run all passes in sequence.
+
+        Each pass receives the directives produced by the previous pass.
+        The first pass receives the initial directives built from
+        the QueryContext's raw paths.
+        """
+        self._last_results = []
+
+        # Build initial directives from raw user intent
+        directives = self._build_initial_directives(graph, context)
+
+        # Run each pass
+        for planning_pass in self._passes:
+            result = planning_pass.execute(graph, context, directives)
+            self._last_results.append(result)
+
+            if result.directives:
+                directives = result.directives
+
+            if result.warnings:
+                for w in result.warnings:
+                    logger.warning(
+                        "planning_warning",
+                        pass_name=planning_pass.name,
+                        detail=w,
+                    )
+
+            if result.mutations:
+                for m in result.mutations:
+                    logger.info(
+                        "planning_mutation",
+                        pass_name=planning_pass.name,
+                        detail=m,
+                    )
+
+        return directives
+
+    def _build_initial_directives(
+        self,
+        graph: JoinGraphReadable,
+        context: QueryContext,
+    ) -> list[LoadDirective]:
+        """
+        Translate raw user paths into initial LoadDirectives.
+
+        This is pure translation — no intelligence yet.
+        The passes add the intelligence.
+        """
+        directives: list[LoadDirective] = []
+
+        for path in sorted(context.select_related_paths):
+            relationships = graph.resolve_relationship_path(path)
+            if relationships is None:
+                raise RelationshipError(f"Invalid select_related path {path!r}.")
+            if any(relationship.uselist for relationship in relationships):
+                raise RelationshipError(
+                    f"select_related path {path!r} crosses a collection relationship; "
+                    "use prefetch_related() instead."
+                )
+            directives.append(
+                LoadDirective(
+                    path=path,
+                    strategy=LoadStrategy.JOINEDLOAD,
+                    reason="User-requested select_related",
+                )
+            )
+
+        for path in sorted(context.prefetch_paths):
+            if path in context.select_related_paths:
+                continue  # select_related takes precedence for the exact path
+
+            relationships = graph.resolve_relationship_path(path)
+            if relationships is None:
+                raise RelationshipError(f"Invalid prefetch_related path {path!r}.")
+            if relationships[-1].uselist:
+                reason = "User-requested prefetch (collection)"
+            else:
+                reason = "User-requested prefetch (scalar relationship)"
+            strategy = LoadStrategy.SELECTINLOAD
+
+            directives.append(
+                LoadDirective(
+                    path=path,
+                    strategy=strategy,
+                    reason=reason,
+                )
+            )
+
+        return directives
+
+    def explain(self) -> str:
+        """
+        Human-readable explanation of what each pass did.
+        Invaluable for debugging slow queries.
+        """
+        lines = ["Query Plan Explanation", "=" * 40]
+
+        for result in self._last_results:
+            lines.append(f"\n── {result.pass_name} ──")
+
+            if not result.had_effect:
+                lines.append("  (no effect)")
+                continue
+
+            for m in result.mutations:
+                lines.append(f"  [MUTATE] {m}")
+            for w in result.warnings:
+                lines.append(f"  [WARN]   {w}")
+
+        lines.append("\n── Final Directives ──")
+        if self._last_results and self._last_results[-1].directives:
+            for d in self._last_results[-1].directives:
+                marker = " ⚠ DEMOTED" if d.was_demoted else ""
+                lines.append(f"  {d.path:40s} → {d.strategy.name:15s}{marker}")
+                lines.append(f"    reason: {d.reason}")
+
+        return "\n".join(lines)

async_orm_kit/query/planning/protocols.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Protocol, runtime_checkable
+
+
+@runtime_checkable
+class JoinNodeReadable(Protocol):
+    """Read-only view of a single join node."""
+
+    @property
+    def path(self) -> str: ...
+
+    @property
+    def join_type(self) -> str: ...
+
+    @property
+    def used_in_filter(self) -> bool: ...
+
+    @property
+    def used_in_or_context(self) -> bool: ...
+
+    @property
+    def has_nullable_check(self) -> bool: ...
+
+    @property
+    def used_in_negation(self) -> bool: ...
+
+    @property
+    def parent(self) -> JoinNodeReadable | None: ...
+
+    @property
+    def relationship(self) -> RelationshipReadable | None: ...
+
+
+@runtime_checkable
+class RelationshipReadable(Protocol):
+    """What the planner needs to know about a relationship."""
+
+    @property
+    def uselist(self) -> bool: ...
+
+    @property
+    def lazy(self) -> str: ...
+
+
+@runtime_checkable
+class JoinGraphReadable(Protocol):
+    """
+    Read-only view of the join graph.
+
+    The planner NEVER mutates the graph through this protocol.
+    Mutation (like join promotion) goes through a separate
+    mutable interface.
+    """
+
+    def has_node(self, path: str) -> bool: ...
+
+    def get_node(self, path: str) -> JoinNodeReadable | None: ...
+
+    @property
+    def join_order(self) -> Sequence[JoinNodeReadable]: ...
+
+    def resolve_relationship(self, path: str) -> RelationshipReadable | None: ...
+
+    def resolve_relationship_path(
+        self, path: str
+    ) -> Sequence[RelationshipReadable] | None: ...
+
+
+@runtime_checkable
+class JoinGraphMutable(Protocol):
+    """
+    Narrow mutable interface — only for passes that need
+    to modify join types (e.g., promotion pass).
+    """
+
+    def set_join_type(self, path: str, join_type: str) -> None: ...

async_orm_kit/query/postgres/optimizer.py

@@ -0,0 +1,189 @@
+from __future__ import annotations
+
+from typing import Any
+
+from sqlalchemy import cast, exists, inspect as sa_inspect, literal, select
+from sqlalchemy.dialects.postgresql import JSONPATH
+
+from async_orm_kit._logging import get_logger
+
+logger = get_logger(__name__)
+
+
+class PostgresOptimizer:
+    """
+    PostgreSQL-specific query optimizations.
+
+    Features:
+    - LATERAL joins for top-N per group
+    - EXISTS optimization for relationship filters
+    - JSONB operator support
+    - Array operations
+    """
+
+    @staticmethod
+    def is_postgres(session) -> bool:
+        """Check if current database is PostgreSQL"""
+        dialect_name = session.bind.dialect.name
+        return dialect_name == "postgresql"
+
+    @staticmethod
+    def lateral_join_top_n(
+        base_model: type,
+        related_model: type,
+        relationship_fk: str,
+        order_by_field: str,
+        limit: int = 1,
+        descending: bool = True,
+    ):
+        """
+        Create LATERAL join for top-N per parent.
+
+        Example: Get latest post per user
+
+        Usage:
+            lateral_subq = PostgresOptimizer.lateral_join_top_n(
+                User, Post, "user_id", "created_at", limit=1
+            )
+            stmt = select(User).join(lateral_subq, isouter=True)
+        """
+        primary_key: Any = sa_inspect(base_model).primary_key
+        if len(primary_key) != 1:
+            raise ValueError("lateral_join_top_n requires a single-column primary key")
+        order_col = getattr(related_model, order_by_field)
+        fk_col = getattr(related_model, relationship_fk)
+
+        direction = order_col.desc() if descending else order_col.asc()
+
+        lateral_subq = (
+            select(related_model)
+            .where(fk_col == getattr(base_model, primary_key[0].key))
+            .order_by(direction)
+            .limit(limit)
+            .lateral()
+        )
+
+        return lateral_subq
+
+    @staticmethod
+    def optimize_exists_filter(base_model: type, relationship_attr, condition):
+        """
+        Use EXISTS instead of JOIN for relationship filters.
+
+        More efficient for large collections.
+
+        Usage:
+            exists_cond = PostgresOptimizer.optimize_exists_filter(
+                User, User.posts, Post.published == True
+            )
+            stmt = select(User).where(exists_cond)
+        """
+        related_model = relationship_attr.property.mapper.class_
+
+        exists_subq = (
+            exists()
+            .where(
+                getattr(
+                    related_model,
+                    next(iter(relationship_attr.property.remote_side)).name,
+                )
+                == getattr(
+                    base_model, relationship_attr.property.local_remote_pairs[0][0].name
+                )
+            )
+            .where(condition)
+        )
+
+        return exists_subq
+
+    @staticmethod
+    def jsonb_contains(column, value: dict):
+        """
+        PostgreSQL JSONB @> operator.
+
+        Usage:
+            User.metadata.op('@>')({"role": "admin"})
+        """
+        return column.op("@>")(value)
+
+    @staticmethod
+    def jsonb_path_exists(column, path: str):
+        """
+        PostgreSQL JSONB path query.
+
+        Usage:
+            jsonb_path_exists(User.data, '$.address.city')
+        """
+        if not isinstance(path, str) or not path.strip():
+            raise ValueError("JSONPath must be a non-empty string.")
+        normalized_path = path.strip()
+        if not normalized_path.startswith("$"):
+            normalized_path = f"$.{normalized_path}"
+        # Bind the path as data. Never interpolate a caller-controlled path into SQL.
+        return column.op("@?")(cast(literal(normalized_path), JSONPATH))
+
+    @staticmethod
+    def array_contains(column, value):
+        """
+        PostgreSQL array @> operator.
+
+        Usage:
+            User.tags.op('@>')(["python", "sql"])
+        """
+        return column.op("@>")(value)
+
+    @staticmethod
+    def explain_analyze(session, stmt, verbose: bool = False) -> str:
+        """Run parameterized EXPLAIN ANALYZE on PostgreSQL.
+
+        Values remain DBAPI parameters. This avoids executing SQL produced by
+        interpolating literal values into a text statement.
+        """
+        if type(verbose) is not bool:
+            raise TypeError("verbose must be a bool.")
+        if not PostgresOptimizer.is_postgres(session):
+            logger.warning("EXPLAIN ANALYZE only available on PostgreSQL")
+            return ""
+
+        bind = session.get_bind()
+        compiled = stmt.compile(
+            dialect=bind.dialect,
+            compile_kwargs={"render_postcompile": True},
+        )
+        options = ["ANALYZE", "BUFFERS"]
+        if verbose:
+            options.append("VERBOSE")
+
+        parameters = (
+            tuple(compiled.params[name] for name in compiled.positiontup)
+            if compiled.positional
+            else compiled.params
+        )
+        result = session.connection().exec_driver_sql(
+            f"EXPLAIN ({', '.join(options)}) {compiled.string}",
+            parameters,
+        )
+        rows = result.fetchall()
+        return "\n".join(str(row[0]) for row in rows)
+
+    @staticmethod
+    def suggest_index(session, stmt) -> list[str]:
+        """
+        Analyze query plan and suggest missing indexes.
+
+        Returns list of suggested CREATE INDEX statements.
+        """
+        if not PostgresOptimizer.is_postgres(session):
+            return []
+
+        plan = PostgresOptimizer.explain_analyze(session, stmt)
+        suggestions = []
+
+        # Simple heuristic: look for Seq Scan in plan
+        if "Seq Scan" in plan:
+            logger.warning("Sequential scan detected - consider adding index")
+            # Parse table and filter conditions from plan
+            # This is simplified - production version would parse EXPLAIN output
+            suggestions.append("-- Analyze EXPLAIN output and add appropriate indexes")
+
+        return suggestions