MINOR: Fix Databricks SDK Breaking Change (#15037)

This commit is contained in:
Mayur Singal 2024-02-06 10:42:53 +05:30 committed by GitHub
parent 22b7ac3a0b
commit d76809801d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 17 additions and 32 deletions

View File

@ -41,7 +41,7 @@ VERSIONS = {
"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",

View File

@ -13,11 +13,10 @@ Databricks Unity Catalog Source source methods.
"""
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 @@ from metadata.utils.logger import ingestion_logger
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 @@ class UnitycatalogSource(DatabaseServiceSource, MultiDBSource):
)
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 @@ class UnitycatalogSource(DatabaseServiceSource, MultiDBSource):
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]: