plain.postgres 0.84.0__py3-none-any.whl

plain/postgres/functions/text.py

@@ -0,0 +1,282 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from plain.postgres.expressions import Func, ResolvableExpression, Value
+from plain.postgres.fields import CharField, IntegerField, TextField
+from plain.postgres.functions import Cast, Coalesce
+from plain.postgres.lookups import Transform
+
+if TYPE_CHECKING:
+    from plain.postgres.connection import DatabaseConnection
+    from plain.postgres.sql.compiler import SQLCompiler
+
+
+class SHAMixin(Transform):
+    """Base class for SHA hashing using PostgreSQL's pgcrypto extension."""
+
+    def as_sql(
+        self,
+        compiler: SQLCompiler,
+        connection: DatabaseConnection,
+        function: str | None = None,
+        template: str | None = None,
+        arg_joiner: str | None = None,
+        **extra_context: Any,
+    ) -> tuple[str, list[Any]]:
+        assert self.function is not None
+        return super().as_sql(
+            compiler,
+            connection,
+            template="ENCODE(DIGEST(%(expressions)s, '%(function)s'), 'hex')",
+            function=self.function.lower(),
+            **extra_context,
+        )
+
+
+class Chr(Transform):
+    function = "CHR"
+    lookup_name = "chr"
+
+
+class ConcatPair(Func):
+    """Concatenate two arguments together."""
+
+    function = "CONCAT"
+
+    def as_sql(
+        self,
+        compiler: SQLCompiler,
+        connection: DatabaseConnection,
+        function: str | None = None,
+        template: str | None = None,
+        arg_joiner: str | None = None,
+        **extra_context: Any,
+    ) -> tuple[str, list[Any]]:
+        # PostgreSQL requires explicit cast to text for CONCAT.
+        copy = self.copy()
+        copy.set_source_expressions(
+            [
+                Cast(expression, TextField())
+                for expression in copy.get_source_expressions()
+            ]
+        )
+        return super(ConcatPair, copy).as_sql(
+            compiler,
+            connection,
+            **extra_context,
+        )
+
+    def coalesce(self) -> ConcatPair:
+        # null on either side results in null for expression, wrap with coalesce
+        c = self.copy()
+        c.set_source_expressions(
+            [
+                Coalesce(expression, Value(""))
+                for expression in c.get_source_expressions()
+            ]
+        )
+        return c
+
+
+class Concat(Func):
+    """
+    Concatenate text fields together. Wraps each argument in coalesce
+    functions to ensure a non-null result.
+    """
+
+    function = None
+    template = "%(expressions)s"
+
+    def __init__(self, *expressions: Any, **extra: Any) -> None:
+        if len(expressions) < 2:
+            raise ValueError("Concat must take at least two expressions")
+        paired = self._paired(expressions)
+        super().__init__(paired, **extra)
+
+    def _paired(self, expressions: tuple[Any, ...]) -> ConcatPair:
+        # wrap pairs of expressions in successive concat functions
+        # exp = [a, b, c, d]
+        # -> ConcatPair(a, ConcatPair(b, ConcatPair(c, d))))
+        if len(expressions) == 2:
+            return ConcatPair(*expressions)
+        return ConcatPair(expressions[0], self._paired(expressions[1:]))
+
+
+class Left(Func):
+    function = "LEFT"
+    arity = 2
+    output_field = CharField()
+
+    def __init__(self, expression: Any, length: Any, **extra: Any) -> None:
+        """
+        expression: the name of a field, or an expression returning a string
+        length: the number of characters to return from the start of the string
+        """
+        if not isinstance(length, ResolvableExpression):
+            if length < 1:
+                raise ValueError("'length' must be greater than 0.")
+        super().__init__(expression, length, **extra)
+
+    def get_substr(self) -> Substr:
+        return Substr(self.source_expressions[0], Value(1), self.source_expressions[1])
+
+
+class Length(Transform):
+    """Return the number of characters in the expression."""
+
+    function = "LENGTH"
+    lookup_name = "length"
+    output_field = IntegerField()
+
+
+class Lower(Transform):
+    function = "LOWER"
+    lookup_name = "lower"
+
+
+class LPad(Func):
+    function = "LPAD"
+    output_field = CharField()
+
+    def __init__(
+        self, expression: Any, length: Any, fill_text: Any = Value(" "), **extra: Any
+    ) -> None:
+        if (
+            not isinstance(length, ResolvableExpression)
+            and length is not None
+            and length < 0
+        ):
+            raise ValueError("'length' must be greater or equal to 0.")
+        super().__init__(expression, length, fill_text, **extra)
+
+
+class LTrim(Transform):
+    function = "LTRIM"
+    lookup_name = "ltrim"
+
+
+class MD5(Transform):
+    function = "MD5"
+    lookup_name = "md5"
+
+
+class Ord(Transform):
+    function = "ASCII"
+    lookup_name = "ord"
+    output_field = IntegerField()
+
+
+class Repeat(Func):
+    function = "REPEAT"
+    output_field = CharField()
+
+    def __init__(self, expression: Any, number: Any, **extra: Any) -> None:
+        if (
+            not isinstance(number, ResolvableExpression)
+            and number is not None
+            and number < 0
+        ):
+            raise ValueError("'number' must be greater or equal to 0.")
+        super().__init__(expression, number, **extra)
+
+
+class Replace(Func):
+    function = "REPLACE"
+
+    def __init__(
+        self, expression: Any, text: Any, replacement: Any = Value(""), **extra: Any
+    ) -> None:
+        super().__init__(expression, text, replacement, **extra)
+
+
+class Reverse(Transform):
+    function = "REVERSE"
+    lookup_name = "reverse"
+
+
+class Right(Left):
+    function = "RIGHT"
+
+    def get_substr(self) -> Substr:
+        return Substr(
+            self.source_expressions[0], self.source_expressions[1] * Value(-1)
+        )
+
+
+class RPad(LPad):
+    function = "RPAD"
+
+
+class RTrim(Transform):
+    function = "RTRIM"
+    lookup_name = "rtrim"
+
+
+class SHA1(SHAMixin, Transform):
+    function = "SHA1"
+    lookup_name = "sha1"
+
+
+class SHA224(SHAMixin, Transform):
+    function = "SHA224"
+    lookup_name = "sha224"
+
+
+class SHA256(SHAMixin, Transform):
+    function = "SHA256"
+    lookup_name = "sha256"
+
+
+class SHA384(SHAMixin, Transform):
+    function = "SHA384"
+    lookup_name = "sha384"
+
+
+class SHA512(SHAMixin, Transform):
+    function = "SHA512"
+    lookup_name = "sha512"
+
+
+class StrIndex(Func):
+    """
+    Return a positive integer corresponding to the 1-indexed position of the
+    first occurrence of a substring inside another string, or 0 if the
+    substring is not found.
+    """
+
+    # PostgreSQL uses STRPOS instead of INSTR.
+    function = "STRPOS"
+    arity = 2
+    output_field = IntegerField()
+
+
+class Substr(Func):
+    function = "SUBSTRING"
+    output_field = CharField()
+
+    def __init__(
+        self, expression: Any, pos: Any, length: Any = None, **extra: Any
+    ) -> None:
+        """
+        expression: the name of a field, or an expression returning a string
+        pos: an integer > 0, or an expression returning an integer
+        length: an optional number of characters to return
+        """
+        if not isinstance(pos, ResolvableExpression):
+            if pos < 1:
+                raise ValueError("'pos' must be greater than 0")
+        expressions = [expression, pos]
+        if length is not None:
+            expressions.append(length)
+        super().__init__(*expressions, **extra)
+
+
+class Trim(Transform):
+    function = "TRIM"
+    lookup_name = "trim"
+
+
+class Upper(Transform):
+    function = "UPPER"
+    lookup_name = "upper"

