fix(ingest/databricks): Updating code to work with Databricks sdk 0.30 (#11158)

This commit is contained in:
Tamas Nemeth 2024-08-13 16:57:31 +02:00 committed by GitHub
parent 9dc85cb5bc
commit 5e9188ca2c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 8 additions and 8 deletions

View File

@ -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"
// ],
],
};

View File

@ -104,6 +104,8 @@ sqlglot_lib = {
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",

View File

@ -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,
)
@ -125,7 +124,7 @@ class UnityCatalogProxyProfilingMixin:
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"
@ -139,7 +138,7 @@ class UnityCatalogProxyProfilingMixin:
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
@ -230,9 +229,7 @@ class UnityCatalogProxyProfilingMixin:
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,