implement second-level type resolution for literals

zzzeek · zzzeek · commit 8a62aa58fa5e · 2022-01-11T09:24:34.000-05:00
Added additional rule to the system that determines ``TypeEngine`` implementations from Python literals to apply a second level of adjustment to the type, so that a Python datetime with or without tzinfo can set the ``timezone=True`` parameter on the returned :class:`.DateTime` object, as well as :class:`.Time`. This helps with some round-trip scenarios on type-sensitive PostgreSQL dialects such as asyncpg, psycopg3 (2.0 only). For 1.4 specifically, the backport improves support for asyncpg handling of TIME WITH TIMEZONE, which was not fully implemented. 2.0's reworked PostgreSQL architecture had this handled already. Fixes: #7537 Change-Id: Icdb07db85af5f7f39f1c1ef855fe27609770094b
diff --git a/doc/build/changelog/unreleased_14/7537.rst b/doc/build/changelog/unreleased_14/7537.rst
@@ -0,0 +1,17 @@
+.. change::
+    :tags: bug, sql, postgresql
+    :tickets: 7537
+
+    Added additional rule to the system that determines ``TypeEngine``
+    implementations from Python literals to apply a second level of adjustment
+    to the type, so that a Python datetime with or without tzinfo can set the
+    ``timezone=True`` parameter on the returned :class:`.DateTime` object, as
+    well as :class:`.Time`. This helps with some round-trip scenarios on
+    type-sensitive PostgreSQL dialects such as asyncpg, psycopg3 (2.0 only).
+
+.. change::
+    :tags: bug, postgresql, asyncpg
+    :tickets: 7537
+
+    Improved support for asyncpg handling of TIME WITH TIMEZONE, which
+    was not fully implemented.
diff --git a/lib/sqlalchemy/sql/sqltypes.py b/lib/sqlalchemy/sql/sqltypes.py
@@ -622,6 +622,13 @@ def __init__(self, timezone=False):
     def get_dbapi_type(self, dbapi):
         return dbapi.DATETIME
 
+    def _resolve_for_literal(self, value):
+        with_timezone = value.tzinfo is not None
+        if with_timezone and not self.timezone:
+            return DATETIME_TIMEZONE
+        else:
+            return self
+
     @property
     def python_type(self):
         return dt.datetime
@@ -692,6 +699,13 @@ def get_dbapi_type(self, dbapi):
     def python_type(self):
         return dt.time
 
+    def _resolve_for_literal(self, value):
+        with_timezone = value.tzinfo is not None
+        if with_timezone and not self.timezone:
+            return TIME_TIMEZONE
+        else:
+            return self
+
     @util.memoized_property
     def _expression_adaptations(self):
         # Based on https://www.postgresql.org/docs/current/\
@@ -2994,6 +3008,8 @@ class MatchType(Boolean):
 INTEGERTYPE = Integer()
 MATCHTYPE = MatchType()
 TABLEVALUE = TableValueType()
+DATETIME_TIMEZONE = DateTime(timezone=True)
+TIME_TIMEZONE = Time(timezone=True)
 
 _type_map = {
     int: Integer(),
@@ -3031,7 +3047,7 @@ def _resolve_value_to_type(value):
             )
         return NULLTYPE
     else:
-        return _result_type
+        return _result_type._resolve_for_literal(value)
 
 
 # back-assign to type_api
diff --git a/lib/sqlalchemy/sql/type_api.py b/lib/sqlalchemy/sql/type_api.py
@@ -592,6 +592,17 @@ def with_variant(
         )
         return new_type
 
+    def _resolve_for_literal(self, value):
+        """adjust this type given a literal Python value that will be
+        stored in a bound parameter.
+
+        Used exclusively by _resolve_value_to_type().
+
+        .. versionadded:: 1.4.30 or 2.0
+
+        """
+        return self
+
     @util.memoized_property
     def _type_affinity(self):
         """Return a rudimental 'affinity' value expressing the general class
diff --git a/lib/sqlalchemy/testing/requirements.py b/lib/sqlalchemy/testing/requirements.py
@@ -753,6 +753,29 @@ def datetime(self):
 
         return exclusions.open()
 
+    @property
+    def datetime_timezone(self):
+        """target dialect supports representation of Python
+        datetime.datetime() with tzinfo with DateTime(timezone=True)."""
+
+        return exclusions.closed()
+
+    @property
+    def time_timezone(self):
+        """target dialect supports representation of Python
+        datetime.time() with tzinfo with Time(timezone=True)."""
+
+        return exclusions.closed()
+
+    @property
+    def datetime_implicit_bound(self):
+        """target dialect when given a datetime object will bind it such
+        that the database server knows the object is a datetime, and not
+        a plain string.
+
+        """
+        return exclusions.open()
+
     @property
     def datetime_microseconds(self):
         """target dialect supports representation of Python
