Docs plain-postgres plain postgres functions datetime.py
Raw LLM
from __future__ import annotations

from datetime import datetime
from typing import TYPE_CHECKING, Any

from plain.postgres.dialect import (
    date_extract_sql,
    date_trunc_sql,
    datetime_cast_date_sql,
    datetime_cast_time_sql,
    datetime_extract_sql,
    datetime_trunc_sql,
    time_extract_sql,
    time_trunc_sql,
)
from plain.postgres.expressions import Func
from plain.postgres.fields import (
    DateField,
    DateTimeField,
    DurationField,
    Field,
    IntegerField,
    TimeField,
)
from plain.postgres.lookups import (
    Transform,
    YearExact,
    YearGt,
    YearGte,
    YearLt,
    YearLte,
)
from plain.utils import timezone

if TYPE_CHECKING:
    from plain.postgres.connection import DatabaseConnection
    from plain.postgres.sql.compiler import SQLCompiler


class TimezoneMixin(Transform):
    tzinfo = None

    def get_tzname(self) -> str | None:
        # Timezone conversions must happen to the input datetime *before*
        # applying a function. 2015-12-31 23:00:00 -02:00 is stored in the
        # database as 2016-01-01 01:00:00 +00:00. Any results should be
        # based on the input datetime not the stored datetime.
        if self.tzinfo is None:
            return timezone.get_current_timezone_name()
        else:
            return timezone._get_timezone_name(self.tzinfo)


class Extract(TimezoneMixin, Transform):
    lookup_name: str | None = None
    output_field = IntegerField()

    def __init__(
        self,
        expression: Any,
        lookup_name: str | None = None,
        tzinfo: Any = None,
        **extra: Any,
    ) -> None:
        if self.lookup_name is None:
            self.lookup_name = lookup_name
        if self.lookup_name is None:
            raise ValueError("lookup_name must be provided")
        self.tzinfo = tzinfo
        super().__init__(expression, **extra)

    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]]:
        # lookup_name is guaranteed to be str after __init__ validation
        assert self.lookup_name is not None
        sql, params = compiler.compile(self.lhs)
        lhs_output_field = self.lhs.output_field
        if isinstance(lhs_output_field, DateTimeField):
            tzname = self.get_tzname()
            sql, params = datetime_extract_sql(
                self.lookup_name, sql, tuple(params), tzname
            )
        elif self.tzinfo is not None:
            raise ValueError("tzinfo can only be used with DateTimeField.")
        elif isinstance(lhs_output_field, DateField):
            sql, params = date_extract_sql(self.lookup_name, sql, tuple(params))
        elif isinstance(lhs_output_field, TimeField):
            sql, params = time_extract_sql(self.lookup_name, sql, tuple(params))
        elif isinstance(lhs_output_field, DurationField):
            # PostgreSQL has native duration (interval) type
            sql, params = time_extract_sql(self.lookup_name, sql, tuple(params))
        else:
            # resolve_expression has already validated the output_field so this
            # assert should never be hit.
            raise ValueError("Tried to Extract from an invalid type.")
        return sql, list(params)

    def resolve_expression(
        self,
        query: Any = None,
        allow_joins: bool = True,
        reuse: Any = None,
        summarize: bool = False,
        for_save: bool = False,
    ) -> Extract:
        copy = super().resolve_expression(
            query, allow_joins, reuse, summarize, for_save
        )
        field = getattr(copy.lhs, "output_field", None)
        if field is None:
            return copy
        if not isinstance(field, DateField | DateTimeField | TimeField | DurationField):
            raise ValueError(
                "Extract input expression must be DateField, DateTimeField, "
                "TimeField, or DurationField."
            )
        # Passing dates to functions expecting datetimes is most likely a mistake.
        if type(field) == DateField and copy.lookup_name in (  # noqa: E721
            "hour",
            "minute",
            "second",
        ):
            raise ValueError(
                f"Cannot extract time component '{copy.lookup_name}' from DateField '{field.name}'."
            )
        if isinstance(field, DurationField) and copy.lookup_name in (
            "year",
            "iso_year",
            "month",
            "week",
            "week_day",
            "iso_week_day",
            "quarter",
        ):
            raise ValueError(
                f"Cannot extract component '{copy.lookup_name}' from DurationField '{field.name}'."
            )
        return copy


