Merge branch 'master' into fix-numeric-precision-recognition-bq-pg

Author: vvkh

.github/workflows/formatter.yml

@@ -21,13 +21,25 @@ jobs:
         uses: actions/checkout@v3
         if: github.event_name == 'workflow_dispatch'
 
-      - name: Check files using the ruff formatter
+      # This is used for forked PRs as write permissions are required to format files
+      - name: Run and commit changes with `ruff format .` locally on your forked branch to fix errors if they appear
+        if: ${{ github.event.pull_request.head.repo.fork == true }}
+        uses: chartboost/ruff-action@v1
+        id: ruff_formatter_suggestions
+        with:
+          args: format --diff
+
+      # This only runs if the PR is NOT from a forked repo
+      - name: Format files using ruff
+        if: ${{ github.event.pull_request.head.repo.fork == false }}
         uses: chartboost/ruff-action@v1
         id: ruff_formatter
         with:
           args: format
 
+      # This only runs if the PR is NOT from a forked repo
       - name: Auto commit ruff formatting
+        if: ${{ github.event.pull_request.head.repo.fork == false }}
         uses: stefanzweifel/git-auto-commit-action@v5
         with:
-          commit_message: 'style fixes by ruff'
+          commit_message: 'style fixes by ruff'

data_diff/__main__.py

@@ -12,8 +12,8 @@
 from rich.logging import RichHandler
 import click
 
-from data_diff import Database
-from data_diff.schema import create_schema
+from data_diff import Database, DbPath
+from data_diff.schema import RawColumnInfo, create_schema
 from data_diff.queries.api import current_timestamp
 
 from data_diff.dbt import dbt_diff
@@ -72,7 +72,7 @@ def _remove_passwords_in_dict(d: dict) -> None:
             d[k] = remove_password_from_url(v)
 
 
-def _get_schema(pair):
+def _get_schema(pair: Tuple[Database, DbPath]) -> Dict[str, RawColumnInfo]:
     db, table_path = pair
     return db.query_table_schema(table_path)
 

data_diff/abcs/database_types.py

@@ -1,6 +1,6 @@
 import decimal
 from abc import ABC, abstractmethod
-from typing import List, Optional, Tuple, Type, TypeVar, Union
+from typing import Collection, List, Optional, Tuple, Type, TypeVar, Union
 from datetime import datetime
 
 import attrs
@@ -15,6 +15,91 @@
 N = TypeVar("N")
 
 
+@attrs.frozen(kw_only=True, eq=False, order=False, unsafe_hash=True)
+class Collation:
+    """
+    A pre-parsed or pre-known record about db collation, per column.
+
+    The "greater" collation should be used as a target collation for textual PKs
+    on both sides of the diff — by coverting the "lesser" collation to self.
+
+    Snowflake easily absorbs the performance losses, so it has a boost to always
+    be greater than any other collation in non-Snowflake databases.
+    Other databases need to negotiate which side absorbs the performance impact.
+    """
+
+    # A boost for special databases that are known to absorb the performance dmaage well.
+    absorbs_damage: bool = False
+
+    # Ordinal soring by ASCII/UTF8 (True), or alphabetic as per locale/country/etc (False).
+    ordinal: Optional[bool] = None
+
+    # Lowercase first (aAbBcC or abcABC). Otherwise, uppercase first (AaBbCc or ABCabc).
+    lower_first: Optional[bool] = None
+
+    # 2-letter lower-case locale and upper-case country codes, e.g. en_US. Ignored for ordinals.
+    language: Optional[str] = None
+    country: Optional[str] = None
+
+    # There are also space-, punctuation-, width-, kana-(in)sensitivity, so on.
+    # Ignore everything not related to xdb alignment. Only case- & accent-sensitivity are common.
+    case_sensitive: Optional[bool] = None
+    accent_sensitive: Optional[bool] = None
+
+    # Purely informational, for debugging:
+    _source: Union[None, str, Collection[str]] = None
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Collation):
+            return NotImplemented
+        if self.ordinal and other.ordinal:
+            # TODO: does it depend on language? what does Albanic_BIN mean in MS SQL?
+            return True
+        return (
+            self.language == other.language
+            and (self.country is None or other.country is None or self.country == other.country)
+            and self.case_sensitive == other.case_sensitive
+            and self.accent_sensitive == other.accent_sensitive
+            and self.lower_first == other.lower_first
+        )
+
+    def __ne__(self, other: object) -> bool:
+        if not isinstance(other, Collation):
+            return NotImplemented
+        return not self.__eq__(other)
+
+    def __gt__(self, other: object) -> bool:
+        if not isinstance(other, Collation):
+            return NotImplemented
+        if self == other:
+            return False
+        if self.absorbs_damage and not other.absorbs_damage:
+            return False
+        if other.absorbs_damage and not self.absorbs_damage:
+            return True  # this one is preferred if it cannot absorb damage as its counterpart can
+        if self.ordinal and not other.ordinal:
+            return True
+        if other.ordinal and not self.ordinal:
+            return False
+        # TODO: try to align the languages & countries?
+        return False
+
+    def __ge__(self, other: object) -> bool:
+        if not isinstance(other, Collation):
+            return NotImplemented
+        return self == other or self.__gt__(other)
+
+    def __lt__(self, other: object) -> bool:
+        if not isinstance(other, Collation):
+            return NotImplemented
+        return self != other and not self.__gt__(other)
+
+    def __le__(self, other: object) -> bool:
+        if not isinstance(other, Collation):
+            return NotImplemented
+        return self == other or not self.__gt__(other)
+
+
 @attrs.define(frozen=True, kw_only=True)
 class ColType:
     # Arbitrary metadata added and fetched at runtime.
