Now tests for unique key constraints (if possible) instead of always actively validating (+ tests)

Author: erezsh

data_diff/databases/base.py

@@ -108,6 +108,7 @@ class Database(AbstractDatabase):
     TYPE_CLASSES: Dict[str, type] = {}
     default_schema: str = None
     SUPPORTS_ALPHANUMS = True
+    SUPPORTS_PRIMARY_KEY = False
 
     _interactive = False
 
@@ -232,6 +233,20 @@ def query_table_schema(self, path: DbPath) -> Dict[str, tuple]:
         assert len(d) == len(rows)
         return d
 
+
+    def select_table_unique_columns(self, path: DbPath) -> str:
+        schema, table = self._normalize_table_path(path)
+
+        return (
+            "SELECT column_name "
+            "FROM information_schema.key_column_usage "
+            f"WHERE table_name = '{table}' AND table_schema = '{schema}'"
+        )
+
+    def query_table_unique_columns(self, path: DbPath) -> List[str]:
+        res = self.query(self.select_table_unique_columns(path), List[str])
+        return list(res)
+
     def _process_table_schema(
         self, path: DbPath, raw_schema: Dict[str, tuple], filter_columns: Sequence[str], where: str = None
     ):

data_diff/databases/bigquery.py

@@ -1,4 +1,4 @@
-from typing import Union
+from typing import List, Union
 from .database_types import Timestamp, Datetime, Integer, Decimal, Float, Text, DbPath, FractionalType, TemporalType
 from .base import Database, import_helper, parse_table_name, ConnectError, apply_query
 from .base import TIMESTAMP_PRECISION_POS, ThreadLocalInterpreter
@@ -78,6 +78,9 @@ def select_table_schema(self, path: DbPath) -> str:
             f"WHERE table_name = '{table}' AND table_schema = '{schema}'"
         )
 
+    def query_table_unique_columns(self, path: DbPath) -> List[str]:
+        return []
+
     def normalize_timestamp(self, value: str, coltype: TemporalType) -> str:
         if coltype.rounds:
             timestamp = f"timestamp_micros(cast(round(unix_micros(cast({value} as timestamp))/1000000, {coltype.precision})*1000000 as int))"

data_diff/databases/database_types.py