class ExtractYear(Extract):
    lookup_name = "year"


class ExtractIsoYear(Extract):
    """Return the ISO-8601 week-numbering year."""

    lookup_name = "iso_year"


class ExtractMonth(Extract):
    lookup_name = "month"


class ExtractDay(Extract):
    lookup_name = "day"


class ExtractWeek(Extract):
    """
    Return 1-52 or 53, based on ISO-8601, i.e., Monday is the first of the
    week.
    """

    lookup_name = "week"


class ExtractWeekDay(Extract):
    """
    Return Sunday=1 through Saturday=7.

    To replicate this in Python: (mydatetime.isoweekday() % 7) + 1
    """

    lookup_name = "week_day"


class ExtractIsoWeekDay(Extract):
    """Return Monday=1 through Sunday=7, based on ISO-8601."""

    lookup_name = "iso_week_day"


class ExtractQuarter(Extract):
    lookup_name = "quarter"


class ExtractHour(Extract):
    lookup_name = "hour"


class ExtractMinute(Extract):
    lookup_name = "minute"


class ExtractSecond(Extract):
    lookup_name = "second"


DateField.register_lookup(ExtractYear)
DateField.register_lookup(ExtractMonth)
DateField.register_lookup(ExtractDay)
DateField.register_lookup(ExtractWeekDay)
DateField.register_lookup(ExtractIsoWeekDay)
DateField.register_lookup(ExtractWeek)
DateField.register_lookup(ExtractIsoYear)
DateField.register_lookup(ExtractQuarter)

DateTimeField.register_lookup(ExtractYear)
DateTimeField.register_lookup(ExtractMonth)
DateTimeField.register_lookup(ExtractDay)
DateTimeField.register_lookup(ExtractWeekDay)
DateTimeField.register_lookup(ExtractIsoWeekDay)
DateTimeField.register_lookup(ExtractWeek)
DateTimeField.register_lookup(ExtractIsoYear)
DateTimeField.register_lookup(ExtractQuarter)
DateTimeField.register_lookup(ExtractHour)
DateTimeField.register_lookup(ExtractMinute)
DateTimeField.register_lookup(ExtractSecond)

TimeField.register_lookup(ExtractHour)
TimeField.register_lookup(ExtractMinute)
TimeField.register_lookup(ExtractSecond)

ExtractYear.register_lookup(YearExact)
ExtractYear.register_lookup(YearGt)
ExtractYear.register_lookup(YearGte)
ExtractYear.register_lookup(YearLt)
ExtractYear.register_lookup(YearLte)

ExtractIsoYear.register_lookup(YearExact)
ExtractIsoYear.register_lookup(YearGt)
ExtractIsoYear.register_lookup(YearGte)
ExtractIsoYear.register_lookup(YearLt)
ExtractIsoYear.register_lookup(YearLte)


class Now(Func):
    # STATEMENT_TIMESTAMP() returns the time at the start of the current statement,
    # as opposed to CURRENT_TIMESTAMP which returns the time at the start of the
    # transaction.
    template = "STATEMENT_TIMESTAMP()"
    output_field = DateTimeField()


