From 37bf91addcae6e35998be2bb136998b1e389ece1 Mon Sep 17 00:00:00 2001
From: Mayur Singal <39544459+ulixius9@users.noreply.github.com>
Date: Tue, 6 Feb 2024 10:42:53 +0530
Subject: [PATCH] MINOR: Fix Databricks SDK Breaking Change (#15037)

---
 ingestion/setup.py                            |  2 +-
 .../source/database/unitycatalog/metadata.py  | 47 +++++++------------
 2 files changed, 17 insertions(+), 32 deletions(-)

diff --git a/ingestion/setup.py b/ingestion/setup.py
index 2447b7b6b6b5..804f864e4e7e 100644
--- 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
index 5a1a0db5228b..d2c404c82524 100644
--- 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,7 +337,7 @@ 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
@@ -359,23 +345,22 @@ def get_table_constraints(
 
         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]: