data-diff now uses database A's now instead of cli's now. (issue #284)

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,19 @@ 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/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,7 @@ def timestamp_value(self, t: datetime) -> str:
         "Provide SQL for the given timestamp value"
         ...
 
+
     @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/queries/api.py

@@ -95,4 +95,8 @@ def insert_rows_in_batches(db, table: TablePath, rows, *, columns=None, batch_si
         db.query(table.insert_rows(batch, columns=columns))
 
 
+def current_timestamp():
+    return CurrentTimestamp()
+
+
 commit = Commit()

data_diff/sqeleton/queries/ast_classes.py

@@ -8,7 +8,7 @@
 from ..abcs import Compilable
 from ..schema import Schema
 
-from .compiler import Compiler, cv_params
+from .compiler import Compiler, cv_params, Root
 from .base import SKIP, CompileError, DbPath, args_as_tuple
 
 
@@ -47,7 +47,7 @@ def cast_to(self, to):
 
 
 @dataclass
-class Code(ExprNode):
+class Code(ExprNode, Root):
     code: str
 
     def compile(self, c: Compiler) -> str:
@@ -434,7 +434,7 @@ def compile(self, c: Compiler) -> str:
 
 
 @dataclass
-class Join(ExprNode, ITable):
+class Join(ExprNode, ITable, Root):
     source_tables: Sequence[ITable]
     op: str = None
     on_exprs: Sequence[Expr] = None
@@ -499,7 +499,7 @@ def compile(self, parent_c: Compiler) -> str:
 
 
 @dataclass
-class GroupBy(ExprNode, ITable):
+class GroupBy(ExprNode, ITable, Root):
     table: ITable
     keys: Sequence[Expr] = None  # IKey?
     values: Sequence[Expr] = None
@@ -540,7 +540,7 @@ def compile(self, c: Compiler) -> str:
 
 
 @dataclass
-class TableOp(ExprNode, ITable):
+class TableOp(ExprNode, ITable, Root):
     op: str
     table1: ITable
     table2: ITable
@@ -571,7 +571,7 @@ def compile(self, parent_c: Compiler) -> str:
 
 
 @dataclass
-class Select(ExprNode, ITable):
+class Select(ExprNode, ITable, Root):
     table: Expr = None
     columns: Sequence[Expr] = None
     where_exprs: Sequence[Expr] = None
@@ -771,7 +771,7 @@ def compile_for_insert(self, c: Compiler):
 
 
 @dataclass
-class Explain(ExprNode):
+class Explain(ExprNode, Root):
     select: Select
 
     type = str
@@ -780,10 +780,16 @@ def compile(self, c: Compiler) -> str:
         return c.dialect.explain_as_text(c.compile(self.select))
 
 
+class CurrentTimestamp(ExprNode):
+    type = datetime
+
+    def compile(self, c: Compiler) -> str:
+        return c.dialect.current_timestamp()
+
 # DDL
 
 
-class Statement(Compilable):
+class Statement(Compilable, Root):
     type = None
 
 

data_diff/sqeleton/queries/compiler.py

@@ -11,6 +11,9 @@
 
 cv_params = contextvars.ContextVar("params")
 
+class Root:
+    "Nodes inheriting from Root can be used as root statements in SQL (e.g. SELECT yes, RANDOM() no)"
+
 
 @dataclass
 class Compiler(AbstractCompiler):
@@ -33,6 +36,10 @@ def compile(self, elem, params=None) -> str:
         if params:
             cv_params.set(params)
 
+        if self.root and isinstance(elem, Compilable) and not isinstance(elem, Root):
+            from .ast_classes import Select
+            elem = Select(columns=[elem])
+
         res = self._compile(elem)
         if self.root and self._subqueries:
             subq = ", ".join(f"\n  {k} AS ({v})" for k, v in self._subqueries.items())

tests/common.py

@@ -13,6 +13,7 @@
 from data_diff import tracking
 from data_diff import connect
 from data_diff.sqeleton.queries import table
+from data_diff.table_segment import TableSegment
 from data_diff.sqeleton.databases import Database
 from data_diff.query_utils import drop_table
 
@@ -86,10 +87,13 @@ def get_git_revision_short_hash() -> str:
 _database_instances = {}
 
 
-def get_conn(cls: type) -> Database:
-    if cls not in _database_instances:
-        _database_instances[cls] = connect(CONN_STRINGS[cls], N_THREADS)
-    return _database_instances[cls]
+def get_conn(cls: type, shared: bool =True) -> Database:
+    if shared:
+        if cls not in _database_instances:
+            _database_instances[cls] = get_conn(cls, shared=False)
+        return _database_instances[cls]
+
+    return connect(CONN_STRINGS[cls], N_THREADS)
 
 
 def _print_used_dbs():
@@ -134,11 +138,12 @@ class DiffTestCase(unittest.TestCase):
     db_cls = None
     src_schema = None
     dst_schema = None
+    shared_connection = True
 
     def setUp(self):
         assert self.db_cls, self.db_cls
 
-        self.connection = get_conn(self.db_cls)
+        self.connection = get_conn(self.db_cls, self.shared_connection)
 
         table_suffix = random_table_suffix()
         self.table_src_name = f"src{table_suffix}"
@@ -150,11 +155,11 @@ def setUp(self):
         drop_table(self.connection, self.table_src_path)
         drop_table(self.connection, self.table_dst_path)
 
+        self.src_table = table(self.table_src_path, schema=self.src_schema)
+        self.dst_table = table(self.table_dst_path, schema=self.dst_schema)
         if self.src_schema:
-            self.src_table = table(self.table_src_path, schema=self.src_schema)
             self.connection.query(self.src_table.create())
         if self.dst_schema:
-            self.dst_table = table(self.table_dst_path, schema=self.dst_schema)
             self.connection.query(self.dst_table.create())
 
         return super().setUp()
@@ -175,3 +180,8 @@ def _test_per_database(cls):
         return _parameterized_class_per_conn(databases)(cls)
 
     return _test_per_database
+
+def table_segment(database, table_path, key_columns, *args, **kw):
+    if isinstance(key_columns, str):
+        key_columns = (key_columns,)
+    return TableSegment(database, table_path, key_columns, *args, **kw)

tests/sqeleton/test_query.py

@@ -33,6 +33,9 @@ def is_distinct_from(self, a: str, b: str) -> str:
     def random(self) -> str:
         return "random()"
 
+    def current_timestamp(self) -> str:
+        return "now()"
+
     def offset_limit(self, offset: Optional[int] = None, limit: Optional[int] = None):
         x = offset and f"OFFSET {offset}", limit and f"LIMIT {limit}"
         return " ".join(filter(None, x))