Merge pull request #30 from datafold/bjoernhaeuser-patch-1

PR #20 with added infrastructure (TimestampTZ repr)

Author: erezsh

sqeleton/abcs/database_types.py

@@ -5,7 +5,7 @@
 
 from runtype import dataclass
 
-from ..utils import ArithAlphanumeric, ArithUUID, Self
+from ..utils import ArithAlphanumeric, ArithUUID, Self, Unknown
 
 
 DbPath = Tuple[str, ...]
@@ -20,7 +20,7 @@ class ColType:
 @dataclass
 class PrecisionType(ColType):
     precision: int
-    rounds: bool
+    rounds: Union[bool, Unknown] = Unknown
 
 
 class Boolean(ColType):

sqeleton/databases/base.py

@@ -30,6 +30,7 @@
     String_VaryingAlphanum,
     TemporalType,
     UnknownColType,
+    TimestampTZ,
     Text,
     DbTime,
     DbPath,
@@ -202,6 +203,8 @@ def constant_values(self, rows) -> str:
     def type_repr(self, t) -> str:
         if isinstance(t, str):
             return t
+        elif isinstance(t, TimestampTZ):
+            return f"TIMESTAMP({min(t.precision, DEFAULT_DATETIME_PRECISION)})"
         return {
             int: "INT",
             str: "VARCHAR",

sqeleton/databases/postgresql.py

@@ -97,6 +97,11 @@ def set_timezone_to_utc(self) -> str:
     def current_timestamp(self) -> str:
         return "current_timestamp"
 
+    def type_repr(self, t) -> str:
+        if isinstance(t, TimestampTZ):
+            return f"timestamp ({t.precision}) with time zone"
+        return super().type_repr(t)
+
 
 class PostgreSQL(ThreadedDatabase):
     dialect = PostgresqlDialect()

sqeleton/databases/presto.py

@@ -141,6 +141,11 @@ def set_timezone_to_utc(self) -> str:
     def current_timestamp(self) -> str:
         return "current_timestamp"
 
+    def type_repr(self, t) -> str:
+        if isinstance(t, TimestampTZ):
+            return f"timestamp with time zone"
+        return super().type_repr(t)
+
 
 class Presto(Database):
     dialect = Dialect()

sqeleton/databases/redshift.py

@@ -1,5 +1,5 @@
 from typing import List, Dict
-from ..abcs.database_types import Float, TemporalType, FractionalType, DbPath
+from ..abcs.database_types import Float, TemporalType, FractionalType, DbPath, TimestampTZ
 from ..abcs.mixins import AbstractMixin_MD5
 from .postgresql import (
     PostgreSQL,
@@ -57,6 +57,11 @@ def concat(self, items: List[str]) -> str:
     def is_distinct_from(self, a: str, b: str) -> str:
         return f"({a} IS NULL != {b} IS NULL) OR ({a}!={b})"
 
+    def type_repr(self, t) -> str:
+        if isinstance(t, TimestampTZ):
+            return f"timestamptz"
+        return super().type_repr(t)
+
 
 class Redshift(PostgreSQL):
     dialect = Dialect()

sqeleton/databases/snowflake.py

@@ -48,9 +48,9 @@ def md5_as_int(self, s: str) -> str:
 class Mixin_NormalizeValue(AbstractMixin_NormalizeValue):
     def normalize_timestamp(self, value: str, coltype: TemporalType) -> str:
         if coltype.rounds:
-            timestamp = f"to_timestamp(round(date_part(epoch_nanosecond, {value}::timestamp(9))/1000000000, {coltype.precision}))"
+            timestamp = f"to_timestamp(round(date_part(epoch_nanosecond, convert_timezone('UTC', {value})::timestamp(9))/1000000000, {coltype.precision}))"
         else:
-            timestamp = f"cast({value} as timestamp({coltype.precision}))"
+            timestamp = f"cast(convert_timezone('UTC', {value}) as timestamp({coltype.precision}))"
 
         return f"to_char({timestamp}, 'YYYY-MM-DD HH24:MI:SS.FF6')"
 
@@ -139,6 +139,11 @@ def set_timezone_to_utc(self) -> str:
     def optimizer_hints(self, hints: str) -> str:
         raise NotImplementedError("Optimizer hints not yet implemented in snowflake")
 
+    def type_repr(self, t) -> str:
+        if isinstance(t, TimestampTZ):
+            return f"timestamp_tz({t.precision})"
+        return super().type_repr(t)
+
 
 class Snowflake(Database):
     dialect = Dialect()

sqeleton/utils.py

@@ -322,3 +322,20 @@ def match_like(pattern: str, strs: Sequence[str]) -> Iterable[str]:
     for s in strs:
         if reo.match(s):
             yield s
+
+
+
+class UnknownMeta(type):
+    def __instancecheck__(self, instance):
+        return instance is Unknown
+
+    def __repr__(self):
+        return "Unknown"
+
+
+class Unknown(metaclass=UnknownMeta):
+    def __nonzero__(self):
+        raise TypeError()
+
+    def __new__(class_, *args, **kwargs):
+        raise RuntimeError("Unknown is a singleton")

tests/common.py

@@ -7,10 +7,12 @@
 import logging
 import subprocess
 
+import sqeleton
 from parameterized import parameterized_class
 
 from sqeleton import databases as db
 from sqeleton import connect
+from sqeleton.abcs.mixins import AbstractMixin_NormalizeValue
 from sqeleton.queries import table
 from sqeleton.databases import Database
 from sqeleton.query_utils import drop_table
@@ -83,7 +85,8 @@ def get_conn(cls: type, shared: bool = True) -> Database:
             _database_instances[cls] = get_conn(cls, shared=False)
         return _database_instances[cls]
 
-    return connect(CONN_STRINGS[cls], N_THREADS)
+    con = sqeleton.connect.load_mixins(AbstractMixin_NormalizeValue)
+    return con(CONN_STRINGS[cls], N_THREADS)
 
 
 def _print_used_dbs():

tests/test_database.py

@@ -1,15 +1,15 @@
-from typing import Callable, List
-from datetime import datetime
 import unittest
+from datetime import datetime
+from typing import Callable, List, Tuple
 
-from .common import str_to_checksum, TEST_MYSQL_CONN_STRING
-from .common import str_to_checksum, test_each_database_in_list, get_conn, random_table_suffix
-
-from sqeleton.queries import table, current_timestamp
+import pytz
 
-from sqeleton import databases as dbs
 from sqeleton import connect
-
+from sqeleton import databases as dbs
+from sqeleton.queries import table, current_timestamp, NormalizeAsString
+from .common import TEST_MYSQL_CONN_STRING
+from .common import str_to_checksum, test_each_database_in_list, get_conn, random_table_suffix
+from sqeleton.abcs.database_types import TimestampTZ
 
 TEST_DATABASES = {
     dbs.MySQL,
@@ -81,6 +81,43 @@ def test_current_timestamp(self):
         res = db.query(current_timestamp(), datetime)
         assert isinstance(res, datetime), (res, type(res))
 
+    def test_correct_timezone(self):
+        name = "tbl_" + random_table_suffix()
+        db = get_conn(self.db_cls)
+        tbl = table(name, schema={
+            "id": int, "created_at": TimestampTZ(9), "updated_at": TimestampTZ(9)
+        })
+
+        db.query(tbl.create())
+
+        tz = pytz.timezone('Europe/Berlin')
+
+        now = datetime.now(tz)
+        if isinstance(db, dbs.Presto):
+            ms = now.microsecond // 1000 * 1000  # Presto max precision is 3
+            now = now.replace(microsecond = ms)
+
+        db.query(table(name).insert_row(1, now, now))
+        db.query(db.dialect.set_timezone_to_utc())
+
+        t = db.table(name).query_schema()
+        t.schema["created_at"] = t.schema["created_at"].replace(precision=t.schema["created_at"].precision)
+
+        tbl = table(name, schema=t.schema)
+
+        results = db.query(tbl.select(NormalizeAsString(tbl[c]) for c in ["created_at", "updated_at"]), List[Tuple])
+
+        created_at = results[0][1]
+        updated_at = results[0][1]
+
+        utc = now.astimezone(pytz.UTC)
+        expected = utc.__format__("%Y-%m-%d %H:%M:%S.%f")
+
+
+        self.assertEqual(created_at, expected)
+        self.assertEqual(updated_at, expected)
+
+        db.query(tbl.drop())
 
 @test_each_database
 class TestThreePartIds(unittest.TestCase):
@@ -104,3 +141,4 @@ def test_three_part_support(self):
             d = db.query_table_schema(part.path)
             assert len(d) == 1
             db.query(part.drop())
+