@@ -97,6 +182,8 @@ def python_type(self) -> type:
         "Return the equivalent Python type of the key"
 
     def make_value(self, value):
+        if isinstance(value, self.python_type):
+            return value
         return self.python_type(value)
 
 
@@ -112,6 +199,7 @@ def python_type(self) -> type:
 @attrs.define(frozen=True)
 class StringType(ColType):
     python_type = str
+    collation: Optional[Collation] = attrs.field(default=None, kw_only=True)
 
 
 @attrs.define(frozen=True)
@@ -131,7 +219,14 @@ class Native_UUID(ColType_UUID):
 
 @attrs.define(frozen=True)
 class String_UUID(ColType_UUID, StringType):
-    pass
+    # Case is important for UUIDs stored as regular string, not native UUIDs stored as numbers.
+    # We slice them internally as numbers, but render them back to SQL as lower/upper case.
+    # None means we do not know for sure, behave as with False, but it might be unreliable.
+    lowercase: Optional[bool] = None
+    uppercase: Optional[bool] = None
+
+    def make_value(self, v: str) -> ArithUUID:
+        return self.python_type(v, lowercase=self.lowercase, uppercase=self.uppercase)
 
 
 @attrs.define(frozen=True)
@@ -144,9 +239,6 @@ def test_value(value: str) -> bool:
         except ValueError:
             return False
 
-    def make_value(self, value):
-        return self.python_type(value)
-
 
 @attrs.define(frozen=True)
 class String_VaryingAlphanum(String_Alphanum):
@@ -158,6 +250,8 @@ class String_FixedAlphanum(String_Alphanum):
     length: int
 
     def make_value(self, value):
+        if isinstance(value, self.python_type):
+            return 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)

data_diff/databases/base.py

@@ -19,7 +19,8 @@
 
 from data_diff.abcs.compiler import AbstractCompiler, Compilable
 from data_diff.queries.extras import ApplyFuncAndNormalizeAsString, Checksum, NormalizeAsString
-from data_diff.utils import ArithString, is_uuid, join_iter, safezip
+from data_diff.schema import RawColumnInfo
+from data_diff.utils import ArithString, ArithUUID, is_uuid, join_iter, safezip
 from data_diff.queries.api import Expr, table, Select, SKIP, Explain, Code, this
 from data_diff.queries.ast_classes import (
     Alias,
@@ -248,6 +249,9 @@ def _compile(self, compiler: Compiler, elem) -> str:
             return self.timestamp_value(elem)
         elif isinstance(elem, bytes):
             return f"b'{elem.decode()}'"
+        elif isinstance(elem, ArithUUID):
+            s = f"'{elem.uuid}'"
+            return s.upper() if elem.uppercase else s.lower() if elem.lowercase else s
         elif isinstance(elem, ArithString):
             return f"'{elem}'"
         assert False, elem
@@ -681,8 +685,10 @@ def _constant_value(self, v):
             return f"'{v}'"
         elif isinstance(v, datetime):
             return self.timestamp_value(v)
-        elif isinstance(v, UUID):
+        elif isinstance(v, UUID):  # probably unused anymore in favour of ArithUUID
             return f"'{v}'"
+        elif isinstance(v, ArithUUID):
+            return f"'{v.uuid}'"
         elif isinstance(v, decimal.Decimal):
             return str(v)
         elif isinstance(v, bytearray):
@@ -708,27 +714,18 @@ def type_repr(self, t) -> str:
             datetime: "TIMESTAMP",
         }[t]
 
-    def _parse_type_repr(self, type_repr: str) -> Optional[Type[ColType]]:
-        return self.TYPE_CLASSES.get(type_repr)
-
-    def parse_type(
-        self,
-        table_path: DbPath,
-        col_name: str,
-        type_repr: str,
-        datetime_precision: int = None,
-        numeric_precision: int = None,
-        numeric_scale: int = None,
-    ) -> ColType:
+    def parse_type(self, table_path: DbPath, info: RawColumnInfo) -> ColType:
         "Parse type info as returned by the database"
 
-        cls = self._parse_type_repr(type_repr)
+        cls = self.TYPE_CLASSES.get(info.data_type)
         if cls is None:
-            return UnknownColType(type_repr)
+            return UnknownColType(info.data_type)
 
         if issubclass(cls, TemporalType):
             return cls(
-                precision=datetime_precision if datetime_precision is not None else DEFAULT_DATETIME_PRECISION,
+                precision=info.datetime_precision
+                if info.datetime_precision is not None
+                else DEFAULT_DATETIME_PRECISION,
                 rounds=self.ROUNDS_ON_PREC_LOSS,
             )
 
