From fce1d954aa57feca9c163f9d8cf66df5e8ce7b65 Mon Sep 17 00:00:00 2001 From: Mike Bayer Date: Thu, 4 Aug 2022 10:27:59 -0400 Subject: implement PG ranges/multiranges agnostically Ranges now work using a new Range object, multiranges as lists of Range objects (this is what asyncpg does. not sure why psycopg has a "Multirange" type). psycopg, psycopg2, and asyncpg are currently supported. It's not clear how to make ranges work with pg8000, likely needs string conversion; this is straightforward with the new archicture and can be added later. Fixes: #8178 Change-Id: Iab8d8382873d5c14199adbe3f09fd0dc17e2b9f1 --- lib/sqlalchemy/dialects/postgresql/__init__.py | 2 + .../dialects/postgresql/_psycopg_common.py | 1 + lib/sqlalchemy/dialects/postgresql/asyncpg.py | 95 +++++++++++++ lib/sqlalchemy/dialects/postgresql/base.py | 151 +++++++++++++++++++++ lib/sqlalchemy/dialects/postgresql/psycopg.py | 94 +++++++++++++ lib/sqlalchemy/dialects/postgresql/psycopg2.py | 61 ++++++++- lib/sqlalchemy/dialects/postgresql/ranges.py | 117 ++++++++++++---- 7 files changed, 494 insertions(+), 27 deletions(-) (limited to 'lib/sqlalchemy/dialects/postgresql') diff --git a/lib/sqlalchemy/dialects/postgresql/__init__.py b/lib/sqlalchemy/dialects/postgresql/__init__.py index baafdb181..104077a17 100644 --- a/lib/sqlalchemy/dialects/postgresql/__init__.py +++ b/lib/sqlalchemy/dialects/postgresql/__init__.py @@ -55,6 +55,7 @@ from .ranges import INT8MULTIRANGE from .ranges import INT8RANGE from .ranges import NUMMULTIRANGE from .ranges import NUMRANGE +from .ranges import Range from .ranges import TSMULTIRANGE from .ranges import TSRANGE from .ranges import TSTZMULTIRANGE @@ -135,6 +136,7 @@ __all__ = ( "NamedType", "CreateEnumType", "ExcludeConstraint", + "Range", "aggregate_order_by", "array_agg", "insert", diff --git a/lib/sqlalchemy/dialects/postgresql/_psycopg_common.py b/lib/sqlalchemy/dialects/postgresql/_psycopg_common.py index efd1dbe41..92341d2da 100644 --- a/lib/sqlalchemy/dialects/postgresql/_psycopg_common.py +++ b/lib/sqlalchemy/dialects/postgresql/_psycopg_common.py @@ -4,6 +4,7 @@ # This module is part of SQLAlchemy and is released under # the MIT License: https://www.opensource.org/licenses/mit-license.php # mypy: ignore-errors +from __future__ import annotations import decimal diff --git a/lib/sqlalchemy/dialects/postgresql/asyncpg.py b/lib/sqlalchemy/dialects/postgresql/asyncpg.py index d6385a5d6..38f8fddee 100644 --- a/lib/sqlalchemy/dialects/postgresql/asyncpg.py +++ b/lib/sqlalchemy/dialects/postgresql/asyncpg.py @@ -119,14 +119,19 @@ client using this setting passed to :func:`_asyncio.create_async_engine`:: """ # noqa +from __future__ import annotations + import collections import collections.abc as collections_abc import decimal import json as _py_json import re import time +from typing import cast +from typing import TYPE_CHECKING from . import json +from . import ranges from .base import _DECIMAL_TYPES from .base import _FLOAT_TYPES from .base import _INT_TYPES @@ -148,6 +153,9 @@ from ...util.concurrency import asyncio from ...util.concurrency import await_fallback from ...util.concurrency import await_only +if TYPE_CHECKING: + from typing import Iterable + class AsyncpgString(sqltypes.String): render_bind_cast = True @@ -278,6 +286,91 @@ class AsyncpgCHAR(sqltypes.CHAR): render_bind_cast = True +class _AsyncpgRange(ranges.AbstractRange): + def bind_processor(self, dialect): + Range = dialect.dbapi.asyncpg.Range + + NoneType = type(None) + + def to_range(value): + if not isinstance(value, (str, NoneType)): + value = Range( + value.lower, + value.upper, + lower_inc=value.bounds[0] == "[", + upper_inc=value.bounds[1] == "]", + empty=value.empty, + ) + return value + + return to_range + + def result_processor(self, dialect, coltype): + def to_range(value): + if value is not None: + empty = value.isempty + value = ranges.Range( + value.lower, + value.upper, + bounds=f"{'[' if empty or value.lower_inc else '('}" # type: ignore # noqa: E501 + f"{']' if not empty and value.upper_inc else ')'}", + empty=empty, + ) + return value + + return to_range + + +class _AsyncpgMultiRange(ranges.AbstractMultiRange): + def bind_processor(self, dialect): + Range = dialect.dbapi.asyncpg.Range + + NoneType = type(None) + + def to_range(value): + if isinstance(value, (str, NoneType)): + return value + + def to_range(value): + if not isinstance(value, (str, NoneType)): + value = Range( + value.lower, + value.upper, + lower_inc=value.bounds[0] == "[", + upper_inc=value.bounds[1] == "]", + empty=value.empty, + ) + return value + + return [ + to_range(element) + for element in cast("Iterable[ranges.Range]", value) + ] + + return to_range + + def result_processor(self, dialect, coltype): + def to_range_array(value): + def to_range(rvalue): + if rvalue is not None: + empty = rvalue.isempty + rvalue = ranges.Range( + rvalue.lower, + rvalue.upper, + bounds=f"{'[' if empty or rvalue.lower_inc else '('}" # type: ignore # noqa: E501 + f"{']' if not empty and rvalue.upper_inc else ')'}", + empty=empty, + ) + return rvalue + + if value is not None: + value = [to_range(elem) for elem in value] + + return value + + return to_range_array + + class PGExecutionContext_asyncpg(PGExecutionContext): def handle_dbapi_exception(self, e): if isinstance( @@ -828,6 +921,8 @@ class PGDialect_asyncpg(PGDialect): OID: AsyncpgOID, REGCLASS: AsyncpgREGCLASS, sqltypes.CHAR: AsyncpgCHAR, + ranges.AbstractRange: _AsyncpgRange, + ranges.AbstractMultiRange: _AsyncpgMultiRange, }, ) is_async = True diff --git a/lib/sqlalchemy/dialects/postgresql/base.py b/lib/sqlalchemy/dialects/postgresql/base.py index efb4dd547..2ee679e8e 100644 --- a/lib/sqlalchemy/dialects/postgresql/base.py +++ b/lib/sqlalchemy/dialects/postgresql/base.py @@ -1445,6 +1445,157 @@ E.g.:: Column('data', CastingArray(JSONB)) ) +Range and Multirange Types +-------------------------- + +PostgreSQL range and multirange types are supported for the psycopg2, +psycopg, and asyncpg dialects. + +Data values being passed to the database may be passed as string +values or by using the :class:`_postgresql.Range` data object. + +.. versionadded:: 2.0 Added the backend-agnostic :class:`_postgresql.Range` + object used to indicate ranges. The ``psycopg2``-specific range classes + are no longer exposed and are only used internally by that particular + dialect. + +E.g. an example of a fully typed model using the +:class:`_postgresql.TSRANGE` datatype:: + + from datetime import datetime + + from sqlalchemy.dialects.postgresql import Range + from sqlalchemy.dialects.postgresql import TSRANGE + from sqlalchemy.orm import DeclarativeBase + from sqlalchemy.orm import Mapped + from sqlalchemy.orm import mapped_column + + class Base(DeclarativeBase): + pass + + class RoomBooking(Base): + + __tablename__ = 'room_booking' + + id: Mapped[int] = mapped_column(primary_key=True) + room: Mapped[str] + during: Mapped[Range[datetime]] = mapped_column(TSRANGE) + +To represent data for the ``during`` column above, the :class:`_postgresql.Range` +type is a simple dataclass that will represent the bounds of the range. +Below illustrates an INSERT of a row into the above ``room_booking`` table:: + + from sqlalchemy import create_engine + from sqlalchemy.orm import Session + + engine = create_engine("postgresql+psycopg://scott:tiger@pg14/dbname") + + Base.metadata.create_all(engine) + + with Session(engine) as session: + booking = RoomBooking( + room="101", + during=Range(datetime(2013, 3, 23), datetime(2013, 3, 25)) + ) + session.add(booking) + session.commit() + +Selecting from any range column will also return :class:`_postgresql.Range` +objects as indicated:: + + from sqlalchemy import select + + with Session(engine) as session: + for row in session.execute(select(RoomBooking.during)): + print(row) + +The available range datatypes are as follows: + +* :class:`_postgresql.INT4RANGE` +* :class:`_postgresql.INT8RANGE` +* :class:`_postgresql.NUMRANGE` +* :class:`_postgresql.DATERANGE` +* :class:`_postgresql.TSRANGE` +* :class:`_postgresql.TSTZRANGE` + +.. autoclass:: sqlalchemy.dialects.postgresql.Range + +Multiranges +^^^^^^^^^^^ + +Multiranges are supported by PostgreSQL 14 and above. SQLAlchemy's +multirange datatypes deal in lists of :class:`_postgresql.Range` types. + +.. versionadded:: 2.0 Added support for MULTIRANGE datatypes. In contrast + to the ``psycopg`` multirange feature, SQLAlchemy's adaptation represents + a multirange datatype as a list of :class:`_postgresql.Range` objects. + +The example below illustrates use of the :class:`_postgresql.TSMULTIRANGE` +datatype:: + + from datetime import datetime + from typing import List + + from sqlalchemy.dialects.postgresql import Range + from sqlalchemy.dialects.postgresql import TSMULTIRANGE + from sqlalchemy.orm import DeclarativeBase + from sqlalchemy.orm import Mapped + from sqlalchemy.orm import mapped_column + + class Base(DeclarativeBase): + pass + + class EventCalendar(Base): + + __tablename__ = 'event_calendar' + + id: Mapped[int] = mapped_column(primary_key=True) + event_name: Mapped[str] + in_session_periods: Mapped[List[Range[datetime]]] = mapped_column(TSMULTIRANGE) + +Illustrating insertion and selecting of a record:: + + from sqlalchemy import create_engine + from sqlalchemy import select + from sqlalchemy.orm import Session + + engine = create_engine("postgresql+psycopg://scott:tiger@pg14/test") + + Base.metadata.create_all(engine) + + with Session(engine) as session: + calendar = EventCalendar( + event_name="SQLAlchemy Tutorial Sessions", + in_session_periods= [ + Range(datetime(2013, 3, 23), datetime(2013, 3, 25)), + Range(datetime(2013, 4, 12), datetime(2013, 4, 15)), + Range(datetime(2013, 5, 9), datetime(2013, 5, 12)), + ] + ) + session.add(calendar) + session.commit() + + for multirange in session.scalars(select(EventCalendar.in_session_periods)): + for range_ in multirange: + print(f"Start: {range_.lower} End: {range_.upper}") + +.. note:: In the above example, the list of :class:`_postgresql.Range` types + as handled by the ORM will not automatically detect in-place changes to + a particular list value; to update list values with the ORM, either re-assign + a new list to the attribute, or use the :class:`.MutableList` + type modifier. See the section :ref:`mutable_toplevel` for background. + + +The available multirange datatypes are as follows: + +* :class:`_postgresql.INT4MULTIRANGE` +* :class:`_postgresql.INT8MULTIRANGE` +* :class:`_postgresql.NUMMULTIRANGE` +* :class:`_postgresql.DATEMULTIRANGE` +* :class:`_postgresql.TSMULTIRANGE` +* :class:`_postgresql.TSTZMULTIRANGE` + + """ # noqa: E501 diff --git a/lib/sqlalchemy/dialects/postgresql/psycopg.py b/lib/sqlalchemy/dialects/postgresql/psycopg.py index 414976a62..633357a74 100644 --- a/lib/sqlalchemy/dialects/postgresql/psycopg.py +++ b/lib/sqlalchemy/dialects/postgresql/psycopg.py @@ -57,9 +57,14 @@ release of SQLAlchemy 2.0, however. Further documentation is available there. """ # noqa +from __future__ import annotations + import logging import re +from typing import cast +from typing import TYPE_CHECKING +from . import ranges from ._psycopg_common import _PGDialect_common_psycopg from ._psycopg_common import _PGExecutionContext_common_psycopg from .base import INTERVAL @@ -75,6 +80,9 @@ from ...sql import sqltypes from ...util.concurrency import await_fallback from ...util.concurrency import await_only +if TYPE_CHECKING: + from typing import Iterable + logger = logging.getLogger("sqlalchemy.dialects.postgresql") @@ -154,6 +162,78 @@ class _PGBoolean(sqltypes.Boolean): render_bind_cast = True +class _PsycopgRange(ranges.AbstractRange): + def bind_processor(self, dialect): + Range = cast(PGDialect_psycopg, dialect)._psycopg_Range + + NoneType = type(None) + + def to_range(value): + if not isinstance(value, (str, NoneType)): + value = Range( + value.lower, value.upper, value.bounds, value.empty + ) + return value + + return to_range + + def result_processor(self, dialect, coltype): + def to_range(value): + if value is not None: + value = ranges.Range( + value._lower, + value._upper, + bounds=value._bounds if value._bounds else "[)", + empty=not value._bounds, + ) + return value + + return to_range + + +class _PsycopgMultiRange(ranges.AbstractMultiRange): + def bind_processor(self, dialect): + Range = cast(PGDialect_psycopg, dialect)._psycopg_Range + Multirange = cast(PGDialect_psycopg, dialect)._psycopg_Multirange + + NoneType = type(None) + + def to_range(value): + if isinstance(value, (str, NoneType)): + return value + + return Multirange( + [ + Range( + element.lower, + element.upper, + element.bounds, + element.empty, + ) + for element in cast("Iterable[ranges.Range]", value) + ] + ) + + return to_range + + def result_processor(self, dialect, coltype): + def to_range(value): + if value is not None: + value = [ + ranges.Range( + elem._lower, + elem._upper, + bounds=elem._bounds if elem._bounds else "[)", + empty=not elem._bounds, + ) + for elem in value + ] + + return value + + return to_range + + class PGExecutionContext_psycopg(_PGExecutionContext_common_psycopg): pass @@ -204,6 +284,8 @@ class PGDialect_psycopg(_PGDialect_common_psycopg): sqltypes.Integer: _PGInteger, sqltypes.SmallInteger: _PGSmallInteger, sqltypes.BigInteger: _PGBigInteger, + ranges.AbstractRange: _PsycopgRange, + ranges.AbstractMultiRange: _PsycopgMultiRange, }, ) @@ -314,6 +396,18 @@ class PGDialect_psycopg(_PGDialect_common_psycopg): return TransactionStatus + @util.memoized_property + def _psycopg_Range(self): + from psycopg.types.range import Range + + return Range + + @util.memoized_property + def _psycopg_Multirange(self): + from psycopg.types.multirange import Multirange + + return Multirange + def _do_isolation_level(self, connection, autocommit, isolation_level): connection.autocommit = autocommit connection.isolation_level = isolation_level diff --git a/lib/sqlalchemy/dialects/postgresql/psycopg2.py b/lib/sqlalchemy/dialects/postgresql/psycopg2.py index 6f78dafdd..5dcd449ca 100644 --- a/lib/sqlalchemy/dialects/postgresql/psycopg2.py +++ b/lib/sqlalchemy/dialects/postgresql/psycopg2.py @@ -474,10 +474,14 @@ place within SQLAlchemy's own marshalling logic, and not that of ``psycopg2`` which may be more performant. """ # noqa +from __future__ import annotations + import collections.abc as collections_abc import logging import re +from typing import cast +from . import ranges from ._psycopg_common import _PGDialect_common_psycopg from ._psycopg_common import _PGExecutionContext_common_psycopg from .base import PGCompiler @@ -490,7 +494,6 @@ from ...engine import cursor as _cursor from ...util import FastIntFlag from ...util import parse_user_argument_for_enum - logger = logging.getLogger("sqlalchemy.dialects.postgresql") @@ -504,6 +507,56 @@ class _PGJSONB(JSONB): return None +class _Psycopg2Range(ranges.AbstractRange): + _psycopg2_range_cls = "none" + + def bind_processor(self, dialect): + Range = getattr( + cast(PGDialect_psycopg2, dialect)._psycopg2_extras, + self._psycopg2_range_cls, + ) + + NoneType = type(None) + + def to_range(value): + if not isinstance(value, (str, NoneType)): + value = Range( + value.lower, value.upper, value.bounds, value.empty + ) + return value + + return to_range + + def result_processor(self, dialect, coltype): + def to_range(value): + if value is not None: + value = ranges.Range( + value._lower, + value._upper, + bounds=value._bounds if value._bounds else "[)", + empty=not value._bounds, + ) + return value + + return to_range + + +class _Psycopg2NumericRange(_Psycopg2Range): + _psycopg2_range_cls = "NumericRange" + + +class _Psycopg2DateRange(_Psycopg2Range): + _psycopg2_range_cls = "DateRange" + + +class _Psycopg2DateTimeRange(_Psycopg2Range): + _psycopg2_range_cls = "DateTimeRange" + + +class _Psycopg2DateTimeTZRange(_Psycopg2Range): + _psycopg2_range_cls = "DateTimeTZRange" + + class PGExecutionContext_psycopg2(_PGExecutionContext_common_psycopg): _psycopg2_fetched_rows = None @@ -589,6 +642,12 @@ class PGDialect_psycopg2(_PGDialect_common_psycopg): JSON: _PGJSON, sqltypes.JSON: _PGJSON, JSONB: _PGJSONB, + ranges.INT4RANGE: _Psycopg2NumericRange, + ranges.INT8RANGE: _Psycopg2NumericRange, + ranges.NUMRANGE: _Psycopg2NumericRange, + ranges.DATERANGE: _Psycopg2DateRange, + ranges.TSRANGE: _Psycopg2DateTimeRange, + ranges.TSTZRANGE: _Psycopg2DateTimeTZRange, }, ) diff --git a/lib/sqlalchemy/dialects/postgresql/ranges.py b/lib/sqlalchemy/dialects/postgresql/ranges.py index 4f010abf1..edbe165d9 100644 --- a/lib/sqlalchemy/dialects/postgresql/ranges.py +++ b/lib/sqlalchemy/dialects/postgresql/ranges.py @@ -5,28 +5,91 @@ # the MIT License: https://www.opensource.org/licenses/mit-license.php # mypy: ignore-errors +from __future__ import annotations + +import dataclasses +from typing import Any +from typing import Generic +from typing import Optional +from typing import TypeVar from ... import types as sqltypes +from ...util import py310 +from ...util.typing import Literal + +_T = TypeVar("_T", bound=Any) + + +if py310: + dc_slots = {"slots": True} + dc_kwonly = {"kw_only": True} +else: + dc_slots = {} + dc_kwonly = {} + +@dataclasses.dataclass(frozen=True, **dc_slots) +class Range(Generic[_T]): + """Represent a PostgreSQL range. -__all__ = ("INT4RANGE", "INT8RANGE", "NUMRANGE") + E.g.:: + r = Range(10, 50, bounds="()") + + The calling style is similar to that of psycopg and psycopg2, in part + to allow easier migration from previous SQLAlchemy versions that used + these objects directly. + + :param lower: Lower bound value, or None + :param upper: Upper bound value, or None + :param bounds: keyword-only, optional string value that is one of + ``"()"``, ``"[)"``, ``"(]"``, ``"[]"``. Defaults to ``"[)"``. + :param empty: keyword-only, optional bool indicating this is an "empty" + range + + .. versionadded:: 2.0 -class RangeOperators: """ - This mixin provides functionality for the Range Operators - listed in the Range Operators table of the `PostgreSQL documentation`__ - for Range Functions and Operators. It is used by all the range types - provided in the ``postgres`` dialect and can likely be used for - any range types you create yourself. - __ https://www.postgresql.org/docs/current/static/functions-range.html + lower: Optional[_T] = None + """the lower bound""" + + upper: Optional[_T] = None + """the upper bound""" - No extra support is provided for the Range Functions listed in the Range - Functions table of the PostgreSQL documentation. For these, the normal - :func:`~sqlalchemy.sql.expression.func` object should be used. + bounds: Literal["()", "[)", "(]", "[]"] = dataclasses.field( + default="[)", **dc_kwonly + ) + empty: bool = dataclasses.field(default=False, **dc_kwonly) + if not py310: + + def __init__( + self, lower=None, upper=None, *, bounds="[)", empty=False + ): + # no __slots__ either so we can update dict + self.__dict__.update( + { + "lower": lower, + "upper": upper, + "bounds": bounds, + "empty": empty, + } + ) + + def __bool__(self) -> bool: + return self.empty + + +class AbstractRange(sqltypes.TypeEngine): """ + Base for PostgreSQL RANGE types. + + .. seealso:: + + `PostgreSQL range functions `_ + + """ # noqa: E501 class comparator_factory(sqltypes.Concatenable.Comparator): """Define comparison operations for range types.""" @@ -34,9 +97,7 @@ class RangeOperators: def __ne__(self, other): "Boolean expression. Returns true if two ranges are not equal" if other is None: - return super(RangeOperators.comparator_factory, self).__ne__( - other - ) + return super().__ne__(other) else: return self.expr.op("<>", is_comparison=True)(other) @@ -104,73 +165,77 @@ class RangeOperators: return self.expr.op("+")(other) -class INT4RANGE(RangeOperators, sqltypes.TypeEngine): +class AbstractMultiRange(AbstractRange): + """base for PostgreSQL MULTIRANGE types""" + + +class INT4RANGE(AbstractRange): """Represent the PostgreSQL INT4RANGE type.""" __visit_name__ = "INT4RANGE" -class INT8RANGE(RangeOperators, sqltypes.TypeEngine): +class INT8RANGE(AbstractRange): """Represent the PostgreSQL INT8RANGE type.""" __visit_name__ = "INT8RANGE" -class NUMRANGE(RangeOperators, sqltypes.TypeEngine): +class NUMRANGE(AbstractRange): """Represent the PostgreSQL NUMRANGE type.""" __visit_name__ = "NUMRANGE" -class DATERANGE(RangeOperators, sqltypes.TypeEngine): +class DATERANGE(AbstractRange): """Represent the PostgreSQL DATERANGE type.""" __visit_name__ = "DATERANGE" -class TSRANGE(RangeOperators, sqltypes.TypeEngine): +class TSRANGE(AbstractRange): """Represent the PostgreSQL TSRANGE type.""" __visit_name__ = "TSRANGE" -class TSTZRANGE(RangeOperators, sqltypes.TypeEngine): +class TSTZRANGE(AbstractRange): """Represent the PostgreSQL TSTZRANGE type.""" __visit_name__ = "TSTZRANGE" -class INT4MULTIRANGE(RangeOperators, sqltypes.TypeEngine): +class INT4MULTIRANGE(AbstractMultiRange): """Represent the PostgreSQL INT4MULTIRANGE type.""" __visit_name__ = "INT4MULTIRANGE" -class INT8MULTIRANGE(RangeOperators, sqltypes.TypeEngine): +class INT8MULTIRANGE(AbstractMultiRange): """Represent the PostgreSQL INT8MULTIRANGE type.""" __visit_name__ = "INT8MULTIRANGE" -class NUMMULTIRANGE(RangeOperators, sqltypes.TypeEngine): +class NUMMULTIRANGE(AbstractMultiRange): """Represent the PostgreSQL NUMMULTIRANGE type.""" __visit_name__ = "NUMMULTIRANGE" -class DATEMULTIRANGE(RangeOperators, sqltypes.TypeEngine): +class DATEMULTIRANGE(AbstractMultiRange): """Represent the PostgreSQL DATEMULTIRANGE type.""" __visit_name__ = "DATEMULTIRANGE" -class TSMULTIRANGE(RangeOperators, sqltypes.TypeEngine): +class TSMULTIRANGE(AbstractMultiRange): """Represent the PostgreSQL TSRANGE type.""" __visit_name__ = "TSMULTIRANGE" -class TSTZMULTIRANGE(RangeOperators, sqltypes.TypeEngine): +class TSTZMULTIRANGE(AbstractMultiRange): """Represent the PostgreSQL TSTZRANGE type.""" __visit_name__ = "TSTZMULTIRANGE" -- cgit v1.2.1