MINOR: Fix Databricks SDK Breaking Change (#15037)

open-metadata · Feb 6, 2024 · 37bf91a · 37bf91a
1 parent ce5938a
commit 37bf91a
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 32 deletions.
diff --git a/ingestion/setup.py b/ingestion/setup.py
@@ -41,7 +41,7 @@
     "azure-storage-blob": "azure-storage-blob~=12.14",
     "azure-identity": "azure-identity~=1.12",
     "sqlalchemy-databricks": "sqlalchemy-databricks~=0.1",
-    "databricks-sdk": "databricks-sdk~=0.1",
+    "databricks-sdk": "databricks-sdk>=0.18.0,<0.20.0",
     "google": "google>=3.0.0",
     "trino": "trino[sqlalchemy]",
     "spacy": "spacy==3.5.0",

diff --git a/ingestion/src/metadata/ingestion/source/database/unitycatalog/metadata.py b/ingestion/src/metadata/ingestion/source/database/unitycatalog/metadata.py
@@ -13,11 +13,10 @@
 """
 import json
 import traceback
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
+from typing import Any, Iterable, List, Optional, Tuple, Union
 
 from databricks.sdk.service.catalog import ColumnInfo
 from databricks.sdk.service.catalog import TableConstraint as DBTableConstraint
-from databricks.sdk.service.catalog import TableConstraintList
 
 from metadata.generated.schema.api.data.createDatabase import CreateDatabaseRequest
 from metadata.generated.schema.api.data.createDatabaseSchema import (
@@ -75,19 +74,6 @@
 logger = ingestion_logger()
 
 
-# pylint: disable=not-callable
-@classmethod
-def from_dict(cls, dct: Dict[str, Any]) -> "TableConstraintList":
-    return cls(
-        table_constraints=[
-            DBTableConstraint.from_dict(constraint) for constraint in dct
-        ]
-    )
-
-
-TableConstraintList.from_dict = from_dict
-
-
 class UnitycatalogSource(DatabaseServiceSource, MultiDBSource):
     """
     Implements the necessary methods to extract
@@ -351,31 +337,30 @@ def yield_table(
             )
 
     def get_table_constraints(
-        self, constraints: TableConstraintList
+        self, constraints: List[DBTableConstraint]
     ) -> Tuple[List[TableConstraint], List[ForeignConstrains]]:
         """
         Function to handle table constraint for the current table and add it to context
         """
 
         primary_constraints = []
         foreign_constraints = []
-        if constraints and constraints.table_constraints:
-            for constraint in constraints.table_constraints:
-                if constraint.primary_key_constraint:
-                    primary_constraints.append(
-                        TableConstraint(
-                            constraintType=ConstraintType.PRIMARY_KEY,
-                            columns=constraint.primary_key_constraint.child_columns,
-                        )
+        for constraint in constraints:
+            if constraint.primary_key_constraint:
+                primary_constraints.append(
+                    TableConstraint(
+                        constraintType=ConstraintType.PRIMARY_KEY,
+                        columns=constraint.primary_key_constraint.child_columns,
                     )
-                if constraint.foreign_key_constraint:
-                    foreign_constraints.append(
-                        ForeignConstrains(
-                            child_columns=constraint.foreign_key_constraint.child_columns,
-                            parent_columns=constraint.foreign_key_constraint.parent_columns,
-                            parent_table=constraint.foreign_key_constraint.parent_table,
-                        )
+                )
+            if constraint.foreign_key_constraint:
+                foreign_constraints.append(
+                    ForeignConstrains(
+                        child_columns=constraint.foreign_key_constraint.child_columns,
+                        parent_columns=constraint.foreign_key_constraint.parent_columns,
+                        parent_table=constraint.foreign_key_constraint.parent_table,
                     )
+                )
         return primary_constraints, foreign_constraints
 
     def _get_foreign_constraints(self, foreign_columns) -> List[TableConstraint]: