Merge pull request #227 from datafold/sep6

Refactors and fixes

Author: erezsh

data_diff/__main__.py

@@ -11,7 +11,8 @@
 
 from .utils import remove_password_from_url, safezip, match_like
 from .diff_tables import TableDiffer, DEFAULT_BISECTION_THRESHOLD, DEFAULT_BISECTION_FACTOR
-from .table_segment import create_schema, TableSegment
+from .table_segment import TableSegment
+from .databases.database_types import create_schema
 from .databases.connect import connect
 from .parse_time import parse_time_before_now, UNITS_STR, ParseError
 from .config import apply_config_from_file

data_diff/databases/base.py

@@ -270,9 +270,13 @@ def offset_limit(self, offset: Optional[int] = None, limit: Optional[int] = None
         return f"LIMIT {limit}"
 
     def concat(self, l: List[str]) -> str:
+        assert len(l) > 1
         joined_exprs = ", ".join(l)
         return f"concat({joined_exprs})"
 
+    def is_distinct_from(self, a: str, b: str) -> str:
+        return f"{a} is distinct from {b}"
+
     def timestamp_value(self, t: DbTime) -> str:
         return f"'{t.isoformat()}'"
 

data_diff/databases/database_types.py

@@ -1,17 +1,20 @@
+import logging
 import decimal
 from abc import ABC, abstractmethod
 from typing import Sequence, Optional, Tuple, Union, Dict, List
 from datetime import datetime
 
 from runtype import dataclass
 
-from data_diff.utils import ArithAlphanumeric, ArithUUID, CaseAwareMapping
+from data_diff.utils import ArithAlphanumeric, ArithUUID, CaseAwareMapping, CaseInsensitiveDict, CaseSensitiveDict
 
 
 DbPath = Tuple[str, ...]
 DbKey = Union[int, str, bytes, ArithUUID, ArithAlphanumeric]
 DbTime = datetime
 
+logger = logging.getLogger("databases")
+
 
 class ColType:
     supported = True
@@ -129,6 +132,8 @@ class UnknownColType(ColType):
 
 
 class AbstractDatabase(ABC):
+    name: str
+
     @abstractmethod
     def quote(self, s: str):
         "Quote SQL name (implementation specific)"
@@ -144,6 +149,11 @@ def concat(self, l: List[str]) -> str:
         "Provide SQL for concatenating a bunch of column into a string"
         ...
 
+    @abstractmethod
+    def is_distinct_from(self, a: str, b: str) -> str:
+        "Provide SQL for a comparison where NULL = NULL is true"
+        ...
+
     @abstractmethod
     def timestamp_value(self, t: DbTime) -> str:
         "Provide SQL for the given timestamp value"
@@ -270,3 +280,15 @@ def _normalize_table_path(self, path: DbPath) -> DbPath:
 
 
 Schema = CaseAwareMapping
+
+
+def create_schema(db: AbstractDatabase, table_path: DbPath, schema: dict, case_sensitive: bool) -> CaseAwareMapping:
+    logger.debug(f"[{db.name}] Schema = {schema}")
+
+    if case_sensitive:
+        return CaseSensitiveDict(schema)
+
+    if len({k.lower() for k in schema}) < len(schema):
+        logger.warning(f'Ambiguous schema for {db}:{".".join(table_path)} | Columns = {", ".join(list(schema))}')
+        logger.warning("We recommend to disable case-insensitivity (set --case-sensitive).")
+    return CaseInsensitiveDict(schema)

data_diff/databases/mysql.py

@@ -69,3 +69,6 @@ def normalize_number(self, value: str, coltype: FractionalType) -> str:
 
     def normalize_uuid(self, value: str, coltype: ColType_UUID) -> str:
         return f"TRIM(CAST({value} AS char))"
+
+    def is_distinct_from(self, a: str, b: str) -> str:
+        return f"not ({a} <=> {b})"

data_diff/table_segment.py

@@ -4,29 +4,17 @@
 
 from runtype import dataclass
 
-from .utils import ArithString, split_space, CaseAwareMapping, CaseInsensitiveDict, CaseSensitiveDict
+from .utils import ArithString, split_space
 
 from .databases.base import Database
-from .databases.database_types import DbPath, DbKey, DbTime, Native_UUID, Schema
+from .databases.database_types import DbPath, DbKey, DbTime, Native_UUID, Schema, create_schema
 from .sql import Select, Checksum, Compare, Count, TableName, Time, Value
 
 logger = logging.getLogger("table_segment")
 
 RECOMMENDED_CHECKSUM_DURATION = 10
 
 
-def create_schema(db: Database, table_path: DbPath, schema: dict, case_sensitive: bool) -> CaseAwareMapping:
-    logger.debug(f"[{db.name}] Schema = {schema}")
-
-    if case_sensitive:
-        return CaseSensitiveDict(schema)
-
-    if len({k.lower() for k in schema}) < len(schema):
-        logger.warning(f'Ambiguous schema for {db}:{".".join(table_path)} | Columns = {", ".join(list(schema))}')
-        logger.warning("We recommend to disable case-insensitivity (remove --any-case).")
-    return CaseInsensitiveDict(schema)
-
-
 @dataclass
 class TableSegment:
     """Signifies a segment of rows (and selected columns) within a table

data_diff/utils.py

@@ -212,6 +212,7 @@ def items(self) -> Iterable[Tuple[str, V]]:
 
 class CaseSensitiveDict(dict, CaseAwareMapping):
     def get_key(self, key):
+        self[key]  # Throw KeyError is key doesn't exist
         return key
 
     def as_insensitive(self):

tests/test_diff_tables.py

@@ -521,6 +521,25 @@ def test_table_segment(self):
 
         self.assertRaises(ValueError, self.table.replace, min_key=10, max_key=0)
 
+    def test_case_awareness(self):
+        # create table
+        self.connection.query(f"create table {self.table_src}(id int, userid int, timestamp timestamp)", None)
+        _commit(self.connection)
+
+        # insert rows
+        cols = "id userid timestamp".split()
+        time = "2022-01-01 00:00:00.000000"
+        time_str = f"timestamp '{time}'"
+        _insert_rows(self.connection, self.table_src, cols, [[1, 9, time_str], [2, 2, time_str]])
+        _commit(self.connection)
+
+        res = tuple(self.table.replace(key_column="Id", case_sensitive=False).with_schema().query_key_range())
+        assert res == ("1", "2")
+
+        self.assertRaises(
+            KeyError, self.table.replace(key_column="Id", case_sensitive=True).with_schema().query_key_range
+        )
+
 
 @test_per_database
 class TestTableUUID(TestPerDatabase):