Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

MINOR: Fix Databricks SDK Breaking Change #15037

Merged
merged 4 commits into from
Feb 6, 2024
Merged

MINOR: Fix Databricks SDK Breaking Change #15037

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
4b76bc4
MINOR: Fix Databricks SDK Breaking Change
ulixius9 Feb 5, 2024
2cd4c00
databrick dependency fix
ulixius9 Feb 5, 2024
944da9b
Fix Version Range
ulixius9 Feb 5, 2024
00c06f9
Merge branch 'main' into databricks_fix_version
ulixius9 Feb 6, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion ingestion/setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
47 changes: 16 additions & 31 deletions ingestion/src/metadata/ingestion/source/database/unitycatalog/metadata.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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]:
Expand Down
Loading