diff --git a/ingestion/src/metadata/core/connections/test_connection/checks/database.py b/ingestion/src/metadata/core/connections/test_connection/checks/database.py index 6792d9efbeb2..342d5a2d3a69 100644 --- a/ingestion/src/metadata/core/connections/test_connection/checks/database.py +++ b/ingestion/src/metadata/core/connections/test_connection/checks/database.py @@ -59,6 +59,13 @@ class DatabaseStep(StepName): GetColumnMetadata = "GetColumnMetadata" GetTableComments = "GetTableComments" GetInformationSchemaColumns = "GetInformationSchemaColumns" + GetViewDefinitions = "GetViewDefinitions" + GetCatalogTags = "GetCatalogTags" + GetSchemaTags = "GetSchemaTags" + GetTableTags = "GetTableTags" + GetColumnTags = "GetColumnTags" + GetTableLineage = "GetTableLineage" + GetColumnLineage = "GetColumnLineage" @contextmanager diff --git a/ingestion/src/metadata/ingestion/source/database/databricks/auth.py b/ingestion/src/metadata/ingestion/source/database/databricks/auth.py index fcef05baa107..d3d0c10bf1c5 100644 --- a/ingestion/src/metadata/ingestion/source/database/databricks/auth.py +++ b/ingestion/src/metadata/ingestion/source/database/databricks/auth.py @@ -34,6 +34,15 @@ ) +def normalize_host_port(host_port: str) -> str: + """Strip a pasted URL scheme and path, leaving ``host:port``.""" + return host_port.split("://", 1)[-1].split("/", 1)[0] + + +def _host(connection: Union[DatabricksConnection, UnityCatalogConnection]) -> str: # noqa: UP007 + return normalize_host_port(connection.hostPort).split(":")[0] + + def get_personal_access_token_auth( connection: Union[DatabricksConnection, UnityCatalogConnection], # noqa: UP007 ) -> dict: @@ -51,7 +60,7 @@ def get_databricks_oauth_auth( """ def credential_provider(): - hostname = connection.hostPort.split(":")[0] + hostname = _host(connection) config = Config( host=f"https://{hostname}", client_id=connection.authType.clientId, @@ -68,7 +77,7 @@ def get_azure_ad_auth(connection: Union[DatabricksConnection, UnityCatalogConnec """ def credential_provider(): - hostname = connection.hostPort.split(":")[0] + hostname = _host(connection) config = Config( host=f"https://{hostname}", azure_client_secret=connection.authType.azureClientSecret.get_secret_value(), diff --git a/ingestion/src/metadata/ingestion/source/database/databricks/connection.py b/ingestion/src/metadata/ingestion/source/database/databricks/connection.py index 3264f98f85ea..2ba2b8f34c2c 100644 --- a/ingestion/src/metadata/ingestion/source/database/databricks/connection.py +++ b/ingestion/src/metadata/ingestion/source/database/databricks/connection.py @@ -13,37 +13,48 @@ Source connection handler """ +from __future__ import annotations + from copy import deepcopy -from functools import partial -from typing import Optional +from typing import TYPE_CHECKING, Optional from urllib.parse import quote_plus from sqlalchemy import text from sqlalchemy.engine import Engine -from sqlalchemy.exc import DatabaseError from sqlalchemy.inspection import inspect -from metadata.generated.schema.entity.automations.workflow import ( - Workflow as AutomationWorkflow, +from metadata.core.connections.test_connection import ( + ErrorPack, + Evidence, + Matchers, + check, + when, +) +from metadata.core.connections.test_connection.check import CheckError +from metadata.core.connections.test_connection.checks.database import ( + DEFAULT_SAMPLE_ROWS, + DatabaseStep, + run_sql, +) +from metadata.core.connections.test_connection.constants import STEP_TIMEOUT_SECONDS +from metadata.core.connections.test_connection.network import ( + NETWORK_ERRORS, + NetworkUnreachableError, + tcp_probe, ) from metadata.generated.schema.entity.services.connections.database.databricksConnection import ( DatabricksConnection as DatabricksConnectionConfig, ) -from metadata.generated.schema.entity.services.connections.testConnectionResult import ( - TestConnectionResult, -) from metadata.ingestion.connections.builders import ( create_generic_db_connection, get_connection_args_common, init_empty_connection_arguments, ) from metadata.ingestion.connections.connection import BaseConnection -from metadata.ingestion.connections.test_connections import ( - test_connection_engine_step, - test_connection_steps, +from metadata.ingestion.source.database.databricks.auth import ( + get_auth_config, + normalize_host_port, ) -from metadata.ingestion.ometa.ometa_api import OpenMetadata -from metadata.ingestion.source.database.databricks.auth import get_auth_config from metadata.ingestion.source.database.databricks.log_filters import ( suppress_user_agent_entry_deprecation_log, ) @@ -58,44 +69,132 @@ TEST_TABLE_TAGS, TEST_VIEW_DEFINITIONS, ) -from metadata.utils.constants import THREE_MIN from metadata.utils.logger import ingestion_logger +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + + from metadata.core.connections.test_connection import ChecksProvider + logger = ingestion_logger() suppress_user_agent_entry_deprecation_log() +# Databricks dials the workspace over HTTPS; the gate TCP-probes this port. +DEFAULT_DATABRICKS_PORT = 443 + +DEFAULT_CATALOG = "main" + +SYSTEM_SCHEMAS = frozenset({"information_schema", "performance_schema", "sys"}) + + +# databricks-sql/thrift reports failures as message tokens, not numeric codes, so +# rules key on tokens; specific ones precede broad ones (first match wins). +DATABRICKS_ERRORS = ErrorPack( + when(Matchers.contains("invalid access token")).diagnose( + "Authentication failed", + fix="Check the access token - the workspace rejected it. Verify the token is valid, " + "not expired, and belongs to a user with access to this workspace.", + ), + when(Matchers.contains("token is expired")).diagnose( + "Access token expired", + fix="The access token has expired. Generate a new token and update the connection.", + ), + when(Matchers.contains("forbidden")).diagnose( + "Access denied", + fix="The workspace returned 403 Forbidden. Verify the token's user is entitled to the " + "workspace and the configured HTTP path / SQL warehouse.", + ), + when(Matchers.contains("malformed_request")).diagnose( + "Invalid HTTP path", + fix="The HTTP Path is malformed. Copy it from the SQL warehouse (or cluster) Connection " + "Details in Databricks - it must look like /sql/1.0/warehouses/ or " + "/sql/1.0/endpoints/.", + ), + when(Matchers.contains("no_such_catalog")).diagnose( + "Catalog not found", + fix="The configured catalog does not exist or is not visible to the token's user. Verify " + "the catalog name and that the user has USE CATALOG on it.", + ), + when(Matchers.contains("no_such_schema")).diagnose( + "Schema not found", + fix="The configured schema does not exist or is not visible. Verify the schema name and " + "that the user has USE SCHEMA on it.", + ), + when(Matchers.contains("table_or_view_not_found")).diagnose( + "Table or view not found", + fix="The referenced table or view does not exist or is not visible to the token's user. " + "Verify the object exists and the user has SELECT on it.", + ), + when(Matchers.contains("schema_not_found")).diagnose( + "Schema not found", + fix="The referenced schema does not exist or is not visible. Verify the schema name and " + "that the user has USAGE on it.", + ), + when(Matchers.contains("permission_denied")).diagnose( + "Insufficient privileges", + fix="Grant the token's user the privileges the failing step needs (USAGE on the catalog / " + "schema and SELECT on the system or information_schema tables it reads).", + ), + when(Matchers.contains("insufficient_permissions")).diagnose( + "Insufficient privileges", + fix="Grant the token's user the privileges the failing step needs (USAGE on the catalog / " + "schema and SELECT on the system or information_schema tables it reads).", + ), + when(Matchers.contains("does not exist")).diagnose( + "Object not found", + fix="Verify the configured catalog, schema, and HTTP path exist and the token's user is " + "authorized to use them.", + ), +).including(NETWORK_ERRORS) + + +def _summarize(rows: Sequence[object], noun: str) -> str: + """``N s enumerated`` (``N+`` at the cap) or ``no s enumerated``.""" + count = len(rows) + if not count: + return f"no {noun}s enumerated" + suffix = "+" if count >= DEFAULT_SAMPLE_ROWS else "" + plural = noun if count == 1 else f"{noun}s" + return f"{count}{suffix} {plural} enumerated" + class DatabricksEngineWrapper: - """Wrapper to store engine and schemas to avoid multiple calls""" + """Engine wrapper caching the resolved catalog and schema. Lookups are lazy, so + constructing it touches no network and stays behind the CheckAccess gate.""" def __init__(self, engine: Engine): self.engine = engine - self.inspector = inspect(engine) + self._inspector = None self.schemas = None self.first_schema = None self.first_catalog = None + @property + def inspector(self): + # Lazy, never at construction: inspect(engine) connects eagerly (the + # dialect's _init_engine runs engine.connect().close()), which would + # bypass the CheckAccess gate. + if self._inspector is None: + self._inspector = inspect(self.engine) + return self._inspector + def get_schemas(self, schema_name: Optional[str] = None): # noqa: UP045 """Get schemas and cache them""" if schema_name is not None: - with self.engine.connect() as connection: - connection.execute(text(f"USE CATALOG `{self.first_catalog}`")) + if self.first_catalog: + with self.engine.connect() as connection: + connection.execute(text(f"USE CATALOG `{self.first_catalog}`")) self.first_schema = schema_name return [schema_name] if self.schemas is None: - self.schemas = self.inspector.get_schema_names(database=self.first_catalog) + # Bound the reflected list (no driver-side cap) to the sample size. + self.schemas = self.inspector.get_schema_names(database=self.first_catalog)[:DEFAULT_SAMPLE_ROWS] if self.schemas: - # Find the first schema that's not a system schema for schema in self.schemas: - if schema.lower() not in ( - "information_schema", - "performance_schema", - "sys", - ): + if schema.lower() not in SYSTEM_SCHEMAS: self.first_schema = schema break - # If no non-system schema found, use the first one if self.first_schema is None and self.schemas: self.first_schema = self.schemas[0] return self.schemas @@ -103,41 +202,40 @@ def get_schemas(self, schema_name: Optional[str] = None): # noqa: UP045 def get_tables(self): """Get tables using the cached first schema""" if self.first_schema is None: - self.get_schemas() # This will set first_schema - if self.first_schema: + self.get_schemas() + if self.first_catalog and self.first_schema: with self.engine.connect() as connection: tables = connection.execute(text(f"SHOW TABLES IN `{self.first_catalog}`.`{self.first_schema}`")) - return tables # noqa: RET504 + return tables.fetchmany(DEFAULT_SAMPLE_ROWS) return [] def get_views(self): """Get views using the cached first schema""" if self.first_schema is None: - self.get_schemas() # This will set first_schema - if self.first_schema: + self.get_schemas() + if self.first_catalog and self.first_schema: with self.engine.connect() as connection: views = connection.execute(text(f"SHOW VIEWS IN `{self.first_catalog}`.`{self.first_schema}`")) - return views # noqa: RET504 + return views.fetchmany(DEFAULT_SAMPLE_ROWS) return [] def get_catalogs(self, catalog_name: Optional[str] = None): # noqa: UP045 """Get catalogs""" - catalogs = [] if catalog_name is not None: self.first_catalog = catalog_name return [catalog_name] with self.engine.connect() as connection: - catalogs = connection.execute(text(DATABRICKS_GET_CATALOGS)).fetchall() - for catalog in catalogs: - if catalog[0] != "__databricks_internal": - self.first_catalog = catalog[0] - break + catalogs = connection.execute(text(DATABRICKS_GET_CATALOGS)).fetchmany(DEFAULT_SAMPLE_ROWS) + for catalog in catalogs: + if catalog[0] != "__databricks_internal": + self.first_catalog = catalog[0] + break return catalogs def get_connection_url(connection: DatabricksConnectionConfig) -> str: scheme = connection.scheme.value if connection.scheme else "databricks" - url = f"{scheme}://{connection.hostPort}" + url = f"{scheme}://{normalize_host_port(connection.hostPort)}" if connection.catalog: url = f"{url}?catalog={quote_plus(connection.catalog)}" return url @@ -170,95 +268,129 @@ def get_connection(connection: DatabricksConnectionConfig) -> Engine: return engine +class DatabricksChecks: + """Test-connection checks for Databricks. Catalog/schema listing goes through + ``DatabricksEngineWrapper``; every statement is built inside its ``@check`` so + nothing connects before the gate.""" + + errors = DATABRICKS_ERRORS + + def __init__(self, client: Engine, service_connection: DatabricksConnectionConfig) -> None: + self.client = client + self.service_connection = service_connection + self._engine_wrapper = DatabricksEngineWrapper(client) + + def _first_catalog(self) -> str: + """Resolve and cache the catalog to scope tag probes to.""" + if self._engine_wrapper.first_catalog is None: + self._engine_wrapper.get_catalogs(catalog_name=self.service_connection.catalog) + return self._engine_wrapper.first_catalog or self.service_connection.catalog or DEFAULT_CATALOG + + def _probe_target(self) -> tuple[str, int]: + host_port = normalize_host_port(self.service_connection.hostPort) + host, _, port = host_port.rpartition(":") + if host and port.isdigit(): + return host, int(port) + return host_port, DEFAULT_DATABRICKS_PORT + + def _list(self, operation: Callable[[], Sequence[object] | None], command: str | None, noun: str) -> Evidence: + """Run a wrapper listing op, reporting the command and a row-count summary; + on failure re-raise as ``CheckError`` carrying the attempted command.""" + try: + rows = operation() + except Exception as cause: + raise CheckError(cause, Evidence(command=command)) from cause + rows = list(rows) if rows is not None else [] + return Evidence(summary=_summarize(rows, noun), command=command) + + @check(DatabaseStep.CheckAccess) + def check_access(self) -> Evidence: + host, port = self._probe_target() + try: + tcp_probe(host, port) + except NetworkUnreachableError as error: + raise CheckError(error, Evidence(command=f"TCP connect {host}:{port}")) from error + return run_sql(self.client, "SELECT 1", lambda _: "connection established") + + @check(DatabaseStep.GetDatabases) + def get_databases(self) -> Evidence: + # A configured catalog is trusted without querying, so no SHOW CATALOGS runs. + configured = self.service_connection.catalog + command = None if configured else DATABRICKS_GET_CATALOGS + return self._list( + lambda: self._engine_wrapper.get_catalogs(catalog_name=configured), + command, + "catalog", + ) + + @check(DatabaseStep.GetSchemas) + def get_schemas(self) -> Evidence: + return self._list( + lambda: self._engine_wrapper.get_schemas(schema_name=self.service_connection.databaseSchema), + "SHOW SCHEMAS", + "schema", + ) + + @check(DatabaseStep.GetTables) + def get_tables(self) -> Evidence: + return self._list(self._engine_wrapper.get_tables, "SHOW TABLES", "table") + + @check(DatabaseStep.GetViews) + def get_views(self) -> Evidence: + return self._list(self._engine_wrapper.get_views, "SHOW VIEWS", "view") + + @check(DatabaseStep.GetQueries) + def get_queries(self) -> Evidence: + statement = DATABRICKS_SQL_STATEMENT_TEST.format(query_history=self.service_connection.queryHistoryTable) + return run_sql(self.client, statement, lambda _: "query history accessible") + + @check(DatabaseStep.GetViewDefinitions) + def get_view_definitions(self) -> Evidence: + return run_sql(self.client, TEST_VIEW_DEFINITIONS, lambda _: "view definitions accessible") + + @check(DatabaseStep.GetCatalogTags) + def get_catalog_tags(self) -> Evidence: + statement = TEST_CATALOG_TAGS.format(database_name=self._first_catalog()) + return run_sql(self.client, statement, lambda _: "catalog tags accessible") + + @check(DatabaseStep.GetSchemaTags) + def get_schema_tags(self) -> Evidence: + statement = TEST_SCHEMA_TAGS.format(database_name=self._first_catalog()) + return run_sql(self.client, statement, lambda _: "schema tags accessible") + + @check(DatabaseStep.GetTableTags) + def get_table_tags(self) -> Evidence: + statement = TEST_TABLE_TAGS.format(database_name=self._first_catalog()) + return run_sql(self.client, statement, lambda _: "table tags accessible") + + @check(DatabaseStep.GetColumnTags) + def get_column_tags(self) -> Evidence: + statement = TEST_COLUMN_TAGS.format(database_name=self._first_catalog()) + return run_sql(self.client, statement, lambda _: "column tags accessible") + + @check(DatabaseStep.GetTableLineage) + def get_table_lineage(self) -> Evidence: + return run_sql(self.client, TEST_TABLE_LINEAGE, lambda _: "table lineage accessible") + + @check(DatabaseStep.GetColumnLineage) + def get_column_lineage(self) -> Evidence: + return run_sql(self.client, TEST_COLUMN_LINEAGE, lambda _: "column lineage accessible") + + class DatabricksConnection(BaseConnection[DatabricksConnectionConfig, Engine]): + def __init__(self, service_connection: DatabricksConnectionConfig) -> None: + super().__init__(service_connection) + # Honor the user-facing connectionTimeout (default 120s) as the per-step + # budget; a cold serverless warehouse can exceed the 60s framework default. + self.step_timeout_seconds = service_connection.connectionTimeout or STEP_TIMEOUT_SECONDS + def _get_client(self) -> Engine: - return get_connection(self.service_connection) - - def test_connection( - self, - metadata: OpenMetadata, - automation_workflow: Optional[AutomationWorkflow] = None, # noqa: UP045 - timeout_seconds: Optional[int] = THREE_MIN, # noqa: UP045 - ) -> TestConnectionResult: - """ - Test connection. This can be executed either as part - of a metadata workflow or during an Automation Workflow - """ - connection = self.client - service_connection = self.service_connection - - def test_database_query(engine: Engine, statement: str): - """ - Method used to execute the given query and fetch a result - to test if user has access to the tables specified - in the sql statement - """ - try: - connection = engine.connect() - connection.execute(text(statement)).fetchone() - except DatabaseError as soe: - logger.debug(f"Failed to fetch catalogs due to: {soe}") - - # Create wrapper to avoid multiple schema calls - engine_wrapper = DatabricksEngineWrapper(connection) - - # Helper function to get first catalog for tag queries - def get_first_catalog(): - catalogs = engine_wrapper.get_catalogs(catalog_name=service_connection.catalog) - return catalogs[0] if catalogs else service_connection.catalog or "main" - - test_fn = { - "CheckAccess": partial(test_connection_engine_step, connection), - "GetSchemas": partial(engine_wrapper.get_schemas, schema_name=service_connection.databaseSchema), - "GetTables": engine_wrapper.get_tables, - "GetViews": engine_wrapper.get_views, - "GetDatabases": partial(engine_wrapper.get_catalogs, catalog_name=service_connection.catalog), - "GetQueries": partial( - test_database_query, - engine=connection, - statement=DATABRICKS_SQL_STATEMENT_TEST.format(query_history=service_connection.queryHistoryTable), - ), - "GetViewDefinitions": partial( - test_database_query, - engine=connection, - statement=TEST_VIEW_DEFINITIONS, - ), - "GetCatalogTags": partial( - test_database_query, - engine=connection, - statement=TEST_CATALOG_TAGS.format(database_name=get_first_catalog()), - ), - "GetSchemaTags": partial( - test_database_query, - engine=connection, - statement=TEST_SCHEMA_TAGS.format(database_name=get_first_catalog()), - ), - "GetTableTags": partial( - test_database_query, - engine=connection, - statement=TEST_TABLE_TAGS.format(database_name=get_first_catalog()), - ), - "GetColumnTags": partial( - test_database_query, - engine=connection, - statement=TEST_COLUMN_TAGS.format(database_name=get_first_catalog()), - ), - "GetTableLineage": partial( - test_database_query, - engine=connection, - statement=TEST_TABLE_LINEAGE, - ), - "GetColumnLineage": partial( - test_database_query, - engine=connection, - statement=TEST_COLUMN_LINEAGE, - ), - } - - return test_connection_steps( - metadata=metadata, - test_fn=test_fn, - service_type=service_connection.type.value, # pyright: ignore[reportOptionalMemberAccess] - automation_workflow=automation_workflow, - timeout_seconds=service_connection.connectionTimeout or timeout_seconds, + engine = get_connection(self.service_connection) + self._on_close(engine.dispose) + return engine + + def checks(self) -> ChecksProvider: + return DatabricksChecks( + client=self.client, + service_connection=self.service_connection, ) diff --git a/ingestion/tests/unit/source/database/databricks/__init__.py b/ingestion/tests/unit/source/database/databricks/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/ingestion/tests/unit/source/database/databricks/test_connection.py b/ingestion/tests/unit/source/database/databricks/test_connection.py new file mode 100644 index 000000000000..f8966e82da0e --- /dev/null +++ b/ingestion/tests/unit/source/database/databricks/test_connection.py @@ -0,0 +1,348 @@ +# Copyright 2025 Collate +# Licensed under the Collate Community License, Version 1.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the Databricks BaseConnection wiring and test-connection checks.""" + +from unittest.mock import MagicMock, patch + +import pytest + +from metadata.core.connections.test_connection.check import CheckError, collect_checks +from metadata.core.connections.test_connection.checks.database import DatabaseStep +from metadata.core.connections.test_connection.network import NetworkUnreachableError +from metadata.generated.schema.entity.services.connections.database.databricks.personalAccessToken import ( + PersonalAccessToken, +) +from metadata.generated.schema.entity.services.connections.database.databricksConnection import ( + DatabricksConnection as DatabricksConnectionConfig, +) +from metadata.ingestion.connections.connection import BaseConnection +from metadata.ingestion.source.database.databricks.connection import ( + DATABRICKS_ERRORS, + DatabricksChecks, + DatabricksConnection, +) + +CONNECTION_MODULE = "metadata.ingestion.source.database.databricks.connection" + + +def _config() -> DatabricksConnectionConfig: + return DatabricksConnectionConfig( + hostPort="my-workspace.cloud.databricks.com:443", + authType=PersonalAccessToken(token="dapi-secret"), + httpPath="/sql/1.0/warehouses/abc123", + catalog="my_catalog", + ) + + +class _SqlAlchemyError(Exception): + """Mirror ``sqlalchemy.exc.DBAPIError``: wraps the driver error on ``.orig``.""" + + def __init__(self, orig: Exception) -> None: + super().__init__(str(orig)) + self.orig = orig + + +def test_databricks_connection_is_base_connection(): + assert issubclass(DatabricksConnection, BaseConnection) + + +def test_get_client_uses_the_databricks_url_builder(): + with patch(f"{CONNECTION_MODULE}.create_generic_db_connection") as mock_connection: + _ = DatabricksConnection(_config()).client + assert mock_connection.call_args.kwargs["get_connection_url_fn"].__name__ == "get_connection_url" + + +def test_close_disposes_the_engine_pool(): + engine = MagicMock() + with patch(f"{CONNECTION_MODULE}.create_generic_db_connection", return_value=engine): + connection = DatabricksConnection(_config()) + assert connection.client is engine + connection.close() + engine.dispose.assert_called_once_with() + + +def test_connection_timeout_drives_the_per_step_budget(): + from metadata.core.connections.test_connection.constants import STEP_TIMEOUT_SECONDS + + # connectionTimeout (default 120) must win over the 60s framework default, + # fed through BaseConnection.step_timeout_seconds. + default_config = _config() + assert DatabricksConnection(default_config).step_timeout_seconds == default_config.connectionTimeout + + slow_config = _config() + slow_config.connectionTimeout = 300 + assert DatabricksConnection(slow_config).step_timeout_seconds == 300 + + unset_config = _config() + unset_config.connectionTimeout = None + assert DatabricksConnection(unset_config).step_timeout_seconds == STEP_TIMEOUT_SECONDS + + +def test_invalid_token_message_is_classified(): + error = _SqlAlchemyError(Exception("Error during request to server: Invalid access token")) + assert DATABRICKS_ERRORS.classify(error).title == "Authentication failed" + + +def test_expired_token_message_is_classified(): + error = _SqlAlchemyError(Exception("the access token is expired, please refresh it")) + assert DATABRICKS_ERRORS.classify(error).title == "Access token expired" + + +def test_forbidden_message_is_classified(): + error = _SqlAlchemyError(Exception("HTTP Response code: 403, Forbidden")) + assert DATABRICKS_ERRORS.classify(error).title == "Access denied" + + +def test_malformed_http_path_is_classified(): + # A bad httpPath fails CheckAccess with MALFORMED_REQUEST; needs a diagnosis. + error = _SqlAlchemyError( + Exception( + "MALFORMED_REQUEST: Path /sql/1.0/warehouses/39c390db3a5e19e must match pattern " + "/sql/1.0/endpoints/ or /sql/1.0/warehouses/" + ) + ) + assert DATABRICKS_ERRORS.classify(error).title == "Invalid HTTP path" + + +def test_permission_denied_is_classified(): + error = _SqlAlchemyError(Exception("[PERMISSION_DENIED] User does not have SELECT on system.access.table_lineage")) + assert DATABRICKS_ERRORS.classify(error).title == "Insufficient privileges" + + +def test_insufficient_permissions_is_classified(): + error = _SqlAlchemyError(Exception("INSUFFICIENT_PERMISSIONS: requires USAGE on catalog")) + assert DATABRICKS_ERRORS.classify(error).title == "Insufficient privileges" + + +def test_table_or_view_not_found_is_classified(): + error = _SqlAlchemyError(Exception("[TABLE_OR_VIEW_NOT_FOUND] The table or view cannot be found")) + assert DATABRICKS_ERRORS.classify(error).title == "Table or view not found" + + +def test_schema_not_found_is_classified(): + error = _SqlAlchemyError(Exception("[SCHEMA_NOT_FOUND] The schema cannot be found")) + assert DATABRICKS_ERRORS.classify(error).title == "Schema not found" + + +def test_catalog_does_not_exist_is_classified(): + error = _SqlAlchemyError(Exception("Catalog 'nope' does not exist")) + assert DATABRICKS_ERRORS.classify(error).title == "Object not found" + + +def test_no_such_catalog_exception_is_classified(): + # NO_SUCH_CATALOG_EXCEPTION says "was not found", not "does not exist". + error = _SqlAlchemyError(Exception("[NO_SUCH_CATALOG_EXCEPTION] Catalog 'batata' was not found. SQLSTATE: 42704")) + assert DATABRICKS_ERRORS.classify(error).title == "Catalog not found" + + +def test_no_such_schema_exception_is_classified(): + error = _SqlAlchemyError(Exception("[NO_SUCH_SCHEMA_EXCEPTION] Schema 'nope' was not found")) + assert DATABRICKS_ERRORS.classify(error).title == "Schema not found" + + +def test_network_errors_classify_through_including(): + error = NetworkUnreachableError("workspace:443 is not reachable") + error.__cause__ = ConnectionRefusedError(61, "Connection refused") + assert DATABRICKS_ERRORS.classify(error).title == "Connection refused" + + +def test_unknown_error_returns_no_diagnosis(): + error = _SqlAlchemyError(Exception("something entirely unexpected")) + assert DATABRICKS_ERRORS.classify(error) is None + + +def test_checks_cover_exactly_the_seeded_steps(): + checks = DatabricksChecks(client=MagicMock(), service_connection=_config()) + collected = collect_checks(checks) + assert set(collected.keys()) == { + DatabaseStep.CheckAccess, + DatabaseStep.GetDatabases, + DatabaseStep.GetSchemas, + DatabaseStep.GetTables, + DatabaseStep.GetViews, + DatabaseStep.GetQueries, + DatabaseStep.GetViewDefinitions, + DatabaseStep.GetCatalogTags, + DatabaseStep.GetSchemaTags, + DatabaseStep.GetTableTags, + DatabaseStep.GetColumnTags, + DatabaseStep.GetTableLineage, + DatabaseStep.GetColumnLineage, + } + + +def test_construction_does_not_touch_the_network(): + # Building the provider must not connect - that belongs in a gated check. + engine = MagicMock() + DatabricksChecks(client=engine, service_connection=_config()) + engine.connect.assert_not_called() + + +def test_construction_does_not_build_the_inspector(): + # Regression: inspect(engine) connects eagerly, so building it in __init__ would + # connect before the gate (auth error escapes as a 500). Keep it lazy. + with patch(f"{CONNECTION_MODULE}.inspect") as mock_inspect: + checks = DatabricksChecks(client=MagicMock(), service_connection=_config()) + mock_inspect.assert_not_called() + _ = checks._engine_wrapper.inspector + mock_inspect.assert_called_once() + + +def test_check_access_probes_the_host_port_then_pings(): + checks = DatabricksChecks(client=MagicMock(), service_connection=_config()) + with ( + patch(f"{CONNECTION_MODULE}.tcp_probe") as mock_probe, + patch(f"{CONNECTION_MODULE}.run_sql", return_value=MagicMock()) as mock_run, + ): + checks.check_access() + mock_probe.assert_called_once_with("my-workspace.cloud.databricks.com", 443) + assert mock_run.call_count == 1 + + +def test_check_access_reports_unreachable_host_as_network_failure(): + checks = DatabricksChecks(client=MagicMock(), service_connection=_config()) + probe_error = NetworkUnreachableError("my-workspace.cloud.databricks.com:443 is not reachable") + probe_error.__cause__ = ConnectionRefusedError(61, "Connection refused") + with ( + patch(f"{CONNECTION_MODULE}.tcp_probe", side_effect=probe_error) as mock_probe, + pytest.raises(CheckError) as exc, + ): + checks.check_access() + mock_probe.assert_called_once_with("my-workspace.cloud.databricks.com", 443) + assert exc.value.cause is probe_error + assert DATABRICKS_ERRORS.classify(exc.value.cause).title == "Connection refused" + + +def test_probe_target_defaults_to_443_when_no_port(): + config = _config() + config.hostPort = "my-workspace.cloud.databricks.com" + checks = DatabricksChecks(client=MagicMock(), service_connection=config) + assert checks._probe_target() == ("my-workspace.cloud.databricks.com", 443) + + +def test_first_catalog_uses_configured_catalog_without_querying(): + engine = MagicMock() + checks = DatabricksChecks(client=engine, service_connection=_config()) + assert checks._first_catalog() == "my_catalog" + engine.connect.assert_not_called() + + +def test_listing_caps_rows_at_the_sample_size_at_the_source(): + # Bound the fetch (fetchmany), not fetchall-then-slice. Test-connection only; + # real ingestion goes through the common framework. + from metadata.core.connections.test_connection.checks.database import ( + DEFAULT_SAMPLE_ROWS, + ) + from metadata.ingestion.source.database.databricks.connection import ( + DatabricksEngineWrapper, + ) + + result = MagicMock() + result.fetchmany.return_value = [("t",)] + conn = MagicMock() + conn.execute.return_value = result + engine = MagicMock() + engine.connect.return_value.__enter__.return_value = conn + + wrapper = DatabricksEngineWrapper(engine) + wrapper.first_catalog = "my_catalog" + wrapper.first_schema = "my_schema" + wrapper.get_tables() + result.fetchmany.assert_called_once_with(DEFAULT_SAMPLE_ROWS) + result.fetchall.assert_not_called() + + +def test_get_databases_reports_catalog_count(): + checks = DatabricksChecks(client=MagicMock(), service_connection=_config()) + with patch.object( + checks._engine_wrapper, "get_catalogs", return_value=[("my_catalog",), ("other",)] + ) as mock_catalogs: + evidence = checks.get_databases() + mock_catalogs.assert_called_once_with(catalog_name="my_catalog") + assert evidence.summary == "2 catalogs enumerated" + + +def test_get_databases_failure_reports_the_attempted_command_as_evidence(): + config = _config() + config.catalog = None + checks = DatabricksChecks(client=MagicMock(), service_connection=config) + with ( + patch.object(checks._engine_wrapper, "get_catalogs", side_effect=Exception("boom")), + pytest.raises(CheckError) as exc, + ): + checks.get_databases() + assert exc.value.evidence.command == "SHOW CATALOGS" + + +def test_get_databases_omits_the_command_when_catalog_is_configured(): + # A configured catalog is trusted without running SHOW CATALOGS. + checks = DatabricksChecks(client=MagicMock(), service_connection=_config()) + with patch.object(checks._engine_wrapper, "get_catalogs", return_value=[("my_catalog",)]): + evidence = checks.get_databases() + assert evidence.command is None + + +def test_probe_target_strips_a_pasted_url_scheme_and_path(): + # A pasted workspace URL must not reach the socket probe as the host, or the + # gate fails with a misleading DNS error before the driver. + config = _config() + config.hostPort = "https://my-workspace.cloud.databricks.com:443/sql/1.0/warehouses/x" + checks = DatabricksChecks(client=MagicMock(), service_connection=config) + assert checks._probe_target() == ("my-workspace.cloud.databricks.com", 443) + + +def test_connection_url_strips_a_pasted_url_scheme(): + from urllib.parse import urlparse + + from metadata.ingestion.source.database.databricks.connection import get_connection_url + + config = _config() + config.hostPort = "https://my-workspace.cloud.databricks.com:443" + parsed = urlparse(get_connection_url(config)) + assert parsed.hostname == "my-workspace.cloud.databricks.com" + assert parsed.port == 443 + + +def test_auth_host_derivation_strips_a_pasted_scheme(): + # OAuth/Azure-AD derive the workspace host from hostPort too; a pasted scheme + # must not leak in as the host or auth fails before the driver. + from metadata.ingestion.source.database.databricks.auth import _host, normalize_host_port + + assert normalize_host_port("https://ws.cloud.databricks.com:443/sql/x") == "ws.cloud.databricks.com:443" + connection = MagicMock() + connection.hostPort = "https://ws.cloud.databricks.com:443" + assert _host(connection) == "ws.cloud.databricks.com" + + +def test_schema_scoped_get_schemas_skips_use_catalog_when_catalog_unresolved(): + # first_catalog None (GetDatabases failed/skipped) must not emit USE CATALOG `None`. + from metadata.ingestion.source.database.databricks.connection import ( + DatabricksEngineWrapper, + ) + + engine = MagicMock() + wrapper = DatabricksEngineWrapper(engine) + wrapper.first_catalog = None + assert wrapper.get_schemas(schema_name="my_schema") == ["my_schema"] + engine.connect.assert_not_called() + + +def test_get_tables_returns_empty_when_catalog_unresolved(): + from metadata.ingestion.source.database.databricks.connection import ( + DatabricksEngineWrapper, + ) + + engine = MagicMock() + wrapper = DatabricksEngineWrapper(engine) + wrapper.first_catalog = None + wrapper.first_schema = "my_schema" + assert wrapper.get_tables() == [] + engine.connect.assert_not_called() diff --git a/ingestion/tests/unit/topology/database/test_databricks.py b/ingestion/tests/unit/topology/database/test_databricks.py index cf0c6efd7d9d..86254b42b959 100644 --- a/ingestion/tests/unit/topology/database/test_databricks.py +++ b/ingestion/tests/unit/topology/database/test_databricks.py @@ -586,7 +586,7 @@ def test_databricks_engine_wrapper_get_tables_with_cached_schema(self): # Mock the connection execute method to return a result mock_result = Mock() - mock_result.fetchall.return_value = [("table1",), ("table2",)] + mock_result.fetchmany.return_value = [("table1",), ("table2",)] mock_connection.execute.return_value = mock_result mock_inspector = Mock() @@ -606,7 +606,7 @@ def test_databricks_engine_wrapper_get_tables_with_cached_schema(self): # Then call get_tables tables = wrapper.get_tables() - self.assertEqual(tables, mock_result) + self.assertEqual(tables, [("table1",), ("table2",)]) mock_connection.execute.assert_called_once() def test_databricks_engine_wrapper_get_tables_without_cached_schema(self): @@ -620,7 +620,7 @@ def test_databricks_engine_wrapper_get_tables_without_cached_schema(self): # Mock the connection execute method to return a result mock_result = Mock() - mock_result.fetchall.return_value = [("table1",), ("table2",)] + mock_result.fetchmany.return_value = [("table1",), ("table2",)] mock_connection.execute.return_value = mock_result mock_inspector = Mock() @@ -638,7 +638,7 @@ def test_databricks_engine_wrapper_get_tables_without_cached_schema(self): # Call get_tables directly without calling get_schemas first tables = wrapper.get_tables() - self.assertEqual(tables, mock_result) + self.assertEqual(tables, [("table1",), ("table2",)]) # Should have called get_schemas internally mock_inspector.get_schema_names.assert_called_once() mock_connection.execute.assert_called_once() @@ -669,7 +669,7 @@ def test_databricks_engine_wrapper_get_views_with_cached_schema(self): # Mock the connection execute method to return a result mock_result = Mock() - mock_result.fetchall.return_value = [("view1",), ("view2",)] + mock_result.fetchmany.return_value = [("view1",), ("view2",)] mock_connection.execute.return_value = mock_result mock_inspector = Mock() @@ -689,7 +689,7 @@ def test_databricks_engine_wrapper_get_views_with_cached_schema(self): # Then call get_views views = wrapper.get_views() - self.assertEqual(views, mock_result) + self.assertEqual(views, [("view1",), ("view2",)]) mock_connection.execute.assert_called_once() def test_databricks_engine_wrapper_get_views_without_cached_schema(self): @@ -703,7 +703,7 @@ def test_databricks_engine_wrapper_get_views_without_cached_schema(self): # Mock the connection execute method to return a result mock_result = Mock() - mock_result.fetchall.return_value = [("view1",), ("view2",)] + mock_result.fetchmany.return_value = [("view1",), ("view2",)] mock_connection.execute.return_value = mock_result mock_inspector = Mock() @@ -721,7 +721,7 @@ def test_databricks_engine_wrapper_get_views_without_cached_schema(self): # Call get_views directly without calling get_schemas first views = wrapper.get_views() - self.assertEqual(views, mock_result) + self.assertEqual(views, [("view1",), ("view2",)]) # Should have called get_schemas internally mock_inspector.get_schema_names.assert_called_once() mock_connection.execute.assert_called_once() @@ -767,100 +767,6 @@ def test_databricks_engine_wrapper_caching_behavior(self): # get_schema_names should only be called once due to caching mock_inspector.get_schema_names.assert_called_once() - # pylint: disable=too-many-locals - @patch("metadata.ingestion.source.database.databricks.connection.DatabricksEngineWrapper") - @patch("metadata.ingestion.source.database.databricks.connection.test_connection_steps") - def test_test_connection_function(self, mock_test_connection_steps, mock_engine_wrapper_class): - """Test the test_connection function""" - from metadata.generated.schema.entity.services.connections.database.databricksConnection import ( - DatabricksConnection, - DatabricksScheme, - ) - from metadata.generated.schema.entity.services.connections.testConnectionResult import ( - StatusType, - TestConnectionResult, - TestConnectionStepResult, - ) - - # Mock the test connection result - mock_result = TestConnectionResult( - status=StatusType.Successful, - steps=[ - TestConnectionStepResult( - name="CheckAccess", - mandatory=True, - passed=True, - ), - TestConnectionStepResult( - name="GetSchemas", - mandatory=True, - passed=True, - ), - TestConnectionStepResult( - name="GetTables", - mandatory=True, - passed=True, - ), - TestConnectionStepResult( - name="GetViews", - mandatory=False, - passed=True, - ), - ], - ) - mock_test_connection_steps.return_value = mock_result - - # Create test connection - service_connection = DatabricksConnection( - scheme=DatabricksScheme.databricks, - hostPort="test-host:443", - authType=PersonalAccessToken(token="test-token"), - httpPath="/sql/1.0/warehouses/test", - queryHistoryTable="test_table", - ) - - # Mock the DatabricksEngineWrapper instance to avoid context manager error - mock_wrapper_instance = Mock() - mock_wrapper_instance.get_catalogs.return_value = ["main"] - mock_engine_wrapper_class.return_value = mock_wrapper_instance - - mock_engine = Mock() - mock_engine.connect.return_value = Mock() - mock_metadata = Mock() - - # Test the function - handler = self.DatabricksConnectionHandler(service_connection) - handler._client = mock_engine - result = handler.test_connection(metadata=mock_metadata) - - # Verify the result - self.assertEqual(result, mock_result) - mock_test_connection_steps.assert_called_once() - - # Verify the test_fn dictionary was created with the correct functions - call_args = mock_test_connection_steps.call_args - test_fn = call_args[1]["test_fn"] - - # Check that the test_fn contains the expected keys - expected_keys = [ - "CheckAccess", - "GetSchemas", - "GetTables", - "GetViews", - "GetDatabases", - "GetQueries", - "GetViewDefinitions", - "GetCatalogTags", - "GetSchemaTags", - "GetTableTags", - "GetColumnTags", - "GetTableLineage", - "GetColumnLineage", - ] - for key in expected_keys: - self.assertIn(key, test_fn) - self.assertIsNotNone(test_fn[key]) - def test_databricks_engine_wrapper_system_schema_filtering(self): """Test that system schemas are properly filtered out""" mock_engine = Mock() diff --git a/openmetadata-service/src/main/resources/json/data/testConnections/database/databricks.json b/openmetadata-service/src/main/resources/json/data/testConnections/database/databricks.json index ff8bf5c6335e..7efafcbe9749 100644 --- a/openmetadata-service/src/main/resources/json/data/testConnections/database/databricks.json +++ b/openmetadata-service/src/main/resources/json/data/testConnections/database/databricks.json @@ -8,6 +8,7 @@ "description": "Validate that we can properly reach the database and authenticate with the given credentials.", "errorMessage": "Failed to connect to databricks, please validate to token, http path & hostport", "shortCircuit": true, + "category": "ConnectionGate", "mandatory": true }, {