@@ -253,6 +253,17 @@ def query_table_schema(self, path: DbPath) -> Dict[str, tuple]:
         """
         ...
 
+    @abstractmethod
+    def select_table_unique_columns(self, path: DbPath) -> str:
+        "Provide SQL for selecting the names of unique columns in the table"
+        ...
+
+    @abstractmethod
+    def query_table_unique_columns(self, path: DbPath) -> List[str]:
+        """Query the table for its unique columns for table in 'path', and return {column}
+        """
+        ...
+
     @abstractmethod
     def _process_table_schema(
         self, path: DbPath, raw_schema: Dict[str, tuple], filter_columns: Sequence[str], where: str = None

data_diff/databases/mysql.py

@@ -39,6 +39,7 @@ class MySQL(ThreadedDatabase):
     }
     ROUNDS_ON_PREC_LOSS = True
     SUPPORTS_ALPHANUMS = False
+    SUPPORTS_PRIMARY_KEY = True
 
     def __init__(self, *, thread_count, **kw):
         self._args = kw

data_diff/databases/oracle.py

@@ -39,6 +39,7 @@ class Oracle(ThreadedDatabase):
         "VARCHAR2": Text,
     }
     ROUNDS_ON_PREC_LOSS = True
+    SUPPORTS_PRIMARY_KEY = True
 
     def __init__(self, *, host, database, thread_count, **kw):
         self.kwargs = dict(dsn=f"{host}/{database}" if database else host, **kw)

data_diff/databases/postgresql.py

@@ -46,6 +46,7 @@ class PostgreSQL(ThreadedDatabase):
         "uuid": Native_UUID,
     }
     ROUNDS_ON_PREC_LOSS = True
+    SUPPORTS_PRIMARY_KEY = True
 
     default_schema = "public"
 

data_diff/databases/snowflake.py

@@ -1,4 +1,4 @@
-from typing import Union
+from typing import Union, List
 import logging
 
 from .database_types import Timestamp, TimestampTZ, Decimal, Float, Text, FractionalType, TemporalType, DbPath
@@ -95,3 +95,6 @@ def is_autocommit(self) -> bool:
 
     def explain_as_text(self, query: str) -> str:
         return f"EXPLAIN USING TEXT {query}"
+
+    def query_table_unique_columns(self, path: DbPath) -> List[str]:
+        return []

data_diff/joindiff_tables.py

@@ -190,18 +190,24 @@ def _diff_segments(
                 if not is_xa:
                     yield "+", tuple(b_row)
 
-    def _test_duplicate_keys(self, table1, table2):
+    def _test_duplicate_keys(self, table1: TableSegment, table2: TableSegment):
         logger.debug("Testing for duplicate keys")
 
         # Test duplicate keys
         for ts in [table1, table2]:
+            unique = ts.database.query_table_unique_columns(ts.table_path)
+
             t = ts.make_select()
             key_columns = ts.key_columns
 
-            q = t.select(total=Count(), total_distinct=Count(Concat(this[key_columns]), distinct=True))
-            total, total_distinct = ts.database.query(q, tuple)
-            if total != total_distinct:
-                raise ValueError("Duplicate primary keys")
+            unvalidated = list(set(key_columns) - set(unique))
+            if unvalidated:
+                # Validate that there are no duplicate keys
+                self.stats['validated_unique_keys'] = self.stats.get('validated_unique_keys', []) + [unvalidated]
+                q = t.select(total=Count(), total_distinct=Count(Concat(this[unvalidated]), distinct=True))
+                total, total_distinct = ts.database.query(q, tuple)
+                if total != total_distinct:
+                    raise ValueError("Duplicate primary keys")
 
     def _test_null_keys(self, table1, table2):
         logger.debug("Testing for null keys")

data_diff/queries/ast_classes.py

@@ -299,12 +299,12 @@ class TablePath(ExprNode, ITable):
     path: DbPath
     schema: Optional[Schema] = field(default=None, repr=False)
 
-    def create(self, source_table: ITable = None, *, if_not_exists=False):
+    def create(self, source_table: ITable = None, *, if_not_exists=False, primary_keys=None):
         if source_table is None and not self.schema:
             raise ValueError("Either schema or source table needed to create table")
         if isinstance(source_table, TablePath):
             source_table = source_table.select()
-        return CreateTable(self, source_table, if_not_exists=if_not_exists)
+        return CreateTable(self, source_table, if_not_exists=if_not_exists, primary_keys=primary_keys)
 
     def drop(self, if_exists=False):
         return DropTable(self, if_exists=if_exists)
@@ -647,14 +647,16 @@ class CreateTable(Statement):
     path: TablePath
     source_table: Expr = None
     if_not_exists: bool = False
+    primary_keys: List[str] = None
 
     def compile(self, c: Compiler) -> str:
         ne = "IF NOT EXISTS " if self.if_not_exists else ""
         if self.source_table:
             return f"CREATE TABLE {ne}{c.compile(self.path)} AS {c.compile(self.source_table)}"
 
         schema = ", ".join(f"{c.database.quote(k)} {c.database.type_repr(v)}" for k, v in self.path.schema.items())
-        return f"CREATE TABLE {ne}{c.compile(self.path)}({schema})"
+        pks = ", PRIMARY KEY (%s)" % ', '.join(self.primary_keys) if self.primary_keys and c.database.SUPPORTS_PRIMARY_KEY else ""
+        return f"CREATE TABLE {ne}{c.compile(self.path)}({schema}{pks})"
 
 
 @dataclass

tests/common.py

@@ -150,6 +150,7 @@ def tearDown(self):
 
 
 def _parameterized_class_per_conn(test_databases):
+    test_databases = set(test_databases)
     names = [(cls.__name__, cls) for cls in CONN_STRINGS if cls in test_databases]
     return parameterized_class(("name", "db_cls"), names)