diff --git a/tests/unit/sqlalchemy/test_compiler.py b/tests/unit/sqlalchemy/test_compiler.py index 07638a20..10dce981 100644 --- a/tests/unit/sqlalchemy/test_compiler.py +++ b/tests/unit/sqlalchemy/test_compiler.py @@ -9,6 +9,9 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import datetime +from zoneinfo import ZoneInfo + import pytest from sqlalchemy import Column from sqlalchemy import ForeignKey @@ -19,12 +22,16 @@ from sqlalchemy import select from sqlalchemy import String from sqlalchemy import Table +from sqlalchemy.exc import CompileError from sqlalchemy.exc import SAWarning from sqlalchemy.schema import CreateTable from sqlalchemy.sql import column from sqlalchemy.sql import table from tests.unit.conftest import sqlalchemy_version +from trino.sqlalchemy.datatype import DATE +from trino.sqlalchemy.datatype import TIME +from trino.sqlalchemy.datatype import TIMESTAMP from trino.sqlalchemy.dialect import TrinoDialect metadata = MetaData() @@ -206,6 +213,110 @@ def test_try_cast(dialect): assert str(query) == 'SELECT try_cast("table".id as VARCHAR) AS id \nFROM "table"' +@pytest.mark.skipif( + sqlalchemy_version() < "1.4", + reason="columns argument to select() must be a Python list or other iterable" +) +@pytest.mark.parametrize( + 'col_type,value,expected', + [ + (TIMESTAMP(), + datetime.datetime(2026, 6, 17, 9, 57, 43, 244000), + "TIMESTAMP '2026-06-17 09:57:43.244000'"), + (TIMESTAMP(), + datetime.datetime(2026, 6, 17, 9, 57, 43), + "TIMESTAMP '2026-06-17 09:57:43.000000'"), + (TIMESTAMP(), + datetime.datetime(5, 6, 17, 9, 57, 43), + "TIMESTAMP '0005-06-17 09:57:43.000000'"), + (TIMESTAMP(timezone=True), + datetime.datetime(2026, 6, 17, 9, 57, 43, 244000, + tzinfo=datetime.timezone(datetime.timedelta(hours=5, minutes=30))), + "TIMESTAMP '2026-06-17 09:57:43.244000 +05:30'"), + (TIMESTAMP(timezone=True), + datetime.datetime(2026, 6, 17, 9, 57, 43, 244000, tzinfo=ZoneInfo("Asia/Kolkata")), + "TIMESTAMP '2026-06-17 09:57:43.244000 Asia/Kolkata'"), + (TIME(), + datetime.time(9, 57, 43, 244000), + "TIME '09:57:43.244000'"), + (TIME(), + datetime.time(9, 57, 43), + "TIME '09:57:43.000000'"), + (TIME(timezone=True), + datetime.time(9, 57, 43, 244000, tzinfo=datetime.timezone(datetime.timedelta(hours=-8))), + "TIME '09:57:43.244000 -08:00'"), + # Asia/Kolkata has kept the same offset since 1945, so resolving the + # zone against the current date stays deterministic. + (TIME(timezone=True), + datetime.time(9, 57, 43, 244000, tzinfo=ZoneInfo("Asia/Kolkata")), + "TIME '09:57:43.244000 +05:30'"), + (DATE(), + datetime.date(2026, 6, 17), + "DATE '2026-06-17'"), + (DATE(), + datetime.date(5, 6, 17), + "DATE '0005-06-17'"), + ] +) +def test_temporal_literal_processor(dialect, col_type, value, expected): + col = column("col", col_type) + tbl = table("t", col) + stmt = select(tbl).where(col == value) + query = stmt.compile(dialect=dialect, compile_kwargs={"literal_binds": True}) + assert expected in str(query) + + +@pytest.mark.parametrize( + 'col_type,value', + [ + (TIMESTAMP(), "not-a-timestamp"), + (TIMESTAMP(), datetime.date(2026, 6, 17)), + (TIME(), "not-a-time"), + (DATE(), "not-a-date"), + (DATE(), datetime.datetime(2026, 6, 17, 9, 57, 43)), + ] +) +def test_temporal_literal_processor_invalid_value(dialect, col_type, value): + process = col_type.literal_processor(dialect) + with pytest.raises(CompileError): + process(value) + + +class _FakeDstZone(datetime.tzinfo): + # Behaves like a named zone: no offset without a date. The offset for a + # date is mutable so a test can simulate crossing a DST transition. + def __init__(self): + self.current_offset = datetime.timedelta(hours=-5) + + def utcoffset(self, dt): + return None if dt is None else self.current_offset + + def dst(self, dt): + return None if dt is None else datetime.timedelta(0) + + def tzname(self, dt): + return "Fake/Eastern" + + +def test_time_literal_processor_named_zone_offset_varies_with_date(dialect): + # The renderer resolves a named zone against the current date, so the + # same time value renders different literals on different days. + # trino.temporal carries a TODO to reject named-zone times instead. + process = TIME().literal_processor(dialect) + tz = _FakeDstZone() + value = datetime.time(9, 57, 43, tzinfo=tz) + assert process(value) == "TIME '09:57:43.000000 -05:00'" + tz.current_offset = datetime.timedelta(hours=-4) + assert process(value) == "TIME '09:57:43.000000 -04:00'" + + +def test_timestamp_literal_processor_sub_minute_offset(dialect): + process = TIMESTAMP().literal_processor(dialect) + tz = datetime.timezone(datetime.timedelta(seconds=30)) + with pytest.raises(CompileError): + process(datetime.datetime(2026, 6, 17, 9, 57, 43, tzinfo=tz)) + + def test_catalogs_create_table_with_pk(dialect): with pytest.warns(SAWarning, match="Trino does not support PRIMARY KEY constraints. Constraint will be ignored."): statement = CreateTable(table_with_pk) diff --git a/tests/unit/sqlalchemy/test_datatype_parse.py b/tests/unit/sqlalchemy/test_datatype_parse.py index 99b4b50d..64caf7bf 100644 --- a/tests/unit/sqlalchemy/test_datatype_parse.py +++ b/tests/unit/sqlalchemy/test_datatype_parse.py @@ -13,13 +13,13 @@ from sqlalchemy.exc import UnsupportedCompilationError from sqlalchemy.sql.sqltypes import ARRAY from sqlalchemy.sql.sqltypes import CHAR -from sqlalchemy.sql.sqltypes import DATE from sqlalchemy.sql.sqltypes import DECIMAL from sqlalchemy.sql.sqltypes import INTEGER from sqlalchemy.sql.sqltypes import VARCHAR from sqlalchemy.sql.type_api import TypeEngine from trino.sqlalchemy import datatype +from trino.sqlalchemy.datatype import DATE from trino.sqlalchemy.datatype import MAP from trino.sqlalchemy.datatype import ROW from trino.sqlalchemy.datatype import TIME diff --git a/trino/dbapi.py b/trino/dbapi.py index 69d3047b..d9f2fbaa 100644 --- a/trino/dbapi.py +++ b/trino/dbapi.py @@ -33,7 +33,6 @@ from typing import Optional from typing import Union from urllib.parse import urlparse -from zoneinfo import ZoneInfo import trino.client import trino.exceptions @@ -52,6 +51,7 @@ from trino.exceptions import OperationalError from trino.exceptions import ProgrammingError from trino.exceptions import Warning +from trino.temporal import format_temporal_literal from trino.transaction import IsolationLevel from trino.transaction import NO_TRANSACTION from trino.transaction import Transaction @@ -569,35 +569,11 @@ def _format_prepared_param(self, param): if isinstance(param, (bytes, bytearray)): return "X'%s'" % param.hex() - if isinstance(param, datetime.datetime) and param.tzinfo is None: - datetime_str = param.strftime("%Y-%m-%d %H:%M:%S.%f") - return "TIMESTAMP '%s'" % datetime_str - - if isinstance(param, datetime.datetime) and param.tzinfo is not None: - datetime_str = param.strftime("%Y-%m-%d %H:%M:%S.%f") - # named timezones - if isinstance(param.tzinfo, ZoneInfo): - return "TIMESTAMP '%s %s'" % (datetime_str, param.tzinfo.key) - # offset-based timezones - return "TIMESTAMP '%s %s'" % (datetime_str, param.tzinfo.tzname(param)) - - # We can't calculate the offset for a time without a point in time - if isinstance(param, datetime.time) and param.tzinfo is None: - time_str = param.strftime("%H:%M:%S.%f") - return "TIME '%s'" % time_str - - if isinstance(param, datetime.time) and param.tzinfo is not None: - time_str = param.strftime("%H:%M:%S.%f") - # named timezones - if isinstance(param.tzinfo, ZoneInfo): - utc_offset = datetime.datetime.now(tz=param.tzinfo).strftime('%z') - return "TIME '%s %s:%s'" % (time_str, utc_offset[:3], utc_offset[3:]) - # offset-based timezones - return "TIME '%s %s'" % (time_str, param.strftime('%Z')[3:]) - - if isinstance(param, datetime.date): - date_str = param.strftime("%Y-%m-%d") - return "DATE '%s'" % date_str + if isinstance(param, (datetime.datetime, datetime.time, datetime.date)): + try: + return format_temporal_literal(param) + except ValueError as e: + raise trino.exceptions.NotSupportedError(str(e)) if isinstance(param, list): return "ARRAY[%s]" % ','.join(map(self._format_prepared_param, param)) diff --git a/trino/sqlalchemy/datatype.py b/trino/sqlalchemy/datatype.py index f5ecf433..87282bf3 100644 --- a/trino/sqlalchemy/datatype.py +++ b/trino/sqlalchemy/datatype.py @@ -9,6 +9,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import datetime import re from collections.abc import Iterator from typing import Any @@ -22,11 +23,16 @@ import sqlalchemy from sqlalchemy import func from sqlalchemy import util +from sqlalchemy.exc import CompileError from sqlalchemy.sql import sqltypes from sqlalchemy.sql.type_api import TypeDecorator from sqlalchemy.sql.type_api import TypeEngine from sqlalchemy.types import JSON +from trino.temporal import format_date_literal +from trino.temporal import format_time_literal +from trino.temporal import format_timestamp_literal + SQLType = Union[TypeEngine, Type[TypeEngine]] @@ -66,6 +72,33 @@ def python_type(self): return list +def _temporal_literal_processor(type_name, render, accepted_type, rejected_type=None): + # The rendered literal carries a zone suffix if and only if the value is + # timezone-aware. The type's timezone flag plays no part. The DBAPI + # parameter path (trino.dbapi) behaves the same way. + def process(value): + if not isinstance(value, accepted_type) or (rejected_type is not None and isinstance(value, rejected_type)): + raise CompileError( + f"Don't know how to literal-quote value {value!r} of type {type(value)} for {type_name}" + ) + try: + return render(value) + except ValueError as e: + raise CompileError(str(e)) + + return process + + +class DATE(sqltypes.DATE): + __visit_name__ = "DATE" + + def literal_processor(self, dialect): + # datetime.datetime subclasses datetime.date, rendering it as a DATE + # literal would silently drop the time part, so reject it. + return _temporal_literal_processor( + "DATE", format_date_literal, datetime.date, rejected_type=datetime.datetime) + + class TIME(sqltypes.TIME): __visit_name__ = "TIME" @@ -73,6 +106,9 @@ def __init__(self, precision=None, timezone=False): super(TIME, self).__init__(timezone=timezone) self.precision = precision + def literal_processor(self, dialect): + return _temporal_literal_processor("TIME", format_time_literal, datetime.time) + class TIMESTAMP(sqltypes.TIMESTAMP): __visit_name__ = "TIMESTAMP" @@ -81,6 +117,9 @@ def __init__(self, precision=None, timezone=False): super(TIMESTAMP, self).__init__(timezone=timezone) self.precision = precision + def literal_processor(self, dialect): + return _temporal_literal_processor("TIMESTAMP", format_timestamp_literal, datetime.datetime) + class JSON(TypeDecorator): impl = JSON @@ -159,7 +198,7 @@ def _format_value(self, value): "varbinary": sqltypes.VARBINARY, "json": JSON, # === Date and time === - "date": sqltypes.DATE, + "date": DATE, "time": TIME, "time with time zone": TIME, "timestamp": TIMESTAMP, diff --git a/trino/temporal.py b/trino/temporal.py new file mode 100644 index 00000000..1c8653bf --- /dev/null +++ b/trino/temporal.py @@ -0,0 +1,90 @@ +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Render Python temporal values as Trino SQL literals. + +The DBAPI parameter formatter and the SQLAlchemy literal processors share +these functions. A value therefore renders to the same SQL on both paths. +""" +import datetime +from typing import Optional + + +def format_temporal_literal(value: datetime.date) -> str: + """Render a datetime, time, or date value as a Trino literal.""" + if isinstance(value, datetime.datetime): + return format_timestamp_literal(value) + if isinstance(value, datetime.time): + return format_time_literal(value) + return format_date_literal(value) + + +def format_timestamp_literal(value: datetime.datetime) -> str: + body = value.replace(tzinfo=None).isoformat(sep=" ", timespec="microseconds") + if value.tzinfo is None: + return f"TIMESTAMP '{body}'" + zone = _zone_id(value.tzinfo) + if zone is None: + # utcoffset() at this datetime preserves the instant. Only zone + # identity is lost. No offset at all means the value is + # naive - reject it rather than render a plain TIMESTAMP. + offset = value.utcoffset() + if offset is None: + raise ValueError(f"tzinfo of {value!r} reports no UTC offset") + zone = _format_offset(offset) + return f"TIMESTAMP '{body} {zone}'" + + +def format_time_literal(value: datetime.time) -> str: + """Render a time value as a Trino TIME literal. + + A Trino TIME WITH TIME ZONE literal accepts only a fixed offset. A time + in a named zone resolves to an offset via the current date, so for DST + zones the rendered offset varies with the day the query is built. Pass a + fixed-offset tzinfo to render an exact literal. + """ + # TODO: reject named-zone times instead of resolving via the current + # date. That is a breaking change to the DBAPI parameter path. + body = value.replace(tzinfo=None).isoformat(timespec="microseconds") + if value.tzinfo is None: + return f"TIME '{body}'" + offset = value.utcoffset() + if offset is None: + # A named zone needs a date to resolve an offset. A time carries + # no date so use today. For DST zones the result varies by date. + offset = datetime.datetime.now(tz=value.tzinfo).utcoffset() + if offset is None: + raise ValueError(f"tzinfo of {value!r} reports no UTC offset") + return f"TIME '{body} {_format_offset(offset)}'" + + +def format_date_literal(value: datetime.date) -> str: + return f"DATE '{value.isoformat()}'" + + +def _zone_id(tzinfo: datetime.tzinfo) -> Optional[str]: + # zoneinfo.ZoneInfo exposes the IANA zone name as .key, pytz zones as + # .zone. Trino keeps the zone name in a TIMESTAMP WITH TIME ZONE value, + # so prefer the name over a fixed offset. + zone = getattr(tzinfo, "key", None) or getattr(tzinfo, "zone", None) + if isinstance(zone, str) and zone: + return zone + return None + + +def _format_offset(offset: datetime.timedelta) -> str: + seconds = round(offset.total_seconds()) + sign = "+" if seconds >= 0 else "-" + minutes, remainder = divmod(abs(seconds), 60) + if remainder: + raise ValueError(f"a Trino literal cannot represent the sub-minute UTC offset {offset!r}") + hours, minutes = divmod(minutes, 60) + return f"{sign}{hours:02d}:{minutes:02d}"