open-metadata--openmetadata/ingestion/tests/unit/source/database/bigquery/test_connection.py

#  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 BigQuery test-connection checks and error classification.

Cover the wiring (steps resolve to checks, the network pack is folded in, nothing
connects at construction) and the error-pack mapping (each GCP auth / api_core
scenario classifies to the intended diagnosis).
"""

from unittest.mock import MagicMock, patch

import pytest
from google.api_core.exceptions import Forbidden, NotFound
from google.auth.exceptions import DefaultCredentialsError, RefreshError

from metadata.core.connections.test_connection import Evidence
from metadata.core.connections.test_connection.check import 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.bigQueryConnection import (
    BigQueryConnection as BigQueryConnectionConfig,
)
from metadata.ingestion.source.database.bigquery.connection import (
    BIGQUERY_ERRORS,
    BigQueryChecks,
    _enumerated,
    probe_table_view_enumeration,
)
from metadata.utils.credentials import InvalidPrivateKeyException

_CONNECTION_MODULE = "metadata.ingestion.source.database.bigquery.connection"

_GCP_CONFIG = {
    "type": "service_account",
    "projectId": "placeholder",
    "privateKeyId": "key-id",
    "privateKey": "private-key",
    "clientEmail": "user@example.com",
    "clientId": "1234",
    "authUri": "https://accounts.google.com/o/oauth2/auth",
    "tokenUri": "https://oauth2.googleapis.com/token",
    "authProviderX509CertUrl": "https://www.googleapis.com/oauth2/v1/certs",
    "clientX509CertUrl": "https://www.googleapis.com/oauth2/v1/certs",
}


def _config(**overrides) -> BigQueryConnectionConfig:
    base = {"type": "BigQuery", "credentials": {"gcpConfig": {**_GCP_CONFIG}}}
    base.update(overrides)
    return BigQueryConnectionConfig.model_validate(base)


class _ApiError(Exception):
    """Mirror how SQLAlchemy surfaces a driver error: the wrapper's message embeds
    the original and its ``__cause__`` chains to it (via ``raise ... from``), which
    is what the classifier walks."""

    def __init__(self, orig: Exception) -> None:
        super().__init__(f"(google.api_core.exceptions.{type(orig).__name__}) {orig}")
        self.orig = orig
        self.__cause__ = orig


def test_invalid_grant_is_classified():
    error = RefreshError("('invalid_grant: Invalid grant: account not found', {})")
    assert BIGQUERY_ERRORS.classify(error).title == "Invalid service account credentials"


def test_default_credentials_error_is_classified():
    error = DefaultCredentialsError("Could not automatically determine credentials.")
    assert BIGQUERY_ERRORS.classify(error).title == "Could not determine GCP credentials"


def test_refresh_error_is_classified():
    error = RefreshError("Unable to acquire impersonated credentials")
    assert BIGQUERY_ERRORS.classify(error).title == "Failed to obtain a GCP access token"


def test_missing_jobs_create_permission_is_classified():
    error = _ApiError(
        Forbidden("403 POST https://bigquery.googleapis.com: User does not have bigquery.jobs.create permission")
    )
    assert BIGQUERY_ERRORS.classify(error).title == "Missing permission to run BigQuery jobs"


def test_jobs_create_beats_generic_forbidden():
    # A bigquery.jobs.create denial is a Forbidden; the sharper token rule is
    # ordered first so the user gets the actionable role, not a generic message.
    error = Forbidden("403 User does not have bigquery.jobs.create permission in project x")
    assert BIGQUERY_ERRORS.classify(error).title != "Permission denied"


def test_access_denied_is_classified():
    error = _ApiError(Forbidden("403 Access Denied: Table project:dataset.INFORMATION_SCHEMA.JOBS_BY_PROJECT"))
    assert BIGQUERY_ERRORS.classify(error).title == "Access denied"


def test_generic_forbidden_is_permission_denied():
    error = Forbidden("403 The caller does not have permission")
    assert BIGQUERY_ERRORS.classify(error).title == "Permission denied"


def test_not_found_is_classified():
    error = _ApiError(NotFound("404 Not found: Dataset project:dataset was not found"))
    assert BIGQUERY_ERRORS.classify(error).title == "Project or dataset not found"


def test_unknown_error_returns_no_diagnosis():
    assert BIGQUERY_ERRORS.classify(ValueError("something unexpected")) is None


def test_network_errors_classify_through_including():
    error = NetworkUnreachableError("bigquery.googleapis.com:443 is not reachable")
    error.__cause__ = TimeoutError("timed out")
    assert BIGQUERY_ERRORS.classify(error).title == "Connection timed out"


def test_malformed_private_key_is_classified():
    # A malformed PEM raises InvalidPrivateKeyException while the engine is built;
    # it must be classified (and, being first, win over any message-token rule).
    error = InvalidPrivateKeyException("Cannot serialise key: Unable to load PEM file.")
    assert BIGQUERY_ERRORS.classify(error).title == "Malformed service account private key"


def _checks(service_connection, client=None) -> BigQueryChecks:
    engine = client if client is not None else MagicMock()
    return BigQueryChecks(get_client=lambda: engine, service_connection=service_connection)


def test_checks_cover_exactly_the_wired_steps():
    checks = _checks(_config())
    collected = collect_checks(checks)
    assert set(collected.keys()) == {
        DatabaseStep.CheckAccess,
        DatabaseStep.GetSchemas,
        DatabaseStep.GetTables,
        DatabaseStep.GetViews,
        DatabaseStep.GetTags,
        DatabaseStep.GetQueries,
    }


def test_construction_does_not_build_the_client():
    # The engine is built lazily inside CheckAccess, never at construction - so
    # credential parsing (and any connect) happens behind the gate, not before it.
    get_client = MagicMock()
    BigQueryChecks(get_client=get_client, service_connection=_config())
    get_client.assert_not_called()


def test_check_access_builds_client_lazily_and_pings():
    get_client = MagicMock()
    checks = BigQueryChecks(get_client=get_client, service_connection=_config())
    with patch(f"{_CONNECTION_MODULE}.ping", return_value=Evidence(summary="connection established")) as mock_ping:
        evidence = checks.check_access()
    get_client.assert_called_once_with()
    mock_ping.assert_called_once_with(get_client.return_value)
    assert evidence.summary == "connection established"


def test_check_access_surfaces_malformed_key_error():
    # A malformed key fails while building the engine; because the engine is built
    # inside CheckAccess, the error reaches the runner (here: the caller) to be
    # classified, instead of escaping before the test starts.
    error = InvalidPrivateKeyException("Cannot serialise key: Unable to load PEM file.")
    checks = BigQueryChecks(get_client=MagicMock(side_effect=error), service_connection=_config())
    with pytest.raises(InvalidPrivateKeyException):
        checks.check_access()


def test_get_queries_formats_statement_lazily():
    # The INFORMATION_SCHEMA.JOBS_BY_PROJECT statement must be built inside the
    # check (behind the gate), scoped to the configured usage region.
    checks = _checks(_config(usageLocation="us-east1"))
    captured = {}

    def fake(client, statement, summarize, *args, **kwargs):
        captured["statement"] = statement
        return Evidence(summary=summarize([]), command=statement)

    with patch(f"{_CONNECTION_MODULE}.run_sql", side_effect=fake):
        evidence = checks.get_queries()
    assert "region-us-east1" in captured["statement"]
    assert "JOBS_BY_PROJECT" in captured["statement"]
    assert evidence.summary == "query history accessible"


def test_get_tags_skips_when_policy_tags_disabled():
    checks = _checks(_config(includePolicyTags=False))
    assert checks.get_tags() is None


def test_get_tags_skips_when_no_taxonomy_project():
    client = MagicMock()
    client.url.host = None
    checks = _checks(_config(includePolicyTags=True), client=client)
    assert checks.get_tags() is None


def test_get_tags_skips_when_no_location():
    client = MagicMock()
    client.url.host = "proj-a"
    checks = _checks(_config(includePolicyTags=True, taxonomyLocation=None), client=client)
    assert checks.get_tags() is None


def test_get_tags_counts_policy_tags():
    client = MagicMock()
    client.url.host = "proj-a"
    checks = _checks(_config(includePolicyTags=True, taxonomyLocation="us"), client=client)
    tag_client = MagicMock()
    tag_client.list_taxonomies.return_value = [MagicMock(name="tax")]
    tag_client.list_policy_tags.return_value = [object(), object()]
    with patch(f"{_CONNECTION_MODULE}.PolicyTagManagerClient", return_value=tag_client):
        evidence = checks.get_tags()
    assert evidence.summary == "2 policy tags enumerated"


def test_probe_table_view_enumeration_tolerates_deleted_dataset():
    # A dataset dropped between listing datasets and listing its tables raises
    # NotFound; the probe skips it and still reports the datasets it enumerated.
    conn = MagicMock()
    engine = MagicMock()
    engine.connect.return_value.__enter__.return_value = conn
    bq_client = conn.connection._client
    dataset = MagicMock()
    bq_client.list_datasets.return_value = [dataset]
    bq_client.list_tables.side_effect = NotFound("404 Not found: Dataset was deleted")
    evidence = probe_table_view_enumeration(engine)
    assert evidence.summary == "1 dataset enumerated"


def test_enumerated_pluralizes_by_count():
    assert _enumerated(1, "dataset") == "1 dataset enumerated"
    assert _enumerated(3, "dataset") == "3 datasets enumerated"
    assert _enumerated(1, "policy tag") == "1 policy tag enumerated"