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" // "docs/_api-guide-template"
// - "metadata-service/services/README" // - "metadata-service/services/README"
// "metadata-ingestion/examples/structured_properties/README" // "metadata-ingestion/examples/structured_properties/README"
// "smoke-test/tests/openapi/README"
// ], // ],
], ],
}; };

View File

@ -104,6 +104,8 @@ sqlglot_lib = {
classification_lib = { classification_lib = {
"acryl-datahub-classify==0.0.11", "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, # 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. # we need pip to be available.
"pip", "pip",

View File

@ -1,15 +1,14 @@
import logging import logging
import time import time
from typing import Optional, Union from typing import Optional
from databricks.sdk import WorkspaceClient from databricks.sdk import WorkspaceClient
from databricks.sdk.core import DatabricksError from databricks.sdk.core import DatabricksError
from databricks.sdk.service._internal import Wait from databricks.sdk.service._internal import Wait
from databricks.sdk.service.catalog import TableInfo from databricks.sdk.service.catalog import TableInfo
from databricks.sdk.service.sql import ( from databricks.sdk.service.sql import (
ExecuteStatementResponse,
GetStatementResponse,
GetWarehouseResponse, GetWarehouseResponse,
StatementResponse,
StatementState, StatementState,
StatementStatus, StatementStatus,
) )
@ -125,7 +124,7 @@ class UnityCatalogProxyProfilingMixin:
def _analyze_table( def _analyze_table(
self, ref: TableReference, include_columns: bool self, ref: TableReference, include_columns: bool
) -> ExecuteStatementResponse: ) -> StatementResponse:
statement = f"ANALYZE TABLE {ref.schema}.{ref.table} COMPUTE STATISTICS" statement = f"ANALYZE TABLE {ref.schema}.{ref.table} COMPUTE STATISTICS"
if include_columns: if include_columns:
statement += " FOR ALL COLUMNS" statement += " FOR ALL COLUMNS"
@ -139,7 +138,7 @@ class UnityCatalogProxyProfilingMixin:
return response return response
def _check_analyze_table_statement_status( def _check_analyze_table_statement_status(
self, execute_response: ExecuteStatementResponse, max_wait_secs: int self, execute_response: StatementResponse, max_wait_secs: int
) -> bool: ) -> bool:
if not execute_response.statement_id or not execute_response.status: if not execute_response.statement_id or not execute_response.status:
return False return False
@ -230,9 +229,7 @@ class UnityCatalogProxyProfilingMixin:
return None return None
@staticmethod @staticmethod
def _raise_if_error( def _raise_if_error(response: StatementResponse, key: str) -> None:
response: Union[ExecuteStatementResponse, GetStatementResponse], key: str
) -> None:
if response.status and response.status.state in [ if response.status and response.status.state in [
StatementState.FAILED, StatementState.FAILED,
StatementState.CANCELED, StatementState.CANCELED,