collate-data-diff 0.11.12__tar.gz → 0.11.15__tar.gz

{collate_data_diff-0.11.12 → collate_data_diff-0.11.15}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: collate-data-diff
-Version: 0.11.12
+Version: 0.11.15
 Summary: Command-line tool and Python library to efficiently diff rows across two different databases.
 License: MIT
 License-File: LICENSE

{collate_data_diff-0.11.12 → collate_data_diff-0.11.15}/data_diff/databases/databricks.py

@@ -37,6 +37,18 @@ def import_databricks():
     return databricks
 
 
+@import_helper(text="You can install it using 'pip install databricks-sdk'")
+def import_databricks_sdk():
+    from databricks.sdk.core import Config, oauth_service_principal
+
+    return Config, oauth_service_principal
+
+
+AUTH_PAT = "pat"
+AUTH_OAUTH_M2M = "oauth-m2m"
+AUTH_AZURE_SP_M2M = "azure-sp-m2m"
+
+
 @attrs.define(frozen=False)
 class Dialect(BaseDialect):
     name = "Databricks"
@@ -133,6 +145,61 @@ class Databricks(ThreadedDatabase):
         self.default_schema = kw.get("schema", "default")
         self.catalog = kw.get("catalog", "hive_metastore")
 
+    def _require(self, key: str) -> str:
+        value = self._args.get(key)
+        if not value:
+            auth_method = self._args.get("auth_method", AUTH_PAT)
+            raise ValueError(f"{self.name}: '{key}' is required for auth_method '{auth_method}'")
+        return value
+
+    def _pat_kwargs(self) -> Dict[str, Any]:
+        return {"access_token": self._require("access_token")}
+
+    def _oauth_m2m_kwargs(self) -> Dict[str, Any]:
+        Config, oauth_service_principal = import_databricks_sdk()
+        hostname = self._args["server_hostname"]
+        client_id = self._require("databricks_client_id")
+        client_secret = self._require("databricks_client_secret")
+
+        def credentials_provider():
+            provider = oauth_service_principal(
+                Config(host=f"https://{hostname}", client_id=client_id, client_secret=client_secret)
+            )
+            if provider is None:
+                raise ConnectionError(f"{self.name}: OAuth is not available on {hostname}")
+            return provider
+
+        return {"credentials_provider": credentials_provider}
+
+    def _azure_sp_m2m_kwargs(self) -> Dict[str, Any]:
+        return {
+            "auth_type": AUTH_AZURE_SP_M2M,
+            "azure_client_id": self._require("azure_client_id"),
+            "azure_client_secret": self._require("azure_client_secret"),
+            "azure_tenant_id": self._require("azure_tenant_id"),
+        }
+
+    def _auth_kwargs(self) -> Dict[str, Any]:
+        """Credential kwargs for ``databricks.sql.connect``.
+
+        Raises:
+            ValueError: on an unknown ``auth_method`` or a missing credential.
+        """
+        # Callables cannot travel in the connection config: `Connect` caches on
+        # json.dumps of it. Providers are built here, from plain strings.
+        auth_method = self._args.get("auth_method", AUTH_PAT)
+        builder = {
+            AUTH_PAT: self._pat_kwargs,
+            AUTH_OAUTH_M2M: self._oauth_m2m_kwargs,
+            AUTH_AZURE_SP_M2M: self._azure_sp_m2m_kwargs,
+        }.get(auth_method)
+        if builder is None:
+            raise ValueError(
+                f"{self.name}: unknown auth_method '{auth_method}'. "
+                f"Expected one of {AUTH_PAT}, {AUTH_OAUTH_M2M}, {AUTH_AZURE_SP_M2M}"
+            )
+        return builder()
+
     def create_connection(self):
         databricks = import_databricks()
 
@@ -140,8 +207,8 @@ class Databricks(ThreadedDatabase):
             return databricks.sql.connect(
                 server_hostname=self._args["server_hostname"],
                 http_path=self._args["http_path"],
-                access_token=self._args["access_token"],
                 catalog=self.catalog,
+                **self._auth_kwargs(),
             )
         except databricks.sql.exc.Error as e:
             raise ConnectionError(*e.args) from e

{collate_data_diff-0.11.12 → collate_data_diff-0.11.15}/data_diff/databases/snowflake.py

@@ -160,6 +160,7 @@ class Snowflake(Database):
                 format=serialization.PrivateFormat.PKCS8,
                 encryption_algorithm=serialization.NoEncryption(),
             )
+            kw.pop("private_key_passphrase", None)
 
         self._conn = snowflake.connector.connect(schema=f'"{schema}"', **kw)
 

{collate_data_diff-0.11.12 → collate_data_diff-0.11.15}/data_diff/diff_tables.py

@@ -11,13 +11,13 @@ from concurrent.futures import ThreadPoolExecutor, as_completed
 
 import attrs
 
-from data_diff.errors import DataDiffMismatchingKeyTypesError
+from data_diff.errors import DataDiffMismatchingKeyTypesError, DataDiffUnsupportedKeyValueError
 from data_diff.info_tree import InfoTree, SegmentInfo
 from data_diff.utils import dbt_diff_string_template, run_as_daemon, safezip, getLogger, truncate_error, Vector
 from data_diff.thread_utils import ThreadedYielder
 from data_diff.table_segment import TableSegment, create_mesh_from_points
 from data_diff.tracking import create_end_event_json, create_start_event_json, send_event_json, is_tracking_enabled