@@ -739,22 +736,22 @@ def parse_type(
             return cls()
 
         elif issubclass(cls, Decimal):
-            if numeric_scale is None:
-                numeric_scale = 0  # Needed for Oracle.
-            return cls(precision=numeric_scale)
+            if info.numeric_scale is None:
+                return cls(precision=0)  # Needed for Oracle.
+            return cls(precision=info.numeric_scale)
 
         elif issubclass(cls, Float):
             # assert numeric_scale is None
             return cls(
                 precision=self._convert_db_precision_to_digits(
-                    numeric_precision if numeric_precision is not None else DEFAULT_NUMERIC_PRECISION
+                    info.numeric_precision if info.numeric_precision is not None else DEFAULT_NUMERIC_PRECISION
                 )
             )
 
         elif issubclass(cls, (JSON, Array, Struct, Text, Native_UUID)):
             return cls()
 
-        raise TypeError(f"Parsing {type_repr} returned an unknown type '{cls}'.")
+        raise TypeError(f"Parsing {info.data_type} returned an unknown type {cls!r}.")
 
     def _convert_db_precision_to_digits(self, p: int) -> int:
         """Convert from binary precision, used by floats, to decimal precision."""
@@ -1019,7 +1016,7 @@ def select_table_schema(self, path: DbPath) -> str:
             f"WHERE table_name = '{name}' AND table_schema = '{schema}'"
         )
 
-    def query_table_schema(self, path: DbPath) -> Dict[str, tuple]:
+    def query_table_schema(self, path: DbPath) -> Dict[str, RawColumnInfo]:
         """Query the table for its schema for table in 'path', and return {column: tuple}
         where the tuple is (table_name, col_name, type_repr, datetime_precision?, numeric_precision?, numeric_scale?)
 
@@ -1030,7 +1027,17 @@ def query_table_schema(self, path: DbPath) -> Dict[str, tuple]:
         if not rows:
             raise RuntimeError(f"{self.name}: Table '{'.'.join(path)}' does not exist, or has no columns")
 
-        d = {r[0]: r for r in rows}
+        d = {
+            r[0]: RawColumnInfo(
+                column_name=r[0],
+                data_type=r[1],
+                datetime_precision=r[2],
+                numeric_precision=r[3],
+                numeric_scale=r[4],
+                collation_name=r[5] if len(r) > 5 else None,
+            )
+            for r in rows
+        }
         assert len(d) == len(rows)
         return d
 
@@ -1052,7 +1059,11 @@ def query_table_unique_columns(self, path: DbPath) -> List[str]:
         return list(res)
 
     def _process_table_schema(
-        self, path: DbPath, raw_schema: Dict[str, tuple], filter_columns: Sequence[str] = None, where: str = None
+        self,
+        path: DbPath,
+        raw_schema: Dict[str, RawColumnInfo],
+        filter_columns: Sequence[str] = None,
+        where: str = None,
     ):
         """Process the result of query_table_schema().
 
@@ -1068,7 +1079,7 @@ def _process_table_schema(
             accept = {i.lower() for i in filter_columns}
             filtered_schema = {name: row for name, row in raw_schema.items() if name.lower() in accept}
 
-        col_dict = {row[0]: self.dialect.parse_type(path, *row) for _name, row in filtered_schema.items()}
+        col_dict = {info.column_name: self.dialect.parse_type(path, info) for info in filtered_schema.values()}
 
         self._refine_coltypes(path, col_dict, where)
 
@@ -1077,15 +1088,15 @@ def _process_table_schema(
 
     def _refine_coltypes(
         self, table_path: DbPath, col_dict: Dict[str, ColType], where: Optional[str] = None, sample_size=64
-    ):
+    ) -> Dict[str, ColType]:
         """Refine the types in the column dict, by querying the database for a sample of their values
 
         'where' restricts the rows to be sampled.
         """
 
         text_columns = [k for k, v in col_dict.items() if isinstance(v, Text)]
         if not text_columns:
-            return
+            return col_dict
 
         fields = [Code(self.dialect.normalize_uuid(self.dialect.quote(c), String_UUID())) for c in text_columns]
 
@@ -1105,7 +1116,10 @@ def _refine_coltypes(
                     )
                 else:
                     assert col_name in col_dict
-                    col_dict[col_name] = String_UUID()
+                    col_dict[col_name] = String_UUID(
+                        lowercase=all(s == s.lower() for s in uuid_samples),
+                        uppercase=all(s == s.upper() for s in uuid_samples),
+                    )
                     continue
 
             if self.SUPPORTS_ALPHANUMS:  # Anything but MySQL (so far)
@@ -1117,7 +1131,9 @@ def _refine_coltypes(
                         )
                     else:
                         assert col_name in col_dict
-                        col_dict[col_name] = String_VaryingAlphanum()
+                        col_dict[col_name] = String_VaryingAlphanum(collation=col_dict[col_name].collation)
+
+        return col_dict
 
     def _normalize_table_path(self, path: DbPath) -> DbPath:
         if len(path) == 1: