Set session timezone to UTC if possible

Author: erezsh

data_diff/__main__.py

@@ -317,6 +317,7 @@ def _main(
         logging.error(e)
         return
 
+
     now: datetime = db1.query(current_timestamp(), datetime)
     now = now.replace(tzinfo=None)
     try:

data_diff/databases/_connect.py

@@ -1,4 +1,5 @@
-from data_diff.sqeleton.databases import Connect
+from data_diff.sqeleton.databases import Connect, Database
+import logging
 
 from .postgresql import PostgreSQL
 from .mysql import MySQL
@@ -29,4 +30,22 @@
     "vertica": Vertica,
 }
 
-connect = Connect(DATABASE_BY_SCHEME)
+
+class Connect_SetUTC(Connect):
+    """Provides methods for connecting to a supported database using a URL or connection dict.
+
+    Ensures all sessions use UTC Timezone, if possible.
+    """
+
+    def _connection_created(self, db):
+        db = super()._connection_created(db)
+        try:
+            db.query(db.dialect.set_timezone_to_utc())
+        except NotImplementedError:
+            logging.debug(
+                f"Database '{db}' does not allow setting timezone. We recommend making sure it's set to 'UTC'."
+            )
+        return db
+
+
+connect = Connect_SetUTC(DATABASE_BY_SCHEME)

data_diff/joindiff_tables.py

@@ -62,7 +62,7 @@ def sample(table_expr):
 
 def create_temp_table(c: Compiler, path: TablePath, expr: Expr) -> str:
     db = c.database
-    c = c.replace(root=False)   # we're compiling fragments, not full queries
+    c = c.replace(root=False)  # we're compiling fragments, not full queries
     if isinstance(db, BigQuery):
         return f"create table {c.compile(path)} OPTIONS(expiration_timestamp=TIMESTAMP_ADD(CURRENT_TIMESTAMP(), INTERVAL 1 DAY)) as {c.compile(expr)}"
     elif isinstance(db, Presto):

data_diff/sqeleton/abcs/database_types.py

@@ -203,6 +203,9 @@ def timestamp_value(self, t: datetime) -> str:
         "Provide SQL for the given timestamp value"
         ...
 
+    @abstractmethod
+    def set_timezone_to_utc(self) -> str:
+        "Provide SQL for setting the session timezone to UTC"
 
     @abstractmethod
     def parse_type(

data_diff/sqeleton/databases/bigquery.py

@@ -101,6 +101,9 @@ def type_repr(self, t) -> str:
         except KeyError:
             return super().type_repr(t)
 
+    def set_timezone_to_utc(self) -> str:
+        raise NotImplementedError()
+
 
 class BigQuery(Database):
     CONNECT_URI_HELP = "bigquery://<project>/<dataset>"

data_diff/sqeleton/databases/clickhouse.py

@@ -150,6 +150,9 @@ def _parse_type_repr(self, type_repr: str) -> Optional[Type[ColType]]:
     #     # return f"'{t}'"
     #     return f"'{str(t)[:19]}'"
 
+    def set_timezone_to_utc(self) -> str:
+        raise NotImplementedError()
+
 
 class Clickhouse(ThreadedDatabase):
     dialect = Dialect()

data_diff/sqeleton/databases/connect.py

@@ -93,6 +93,8 @@ def match_path(self, dsn):
 
 
 class Connect:
+    """Provides methods for connecting to a supported database using a URL or connection dict."""
+
     def __init__(self, database_by_scheme: Dict[str, Database]):
         self.database_by_scheme = database_by_scheme
         self.match_uri_path = {
@@ -172,9 +174,11 @@ def connect_to_uri(self, db_uri: str, thread_count: Optional[int] = 1) -> Databa
         kw = {k: v for k, v in kw.items() if v is not None}
 
         if issubclass(cls, ThreadedDatabase):
-            return cls(thread_count=thread_count, **kw)
+            db = cls(thread_count=thread_count, **kw)
+        else:
+            db = cls(**kw)
 
-        return cls(**kw)
+        return self._connection_created(db)
 
     def connect_with_dict(self, d, thread_count):
         d = dict(d)
@@ -186,9 +190,15 @@ def connect_with_dict(self, d, thread_count):
 
         cls = matcher.database_cls
         if issubclass(cls, ThreadedDatabase):
-            return cls(thread_count=thread_count, **d)
+            db = cls(thread_count=thread_count, **d)
+        else:
+            db = cls(**d)
+
+        return self._connection_created(db)
 
-        return cls(**d)
+    def _connection_created(self, db):
+        "Nop function to be overridden by subclasses."
+        return db
 
     def __call__(self, db_conf: Union[str, dict], thread_count: Optional[int] = 1) -> Database:
         """Connect to a database using the given database configuration.

data_diff/sqeleton/databases/databricks.py

@@ -82,6 +82,9 @@ def _convert_db_precision_to_digits(self, p: int) -> int:
         # Subtracting 2 due to wierd precision issues
         return max(super()._convert_db_precision_to_digits(p) - 2, 0)
 
+    def set_timezone_to_utc(self) -> str:
+        return "SET TIME ZONE 'UTC'"
+
 
 class Databricks(ThreadedDatabase):
     dialect = Dialect()

data_diff/sqeleton/databases/duckdb.py

@@ -108,6 +108,9 @@ def parse_type(
 
         return super().parse_type(table_path, col_name, type_repr, datetime_precision, numeric_precision, numeric_scale)
 
+    def set_timezone_to_utc(self) -> str:
+        return "SET GLOBAL TimeZone='UTC'"
+
 
 class DuckDB(Database):
     dialect = Dialect()

data_diff/sqeleton/databases/mysql.py

@@ -94,6 +94,9 @@ def type_repr(self, t) -> str:
     def explain_as_text(self, query: str) -> str:
         return f"EXPLAIN FORMAT=TREE {query}"
 
+    def set_timezone_to_utc(self) -> str:
+        return "SET @@session.time_zone='+00:00'"
+
 
 class MySQL(ThreadedDatabase):
     dialect = Dialect()