plain/postgres/functions/window.py

@@ -0,0 +1,125 @@
+from __future__ import annotations
+
+from typing import Any
+
+from plain.postgres.expressions import Func
+from plain.postgres.fields import Field, FloatField, IntegerField
+
+__all__ = [
+    "CumeDist",
+    "DenseRank",
+    "FirstValue",
+    "Lag",
+    "LastValue",
+    "Lead",
+    "NthValue",
+    "Ntile",
+    "PercentRank",
+    "Rank",
+    "RowNumber",
+]
+
+
+class CumeDist(Func):
+    function = "CUME_DIST"
+    output_field = FloatField()
+    window_compatible = True
+
+
+class DenseRank(Func):
+    function = "DENSE_RANK"
+    output_field = IntegerField()
+    window_compatible = True
+
+
+class FirstValue(Func):
+    arity = 1
+    function = "FIRST_VALUE"
+    window_compatible = True
+
+
+class LagLeadFunction(Func):
+    window_compatible = True
+
+    def __init__(
+        self, expression: Any, offset: int = 1, default: Any = None, **extra: Any
+    ) -> None:
+        if expression is None:
+            raise ValueError(
+                f"{self.__class__.__name__} requires a non-null source expression."
+            )
+        if offset is None or offset <= 0:
+            raise ValueError(
+                f"{self.__class__.__name__} requires a positive integer for the offset."
+            )
+        args = (expression, offset)
+        if default is not None:
+            args += (default,)
+        super().__init__(*args, **extra)
+
+    def _resolve_output_field(self) -> Field:
+        sources = self.get_source_expressions()
+        return sources[0].output_field
+
+
+class Lag(LagLeadFunction):
+    function = "LAG"
+
+
+class LastValue(Func):
+    arity = 1
+    function = "LAST_VALUE"
+    window_compatible = True
+
+
+class Lead(LagLeadFunction):
+    function = "LEAD"
+
+
+class NthValue(Func):
+    function = "NTH_VALUE"
+    window_compatible = True
+
+    def __init__(self, expression: Any, nth: int = 1, **extra: Any) -> None:
+        if expression is None:
+            raise ValueError(
+                f"{self.__class__.__name__} requires a non-null source expression."
+            )
+        if nth is None or nth <= 0:
+            raise ValueError(
+                f"{self.__class__.__name__} requires a positive integer as for nth."
+            )
+        super().__init__(expression, nth, **extra)
+
+    def _resolve_output_field(self) -> Field:
+        sources = self.get_source_expressions()
+        return sources[0].output_field
+
+
+class Ntile(Func):
+    function = "NTILE"
+    output_field = IntegerField()
+    window_compatible = True
+
+    def __init__(self, num_buckets: int = 1, **extra: Any) -> None:
+        if num_buckets <= 0:
+            raise ValueError("num_buckets must be greater than 0.")
+        super().__init__(num_buckets, **extra)
+
+
+class PercentRank(Func):
+    function = "PERCENT_RANK"
+    output_field = FloatField()
+    window_compatible = True
+
+
+class Rank(Func):
+    function = "RANK"
+    output_field = IntegerField()
+    window_compatible = True
+
+
+class RowNumber(Func):
+    function = "ROW_NUMBER"
+    output_field = IntegerField()
+    window_compatible = True

plain/postgres/indexes.py

@@ -0,0 +1,286 @@
+from __future__ import annotations
+
+from types import NoneType
+from typing import TYPE_CHECKING, Any, Self
+
+from plain.postgres.expressions import Col, ExpressionList, F, Func, OrderBy
+from plain.postgres.query_utils import Q
+from plain.postgres.sql.query import Query
+from plain.postgres.utils import names_digest, split_identifier
+from plain.utils.functional import partition
+
+if TYPE_CHECKING:
+    from plain.postgres.base import Model
+    from plain.postgres.expressions import Expression
+    from plain.postgres.schema import DatabaseSchemaEditor, Statement
+
+__all__ = ["Index"]
+
+
+class Index:
+    suffix = "idx"
+    # The max length of the name of the index
+    max_name_length = 30
+
+    def __init__(
+        self,
+        *expressions: Any,
+        fields: tuple[str, ...] | list[str] = (),
+        name: str | None = None,
+        opclasses: tuple[str, ...] | list[str] = (),
+        condition: Q | None = None,
+        include: tuple[str, ...] | list[str] | None = None,
+    ) -> None:
+        if opclasses and not name:
+            raise ValueError("An index must be named to use opclasses.")
+        if not isinstance(condition, NoneType | Q):
+            raise ValueError("Index.condition must be a Q instance.")
+        if condition and not name:
+            raise ValueError("An index must be named to use condition.")
+        if not isinstance(fields, list | tuple):
+            raise ValueError("Index.fields must be a list or tuple.")
+        if not isinstance(opclasses, list | tuple):
+            raise ValueError("Index.opclasses must be a list or tuple.")
+        if not expressions and not fields:
+            raise ValueError(
+                "At least one field or expression is required to define an index."
+            )
+        if expressions and fields:
+            raise ValueError(
+                "Index.fields and expressions are mutually exclusive.",
+            )
+        if expressions and not name:
+            raise ValueError("An index must be named to use expressions.")
+        if expressions and opclasses:
+            raise ValueError(
+                "Index.opclasses cannot be used with expressions. Use "
+                "a custom OpClass() instead."
+            )
+        if opclasses and len(fields) != len(opclasses):
+            raise ValueError(
+                "Index.fields and Index.opclasses must have the same number of "
+                "elements."
+            )
+        if fields and not all(isinstance(field, str) for field in fields):
+            raise ValueError("Index.fields must contain only strings with field names.")
+        if include and not name:
+            raise ValueError("A covering index must be named.")
+        if not isinstance(include, NoneType | list | tuple):
+            raise ValueError("Index.include must be a list or tuple.")
+        self.fields = list(fields)
+        # A list of 2-tuple with the field name and ordering ('' or 'DESC').
+        self.fields_orders = [
+            (field_name.removeprefix("-"), "DESC" if field_name.startswith("-") else "")
+            for field_name in self.fields
+        ]
+        self.name = name or ""
+        self.opclasses: tuple[str, ...] = tuple(opclasses)
+        self.condition = condition
+        self.include = tuple(include) if include else ()
+        self.expressions: tuple[Expression, ...] = tuple(  # type: ignore[assignment]
+            F(expression) if isinstance(expression, str) else expression
+            for expression in expressions
+        )
+
+    @property
+    def contains_expressions(self) -> bool:
+        return bool(self.expressions)
+
+    def _get_condition_sql(
+        self, model: type[Model], schema_editor: DatabaseSchemaEditor
+    ) -> str | None:
+        if self.condition is None:
+            return None
+        query = Query(model=model, alias_cols=False)
+        where = query.build_where(self.condition)
+        compiler = query.get_compiler()
+        sql, params = where.as_sql(compiler, schema_editor.connection)
+        return sql % tuple(schema_editor.quote_value(p) for p in params)
+
+    def create_sql(
+        self, model: type[Model], schema_editor: DatabaseSchemaEditor, **kwargs: Any
+    ) -> Statement:
+        include = [
+            model._model_meta.get_forward_field(field_name).column
+            for field_name in self.include
+        ]
+        condition = self._get_condition_sql(model, schema_editor)
+        if self.expressions:
+            index_expressions = []
+            for expression in self.expressions:
+                index_expression = IndexExpression(expression)
+                index_expressions.append(index_expression)
+            expressions = ExpressionList(*index_expressions).resolve_expression(
+                Query(model, alias_cols=False),
+            )
+            fields = None
+            col_suffixes = ()
+        else:
+            fields = [
+                model._model_meta.get_forward_field(field_name)
+                for field_name, _ in self.fields_orders
+            ]
+            # Support index column ordering (ASC/DESC)
+            col_suffixes = tuple(order[1] for order in self.fields_orders)
+            expressions = None
+        return schema_editor._create_index_sql(
+            model,
+            fields=fields,
+            name=self.name,
+            col_suffixes=col_suffixes,
+            opclasses=self.opclasses,
+            condition=condition,
+            include=include,
+            expressions=expressions,
+            **kwargs,
+        )
+
+    def remove_sql(
+        self, model: type[Model], schema_editor: DatabaseSchemaEditor, **kwargs: Any
+    ) -> Statement:
+        return schema_editor._delete_index_sql(model, self.name, **kwargs)
+
+    def deconstruct(self) -> tuple[str, tuple[Expression, ...], dict[str, Any]]:
+        path = f"{self.__class__.__module__}.{self.__class__.__name__}"
+        path = path.replace("plain.postgres.indexes", "plain.postgres")
+        kwargs = {"name": self.name}
+        if self.fields:
+            kwargs["fields"] = self.fields
+        if self.opclasses:
+            kwargs["opclasses"] = self.opclasses
+        if self.condition:
+            kwargs["condition"] = self.condition
+        if self.include:
+            kwargs["include"] = self.include
+        return (path, self.expressions, kwargs)
+
+    def clone(self) -> Index:
+        """Create a copy of this Index."""
+        _, args, kwargs = self.deconstruct()
+        return self.__class__(*args, **kwargs)
+
+    def set_name_with_model(self, model: type[Model]) -> None:
+        """
+        Generate a unique name for the index.
+
+        The name is divided into 3 parts - table name (12 chars), field name
+        (8 chars) and unique hash + suffix (10 chars). Each part is made to
+        fit its size by truncating the excess length.
+        """
+        _, table_name = split_identifier(model.model_options.db_table)
+        column_names = [
+            model._model_meta.get_forward_field(field_name).column
+            for field_name, order in self.fields_orders
+        ]
+        column_names_with_order = [
+            (("-%s" if order else "%s") % column_name)
+            for column_name, (field_name, order) in zip(
+                column_names, self.fields_orders
+            )
+        ]
+        # The length of the parts of the name is based on the default max
+        # length of 30 characters.
+        hash_data = [table_name] + column_names_with_order + [self.suffix]
+        self.name = "{}_{}_{}".format(
+            table_name[:11],
+            column_names[0][:7],
+            f"{names_digest(*hash_data, length=6)}_{self.suffix}",
+        )
+        if len(self.name) > self.max_name_length:
+            raise ValueError(
+                "Index name too long. Is self.suffix longer than 3 characters?"
+            )
+        if self.name[0] == "_" or self.name[0].isdigit():
+            self.name = f"D{self.name[1:]}"
+
+    def __repr__(self) -> str:
+        return "<{}:{}{}{}{}{}{}>".format(
+            self.__class__.__qualname__,
+            "" if not self.fields else f" fields={repr(self.fields)}",
+            "" if not self.expressions else f" expressions={repr(self.expressions)}",
+            "" if not self.name else f" name={repr(self.name)}",
+            "" if self.condition is None else f" condition={self.condition}",
+            "" if not self.include else f" include={repr(self.include)}",
+            "" if not self.opclasses else f" opclasses={repr(self.opclasses)}",
+        )
+
+    def __eq__(self, other: object) -> bool:
+        if isinstance(other, Index):
+            return self.deconstruct() == other.deconstruct()
+        return NotImplemented
+
+
+class IndexExpression(Func):
+    """Order and wrap expressions for CREATE INDEX statements."""
+
+    template = "%(expressions)s"
+    wrapper_classes = (OrderBy,)
+
+    def resolve_expression(
+        self,
+        query: Any = None,
+        allow_joins: bool = True,
+        reuse: Any = None,
+        summarize: bool = False,
+        for_save: bool = False,
+    ) -> Self:
+        expressions = list(self.flatten())
+        # Split expressions and wrappers.
+        index_expressions, wrappers = partition(
+            lambda e: isinstance(e, self.wrapper_classes),
+            expressions,
+        )
+        wrapper_types = [type(wrapper) for wrapper in wrappers]
+        if len(wrapper_types) != len(set(wrapper_types)):
+            raise ValueError(
+                "Multiple references to {} can't be used in an indexed "
+                "expression.".format(
+                    ", ".join(
+                        [
+                            wrapper_cls.__qualname__
+                            for wrapper_cls in self.wrapper_classes
+                        ]
+                    )
+                )
+            )
+        if expressions[1 : len(wrappers) + 1] != wrappers:
+            raise ValueError(
+                "{} must be topmost expressions in an indexed expression.".format(
+                    ", ".join(
+                        [
+                            wrapper_cls.__qualname__
+                            for wrapper_cls in self.wrapper_classes
+                        ]
+                    )
+                )
+            )
+        # Wrap expressions in parentheses if they are not column references.
+        root_expression = index_expressions[1]
+        resolve_root_expression = root_expression.resolve_expression(
+            query,
+            allow_joins,
+            reuse,
+            summarize,
+            for_save,
+        )
+        if not isinstance(resolve_root_expression, Col):
+            root_expression = Func(root_expression, template="(%(expressions)s)")
+
+        if wrappers:
+            # Order wrappers and set their expressions.
+            wrappers = sorted(
+                wrappers,
+                key=lambda w: self.wrapper_classes.index(type(w)),
+            )
+            wrappers = [wrapper.copy() for wrapper in wrappers]
+            for i, wrapper in enumerate(wrappers[:-1]):
+                wrapper.set_source_expressions([wrappers[i + 1]])
+            # Set the root expression on the deepest wrapper.
+            wrappers[-1].set_source_expressions([root_expression])
+            self.set_source_expressions([wrappers[0]])
+        else:
+            # Use the root expression, if there are no wrappers.
+            self.set_source_expressions([root_expression])
+        return super().resolve_expression(
+            query, allow_joins, reuse, summarize, for_save
+        )