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