plain.postgres 0.84.0__py3-none-any.whl

plain/postgres/sql/where.py

@@ -0,0 +1,374 @@
+"""
+Code to manage the creation and SQL rendering of 'where' constraints.
+"""
+
+from __future__ import annotations
+
+import operator
+from functools import cached_property, reduce
+from typing import TYPE_CHECKING, Any
+
+from plain.postgres.exceptions import EmptyResultSet, FullResultSet
+from plain.postgres.expressions import Case, ResolvableExpression, When
+from plain.postgres.lookups import Exact
+from plain.utils import tree
+
+if TYPE_CHECKING:
+    from plain.postgres.connection import DatabaseConnection
+    from plain.postgres.lookups import Lookup
+    from plain.postgres.sql.compiler import SQLCompiler
+
+# Connection types
+AND = "AND"
+OR = "OR"
+XOR = "XOR"
+
+
+class WhereNode(tree.Node):
+    """
+    An SQL WHERE clause.
+
+    The class is tied to the Query class that created it (in order to create
+    the correct SQL).
+
+    A child is usually an expression producing boolean values. Most likely the
+    expression is a Lookup instance.
+
+    However, a child could also be any class with as_sql() and either
+    relabeled_clone() method or relabel_aliases() and clone() methods and
+    contains_aggregate attribute.
+    """
+
+    default = AND
+    resolved = False
+    conditional = True
+
+    def split_having_qualify(
+        self, negated: bool = False, must_group_by: bool = False
+    ) -> tuple[WhereNode | None, WhereNode | None, WhereNode | None]:
+        """
+        Return three possibly None nodes: one for those parts of self that
+        should be included in the WHERE clause, one for those parts of self
+        that must be included in the HAVING clause, and one for those parts
+        that refer to window functions.
+        """
+        if not self.contains_aggregate and not self.contains_over_clause:
+            return self, None, None
+        in_negated = negated ^ self.negated
+        # Whether or not children must be connected in the same filtering
+        # clause (WHERE > HAVING > QUALIFY) to maintain logical semantic.
+        must_remain_connected = (
+            (in_negated and self.connector == AND)
+            or (not in_negated and self.connector == OR)
+            or self.connector == XOR
+        )
+        if (
+            must_remain_connected
+            and self.contains_aggregate
+            and not self.contains_over_clause
+        ):
+            # It's must cheaper to short-circuit and stash everything in the
+            # HAVING clause than split children if possible.
+            return None, self, None
+        where_parts = []
+        having_parts = []
+        qualify_parts = []
+        for c in self.children:
+            if hasattr(c, "split_having_qualify"):
+                where_part, having_part, qualify_part = c.split_having_qualify(
+                    in_negated, must_group_by
+                )
+                if where_part is not None:
+                    where_parts.append(where_part)
+                if having_part is not None:
+                    having_parts.append(having_part)
+                if qualify_part is not None:
+                    qualify_parts.append(qualify_part)
+            elif c.contains_over_clause:
+                qualify_parts.append(c)
+            elif c.contains_aggregate:
+                having_parts.append(c)
+            else:
+                where_parts.append(c)
+        if must_remain_connected and qualify_parts:
+            # Disjunctive heterogeneous predicates can be pushed down to
+            # qualify as long as no conditional aggregation is involved.
+            if not where_parts or (where_parts and not must_group_by):
+                return None, None, self
+            elif where_parts:
+                # In theory this should only be enforced when dealing with
+                # where_parts containing predicates against multi-valued
+                # relationships that could affect aggregation results but this
+                # is complex to infer properly.
+                raise NotImplementedError(
+                    "Heterogeneous disjunctive predicates against window functions are "
+                    "not implemented when performing conditional aggregation."
+                )
+        where_node = (
+            self.create(where_parts, self.connector, self.negated)
+            if where_parts
+            else None
+        )
+        having_node = (
+            self.create(having_parts, self.connector, self.negated)
+            if having_parts
+            else None
+        )
+        qualify_node = (
+            self.create(qualify_parts, self.connector, self.negated)
+            if qualify_parts
+            else None
+        )
+        return where_node, having_node, qualify_node
+
+    def as_sql(
+        self, compiler: SQLCompiler, connection: DatabaseConnection
+    ) -> tuple[str, list[Any]]:
+        """
+        Return the SQL version of the where clause and the value to be
+        substituted in. Return '', [] if this node matches everything,
+        None, [] if this node is empty, and raise EmptyResultSet if this
+        node can't match anything.
+        """
+        result = []
+        result_params = []
+        if self.connector == AND:
+            full_needed, empty_needed = len(self.children), 1
+        else:
+            full_needed, empty_needed = 1, len(self.children)
+
+        if self.connector == XOR:
+            # PostgreSQL doesn't have a native XOR operator, so convert:
+            #   a XOR b XOR c XOR ...
+            # to:
+            #   (a OR b OR c OR ...) AND (a + b + c + ...) == 1
+            lhs = self.__class__(self.children, OR)
+            rhs_sum = reduce(
+                operator.add,
+                (Case(When(c, then=1), default=0) for c in self.children),
+            )
+            rhs = Exact(1, rhs_sum)
+            return self.__class__([lhs, rhs], AND, self.negated).as_sql(
+                compiler, connection
+            )
+
+        for child in self.children:
+            try:
+                sql, params = compiler.compile(child)
+            except EmptyResultSet:
+                empty_needed -= 1
+            except FullResultSet:
+                full_needed -= 1
+            else:
+                if sql:
+                    result.append(sql)
+                    result_params.extend(params)
+                else:
+                    full_needed -= 1
+            # Check if this node matches nothing or everything.
+            # First check the amount of full nodes and empty nodes
+            # to make this node empty/full.
+            # Now, check if this node is full/empty using the
+            # counts.
+            if empty_needed == 0:
+                if self.negated:
+                    raise FullResultSet
+                else:
+                    raise EmptyResultSet
+            if full_needed == 0:
+                if self.negated:
+                    raise EmptyResultSet
+                else:
+                    raise FullResultSet
+        conn = f" {self.connector} "
+        sql_string = conn.join(result)
+        if not sql_string:
+            raise FullResultSet
+        if self.negated:
+            sql_string = f"NOT ({sql_string})"
+        elif len(result) > 1 or self.resolved:
+            sql_string = f"({sql_string})"
+        return sql_string, result_params
+
+    def get_group_by_cols(self) -> list[Any]:
+        cols = []
+        for child in self.children:
+            cols.extend(child.get_group_by_cols())
+        return cols
+
+    def get_source_expressions(self) -> list[Any]:
+        return self.children[:]
+
+    def set_source_expressions(self, children: list[Any]) -> None:
+        assert len(children) == len(self.children)
+        self.children = children
+
+    def relabel_aliases(self, change_map: dict[str, str]) -> None:
+        """
+        Relabel the alias values of any children. 'change_map' is a dictionary
+        mapping old (current) alias values to the new values.
+        """
+        for pos, child in enumerate(self.children):
+            if hasattr(child, "relabel_aliases"):
+                # For example another WhereNode
+                child.relabel_aliases(change_map)
+            elif hasattr(child, "relabeled_clone"):
+                self.children[pos] = child.relabeled_clone(change_map)
+
+    def clone(self) -> WhereNode:
+        clone = self.create(connector=self.connector, negated=self.negated)
+        for child in self.children:
+            if hasattr(child, "clone"):
+                child = child.clone()
+            clone.children.append(child)
+        return clone
+
+    def relabeled_clone(self, change_map: dict[str, str]) -> WhereNode:
+        clone = self.clone()
+        clone.relabel_aliases(change_map)
+        return clone
+
+    def replace_expressions(self, replacements: dict[Any, Any]) -> WhereNode:
+        if replacement := replacements.get(self):
+            return replacement
+        clone = self.create(connector=self.connector, negated=self.negated)
+        for child in self.children:
+            clone.children.append(child.replace_expressions(replacements))
+        return clone
+
+    def get_refs(self) -> set[Any]:
+        refs = set()
+        for child in self.children:
+            refs |= child.get_refs()
+        return refs
+
+    @classmethod
+    def _contains_aggregate(cls, obj: Any) -> bool:
+        if isinstance(obj, tree.Node):
+            return any(cls._contains_aggregate(c) for c in obj.children)
+        return obj.contains_aggregate
+
+    @cached_property
+    def contains_aggregate(self) -> bool:
+        return self._contains_aggregate(self)
+
+    @classmethod
+    def _contains_over_clause(cls, obj: Any) -> bool:
+        if isinstance(obj, tree.Node):
+            return any(cls._contains_over_clause(c) for c in obj.children)
+        return obj.contains_over_clause
+
+    @cached_property
+    def contains_over_clause(self) -> bool:
+        return self._contains_over_clause(self)
+
+    @property
+    def is_summary(self) -> bool:
+        return any(child.is_summary for child in self.children)
+
+    @staticmethod
+    def _resolve_leaf(expr: Any, query: Any, *args: Any, **kwargs: Any) -> Any:
+        if isinstance(expr, ResolvableExpression):
+            expr = expr.resolve_expression(query, *args, **kwargs)
+        return expr
+
+    @classmethod
+    def _resolve_node(cls, node: Any, query: Any, *args: Any, **kwargs: Any) -> None:
+        if hasattr(node, "children"):
+            for child in node.children:
+                cls._resolve_node(child, query, *args, **kwargs)
+        if hasattr(node, "lhs"):
+            node.lhs = cls._resolve_leaf(node.lhs, query, *args, **kwargs)
+        if hasattr(node, "rhs"):
+            node.rhs = cls._resolve_leaf(node.rhs, query, *args, **kwargs)
+
+    def resolve_expression(self, *args: Any, **kwargs: Any) -> WhereNode:
+        clone = self.clone()
+        clone._resolve_node(clone, *args, **kwargs)
+        clone.resolved = True
+        return clone
+
+    @cached_property
+    def output_field(self) -> Any:
+        from plain.postgres.fields import BooleanField
+
+        return BooleanField()
+
+    @property
+    def _output_field_or_none(self) -> Any:
+        return self.output_field
+
+    def select_format(
+        self, compiler: SQLCompiler, sql: str, params: list[Any]
+    ) -> tuple[str, list[Any]]:
+        # Boolean expressions work directly in SELECT
+        return sql, params
+
+    def get_db_converters(self, connection: DatabaseConnection) -> list[Any]:
+        return self.output_field.get_db_converters(connection)
+
+    def get_lookup(self, lookup: str) -> type[Lookup] | None:
+        return self.output_field.get_lookup(lookup)
+
+    def leaves(self) -> Any:
+        for child in self.children:
+            if isinstance(child, WhereNode):
+                yield from child.leaves()
+            else:
+                yield child
+
+
+class NothingNode:
+    """A node that matches nothing."""
+
+    contains_aggregate = False
+    contains_over_clause = False
+
+    def as_sql(
+        self,
+        compiler: SQLCompiler | None = None,
+        connection: DatabaseConnection | None = None,
+    ) -> tuple[str, list[Any]]:
+        raise EmptyResultSet
+
+
+class ExtraWhere:
+    # The contents are a black box - assume no aggregates or windows are used.
+    contains_aggregate = False
+    contains_over_clause = False
+
+    def __init__(self, sqls: list[str], params: list[Any] | None):
+        self.sqls = sqls
+        self.params = params
+
+    def as_sql(
+        self,
+        compiler: SQLCompiler | None = None,
+        connection: DatabaseConnection | None = None,
+    ) -> tuple[str, list[Any]]:
+        sqls = [f"({sql})" for sql in self.sqls]
+        return " AND ".join(sqls), list(self.params or ())
+
+
+class SubqueryConstraint:
+    # Even if aggregates or windows would be used in a subquery,
+    # the outer query isn't interested about those.
+    contains_aggregate = False
+    contains_over_clause = False
+
+    def __init__(
+        self, alias: str, columns: list[str], targets: list[Any], query_object: Any
+    ):
+        self.alias = alias
+        self.columns = columns
+        self.targets = targets
+        query_object.clear_ordering(clear_default=True)
+        self.query_object = query_object
+
+    def as_sql(
+        self, compiler: SQLCompiler, connection: DatabaseConnection
+    ) -> tuple[str, list[Any]]:
+        query = self.query_object
+        query.set_values(self.targets)
+        query_compiler = query.get_compiler()
+        return query_compiler.as_subquery_condition(self.alias, self.columns, compiler)

plain/postgres/test/pytest.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+import re
+from collections.abc import Generator
+from typing import Any
+
+import pytest
+
+from plain.postgres.otel import suppress_db_tracing
+from plain.signals import request_finished, request_started
+
+from .. import transaction
+from ..connection import DatabaseConnection
+from ..db import close_old_connections, get_connection
+from .utils import (
+    setup_database,
+    teardown_database,
+)
+
+
+@pytest.fixture(autouse=True)
+def _db_disabled() -> Generator[None]:
+    """
+    Every test should use this fixture by default to prevent
+    access to the normal database.
+    """
+
+    def cursor_disabled(self: Any) -> None:
+        pytest.fail("Database access not allowed without the `db` fixture")  # type: ignore[invalid-argument-type]
+
+    # Save original cursor method and replace with disabled version
+    setattr(DatabaseConnection, "_enabled_cursor", DatabaseConnection.cursor)
+    DatabaseConnection.cursor = cursor_disabled  # type: ignore[assignment]
+
+    yield
+
+    # Restore original cursor method
+    DatabaseConnection.cursor = getattr(DatabaseConnection, "_enabled_cursor")
+
+
+@pytest.fixture(scope="session")
+def setup_db(request: Any) -> Generator[None]:
+    """
+    This fixture is called automatically by `db`,
+    so a test database will only be setup if the `db` fixture is used.
+    """
+    verbosity = request.config.option.verbose
+
+    # Set up the test db across the entire session
+    _old_db_name = setup_database(verbosity=verbosity)
+
+    # Keep connections open during request client / testing
+    request_started.disconnect(close_old_connections)
+    request_finished.disconnect(close_old_connections)
+
+    yield
+
+    # Put the signals back...
+    request_started.connect(close_old_connections)
+    request_finished.connect(close_old_connections)
+
+    # When the test session is done, tear down the test db
+    teardown_database(_old_db_name, verbosity=verbosity)
+
+
+@pytest.fixture
+def db(setup_db: Any, request: Any) -> Generator[None]:
+    if "isolated_db" in request.fixturenames:
+        pytest.fail("The 'db' and 'isolated_db' fixtures cannot be used together")  # type: ignore[invalid-argument-type]
+
+    # Set .cursor() back to the original implementation to unblock it
+    DatabaseConnection.cursor = getattr(DatabaseConnection, "_enabled_cursor")
+
+    with suppress_db_tracing():
+        atomic = transaction.atomic()
+        atomic._from_testcase = True
+        atomic.__enter__()
+
+    yield
+
+    with suppress_db_tracing():
+        conn = get_connection()
+        # PostgreSQL can defer constraint checks
+        if not conn.needs_rollback and conn.is_usable():
+            conn.check_constraints()
+
+        conn.set_rollback(True)
+        atomic.__exit__(None, None, None)
+
+        conn.close()
+
+
+@pytest.fixture
+def isolated_db(request: Any) -> Generator[None]:
+    """
+    Create and destroy a unique test database for each test, using a prefix
+    derived from the test function name to ensure isolation from the default
+    test database.
+    """
+    if "db" in request.fixturenames:
+        pytest.fail("The 'db' and 'isolated_db' fixtures cannot be used together")  # type: ignore[invalid-argument-type]
+    # Set .cursor() back to the original implementation to unblock it
+    DatabaseConnection.cursor = getattr(DatabaseConnection, "_enabled_cursor")
+
+    verbosity = 1
+
+    # Derive a safe prefix from the test function name
+    raw_name = request.node.name
+    prefix = re.sub(r"[^0-9A-Za-z_]+", "_", raw_name)
+
+    # Set up a fresh test database for this test, using the prefix
+    _old_db_name = setup_database(verbosity=verbosity, prefix=prefix)
+
+    yield
+
+    # Tear down the test database created for this test
+    teardown_database(_old_db_name, verbosity=verbosity)

plain/postgres/test/utils.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from plain.postgres.db import get_connection
+from plain.postgres.otel import suppress_db_tracing
+
+
+def setup_database(*, verbosity: int, prefix: str = "") -> str:
+    conn = get_connection()
+    old_name = conn.settings_dict["DATABASE"]
+    assert old_name is not None, "DATABASE setting must be set before creating test db"
+    with suppress_db_tracing():
+        conn.create_test_db(verbosity=verbosity, prefix=prefix)
+    return old_name
+
+
+def teardown_database(old_name: str, verbosity: int) -> None:
+    with suppress_db_tracing():
+        get_connection().destroy_test_db(old_name, verbosity)