-from data_diff.abcs.database_types import IKey
+from data_diff.abcs.database_types import IKey, ColType_Alphanum
 
 logger = getLogger(__name__)
 
@@ -309,7 +309,7 @@ class TableDiffer(ThreadBase, ABC):
         key_ranges = self._threaded_call_as_completed("query_key_range", [table1, table2])
 
         # Start with the first completed value, so we don't waste time waiting
-        min_key1, max_key1 = self._parse_key_range_result(key_types1, next(key_ranges))
+        min_key1, max_key1 = self._parse_key_range_result(key_types1, next(key_ranges), table1.key_columns)
 
         btable1 = table1.new_key_bounds(min_key=min_key1, max_key=max_key1, key_types=key_types1)
         btable2 = table2.new_key_bounds(min_key=min_key1, max_key=max_key1, key_types=key_types2)
@@ -337,7 +337,7 @@ class TableDiffer(ThreadBase, ABC):
         # Overall, the max number of new regions in this 2nd pass is 3^|k| - 1
 
         # Note: python types can be the same, but the rendering parameters (e.g. casing) can differ.
-        min_key2, max_key2 = self._parse_key_range_result(key_types2, next(key_ranges))
+        min_key2, max_key2 = self._parse_key_range_result(key_types2, next(key_ranges), table2.key_columns)
 
         points = [list(sorted(p)) for p in safezip(min_key1, min_key2, max_key1, max_key2)]
         box_mesh = create_mesh_from_points(*points)
@@ -351,18 +351,48 @@ class TableDiffer(ThreadBase, ABC):
 
         return ti
 
-    def _parse_key_range_result(self, key_types, key_range) -> Tuple[Vector, Vector]:
+    def _parse_key_range_result(self, key_types, key_range, key_columns=None) -> Tuple[Vector, Vector]:
         min_key_values, max_key_values = key_range
 
         # We add 1 because our ranges are exclusive of the end (like in Python)
         try:
             min_key = Vector(key_type.make_value(mn) for key_type, mn in safezip(key_types, min_key_values))
             max_key = Vector(key_type.make_value(mx) + 1 for key_type, mx in safezip(key_types, max_key_values))
-        except (TypeError, ValueError) as e:
+        except ValueError as e:
+            # A key value could not be mapped to a bisectable range. For an
+            # alphanumeric string key this means its min/max contain characters
+            # outside data-diff's bisection alphabet (e.g. '.', accents, other
+            # non-ASCII) — such keys can't be bisected safely, so we surface an
+            # actionable message. Other value errors (e.g. a malformed UUID) keep
+            # the original generic error.
+            raise self._key_range_error(key_types, min_key_values, max_key_values, key_columns) from e
+        except TypeError as e:
             raise type(e)(f"Cannot apply {key_types} to '{min_key_values}', '{max_key_values}'.") from e
 
         return min_key, max_key
 
+    @staticmethod
+    def _key_range_error(key_types, min_key_values, max_key_values, key_columns) -> ValueError:
+        names = list(key_columns) if key_columns else [f"#{i}" for i in range(len(key_types))]
+        offenders = []
+        for name, key_type, mn, mx in safezip(names, key_types, min_key_values, max_key_values):
+            if not isinstance(key_type, ColType_Alphanum):
+                continue  # only alphanum keys get the tailored, actionable guidance
+            for bound, value in (("min", mn), ("max", mx)):
+                try:
+                    key_type.make_value(value)
+                except ValueError as inner:
+                    offenders.append(f"column '{name}' ({bound}={value!r}): {inner}")
+        if offenders:
+            return DataDiffUnsupportedKeyValueError(
+                f"Cannot bisect on key {'; '.join(offenders)}. data-diff can only bisect string keys built "
+                "from spaces, hyphens, underscores, digits and unaccented ASCII letters (a-z, A-Z); keys "
+                "containing other characters (e.g. '.', accented or non-ASCII characters) cannot be mapped to "
+                "a comparable range. Use a numeric or UUID key column for this diff instead."
+            )
+        # Non-alphanum failure (e.g. a malformed UUID): preserve the original error.
+        return ValueError(f"Cannot apply {key_types} to '{min_key_values}', '{max_key_values}'.")
+
     def _bisect_and_diff_segments(
         self,
         ti: ThreadedYielder,

{collate_data_diff-0.11.12 → collate_data_diff-0.11.15}/data_diff/errors.py

@@ -72,3 +72,9 @@ class DataDiffSimpleSelectNotFound(Exception):
 
 class DataDiffMismatchingKeyTypesError(Exception):
     "Raised when the key types of two tables do not match, like VARCHAR and INT."
+
+
+# Subclasses ValueError because it replaces the ValueError previously raised from
+# _parse_key_range_result — keeps the exception contract for existing callers.
+class DataDiffUnsupportedKeyValueError(ValueError):
+    "Raised when a key column holds values that cannot be mapped to a bisectable range (e.g. text keys with '.', accents or other non-ASCII characters)."

{collate_data_diff-0.11.12 → collate_data_diff-0.11.15}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "collate-data-diff"
-version = "0.11.12"
+version = "0.11.15"
 description = "Command-line tool and Python library to efficiently diff rows across two different databases."
 authors = ["Collate <info@getcollate.io>", "Datafold <data-diff@datafold.com>"]
 license = "MIT"