class TruncBase(TimezoneMixin, Transform):
    kind: str | None = None

    def __init__(
        self,
        expression: Any,
        output_field: Field | None = None,
        tzinfo: Any = None,
        **extra: Any,
    ) -> None:
        self.tzinfo = tzinfo
        super().__init__(expression, output_field=output_field, **extra)

    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]]:
        # kind is guaranteed to be str in subclasses
        assert self.kind is not None
        sql, params = compiler.compile(self.lhs)
        tzname = None
        if isinstance(self.lhs.output_field, DateTimeField):
            tzname = self.get_tzname()
        elif self.tzinfo is not None:
            raise ValueError("tzinfo can only be used with DateTimeField.")
        if isinstance(self.output_field, DateTimeField):
            sql, params = datetime_trunc_sql(self.kind, sql, tuple(params), tzname)
        elif isinstance(self.output_field, DateField):
            sql, params = date_trunc_sql(self.kind, sql, tuple(params), tzname)
        elif isinstance(self.output_field, TimeField):
            sql, params = time_trunc_sql(self.kind, sql, tuple(params), tzname)
        else:
            raise ValueError(
                "Trunc only valid on DateField, TimeField, or DateTimeField."
            )
        return sql, list(params)

    def resolve_expression(
        self,
        query: Any = None,
        allow_joins: bool = True,
        reuse: Any = None,
        summarize: bool = False,
        for_save: bool = False,
    ) -> TruncBase:
        copy = super().resolve_expression(
            query, allow_joins, reuse, summarize, for_save
        )
        field = copy.lhs.output_field
        if not isinstance(field, DateField | DateTimeField | TimeField):
            raise TypeError(
                f"{field.name!r} isn't a DateField, TimeField, or DateTimeField."
            )
        # If self.output_field was None, then accessing the field will trigger
        # the resolver to assign it to self.lhs.output_field.
        if not isinstance(copy.output_field, DateField | DateTimeField | TimeField):
            raise ValueError(
                "output_field must be either DateField, TimeField, or DateTimeField"
            )
        # Passing dates or times to functions expecting datetimes is most
        # likely a mistake.
        class_output_field = (
            self.__class__.output_field
            if isinstance(self.__class__.output_field, Field)
            else None
        )
        output_field = class_output_field or copy.output_field
        has_explicit_output_field = (
            class_output_field or field.__class__ is not copy.output_field.__class__
        )
        if type(field) == DateField and (  # noqa: E721
            isinstance(output_field, DateTimeField)
            or copy.kind in ("hour", "minute", "second", "time")
        ):
            raise ValueError(
                "Cannot truncate DateField '{}' to {}.".format(
                    field.name,
                    output_field.__class__.__name__
                    if has_explicit_output_field
                    else "DateTimeField",
                )
            )
        elif isinstance(field, TimeField) and (
            isinstance(output_field, DateTimeField)
            or copy.kind in ("year", "quarter", "month", "week", "day", "date")
        ):
            raise ValueError(
                "Cannot truncate TimeField '{}' to {}.".format(
                    field.name,
                    output_field.__class__.__name__
                    if has_explicit_output_field
                    else "DateTimeField",
                )
            )
        return copy

    def convert_value(
        self, value: Any, expression: Any, connection: DatabaseConnection
    ) -> Any:
        if isinstance(self.output_field, DateTimeField):
            if value is not None:
                value = value.replace(tzinfo=None)
                value = timezone.make_aware(value, self.tzinfo)
        elif isinstance(value, datetime):
            if value is None:
                pass
            elif isinstance(self.output_field, DateField):
                value = value.date()
            elif isinstance(self.output_field, TimeField):
                value = value.time()
        return value


class Trunc(TruncBase):
    def __init__(
        self,
        expression: Any,
        kind: str,
        output_field: Field | None = None,
        tzinfo: Any = None,
        **extra: Any,
    ) -> None:
        self.kind = kind
        super().__init__(expression, output_field=output_field, tzinfo=tzinfo, **extra)


class TruncYear(TruncBase):
    kind = "year"


class TruncQuarter(TruncBase):
    kind = "quarter"


class TruncMonth(TruncBase):
    kind = "month"


class TruncWeek(TruncBase):
    """Truncate to midnight on the Monday of the week."""

    kind = "week"


class TruncDay(TruncBase):
    kind = "day"


class TruncDate(TruncBase):
    kind = "date"
    lookup_name = "date"
    output_field = DateField()

    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]]:
        # Cast to date rather than truncate to date.
        sql, params = compiler.compile(self.lhs)
        tzname = self.get_tzname()
        sql, params = datetime_cast_date_sql(sql, tuple(params), tzname)
        return sql, list(params)


class TruncTime(TruncBase):
    kind = "time"
    lookup_name = "time"
    output_field = TimeField()

    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]]:
        # Cast to time rather than truncate to time.
        sql, params = compiler.compile(self.lhs)
        tzname = self.get_tzname()
        sql, params = datetime_cast_time_sql(sql, tuple(params), tzname)
        return sql, list(params)


class TruncHour(TruncBase):
    kind = "hour"


class TruncMinute(TruncBase):
    kind = "minute"


class TruncSecond(TruncBase):
    kind = "second"


DateTimeField.register_lookup(TruncDate)
DateTimeField.register_lookup(TruncTime)
On this page