Skip to content

Commit

Permalink
fix(ingest/databricks): Updating code to work with Databricks sdk 0.30 (
Browse files Browse the repository at this point in the history
  • Loading branch information
treff7es authored Aug 13, 2024
1 parent 9dc85cb commit 5e9188c
Show file tree
Hide file tree
Showing 3 changed files with 8 additions and 8 deletions.
1 change: 1 addition & 0 deletions docs-website/sidebars.js
Original file line number Diff line number Diff line change
Expand Up @@ -924,6 +924,7 @@ module.exports = {
// "docs/_api-guide-template"
// - "metadata-service/services/README"
// "metadata-ingestion/examples/structured_properties/README"
// "smoke-test/tests/openapi/README"
// ],
],
};
2 changes: 2 additions & 0 deletions metadata-ingestion/setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,8 @@

classification_lib = {
"acryl-datahub-classify==0.0.11",
# schwifty is needed for the classify plugin but in 2024.08.0 they broke the python 3.8 compatibility
"schwifty<2024.08.0",
# This is a bit of a hack. Because we download the SpaCy model at runtime in the classify plugin,
# we need pip to be available.
"pip",
Expand Down
Original file line number Diff line number Diff line change
@@ -1,15 +1,14 @@
import logging
import time
from typing import Optional, Union
from typing import Optional

from databricks.sdk import WorkspaceClient
from databricks.sdk.core import DatabricksError
from databricks.sdk.service._internal import Wait
from databricks.sdk.service.catalog import TableInfo
from databricks.sdk.service.sql import (
ExecuteStatementResponse,
GetStatementResponse,
GetWarehouseResponse,
StatementResponse,
StatementState,
StatementStatus,
)
Expand Down Expand Up @@ -125,7 +124,7 @@ def _should_retry_unsupported_column(

def _analyze_table(
self, ref: TableReference, include_columns: bool
) -> ExecuteStatementResponse:
) -> StatementResponse:
statement = f"ANALYZE TABLE {ref.schema}.{ref.table} COMPUTE STATISTICS"
if include_columns:
statement += " FOR ALL COLUMNS"
Expand All @@ -139,7 +138,7 @@ def _analyze_table(
return response

def _check_analyze_table_statement_status(
self, execute_response: ExecuteStatementResponse, max_wait_secs: int
self, execute_response: StatementResponse, max_wait_secs: int
) -> bool:
if not execute_response.statement_id or not execute_response.status:
return False
Expand Down Expand Up @@ -230,9 +229,7 @@ def _get_int(self, table_info: TableInfo, field: str) -> Optional[int]:
return None

@staticmethod
def _raise_if_error(
response: Union[ExecuteStatementResponse, GetStatementResponse], key: str
) -> None:
def _raise_if_error(response: StatementResponse, key: str) -> None:
if response.status and response.status.state in [
StatementState.FAILED,
StatementState.CANCELED,
Expand Down

0 comments on commit 5e9188c

Please sign in to comment.