Merge pull request #306 from datafold/issue284

data-diff now uses database A's now instead of cli's now.

Author: erezsh

data_diff/__main__.py

@@ -1,4 +1,5 @@
 from copy import deepcopy
+from datetime import datetime
 import sys
 import time
 import json
@@ -15,8 +16,9 @@
 from .joindiff_tables import TABLE_WRITE_LIMIT, JoinDiffer
 from .table_segment import TableSegment
 from .sqeleton.schema import create_schema
+from .sqeleton.queries.api import current_timestamp
 from .databases import connect
-from .parse_time import parse_time_before_now, UNITS_STR, ParseError
+from .parse_time import parse_time_before, UNITS_STR, ParseError
 from .config import apply_config_from_file
 from .tracking import disable_tracking
 from . import __version__
@@ -299,17 +301,6 @@ def _main(
 
     start = time.monotonic()
 
-    try:
-        options = dict(
-            min_update=max_age and parse_time_before_now(max_age),
-            max_update=min_age and parse_time_before_now(min_age),
-            case_sensitive=case_sensitive,
-            where=where,
-        )
-    except ParseError as e:
-        logging.error(f"Error while parsing age expression: {e}")
-        return
-
     if database1 is None or database2 is None:
         logging.error(
             f"Error: Databases not specified. Got {database1} and {database2}. Use --help for more information."
@@ -326,6 +317,20 @@ def _main(
         logging.error(e)
         return
 
+
+    now: datetime = db1.query(current_timestamp(), datetime)
+    now = now.replace(tzinfo=None)
+    try:
+        options = dict(
+            min_update=max_age and parse_time_before(now, max_age),
+            max_update=min_age and parse_time_before(now, min_age),
+            case_sensitive=case_sensitive,
+            where=where,
+        )
+    except ParseError as e:
+        logging.error(f"Error while parsing age expression: {e}")
+        return
+
     dbs = db1, db2
 
     if interactive:

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,6 +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
     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/parse_time.py

@@ -70,5 +70,5 @@ def parse_time_delta(t: str):
     return timedelta(**time_dict)
 
 
-def parse_time_before_now(t: str):
-    return datetime.now() - parse_time_delta(t)
+def parse_time_before(time: datetime, delta: str):
+    return time - parse_time_delta(delta)

data_diff/sqeleton/abcs/database_types.py

@@ -184,6 +184,10 @@ def to_string(self, s: str) -> str:
     def random(self) -> str:
         "Provide SQL for generating a random number betweein 0..1"
 
+    @abstractmethod
+    def current_timestamp(self) -> str:
+        "Provide SQL for returning the current timestamp, aka now"
+
     @abstractmethod
     def offset_limit(self, offset: Optional[int] = None, limit: Optional[int] = None):
         "Provide SQL fragment for limit and offset inside a select"
@@ -199,6 +203,10 @@ 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(
         self,

data_diff/sqeleton/databases/base.py

@@ -142,7 +142,10 @@ def timestamp_value(self, t: DbTime) -> str:
         return f"'{t.isoformat()}'"
 
     def random(self) -> str:
-        return "RANDOM()"
+        return "random()"
+
+    def current_timestamp(self) -> str:
+        return "current_timestamp()"
 
     def explain_as_text(self, query: str) -> str:
         return f"EXPLAIN {query}"

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()