Merge remote-tracking branch 'upstream/master' into bigquery-dbt-impersonation

Author: ivan-toriya

.github/workflows/ci.yml

@@ -18,10 +18,10 @@ jobs:
       matrix:
         os: [ubuntu-latest]
         python-version:
-          - "3.7"
           - "3.8"
           - "3.9"
           - "3.10"
+          - "3.11"
 
     name: Check Python ${{ matrix.python-version }} on ${{ matrix.os }}
     runs-on: ${{ matrix.os }}

.github/workflows/ci_full.yml

@@ -21,7 +21,7 @@ jobs:
       matrix:
         os: [ubuntu-latest]
         python-version:
-          - "3.10"
+          - "3.11"
 
     name: Check Python ${{ matrix.python-version }} on ${{ matrix.os }}
     runs-on: ${{ matrix.os }}

data_diff/__init__.py

@@ -1,14 +1,14 @@
 from typing import Sequence, Tuple, Iterator, Optional, Union
 
-from data_diff.sqeleton.abcs import DbTime, DbPath
-
-from .tracking import disable_tracking
-from .databases import connect
-from .diff_tables import Algorithm
-from .hashdiff_tables import HashDiffer, DEFAULT_BISECTION_THRESHOLD, DEFAULT_BISECTION_FACTOR
-from .joindiff_tables import JoinDiffer, TABLE_WRITE_LIMIT
-from .table_segment import TableSegment
-from .utils import eval_name_template, Vector
+from data_diff.abcs.database_types import DbTime, DbPath
+from data_diff.databases import Database
+from data_diff.tracking import disable_tracking
+from data_diff.databases._connect import connect
+from data_diff.diff_tables import Algorithm
+from data_diff.hashdiff_tables import HashDiffer, DEFAULT_BISECTION_THRESHOLD, DEFAULT_BISECTION_FACTOR
+from data_diff.joindiff_tables import JoinDiffer, TABLE_WRITE_LIMIT
+from data_diff.table_segment import TableSegment
+from data_diff.utils import eval_name_template, Vector
 
 
 def connect_to_table(
@@ -32,10 +32,10 @@ def connect_to_table(
     if isinstance(key_columns, str):
         key_columns = (key_columns,)
 
-    db = connect(db_info, thread_count=thread_count)
+    db: Database = connect(db_info, thread_count=thread_count)
 
     if isinstance(table_name, str):
-        table_name = db.parse_table_name(table_name)
+        table_name = db.dialect.parse_table_name(table_name)
 
     return TableSegment(db, table_name, key_columns, **kwargs)
 
@@ -162,7 +162,8 @@ def diff_tables(
         )
     elif algorithm == Algorithm.JOINDIFF:
         if isinstance(materialize_to_table, str):
-            materialize_to_table = table1.database.parse_table_name(eval_name_template(materialize_to_table))
+            table_name = eval_name_template(materialize_to_table)
+            materialize_to_table = table1.database.dialect.parse_table_name(table_name)
         differ = JoinDiffer(
             threaded=threaded,
             max_threadpool_size=max_threadpool_size,

data_diff/__main__.py

@@ -6,26 +6,27 @@
 import json
 import logging
 from itertools import islice
-from typing import Dict, Optional
+from typing import Dict, Optional, Tuple
 
 import rich
 from rich.logging import RichHandler
 import click
 
-from data_diff.sqeleton.schema import create_schema
-from data_diff.sqeleton.queries.api import current_timestamp
+from data_diff import Database
+from data_diff.schema import create_schema
+from data_diff.queries.api import current_timestamp
 
-from .dbt import dbt_diff
-from .utils import eval_name_template, remove_password_from_url, safezip, match_like, LogStatusHandler
-from .diff_tables import Algorithm
-from .hashdiff_tables import HashDiffer, DEFAULT_BISECTION_THRESHOLD, DEFAULT_BISECTION_FACTOR
-from .joindiff_tables import TABLE_WRITE_LIMIT, JoinDiffer
-from .table_segment import TableSegment
-from .databases import connect
-from .parse_time import parse_time_before, UNITS_STR, ParseError
-from .config import apply_config_from_file
-from .tracking import disable_tracking, set_entrypoint_name
-from .version import __version__
+from data_diff.dbt import dbt_diff
+from data_diff.utils import eval_name_template, remove_password_from_url, safezip, match_like, LogStatusHandler
+from data_diff.diff_tables import Algorithm
+from data_diff.hashdiff_tables import HashDiffer, DEFAULT_BISECTION_THRESHOLD, DEFAULT_BISECTION_FACTOR
+from data_diff.joindiff_tables import TABLE_WRITE_LIMIT, JoinDiffer
+from data_diff.table_segment import TableSegment
+from data_diff.databases._connect import connect
+from data_diff.parse_time import parse_time_before, UNITS_STR, ParseError
+from data_diff.config import apply_config_from_file
+from data_diff.tracking import disable_tracking, set_entrypoint_name
+from data_diff.version import __version__
 
 
 COLOR_SCHEME = {
@@ -326,8 +327,7 @@ def main(conf, run, **kw):
             )
     except Exception as e:
         logging.error(e)
-        if kw["debug"]:
-            raise
+        raise
 
 
 def _data_diff(
@@ -425,7 +425,7 @@ def _data_diff(
             logging.error(f"Error while parsing age expression: {e}")
             return
 
-    dbs = db1, db2
+    dbs: Tuple[Database, Database] = db1, db2
 
     if interactive:
         for db in dbs:
@@ -444,7 +444,7 @@ def _data_diff(
             materialize_all_rows=materialize_all_rows,
             table_write_limit=table_write_limit,
             materialize_to_table=materialize_to_table
-            and db1.parse_table_name(eval_name_template(materialize_to_table)),
+            and db1.dialect.parse_table_name(eval_name_template(materialize_to_table)),
         )
     else:
         assert algorithm == Algorithm.HASHDIFF
@@ -456,11 +456,11 @@ def _data_diff(
         )
 
     table_names = table1, table2
-    table_paths = [db.parse_table_name(t) for db, t in safezip(dbs, table_names)]
+    table_paths = [db.dialect.parse_table_name(t) for db, t in safezip(dbs, table_names)]
 
     schemas = list(differ._thread_map(_get_schema, safezip(dbs, table_paths)))
     schema1, schema2 = schemas = [
-        create_schema(db, table_path, schema, case_sensitive)
+        create_schema(db.name, table_path, schema, case_sensitive)
         for db, table_path, schema in safezip(dbs, table_paths, schemas)
     ]
 

tests/sqeleton/__init__.py data_diff/abcs/__init__.pytests/sqeleton/__init__.py renamed to data_diff/abcs/__init__.py

@@ -0,0 +1,13 @@
+from abc import ABC
+
+import attrs
+
+
+@attrs.define(frozen=False)
+class AbstractCompiler(ABC):
+    pass
+
+
+@attrs.define(frozen=False, eq=False)
+class Compilable(ABC):
+    pass

data_diff/abcs/compiler.py

@@ -0,0 +1,194 @@
+import decimal
+from abc import ABC, abstractmethod
+from typing import Tuple, Union
+from datetime import datetime
+
+import attrs
+
+from data_diff.utils import ArithAlphanumeric, ArithUUID, Unknown
+
+
+DbPath = Tuple[str, ...]
+DbKey = Union[int, str, bytes, ArithUUID, ArithAlphanumeric]
+DbTime = datetime
+
+
+@attrs.define(frozen=True)
+class ColType:
+    @property
+    def supported(self) -> bool:
+        return True
+
+
+@attrs.define(frozen=True)
+class PrecisionType(ColType):
+    precision: int
+    rounds: Union[bool, Unknown] = Unknown
+
+
+@attrs.define(frozen=True)
+class Boolean(ColType):
+    precision = 0
+
+
+@attrs.define(frozen=True)
+class TemporalType(PrecisionType):
+    pass
+
+
+@attrs.define(frozen=True)
+class Timestamp(TemporalType):
+    pass
+
+
+@attrs.define(frozen=True)
+class TimestampTZ(TemporalType):
+    pass
+
+
+@attrs.define(frozen=True)
+class Datetime(TemporalType):
+    pass
+
+
+@attrs.define(frozen=True)
+class Date(TemporalType):
+    pass
+
+
+@attrs.define(frozen=True)
+class NumericType(ColType):
+    # 'precision' signifies how many fractional digits (after the dot) we want to compare
+    precision: int
+
+
+@attrs.define(frozen=True)
+class FractionalType(NumericType):
+    pass
+
+
+@attrs.define(frozen=True)
+class Float(FractionalType):
+    python_type = float
+
+
+@attrs.define(frozen=True)
+class IKey(ABC):
+    "Interface for ColType, for using a column as a key in table."
+
+    @property
+    @abstractmethod
+    def python_type(self) -> type:
+        "Return the equivalent Python type of the key"
+
+    def make_value(self, value):
+        return self.python_type(value)
+
+
+@attrs.define(frozen=True)
+class Decimal(FractionalType, IKey):  # Snowflake may use Decimal as a key
+    @property
+    def python_type(self) -> type:
+        if self.precision == 0:
+            return int
+        return decimal.Decimal
+
+
+@attrs.define(frozen=True)
+class StringType(ColType):
+    python_type = str
+
+
+@attrs.define(frozen=True)
+class ColType_UUID(ColType, IKey):
+    python_type = ArithUUID
+
+
+@attrs.define(frozen=True)
+class ColType_Alphanum(ColType, IKey):
+    python_type = ArithAlphanumeric
+
+
+@attrs.define(frozen=True)
+class Native_UUID(ColType_UUID):
+    pass
+
+
+@attrs.define(frozen=True)
+class String_UUID(ColType_UUID, StringType):
+    pass
+
+
+@attrs.define(frozen=True)
+class String_Alphanum(ColType_Alphanum, StringType):
+    @staticmethod
+    def test_value(value: str) -> bool:
+        try:
+            ArithAlphanumeric(value)
+            return True
+        except ValueError:
+            return False
+
+    def make_value(self, value):
+        return self.python_type(value)
+
+
+@attrs.define(frozen=True)
+class String_VaryingAlphanum(String_Alphanum):
+    pass
+
+
+@attrs.define(frozen=True)
+class String_FixedAlphanum(String_Alphanum):
+    length: int
+
+    def make_value(self, value):
+        if len(value) != self.length:
+            raise ValueError(f"Expected alphanumeric value of length {self.length}, but got '{value}'.")
+        return self.python_type(value, max_len=self.length)
+
+
+@attrs.define(frozen=True)
+class Text(StringType):
+    @property
+    def supported(self) -> bool:
+        return False
+
+
+# In majority of DBMSes, it is called JSON/JSONB. Only in Snowflake, it is OBJECT.
+@attrs.define(frozen=True)
+class JSON(ColType):
+    pass
+
+
+@attrs.define(frozen=True)
+class Array(ColType):
+    item_type: ColType
+
+
+# Unlike JSON, structs are not free-form and have a very specific set of fields and their types.
+# We do not parse & use those fields now, but we can do this later.
+# For example, in BigQuery:
+# - https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#struct_type
+# - https://cloud.google.com/bigquery/docs/reference/standard-sql/lexical#struct_literals
+@attrs.define(frozen=True)
+class Struct(ColType):
+    pass
+
+
+@attrs.define(frozen=True)
+class Integer(NumericType, IKey):
+    precision: int = 0
+    python_type: type = int
+
+    def __attrs_post_init__(self):
+        assert self.precision == 0
+
+
+@attrs.define(frozen=True)
+class UnknownColType(ColType):
+    text: str
+
+    @property
+    def supported(self) -> bool:
+        return False