@@ -767,6 +790,16 @@ def timestamp_microseconds(self):
         if TIMESTAMP is used."""
         return exclusions.closed()
 
+    @property
+    def timestamp_microseconds_implicit_bound(self):
+        """target dialect when given a datetime object which also includes
+        a microseconds portion when using the TIMESTAMP data type
+        will bind it such that the database server knows
+        the object is a datetime with microseconds, and not a plain string.
+
+        """
+        return self.timestamp_microseconds
+
     @property
     def datetime_historic(self):
         """target dialect supports representation of Python
diff --git a/lib/sqlalchemy/testing/suite/test_types.py b/lib/sqlalchemy/testing/suite/test_types.py
@@ -302,6 +302,11 @@ class Decorated(TypeDecorator):
             Column("decorated_date_data", Decorated),
         )
 
+    @testing.requires.datetime_implicit_bound
+    def test_select_direct(self, connection):
+        result = connection.scalar(select(literal(self.data)))
+        eq_(result, self.data)
+
     def test_round_trip(self, connection):
         date_table = self.tables.date_table
 
@@ -376,6 +381,15 @@ class DateTimeTest(_DateFixture, fixtures.TablesTest):
     data = datetime.datetime(2012, 10, 15, 12, 57, 18)
 
 
+class DateTimeTZTest(_DateFixture, fixtures.TablesTest):
+    __requires__ = ("datetime_timezone",)
+    __backend__ = True
+    datatype = DateTime(timezone=True)
+    data = datetime.datetime(
+        2012, 10, 15, 12, 57, 18, tzinfo=datetime.timezone.utc
+    )
+
+
 class DateTimeMicrosecondsTest(_DateFixture, fixtures.TablesTest):
     __requires__ = ("datetime_microseconds",)
     __backend__ = True
@@ -389,6 +403,11 @@ class TimestampMicrosecondsTest(_DateFixture, fixtures.TablesTest):
     datatype = TIMESTAMP
     data = datetime.datetime(2012, 10, 15, 12, 57, 18, 396)
 
+    @testing.requires.timestamp_microseconds_implicit_bound
+    def test_select_direct(self, connection):
+        result = connection.scalar(select(literal(self.data)))
+        eq_(result, self.data)
+
 
 class TimeTest(_DateFixture, fixtures.TablesTest):
     __requires__ = ("time",)
@@ -397,6 +416,13 @@ class TimeTest(_DateFixture, fixtures.TablesTest):
     data = datetime.time(12, 57, 18)
 
 
+class TimeTZTest(_DateFixture, fixtures.TablesTest):
+    __requires__ = ("time_timezone",)
+    __backend__ = True
+    datatype = Time(timezone=True)
+    data = datetime.time(12, 57, 18, tzinfo=datetime.timezone.utc)
+
+
 class TimeMicrosecondsTest(_DateFixture, fixtures.TablesTest):
     __requires__ = ("time_microseconds",)
     __backend__ = True
@@ -1515,6 +1541,7 @@ def test_string_cast_crit_against_string_basic(self):
     "JSONLegacyStringCastIndexTest",
     "DateTest",
     "DateTimeTest",
+    "DateTimeTZTest",
     "TextTest",
     "NumericTest",
     "IntegerTest",
@@ -1524,6 +1551,7 @@ def test_string_cast_crit_against_string_basic(self):
     "TimeMicrosecondsTest",
     "TimestampMicrosecondsTest",
     "TimeTest",
+    "TimeTZTest",
     "TrueDivTest",
     "DateTimeMicrosecondsTest",
     "DateHistoricTest",
diff --git a/test/requirements.py b/test/requirements.py
@@ -1046,6 +1046,23 @@ def datetime(self):
 
         return exclusions.open()
 
+    @property
+    def datetime_implicit_bound(self):
+        """target dialect when given a datetime object will bind it such
+        that the database server knows the object is a datetime, and not
+        a plain string.
+
+        """
+        return exclusions.fails_on(["mysql", "mariadb"])
+
+    @property
+    def datetime_timezone(self):
+        return exclusions.only_on("postgresql")
+
+    @property
+    def time_timezone(self):
+        return exclusions.only_on("postgresql") + exclusions.skip_if("+pg8000")
+
     @property
     def datetime_microseconds(self):
         """target dialect supports representation of Python
@@ -1061,6 +1078,10 @@ def timestamp_microseconds(self):
 
         return only_on(["oracle"])
 
+    @property
+    def timestamp_microseconds_implicit_bound(self):
+        return self.timestamp_microseconds + exclusions.fails_on(["oracle"])
+
     @property
     def datetime_historic(self):
         """target dialect supports representation of Python
