Many fixes; Added materialize tests;

Now works for : postgresql, mysql, bigquery, presto, trino, snowflake, oracle, redshift

Author: erezsh

data_diff/__main__.py

@@ -9,8 +9,6 @@
 import rich
 import click
 
-from data_diff.databases.base import parse_table_name
-
 from .utils import eval_name_template, remove_password_from_url, safezip, match_like
 from .diff_tables import Algorithm
 from .hashdiff_tables import HashDiffer, DEFAULT_BISECTION_THRESHOLD, DEFAULT_BISECTION_FACTOR
@@ -269,22 +267,6 @@ def _main(
         logging.error(f"Error while parsing age expression: {e}")
         return
 
-    if algorithm == Algorithm.JOINDIFF:
-        differ = JoinDiffer(
-            threaded=threaded,
-            max_threadpool_size=threads and threads * 2,
-            validate_unique_key=not assume_unique_key,
-            materialize_to_table=materialize and parse_table_name(eval_name_template(materialize)),
-        )
-    else:
-        assert algorithm == Algorithm.HASHDIFF
-        differ = HashDiffer(
-            bisection_factor=bisection_factor,
-            bisection_threshold=bisection_threshold,
-            threaded=threaded,
-            max_threadpool_size=threads and threads * 2,
-        )
-
     if database1 is None or database2 is None:
         logging.error(
             f"Error: Databases not specified. Got {database1} and {database2}. Use --help for more information."
@@ -307,6 +289,22 @@ def _main(
         for db in dbs:
             db.enable_interactive()
 
+    if algorithm == Algorithm.JOINDIFF:
+        differ = JoinDiffer(
+            threaded=threaded,
+            max_threadpool_size=threads and threads * 2,
+            validate_unique_key=not assume_unique_key,
+            materialize_to_table=materialize and db1.parse_table_name(eval_name_template(materialize)),
+        )
+    else:
+        assert algorithm == Algorithm.HASHDIFF
+        differ = HashDiffer(
+            bisection_factor=bisection_factor,
+            bisection_threshold=bisection_threshold,
+            threaded=threaded,
+            max_threadpool_size=threads and threads * 2,
+        )
+
     table_names = table1, table2
     table_paths = [db.parse_table_name(t) for db, t in safezip(dbs, table_names)]
 

data_diff/databases/base.py

@@ -1,8 +1,8 @@
 import math
 import sys
 import logging
-from typing import Dict, Generator, Tuple, Optional, Sequence, Type, List, Union
-from functools import wraps
+from typing import Any, Callable, Dict, Generator, Tuple, Optional, Sequence, Type, List, Union
+from functools import partial, wraps
 from concurrent.futures import ThreadPoolExecutor
 import threading
 from abc import abstractmethod
@@ -27,7 +27,7 @@
     DbPath,
 )
 
-from data_diff.queries import Expr, Compiler, table, Select, SKIP, ThreadLocalInterpreter
+from data_diff.queries import Expr, Compiler, table, Select, SKIP
 
 logger = logging.getLogger("database")
 
@@ -66,30 +66,39 @@ def _one(seq):
     return x
 
 
-def _query_cursor(c, sql_code):
-    try:
-        c.execute(sql_code)
-        if sql_code.lower().startswith("select"):
-            return c.fetchall()
-    except Exception as e:
-        logger.exception(e)
-        raise
+class ThreadLocalInterpreter:
+    """An interpeter used to execute a sequence of queries within the same thread.
 
+    Useful for cursor-sensitive operations, such as creating a temporary table.
+    """
 
-def _query_conn(conn, sql_code: Union[str, ThreadLocalInterpreter]) -> list:
-    c = conn.cursor()
+    def __init__(self, compiler: Compiler, gen: Generator):
+        self.gen = gen
+        self.compiler = compiler
 
-    if isinstance(sql_code, ThreadLocalInterpreter):
-        g = sql_code.interpret()
-        q = next(g)
+    def apply_queries(self, callback: Callable[[str], Any]):
+        q: Expr = next(self.gen)
         while True:
-            res = _query_cursor(c, q)
+            sql = self.compiler.compile(q)
             try:
-                q = g.send(res)
+                try:
+                    res = callback(sql) if sql is not SKIP else SKIP
+                except Exception as e:
+                    q = self.gen.throw(type(e), e)
+                else:
+                    q = self.gen.send(res)
             except StopIteration:
                 break
+
+
+def apply_query(callback: Callable[[str], Any], sql_code: Union[str, ThreadLocalInterpreter]) -> list:
+    if isinstance(sql_code, ThreadLocalInterpreter):
+        return sql_code.apply_queries(callback)
     else:
-        return _query_cursor(c, sql_code)
+        return callback(sql_code)
+
+
+
 
 
 class Database(AbstractDatabase):
@@ -108,11 +117,17 @@ class Database(AbstractDatabase):
     def name(self):
         return type(self).__name__
 
-    def query(self, sql_ast: Expr, res_type: type = None):
+    def query(self, sql_ast: Union[Expr, Generator], res_type: type = None):
         "Query the given SQL code/AST, and attempt to convert the result to type 'res_type'"
 
         compiler = Compiler(self)
-        sql_code = compiler.compile(sql_ast)
+        if isinstance(sql_ast, Generator):
+            sql_code = ThreadLocalInterpreter(compiler, sql_ast)
+        else:
+            sql_code = compiler.compile(sql_ast)
+            if sql_code is SKIP:
+                return SKIP
+
         logger.debug("Running SQL (%s): %s", type(self).__name__, sql_code)
         if getattr(self, "_interactive", False) and isinstance(sql_ast, Select):
             explained_sql = compiler.compile(Explain(sql_ast))
@@ -311,6 +326,34 @@ def normalize_uuid(self, value: str, coltype: ColType_UUID) -> str:
     def random(self) -> str:
         return "RANDOM()"
 
+    def type_repr(self, t) -> str:
+        if isinstance(t, str):
+            return t
+        return {
+            int: "INT",
+            str: "VARCHAR",
+            bool: "BOOLEAN",
+            float: "FLOAT",
+        }[t]
+
+    def _query_cursor(self, c, sql_code: str):
+        assert isinstance(sql_code, str), sql_code
+        try:
+            c.execute(sql_code)
+            if sql_code.lower().startswith("select"):
+                return c.fetchall()
+        except Exception as e:
+            # logger.exception(e)
+            # logger.error(f'Caused by SQL: {sql_code}')
+            raise
+
+    def _query_conn(self, conn, sql_code: Union[str, ThreadLocalInterpreter]) -> list:
+        c = conn.cursor()
+        callback = partial(self._query_cursor, c)
+        return apply_query(callback, sql_code)
+
+
+
 
 class ThreadedDatabase(Database):
     """Access the database through singleton threads.
@@ -339,7 +382,7 @@ def _query_in_worker(self, sql_code: Union[str, ThreadLocalInterpreter]):
         "This method runs in a worker thread"
         if self._init_error:
             raise self._init_error
-        return _query_conn(self.thread_local.conn, sql_code)
+        return self._query_conn(self.thread_local.conn, sql_code)
 
     @abstractmethod
     def create_connection(self):
@@ -348,6 +391,10 @@ def create_connection(self):
     def close(self):
         self._queue.shutdown()
 
+    @property
+    def is_autocommit(self) -> bool:
+        return False
+
 
 CHECKSUM_HEXDIGITS = 15  # Must be 15 or lower
 MD5_HEXDIGITS = 32

data_diff/databases/bigquery.py

@@ -1,6 +1,6 @@
 from .database_types import *
-from .base import Database, import_helper, parse_table_name, ConnectError
-from .base import TIMESTAMP_PRECISION_POS
+from .base import Database, import_helper, parse_table_name, ConnectError, apply_query
+from .base import TIMESTAMP_PRECISION_POS, ThreadLocalInterpreter
 
 
 @import_helper(text="Please install BigQuery and configure your google-cloud access.")
@@ -47,7 +47,7 @@ def _normalize_returned_value(self, value):
             return value.decode()
         return value
 
-    def _query(self, sql_code: str):
+    def _query_atom(self, sql_code: str):
         from google.cloud import bigquery
 
         try:
@@ -60,6 +60,9 @@ def _query(self, sql_code: str):
             res = [tuple(self._normalize_returned_value(v) for v in row.values()) for row in res]
         return res
 
+    def _query(self, sql_code: Union[str, ThreadLocalInterpreter]):
+        return apply_query(self._query_atom, sql_code)
+
     def to_string(self, s: str):
         return f"cast({s} as string)"
 
@@ -98,3 +101,15 @@ def parse_table_name(self, name: str) -> DbPath:
 
     def random(self) -> str:
         return "RAND()"
+
+    @property
+    def is_autocommit(self) -> bool:
+        return True
+
+    def type_repr(self, t) -> str:
+        try:
+            return {
+                str: "STRING",
+            }[t]
+        except KeyError:
+            return super().type_repr(t)

data_diff/databases/database_types.py

@@ -1,6 +1,6 @@
 import logging
 import decimal
-from abc import ABC, abstractmethod
+from abc import ABC, abstractmethod, abstractproperty
 from typing import Sequence, Optional, Tuple, Union, Dict, List
 from datetime import datetime
 
@@ -293,6 +293,10 @@ def normalize_value_by_type(self, value: str, coltype: ColType) -> str:
     def _normalize_table_path(self, path: DbPath) -> DbPath:
         ...
 
+    @abstractproperty
+    def is_autocommit(self) -> bool:
+        ...
+
 
 Schema = CaseAwareMapping
 

data_diff/databases/databricks.py

@@ -1,7 +1,7 @@
 import logging
 
 from .database_types import *
-from .base import MD5_HEXDIGITS, CHECKSUM_HEXDIGITS, Database, import_helper, _query_conn, parse_table_name
+from .base import MD5_HEXDIGITS, CHECKSUM_HEXDIGITS, Database, import_helper, parse_table_name
 
 
 @import_helper(text="You can install it using 'pip install databricks-sql-connector'")
@@ -52,7 +52,7 @@ def __init__(
 
     def _query(self, sql_code: str) -> list:
         "Uses the standard SQL cursor interface"
-        return _query_conn(self._conn, sql_code)
+        return self._query_conn(self._conn, sql_code)
 
     def quote(self, s: str):
         return f"`{s}`"

data_diff/databases/mysql.py

@@ -76,3 +76,11 @@ def is_distinct_from(self, a: str, b: str) -> str:
 
     def random(self) -> str:
         return "RAND()"
+
+    def type_repr(self, t) -> str:
+        try:
+            return {
+                str: "VARCHAR(1024)",
+            }[t]
+        except KeyError:
+            return super().type_repr(t)

data_diff/databases/oracle.py

@@ -43,9 +43,9 @@ def create_connection(self):
         except Exception as e:
             raise ConnectError(*e.args) from e
 
-    def _query(self, sql_code: str):
+    def _query_cursor(self, c, sql_code: str):
         try:
-            return super()._query(sql_code)
+            return super()._query_cursor(c, sql_code)
         except self._oracle.DatabaseError as e:
             raise QueryError(e)
 
@@ -130,3 +130,11 @@ def random(self) -> str:
 
     def is_distinct_from(self, a: str, b: str) -> str:
         return f"DECODE({a}, {b}, 1, 0) = 0"
+
+    def type_repr(self, t) -> str:
+        try:
+            return {
+                str: "VARCHAR(1024)",
+            }[t]
+        except KeyError:
+            return super().type_repr(t)

data_diff/databases/presto.py

@@ -1,10 +1,10 @@
+from functools import partial
 import re
 
 from data_diff.utils import match_regexps
-from data_diff.queries import ThreadLocalInterpreter
 
 from .database_types import *
-from .base import Database, import_helper
+from .base import Database, import_helper, ThreadLocalInterpreter
 from .base import (
     MD5_HEXDIGITS,
     CHECKSUM_HEXDIGITS,
@@ -15,7 +15,7 @@
 def query_cursor(c, sql_code):
     c.execute(sql_code)
     if sql_code.lower().startswith("select"):
-        return c.fetchall()
+        return [tuple(x) for x in c.fetchall()]
     # Required for the query to actually run 🤯
     if re.match(r"(insert|create|truncate|drop)", sql_code, re.IGNORECASE):
         return c.fetchone()
@@ -75,16 +75,7 @@ def _query(self, sql_code: str) -> list:
         c = self._conn.cursor()
 
         if isinstance(sql_code, ThreadLocalInterpreter):
-            # TODO reuse code from base.py
-            g = sql_code.interpret()
-            q = next(g)
-            while True:
-                res = query_cursor(c, q)
-                try:
-                    q = g.send(res)
-                except StopIteration:
-                    break
-            return
+            return sql_code.apply_queries(partial(query_cursor, c))
 
         return query_cursor(c, sql_code)
 
@@ -142,3 +133,6 @@ def _parse_type(
     def normalize_uuid(self, value: str, coltype: ColType_UUID) -> str:
         # Trim doesn't work on CHAR type
         return f"TRIM(CAST({value} AS VARCHAR))"
+
+    def is_autocommit(self) -> bool:
+        return False

data_diff/databases/snowflake.py

@@ -1,7 +1,7 @@
 import logging
 
 from .database_types import *
-from .base import ConnectError, Database, import_helper, _query_conn, CHECKSUM_MASK
+from .base import ConnectError, Database, import_helper, CHECKSUM_MASK, ThreadLocalInterpreter
 
 
 @import_helper("snowflake")
@@ -60,9 +60,9 @@ def __init__(self, *, schema: str, **kw):
     def close(self):
         self._conn.close()
 
-    def _query(self, sql_code: str) -> list:
+    def _query(self, sql_code: Union[str, ThreadLocalInterpreter]):
         "Uses the standard SQL cursor interface"
-        return _query_conn(self._conn, sql_code)
+        return self._query_conn(self._conn, sql_code)
 
     def quote(self, s: str):
         return f'"{s}"'
@@ -87,3 +87,6 @@ def normalize_timestamp(self, value: str, coltype: TemporalType) -> str:
 
     def normalize_number(self, value: str, coltype: FractionalType) -> str:
         return self.to_string(f"cast({value} as decimal(38, {coltype.precision}))")
+
+    def is_autocommit(self) -> bool:
+        return True