MINOR: Fix Databricks SDK Breaking Change (#15037)

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",

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]: