chore: import upstream snapshot with attribution

ingestion/tests/unit/source/database/athena/test_connection.py

@@ -0,0 +1,381 @@
+#  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 Athena connection handling."""
+
+import socket
+from unittest.mock import MagicMock, patch
+
+import pytest
+from botocore.exceptions import ClientError
+
+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.records import Evidence
+from metadata.generated.schema.entity.services.connections.database.athenaConnection import (
+    AthenaConnection as AthenaConnectionConfig,
+)
+from metadata.generated.schema.entity.services.connections.database.athenaConnection import (
+    AthenaScheme,
+)
+from metadata.generated.schema.security.credentials.awsCredentials import AWSCredentials
+from metadata.generated.schema.type.filterPattern import FilterPattern
+from metadata.ingestion.connections.connection import BaseConnection
+from metadata.ingestion.source.database.athena import connection as athena_connection
+from metadata.ingestion.source.database.athena.connection import (
+    ATHENA_ERRORS,
+    AthenaChecks,
+    AthenaConnection,
+)
+
+CONNECTION_MODULE = "metadata.ingestion.source.database.athena.connection"
+
+
+def _config(**kwargs) -> AthenaConnectionConfig:
+    base = {
+        "awsConfig": AWSCredentials(awsAccessKeyId="key", awsRegion="us-east-2", awsSecretAccessKey="secret_key"),
+        "s3StagingDir": "s3://postgres/input/",
+        "workgroup": "primary",
+        "scheme": AthenaScheme.awsathena_rest,
+    }
+    base.update(kwargs)
+    return AthenaConnectionConfig(**base)
+
+
+def _checks(**config_kwargs) -> AthenaChecks:
+    conn = AthenaConnection(_config(**config_kwargs))
+    conn._client = MagicMock()
+    return conn.checks()
+
+
+def _client_error(code: str, message: str = "denied") -> ClientError:
+    return ClientError({"Error": {"Code": code, "Message": message}}, "StartQueryExecution")
+
+
+def _raising_factory(exc: BaseException):
+    def factory():
+        raise exc
+
+    return factory
+
+
+def test_athena_connection_is_base_connection():
+    assert issubclass(AthenaConnection, BaseConnection)
+
+
+def test_get_client_uses_the_class_url_builder():
+    with patch(f"{CONNECTION_MODULE}.create_generic_db_connection") as mock_connection:
+        _ = AthenaConnection(_config()).client
+    assert mock_connection.call_args.kwargs["get_connection_url_fn"].__name__ == "get_connection_url"
+
+
+def test_get_client_registers_engine_disposal():
+    conn = AthenaConnection(_config())
+    with patch(f"{CONNECTION_MODULE}.create_generic_db_connection") as mock_connection:
+        engine = mock_connection.return_value
+        _ = conn.client
+        conn.close()
+    engine.dispose.assert_called_once_with()
+
+
+def test_checks_returns_provider_over_the_client():
+    conn = AthenaConnection(_config())
+    conn._client = MagicMock()
+    provider = conn.checks()
+    assert isinstance(provider, AthenaChecks)
+    assert provider.client is conn._client
+
+
+def test_checks_wires_filter_pattern_and_catalog_from_service_connection():
+    pattern = FilterPattern(includes=["db2"])
+    conn = AthenaConnection(_config(catalogId="my_catalog", schemaFilterPattern=pattern))
+    conn._client = MagicMock()
+    provider = conn.checks()
+    assert provider.schema_filter_pattern == pattern
+    assert provider.catalog_id == "my_catalog"
+
+
+def test_every_athena_step_resolves_to_a_check():
+    provider = AthenaChecks(client_factory=MagicMock)
+    resolved = collect_checks(provider)
+    assert set(resolved) == {
+        DatabaseStep.CheckAccess,
+        DatabaseStep.GetSchemas,
+        DatabaseStep.GetTables,
+        DatabaseStep.GetViews,
+    }
+
+
+def test_every_check_reports_the_command_it_ran():
+    # pyathena reflects via the AWS API, so the shared SQL capture yields nothing;
+    # each reflection step names its API operation instead.
+    provider = _checks()
+    inspector = MagicMock()
+    inspector.get_schema_names.return_value = ["ecommerce"]
+    inspector.get_table_names.return_value = ["orders"]
+    inspector.get_view_names.return_value = ["v_orders"]
+
+    with (
+        patch(f"{CONNECTION_MODULE}.list_schemas", return_value=Evidence(summary="1 schema enumerated")),
+        patch(f"{CONNECTION_MODULE}.inspect", return_value=inspector),
+    ):
+        assert provider.get_schemas().command == "athena:ListDatabases (CatalogName=AwsDataCatalog)"
+        assert provider.get_tables().command == "athena:ListTableMetadata (CatalogName=AwsDataCatalog)"
+        assert provider.get_views().command == "athena:ListTableMetadata (CatalogName=AwsDataCatalog)"
+
+
+def test_command_names_a_configured_catalog():
+    provider = _checks(catalogId="my_catalog")
+
+    with patch(f"{CONNECTION_MODULE}.list_schemas", return_value=Evidence(summary="ok")):
+        assert provider.get_schemas().command == "athena:ListDatabases (CatalogName=my_catalog)"
+
+
+def test_error_pack_classifies_access_denied_as_not_authorized():
+    diagnosis = ATHENA_ERRORS.classify(_client_error("AccessDeniedException"))
+    assert diagnosis is not None
+    assert diagnosis.title == "Not authorized"
+
+
+def test_error_pack_classifies_unrecognized_client_as_unknown_access_key():
+    # Athena is json-protocol: an unknown access key ID arrives as this.
+    diagnosis = ATHENA_ERRORS.classify(_client_error("UnrecognizedClientException"))
+    assert diagnosis is not None
+    assert diagnosis.title == "AWS access key not recognized"
+
+
+def test_error_pack_classifies_a_signature_error_through_a_wrapping_cause():
+    wrapped = RuntimeError("query failed")
+    wrapped.__cause__ = _client_error("InvalidSignatureException")
+    diagnosis = ATHENA_ERRORS.classify(wrapped)
+    assert diagnosis is not None
+    assert diagnosis.title == "AWS secret key does not match"
+
+
+def test_error_pack_classifies_missing_workgroup():
+    error = RuntimeError("WorkGroup primary is not found")
+    diagnosis = ATHENA_ERRORS.classify(error)
+    assert diagnosis is not None
+    assert diagnosis.title == "Workgroup not found"
+
+
+def test_error_pack_classifies_missing_result_location():
+    error = RuntimeError("No output location provided for query")
+    diagnosis = ATHENA_ERRORS.classify(error)
+    assert diagnosis is not None
+    assert diagnosis.title == "Query result location not configured"
+
+
+def test_error_pack_classifies_unwritable_result_location():
+    error = RuntimeError("Access denied when writing to location: s3://my-bucket/query-results/abc.csv")
+    diagnosis = ATHENA_ERRORS.classify(error)
+    assert diagnosis is not None
+    assert diagnosis.title == "Cannot write query results"
+
+
+def test_error_pack_classifies_an_unusable_result_bucket():
+    error = RuntimeError("An error occurred (InvalidRequestException): Unable to verify/create output bucket my-bucket")
+    diagnosis = ATHENA_ERRORS.classify(error)
+    assert diagnosis is not None
+    assert diagnosis.title == "Query result bucket not usable"
+
+
+def test_error_pack_classifies_unreachable_endpoint():
+    error = RuntimeError('Could not connect to the endpoint URL: "https://athena.bad.amazonaws.com/"')
+    diagnosis = ATHENA_ERRORS.classify(error)
+    assert diagnosis is not None
+    assert diagnosis.title == "Cannot reach the AWS Athena endpoint"
+
+
+def test_error_pack_classifies_not_authorized():
+    diagnosis = ATHENA_ERRORS.classify(
+        _client_error("SomeOtherException", "User is not authorized to perform: glue:GetTables")
+    )
+    assert diagnosis is not None
+    assert diagnosis.title == "Not authorized"
+
+
+def test_error_pack_folds_in_the_network_pack():
+    diagnosis = ATHENA_ERRORS.classify(socket.gaierror("name resolution failed"))
+    assert diagnosis is not None
+    assert diagnosis.title == "Host could not be resolved"
+
+
+def test_error_pack_classifies_sts_assume_role_denied_as_not_authorized():
+    diagnosis = ATHENA_ERRORS.classify(_client_error("AccessDenied", "not authorized to perform: sts:AssumeRole"))
+    assert diagnosis is not None
+    assert diagnosis.title == "Not authorized"
+
+
+def test_error_pack_classifies_invalid_client_token_as_unknown_access_key():
+    # STS's code, from the assume-role leg.
+    diagnosis = ATHENA_ERRORS.classify(_client_error("InvalidClientTokenId"))
+    assert diagnosis is not None
+    assert diagnosis.title == "AWS access key not recognized"
+
+
+def test_error_pack_classifies_signature_mismatch_as_wrong_secret():
+    diagnosis = ATHENA_ERRORS.classify(_client_error("SignatureDoesNotMatch"))
+    assert diagnosis is not None
+    assert diagnosis.title == "AWS secret key does not match"
+
+
+def test_error_pack_classifies_expired_token():
+    diagnosis = ATHENA_ERRORS.classify(_client_error("ExpiredToken"))
+    assert diagnosis is not None
+    assert diagnosis.title == "AWS session token expired"
+
+
+def test_error_pack_prefers_an_authentication_code_over_not_authorized_text():
+    # The authorization rule sits above AWS_ERRORS; its message fallback must not
+    # swallow a rejected identity.
+    diagnosis = ATHENA_ERRORS.classify(_client_error("ExpiredToken", "not authorized: token expired"))
+    assert diagnosis is not None
+    assert diagnosis.title == "AWS session token expired"
+
+
+def test_error_pack_returns_none_for_unknown_error():
+    assert ATHENA_ERRORS.classify(RuntimeError("something unrelated")) is None
+
+
+def test_client_is_built_lazily_not_at_construction():
+    calls = []
+    provider = AthenaChecks(client_factory=lambda: calls.append(1) or MagicMock())
+    assert calls == []
+    _ = provider.client
+    _ = provider.client
+    assert calls == [1]
+
+
+def test_check_access_surfaces_client_build_failure_for_classification():
+    # An assume-role / credential failure while building the engine now happens
+    # inside CheckAccess (the gate), so its exception propagates from the check and
+    # the runner can classify it - instead of raising at provider construction.
+    error = _client_error("InvalidClientTokenId", "The security token included in the request is invalid")
+    provider = AthenaChecks(client_factory=_raising_factory(error))
+    with pytest.raises(ClientError) as exc_info:
+        provider.check_access()
+    assert ATHENA_ERRORS.classify(exc_info.value).title == "AWS access key not recognized"
+
+
+def test_athena_url():
+    expected = (
+        "awsathena+rest://key:secret_key@athena.us-east-2.amazonaws.com:443"
+        "?s3_staging_dir=s3%3A%2F%2Fpostgres%2Finput%2F&work_group=primary"
+    )
+    assert AthenaConnection.get_connection_url(_config()) == expected
+
+
+def test_athena_url_other_staging_dir():
+    expected = (
+        "awsathena+rest://key:secret_key@athena.us-east-2.amazonaws.com:443"
+        "?s3_staging_dir=s3%3A%2F%2Fpostgres%2Fintput%2F&work_group=primary"
+    )
+    assert AthenaConnection.get_connection_url(_config(s3StagingDir="s3://postgres/intput/")) == expected
+
+
+def test_get_tables_caveats_when_all_targeted_schemas_empty():
+    checks = _checks(catalogId="my_catalog")
+    inspector = MagicMock()
+    inspector.get_schema_names.return_value = ["db1", "db2"]
+    inspector.get_table_names.return_value = []
+    with patch(f"{CONNECTION_MODULE}.inspect", return_value=inspector):
+        evidence = checks.get_tables()
+    assert evidence.caveat is not None
+    assert evidence.caveat.title == "No readable tables"
+
+
+def test_get_tables_passes_without_caveat_when_a_schema_has_tables():
+    checks = _checks()
+    inspector = MagicMock()
+    inspector.get_schema_names.return_value = ["db1", "db2"]
+    inspector.get_table_names.side_effect = lambda schema: ["t1"] if schema == "db2" else []
+    with patch(f"{CONNECTION_MODULE}.inspect", return_value=inspector):
+        evidence = checks.get_tables()
+    assert evidence.caveat is None
+
+
+def test_get_tables_honors_schema_filter_pattern():
+    checks = _checks(schemaFilterPattern=FilterPattern(includes=["db2"]))
+    inspector = MagicMock()
+    inspector.get_schema_names.return_value = ["db1", "db2"]
+    inspector.get_table_names.side_effect = lambda schema: ["t1"] if schema == "db2" else []
+    with patch(f"{CONNECTION_MODULE}.inspect", return_value=inspector):
+        checks.get_tables()
+    probed = {call.args[0] for call in inspector.get_table_names.call_args_list}
+    assert probed == {"db2"}
+
+
+def test_get_tables_caveats_when_filter_matches_nothing():
+    checks = _checks(catalogId="my_catalog", schemaFilterPattern=FilterPattern(includes=["nope"]))
+    inspector = MagicMock()
+    inspector.get_schema_names.return_value = ["db1", "db2"]
+    with patch(f"{CONNECTION_MODULE}.inspect", return_value=inspector):
+        evidence = checks.get_tables()
+    assert evidence.caveat is not None
+    assert evidence.caveat.title == "No schemas visible"
+    assert "my_catalog" in evidence.caveat.remediation
+    inspector.get_table_names.assert_not_called()
+
+
+def test_get_views_does_not_raise_or_caveat_when_empty():
+    checks = _checks()
+    inspector = MagicMock()
+    inspector.get_schema_names.return_value = ["db1"]
+    inspector.get_view_names.return_value = []
+    with patch(f"{CONNECTION_MODULE}.inspect", return_value=inspector):
+        evidence = checks.get_views()
+    assert evidence.caveat is None
+
+
+def test_get_views_reports_visible_views():
+    checks = _checks()
+    inspector = MagicMock()
+    inspector.get_schema_names.return_value = ["db1"]
+    inspector.get_view_names.return_value = ["v1"]
+    with patch(f"{CONNECTION_MODULE}.inspect", return_value=inspector):
+        evidence = checks.get_views()
+    assert evidence.caveat is None
+    assert "views visible" in evidence.summary
+
+
+def test_get_tables_caveat_is_actionable():
+    checks = _checks(catalogId="my_catalog")
+    inspector = MagicMock()
+    inspector.get_schema_names.return_value = ["db1"]
+    inspector.get_table_names.return_value = []
+    with patch(f"{CONNECTION_MODULE}.inspect", return_value=inspector):
+        evidence = checks.get_tables()
+    remediation = evidence.caveat.remediation
+    assert "Lake Formation" in remediation
+    assert any(keyword in remediation.lower() for keyword in ("describe", "select", "grant"))
+
+
+def test_get_tables_caps_probing_at_max_schemas():
+    checks = _checks(catalogId="my_catalog")
+    inspector = MagicMock()
+    inspector.get_schema_names.return_value = [f"db{i}" for i in range(500)]
+    inspector.get_table_names.return_value = []
+    with patch(f"{CONNECTION_MODULE}.inspect", return_value=inspector):
+        evidence = checks.get_tables()
+    assert evidence.caveat is not None
+    assert inspector.get_table_names.call_count <= athena_connection.MAX_SCHEMAS_TO_PROBE
+
+
+def test_targeted_schemas_listed_once_across_table_and_view_steps():
+    checks = _checks()
+    inspector = MagicMock()
+    inspector.get_schema_names.return_value = ["db1"]
+    inspector.get_table_names.return_value = ["t1"]
+    inspector.get_view_names.return_value = []
+    with patch(f"{CONNECTION_MODULE}.inspect", return_value=inspector):
+        checks.get_tables()
+        checks.get_views()
+    assert inspector.get_schema_names.call_count == 1

ingestion/tests/unit/source/database/azuresql/test_connection.py

@@ -0,0 +1,45 @@
+#  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 AzureSQL BaseConnection wiring.
+
+The URL building (Active Directory ODBC and standard) is covered in
+tests/unit/test_build_connection_url.py.
+"""
+
+from unittest.mock import patch
+
+from metadata.generated.schema.entity.services.connections.database.azureSQLConnection import (
+    AzureSQLConnection as AzureSQLConnectionConfig,
+)
+from metadata.ingestion.connections.connection import BaseConnection
+from metadata.ingestion.source.database.azuresql.connection import AzureSQLConnection
+
+CONNECTION_MODULE = "metadata.ingestion.source.database.azuresql.connection"
+
+
+def _config() -> AzureSQLConnectionConfig:
+    return AzureSQLConnectionConfig(
+        username="user",
+        password="pass",
+        hostPort="myhost:1433",
+        database="mydb",
+        driver="ODBC Driver 18 for SQL Server",
+    )
+
+
+def test_azuresql_connection_is_base_connection():
+    assert issubclass(AzureSQLConnection, BaseConnection)
+
+
+def test_get_client_uses_the_module_url_builder():
+    with patch(f"{CONNECTION_MODULE}.create_generic_db_connection") as mock_connection:
+        _ = AzureSQLConnection(_config()).client
+    assert mock_connection.call_args.kwargs["get_connection_url_fn"].__name__ == "get_connection_url"

ingestion/tests/unit/source/database/bigquery/test_connection.py

@@ -0,0 +1,244 @@
+#  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"

ingestion/tests/unit/source/database/bigtable/test_connection.py

@@ -0,0 +1,22 @@
+#  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 BigTable BaseConnection wiring (non-Engine: MultiProjectClient).
+
+Client building requires a full GCP service-account credential; the end-to-end
+behaviour is exercised by tests/unit/topology/database/test_bigtable.py.
+"""
+
+from metadata.ingestion.connections.connection import BaseConnection
+from metadata.ingestion.source.database.bigtable.connection import BigTableConnection
+
+
+def test_bigtable_connection_is_base_connection():
+    assert issubclass(BigTableConnection, BaseConnection)

ingestion/tests/unit/source/database/burstiq/test_connection.py

@@ -0,0 +1,22 @@
+#  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 BurstIQConnection BaseConnection wiring (non-Engine client)."""
+
+import pytest
+
+from metadata.ingestion.connections.connection import BaseConnection
+
+# The connector SDK may be an optional dependency absent from the unit env.
+connection = pytest.importorskip("metadata.ingestion.source.database.burstiq.connection")
+
+
+def test_burstiq_connection_is_base_connection():
+    assert issubclass(connection.BurstIQConnection, BaseConnection)

ingestion/tests/unit/source/database/cassandra/test_connection.py

@@ -0,0 +1,45 @@
+#  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 Cassandra BaseConnection wiring (non-Engine: driver Session)."""
+
+from unittest.mock import patch
+
+from metadata.generated.schema.entity.services.connections.database.cassandraConnection import (
+    CassandraConnection as CassandraConnectionConfig,
+)
+from metadata.ingestion.connections.connection import BaseConnection
+from metadata.ingestion.source.database.cassandra.connection import CassandraConnection
+
+CONNECTION_MODULE = "metadata.ingestion.source.database.cassandra.connection"
+
+
+def _config() -> CassandraConnectionConfig:
+    return CassandraConnectionConfig(hostPort="localhost:9042")
+
+
+def test_cassandra_connection_is_base_connection():
+    assert issubclass(CassandraConnection, BaseConnection)
+
+
+def test_get_client_connects_a_cluster_session():
+    with patch(f"{CONNECTION_MODULE}.Cluster") as mock_cluster:
+        session = CassandraConnection(_config()).client
+    mock_cluster.assert_called_once()
+    mock_cluster.return_value.connect.assert_called_once_with()
+    assert session is mock_cluster.return_value.connect.return_value
+
+
+def test_close_shuts_down_the_cluster():
+    with patch(f"{CONNECTION_MODULE}.Cluster") as mock_cluster:
+        connection = CassandraConnection(_config())
+        _ = connection.client
+        connection.close()
+    mock_cluster.return_value.shutdown.assert_called_once_with()

ingestion/tests/unit/source/database/clickhouse/test_connection.py

@@ -0,0 +1,63 @@
+#  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 Clickhouse connection handling."""
+
+from unittest.mock import patch
+
+from metadata.generated.schema.entity.services.connections.database.clickhouseConnection import (
+    ClickhouseConnection as ClickhouseConnectionConfig,
+)
+from metadata.generated.schema.entity.services.connections.database.clickhouseConnection import (
+    ClickhouseScheme,
+)
+from metadata.ingestion.connections.builders import get_connection_url_common
+from metadata.ingestion.connections.connection import BaseConnection
+from metadata.ingestion.source.database.clickhouse.connection import ClickhouseConnection
+
+CONNECTION_MODULE = "metadata.ingestion.source.database.clickhouse.connection"
+
+
+def test_clickhouse_connection_is_base_connection():
+    assert issubclass(ClickhouseConnection, BaseConnection)
+
+
+def test_url_without_database_schema():
+    connection = ClickhouseConnectionConfig(
+        username="username",
+        hostPort="localhost:8123",
+        scheme=ClickhouseScheme.clickhouse_http,
+        databaseSchema=None,
+    )
+    assert get_connection_url_common(connection) == "clickhouse+http://username:@localhost:8123"
+
+
+def test_url_with_database_schema():
+    connection = ClickhouseConnectionConfig(
+        username="username",
+        hostPort="localhost:8123",
+        scheme=ClickhouseScheme.clickhouse_http,
+        databaseSchema="default",
+    )
+    assert get_connection_url_common(connection) == "clickhouse+http://username:@localhost:8123/default"
+
+
+def test_secure_and_keyfile_move_to_connection_arguments():
+    connection = ClickhouseConnectionConfig(
+        username="username",
+        hostPort="localhost:8123",
+        scheme=ClickhouseScheme.clickhouse_http,
+        secure=True,
+        keyfile="/path/to/key.pem",
+    )
+    with patch(f"{CONNECTION_MODULE}.create_generic_db_connection"):
+        _ = ClickhouseConnection(connection).client
+    assert connection.connectionArguments.root["secure"] is True
+    assert connection.connectionArguments.root["keyfile"] == "/path/to/key.pem"

ingestion/tests/unit/source/database/cockroach/test_connection.py

@@ -0,0 +1,37 @@
+#  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 CockroachDB connection handling."""
+
+from metadata.generated.schema.entity.services.connections.database.cockroachConnection import (
+    CockroachConnection as CockroachConnectionConfig,
+)
+from metadata.generated.schema.entity.services.connections.database.cockroachConnection import (
+    CockroachScheme,
+)
+from metadata.generated.schema.entity.services.connections.database.common.basicAuth import (
+    BasicAuth,
+)
+from metadata.ingestion.source.database.cockroach.connection import CockroachConnection
+
+
+def test_basic_auth_builds_expected_url():
+    connection = CockroachConnectionConfig(
+        username="openmetadata_user",
+        authType=BasicAuth(password="openmetadata_password"),
+        hostPort="localhost:26257",
+        database="openmetadata_db",
+        scheme=CockroachScheme.cockroachdb_psycopg2,
+    )
+    engine = CockroachConnection(connection).client
+    assert (
+        engine.url.render_as_string(hide_password=False)
+        == "cockroachdb+psycopg2://openmetadata_user:openmetadata_password@localhost:26257/openmetadata_db"
+    )

ingestion/tests/unit/source/database/couchbase/test_connection.py

@@ -0,0 +1,23 @@
+#  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 Couchbase BaseConnection wiring (non-Engine: SDK Cluster).
+
+The couchbase SDK is an optional dependency imported lazily inside
+``_get_client``/``test_connection``, so this module only asserts the
+BaseConnection wiring without importing the SDK.
+"""
+
+from metadata.ingestion.connections.connection import BaseConnection
+from metadata.ingestion.source.database.couchbase.connection import CouchbaseConnection
+
+
+def test_couchbase_connection_is_base_connection():
+    assert issubclass(CouchbaseConnection, BaseConnection)

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/<endpointId> or /sql/1.0/warehouses/<warehouseId>"
+        )
+    )
+    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()

ingestion/tests/unit/source/database/db2/test_connection.py

@@ -0,0 +1,92 @@
+#  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 DB2 connection handling."""
+
+from unittest.mock import patch
+
+import pytest
+
+from metadata.generated.schema.entity.services.connections.database.db2Connection import (
+    Db2Connection as Db2ConnectionConfig,
+)
+from metadata.generated.schema.entity.services.connections.database.db2Connection import (
+    Db2Scheme,
+)
+from metadata.ingestion.connections.connection import BaseConnection
+from metadata.ingestion.source.database.db2.connection import (
+    Db2Connection,
+    Db2IbmiStrategy,
+)
+
+CONNECTION_MODULE = "metadata.ingestion.source.database.db2.connection"
+
+
+def _ibmi_config() -> Db2ConnectionConfig:
+    return Db2ConnectionConfig(
+        scheme=Db2Scheme.ibmi,
+        username="openmetadata_user",
+        password="openmetadata_password",
+        hostPort="myhost:8471",
+        database="MYDB",
+    )
+
+
+def _db2_config() -> Db2ConnectionConfig:
+    return Db2ConnectionConfig(
+        scheme=Db2Scheme.db2_ibm_db,
+        username="openmetadata_user",
+        password="openmetadata_password",
+        hostPort="localhost:50000",
+        database="testdb",
+    )
+
+
+def test_db2_connection_is_base_connection():
+    assert issubclass(Db2Connection, BaseConnection)
+
+
+def test_ibmi_url_drops_the_port():
+    config = _ibmi_config()
+    assert (
+        Db2IbmiStrategy(config).get_connection_url(config)
+        == "ibmi://openmetadata_user:openmetadata_password@myhost/MYDB"
+    )
+
+
+def test_ibmi_args_pass_the_port_via_connect_args():
+    config = _ibmi_config()
+    assert Db2IbmiStrategy(config).get_connection_args(config)["port"] == 8471
+
+
+def test_ibmi_args_reject_a_non_numeric_port():
+    config = _ibmi_config()
+    config.hostPort = "myhost:not-a-port"
+    with pytest.raises(ValueError, match="Invalid port"):
+        Db2IbmiStrategy(config).get_connection_args(config)
+
+
+def test_get_client_dispatches_ibmi_scheme_to_ibmi_strategy():
+    config = _ibmi_config()
+    with patch(f"{CONNECTION_MODULE}.create_generic_db_connection") as mock_connection:
+        _ = Db2Connection(config).client
+    url_fn = mock_connection.call_args.kwargs["get_connection_url_fn"]
+    assert url_fn.__name__ == "get_connection_url"
+    assert url_fn(config) == "ibmi://openmetadata_user:openmetadata_password@myhost/MYDB"
+
+
+def test_get_client_dispatches_db2_scheme_to_standard_strategy():
+    config = _db2_config()
+    with (
+        patch(f"{CONNECTION_MODULE}.create_generic_db_connection") as mock_connection,
+        patch(f"{CONNECTION_MODULE}.check_ssl_and_init", return_value=None),
+    ):
+        _ = Db2Connection(config).client
+    assert mock_connection.call_args.kwargs["get_connection_url_fn"].__name__ == "get_connection_url_common"

ingestion/tests/unit/source/database/deltalake/test_connection.py

@@ -0,0 +1,22 @@
+#  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 DeltaLakeConnection BaseConnection wiring (non-Engine client)."""
+
+import pytest
+
+from metadata.ingestion.connections.connection import BaseConnection
+
+# The connector SDK may be an optional dependency absent from the unit env.
+connection = pytest.importorskip("metadata.ingestion.source.database.deltalake.connection")
+
+
+def test_deltalake_connection_is_base_connection():
+    assert issubclass(connection.DeltaLakeConnection, BaseConnection)

ingestion/tests/unit/source/database/domodatabase/test_connection.py

@@ -0,0 +1,22 @@
+#  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 DomoDatabaseConnection BaseConnection wiring (non-Engine client)."""
+
+import pytest
+
+from metadata.ingestion.connections.connection import BaseConnection
+
+# The connector SDK may be an optional dependency absent from the unit env.
+connection = pytest.importorskip("metadata.ingestion.source.database.domodatabase.connection")
+
+
+def test_domodatabase_connection_is_base_connection():
+    assert issubclass(connection.DomoDatabaseConnection, BaseConnection)

ingestion/tests/unit/source/database/doris/test_connection.py

@@ -0,0 +1,42 @@
+#  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 Doris connection handling."""
+
+from metadata.generated.schema.entity.services.connections.database.dorisConnection import (
+    DorisConnection as DorisConnectionConfig,
+)
+from metadata.generated.schema.entity.services.connections.database.dorisConnection import (
+    DorisScheme,
+)
+from metadata.ingestion.connections.builders import get_connection_url_common
+from metadata.ingestion.connections.connection import BaseConnection
+from metadata.ingestion.source.database.doris.connection import DorisConnection
+
+
+def test_doris_connection_is_base_connection():
+    assert issubclass(DorisConnection, BaseConnection)
+
+
+def test_basic_auth_builds_expected_url():
+    connection = DorisConnectionConfig(
+        username="openmetadata_user",
+        password="openmetadata_password",
+        hostPort="localhost:9030",
+        databaseSchema="openmetadata_db",
+        scheme=DorisScheme.doris,
+    )
+    # Assert via the URL builder, not .client: the doris dialect (pydoris-custom,
+    # mysqlclient DBAPI) is installed --no-deps in runtime images only and is
+    # absent from the unit-test environment.
+    assert (
+        get_connection_url_common(connection)
+        == "doris://openmetadata_user:openmetadata_password@localhost:9030/openmetadata_db"
+    )

ingestion/tests/unit/source/database/druid/test_connection.py

@@ -0,0 +1,50 @@
+#  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 Druid connection handling.
+
+The Druid dialect (pydruid) is an optional extra absent from the unit-test
+environment, so URL parity is asserted via the connector's get_connection_url
+rather than by instantiating an engine.
+"""
+
+from metadata.generated.schema.entity.services.connections.database.druidConnection import (
+    DruidConnection as DruidConnectionConfig,
+)
+from metadata.generated.schema.entity.services.connections.database.druidConnection import (
+    DruidScheme,
+)
+from metadata.ingestion.connections.connection import BaseConnection
+from metadata.ingestion.source.database.druid.connection import DruidConnection
+
+
+def test_druid_connection_is_base_connection():
+    assert issubclass(DruidConnection, BaseConnection)
+
+
+def test_get_connection_url_appends_druid_sql_path():
+    connection = DruidConnectionConfig(
+        scheme=DruidScheme.druid,
+        hostPort="localhost:8082",
+    )
+    assert DruidConnection.get_connection_url(connection) == "druid://localhost:8082/druid/v2/sql"
+
+
+def test_get_connection_url_with_basic_auth():
+    connection = DruidConnectionConfig(
+        scheme=DruidScheme.druid,
+        username="openmetadata_user",
+        password="openmetadata_password",
+        hostPort="localhost:8082",
+    )
+    assert (
+        DruidConnection.get_connection_url(connection)
+        == "druid://openmetadata_user:openmetadata_password@localhost:8082/druid/v2/sql"
+    )

ingestion/tests/unit/source/database/dynamodb/test_connection.py

@@ -0,0 +1,37 @@
+#  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 DynamoDB BaseConnection wiring (non-Engine: boto3 client)."""
+
+from unittest.mock import patch
+
+from metadata.generated.schema.entity.services.connections.database.dynamoDBConnection import (
+    DynamoDBConnection as DynamoDBConnectionConfig,
+)
+from metadata.generated.schema.security.credentials.awsCredentials import AWSCredentials
+from metadata.ingestion.connections.connection import BaseConnection
+from metadata.ingestion.source.database.dynamodb.connection import DynamoDBConnection
+
+CONNECTION_MODULE = "metadata.ingestion.source.database.dynamodb.connection"
+
+
+def _config() -> DynamoDBConnectionConfig:
+    return DynamoDBConnectionConfig(awsConfig=AWSCredentials(awsRegion="us-east-1"))
+
+
+def test_dynamodb_connection_is_base_connection():
+    assert issubclass(DynamoDBConnection, BaseConnection)
+
+
+def test_get_client_builds_the_dynamo_client():
+    with patch(f"{CONNECTION_MODULE}.AWSClient") as mock_aws:
+        client = DynamoDBConnection(_config()).client
+    mock_aws.return_value.get_dynamo_client.assert_called_once_with()
+    assert client is mock_aws.return_value.get_dynamo_client.return_value

ingestion/tests/unit/source/database/exasol/test_connection.py

@@ -0,0 +1,56 @@
+#  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 Exasol connection handling."""
+
+from metadata.generated.schema.entity.services.connections.database.exasolConnection import (
+    ExasolConnection as ExasolConnectionConfig,
+)
+from metadata.generated.schema.entity.services.connections.database.exasolConnection import (
+    ExasolScheme,
+    Tls,
+)
+from metadata.ingestion.connections.connection import BaseConnection
+from metadata.ingestion.source.database.exasol.connection import ExasolConnection
+
+
+def _config(tls: Tls) -> ExasolConnectionConfig:
+    return ExasolConnectionConfig(
+        scheme=ExasolScheme.exa_websocket,
+        username="admin",
+        password="password",
+        hostPort="localhost:8563",
+        tls=tls,
+    )
+
+
+def test_exasol_connection_is_base_connection():
+    assert issubclass(ExasolConnection, BaseConnection)
+
+
+def test_url_validate_certificate_has_no_tls_params():
+    assert (
+        ExasolConnection.get_connection_url(_config(Tls.validate_certificate))
+        == "exa+websocket://admin:password@localhost:8563"
+    )
+
+
+def test_url_ignore_certificate():
+    assert (
+        ExasolConnection.get_connection_url(_config(Tls.ignore_certificate))
+        == "exa+websocket://admin:password@localhost:8563?SSLCertificate=SSL_VERIFY_NONE"
+    )
+
+
+def test_url_disable_tls():
+    assert (
+        ExasolConnection.get_connection_url(_config(Tls.disable_tls)) == "exa+websocket://admin:password@localhost:8563"
+        "?SSLCertificate=SSL_VERIFY_NONE&ENCRYPTION=no"
+    )

ingestion/tests/unit/source/database/glue/test_connection.py

@@ -0,0 +1,281 @@
+#  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.
+"""Glue test-connection: its checks match its shipped definition (non-Engine: boto3 client)."""
+
+import json
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+from botocore.exceptions import ClientError
+
+from metadata.core.connections.test_connection.check import CheckError, collect_checks
+from metadata.generated.schema.entity.services.connections.database.glueConnection import (
+    GlueConnection as GlueConnectionConfig,
+)
+from metadata.generated.schema.security.credentials.awsCredentials import AWSCredentials
+from metadata.ingestion.connections.connection import BaseConnection
+from metadata.ingestion.source.database.glue.connection import (
+    GLUE_ERRORS,
+    GlueChecks,
+    GlueConnection,
+    list_databases,
+)
+
+CONNECTION_MODULE = "metadata.ingestion.source.database.glue.connection"
+
+_SEED = (
+    Path(__file__).parents[6]
+    / "openmetadata-service/src/main/resources/json/data"
+    / "testConnections/database/glue.json"
+)
+
+
+def _config() -> GlueConnectionConfig:
+    return GlueConnectionConfig(awsConfig=AWSCredentials(awsRegion="us-east-1"))
+
+
+def _client(pages: dict[str, list[dict]] | None = None) -> MagicMock:
+    """A Glue client whose paginators yield one page per operation."""
+    responses = pages or {}
+    client = MagicMock()
+    client.get_paginator.side_effect = lambda operation: _paginator(responses.get(operation, [{}]))
+    return client
+
+
+def _paginator(pages: list[dict]) -> MagicMock:
+    paginator = MagicMock()
+    paginator.paginate.return_value = pages
+    return paginator
+
+
+def _checks(client=None) -> GlueChecks:
+    return GlueChecks(connect=lambda: client)
+
+
+def _check_names() -> set[str]:
+    return {step.value for step in collect_checks(_checks())}
+
+
+def _client_error(code: str, operation: str = "GetDatabases") -> ClientError:
+    return ClientError({"Error": {"Code": code, "Message": "denied"}}, operation)
+
+
+def test_glue_connection_is_base_connection():
+    assert issubclass(GlueConnection, BaseConnection)
+
+
+def test_get_client_builds_the_glue_client():
+    with patch(f"{CONNECTION_MODULE}.AWSClient") as mock_aws:
+        client = GlueConnection(_config()).client
+
+    mock_aws.return_value.get_glue_client.assert_called_once_with()
+    assert client is mock_aws.return_value.get_glue_client.return_value
+
+
+def test_glue_checks_cover_expected_steps():
+    assert _check_names() == {"GetDatabases", "GetTables"}
+
+
+def test_glue_checks_match_definition_seed():
+    definition_steps = {step["name"] for step in json.loads(_SEED.read_text())["steps"]}
+    assert _check_names() == definition_steps
+
+
+def test_get_databases_is_the_connection_gate():
+    steps = json.loads(_SEED.read_text())["steps"]
+    assert steps[0]["name"] == "GetDatabases"
+    assert steps[0]["category"] == "ConnectionGate"
+
+
+def test_building_checks_does_not_build_the_client():
+    with patch.object(GlueConnection, "_get_client") as mock_build:
+        GlueConnection(_config()).checks()
+
+    mock_build.assert_not_called()
+
+
+def test_get_databases_enumerates_the_catalog():
+    client = _client({"get_databases": [{"DatabaseList": [{"Name": "a"}, {"Name": "b"}]}]})
+
+    evidence = _checks(client).get_databases()
+
+    assert evidence.summary == "2 databases enumerated"
+    assert evidence.command == "glue:GetDatabases"
+    assert evidence.caveat is None
+
+
+def test_get_databases_caps_reported_count_and_flags_more():
+    paginator = _paginator([{"DatabaseList": [{"Name": "a"}, {"Name": "b"}, {"Name": "c"}]}])
+    client = MagicMock()
+    client.get_paginator.return_value = paginator
+
+    evidence = list_databases(client, limit=2)
+
+    assert evidence.summary == "2 databases enumerated (showing first 2; more exist)"
+    client.get_paginator.assert_called_once_with("get_databases")
+    paginator.paginate.assert_called_once_with(PaginationConfig={"MaxItems": 3})
+
+
+def test_get_databases_warns_when_catalog_is_empty():
+    client = _client({"get_databases": [{"DatabaseList": []}]})
+
+    evidence = _checks(client).get_databases()
+
+    assert evidence.summary == "0 databases enumerated"
+    assert evidence.caveat is not None
+    assert "No databases visible" in evidence.caveat.title
+
+
+def test_get_databases_failure_reports_the_command_it_ran():
+    client = _client()
+    cause = _client_error("AccessDeniedException")
+    client.get_paginator.side_effect = cause
+
+    with pytest.raises(CheckError) as failure:
+        _checks(client).get_databases()
+
+    assert failure.value.cause is cause
+    assert failure.value.evidence.command == "glue:GetDatabases"
+
+
+def test_get_tables_probes_the_first_database():
+    client = _client(
+        {
+            "get_databases": [{"DatabaseList": [{"Name": "sales"}]}],
+            "get_tables": [{"TableList": [{"Name": "orders"}]}],
+        }
+    )
+
+    evidence = _checks(client).get_tables()
+
+    assert evidence.summary == "1 table in database 'sales'"
+    assert evidence.command == "glue:GetTables (DatabaseName=sales)"
+
+
+def test_get_tables_reports_a_caveat_when_no_database_can_be_probed():
+    client = _client({"get_databases": [{"DatabaseList": []}]})
+
+    evidence = _checks(client).get_tables()
+
+    assert evidence.command == "glue:GetDatabases"
+    assert evidence.caveat is not None
+    assert evidence.caveat.title == "No tables visible"
+
+
+def test_get_tables_skips_an_empty_leading_database():
+    # A catalog commonly leads with an empty 'default'; the probe must not stop there.
+    client = MagicMock()
+    tables = {"default": [{"TableList": []}], "ecommerce": [{"TableList": [{"Name": "orders"}]}]}
+    calls = {"n": 0}
+
+    def paginator(operation):
+        if operation == "get_databases":
+            return _paginator([{"DatabaseList": [{"Name": "default"}, {"Name": "ecommerce"}]}])
+        calls["n"] += 1
+        return _paginator(tables["default" if calls["n"] == 1 else "ecommerce"])
+
+    client.get_paginator.side_effect = paginator
+
+    evidence = _checks(client).get_tables()
+
+    assert evidence.summary == "1 table in database 'ecommerce'"
+    assert evidence.command == "glue:GetTables (DatabaseName=ecommerce)"
+    assert evidence.caveat is None
+
+
+def test_get_tables_warns_when_no_database_exposes_a_table():
+    client = _client(
+        {
+            "get_databases": [{"DatabaseList": [{"Name": "a"}, {"Name": "b"}]}],
+            "get_tables": [{"TableList": []}],
+        }
+    )
+
+    evidence = _checks(client).get_tables()
+
+    assert evidence.summary == "no tables in the first 2 databases"
+    assert evidence.caveat is not None
+    assert evidence.caveat.title == "No tables visible"
+
+
+def test_get_tables_failure_names_the_database_it_probed():
+    cause = _client_error("EntityNotFoundException", "GetTables")
+    client = _client({"get_databases": [{"DatabaseList": [{"Name": "sales"}]}]})
+    tables_paginator = MagicMock()
+    tables_paginator.paginate.side_effect = cause
+    client.get_paginator.side_effect = lambda operation: (
+        tables_paginator if operation == "get_tables" else _paginator([{"DatabaseList": [{"Name": "sales"}]}])
+    )
+
+    with pytest.raises(CheckError) as failure:
+        _checks(client).get_tables()
+
+    assert failure.value.cause is cause
+    assert failure.value.evidence.command == "glue:GetTables (DatabaseName=sales)"
+
+
+def test_error_pack_access_denied():
+    diagnosis = GLUE_ERRORS.classify(
+        ClientError(
+            {
+                "Error": {
+                    "Code": "AccessDeniedException",
+                    "Message": "User is not authorized to perform: glue:GetDatabases",
+                }
+            },
+            "GetDatabases",
+        )
+    )
+    assert diagnosis is not None
+    assert "authorized" in diagnosis.title.lower()
+
+
+def test_error_pack_entity_not_found():
+    diagnosis = GLUE_ERRORS.classify(_client_error("EntityNotFoundException", "GetTables"))
+    assert diagnosis is not None
+    assert diagnosis.title == "Glue database not found"
+
+
+def test_error_pack_inherits_the_shared_aws_diagnoses():
+    # Glue is json-protocol: expired creds arrive as ExpiredTokenException.
+    diagnosis = GLUE_ERRORS.classify(_client_error("ExpiredTokenException"))
+    assert diagnosis is not None
+    assert diagnosis.title == "AWS session token expired"
+
+
+def test_error_pack_classifies_an_unknown_access_key():
+    diagnosis = GLUE_ERRORS.classify(_client_error("UnrecognizedClientException"))
+    assert diagnosis is not None
+    assert diagnosis.title == "AWS access key not recognized"
+
+
+def test_error_pack_classifies_a_wrong_secret():
+    diagnosis = GLUE_ERRORS.classify(_client_error("InvalidSignatureException"))
+    assert diagnosis is not None
+    assert diagnosis.title == "AWS secret key does not match"
+
+
+def test_error_pack_classifies_an_sts_assume_role_denial():
+    # The assume-role leg denies with the suffix-less code.
+    diagnosis = GLUE_ERRORS.classify(_client_error("AccessDenied", "AssumeRole"))
+    assert diagnosis is not None
+    assert diagnosis.title == "Not authorized"
+
+
+def test_error_pack_inherits_the_network_pack_through_aws_errors():
+    diagnosis = GLUE_ERRORS.classify(ConnectionRefusedError("refused"))
+    assert diagnosis is not None
+    assert "refused" in diagnosis.title.lower()
+
+
+def test_error_pack_unmatched_returns_none():
+    assert GLUE_ERRORS.classify(Exception("novel error")) is None

ingestion/tests/unit/source/database/greenplum/test_connection.py

@@ -0,0 +1,37 @@
+#  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 Greenplum connection handling."""
+
+from metadata.generated.schema.entity.services.connections.database.common.basicAuth import (
+    BasicAuth,
+)
+from metadata.generated.schema.entity.services.connections.database.greenplumConnection import (
+    GreenplumConnection as GreenplumConnectionConfig,
+)
+from metadata.generated.schema.entity.services.connections.database.greenplumConnection import (
+    GreenplumScheme,
+)
+from metadata.ingestion.source.database.greenplum.connection import GreenplumConnection
+
+
+def test_basic_auth_builds_expected_url():
+    connection = GreenplumConnectionConfig(
+        username="openmetadata_user",
+        authType=BasicAuth(password="openmetadata_password"),
+        hostPort="localhost:5432",
+        database="openmetadata_db",
+        scheme=GreenplumScheme.postgresql_psycopg2,
+    )
+    engine = GreenplumConnection(connection).client
+    assert (
+        engine.url.render_as_string(hide_password=False)
+        == "postgresql+psycopg2://openmetadata_user:openmetadata_password@localhost:5432/openmetadata_db"
+    )

ingestion/tests/unit/source/database/hive/test_connection.py

@@ -0,0 +1,167 @@
+#  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 Hive connection handling."""
+
+from unittest.mock import patch
+
+from metadata.generated.schema.entity.services.connections.database.hiveConnection import (
+    Auth,
+    HiveScheme,
+)
+from metadata.generated.schema.entity.services.connections.database.hiveConnection import (
+    HiveConnection as HiveConnectionConfig,
+)
+from metadata.ingestion.connections.connection import BaseConnection
+from metadata.ingestion.source.database.hive.connection import HiveConnection
+
+CONNECTION_MODULE = "metadata.ingestion.source.database.hive.connection"
+
+
+def _config(**kwargs) -> HiveConnectionConfig:
+    base = {
+        "scheme": HiveScheme.hive,
+        "hostPort": "localhost:10000",
+        "databaseSchema": "default",
+    }
+    base.update(kwargs)
+    return HiveConnectionConfig(**base)
+
+
+def test_hive_connection_is_base_connection():
+    assert issubclass(HiveConnection, BaseConnection)
+
+
+def test_get_client_injects_auth_and_kerberos_args():
+    config = _config(auth=Auth.LDAP, kerberosServiceName="hive")
+    with (
+        patch(f"{CONNECTION_MODULE}.create_generic_db_connection") as mock_connection,
+        patch(f"{CONNECTION_MODULE}.check_ssl_and_init", return_value=None),
+    ):
+        _ = HiveConnection(config).client
+    root = mock_connection.call_args.kwargs["connection"].connectionArguments.root
+    assert root["auth"] == "LDAP"
+    assert root["kerberos_service_name"] == "hive"
+
+
+def test_get_client_runs_ssl_setup_when_manager_present():
+    config = _config(useSSL=True)
+    ssl_manager = type("M", (), {"setup_ssl": staticmethod(lambda conn: conn)})()
+    with (
+        patch(f"{CONNECTION_MODULE}.create_generic_db_connection") as mock_connection,
+        patch(f"{CONNECTION_MODULE}.check_ssl_and_init", return_value=ssl_manager),
+    ):
+        _ = HiveConnection(config).client
+    assert mock_connection.call_args.kwargs["connection"].connectionArguments.root["use_ssl"] is True
+
+
+def test_hive_url():
+    assert (
+        HiveConnection.get_connection_url(HiveConnectionConfig(scheme=HiveScheme.hive, hostPort="localhost:10000"))
+        == "hive://localhost:10000"
+    )
+    assert (
+        HiveConnection.get_connection_url(HiveConnectionConfig(scheme=HiveScheme.hive_http, hostPort="localhost:1000"))
+        == "hive+http://localhost:1000"
+    )
+    assert (
+        HiveConnection.get_connection_url(HiveConnectionConfig(scheme=HiveScheme.hive_https, hostPort="localhost:1000"))
+        == "hive+https://localhost:1000"
+    )
+
+
+def test_hive_url_custom_auth():
+    config = HiveConnectionConfig(
+        scheme=HiveScheme.hive.value,
+        username="username",
+        password="password",
+        hostPort="localhost:10000",
+        connectionArguments={"auth": "CUSTOM"},
+    )
+    assert HiveConnection.get_connection_url(config) == "hive://username:password@localhost:10000"
+
+    config = HiveConnectionConfig(
+        scheme=HiveScheme.hive.value,
+        username="username@444",
+        password="password@333",
+        hostPort="localhost:10000",
+        connectionArguments={"auth": "CUSTOM"},
+    )
+    assert HiveConnection.get_connection_url(config) == "hive://username%40444:password%40333@localhost:10000"
+
+
+def test_hive_url_conn_options_with_db():
+    config = HiveConnectionConfig(
+        hostPort="localhost:10000",
+        databaseSchema="test_db",
+        connectionOptions={"Key": "Value"},
+    )
+    assert HiveConnection.get_connection_url(config) == "hive://localhost:10000/test_db?Key=Value"
+
+
+def test_hive_url_conn_options_without_db():
+    config = HiveConnectionConfig(
+        hostPort="localhost:10000",
+        connectionOptions={"Key": "Value"},
+    )
+    assert HiveConnection.get_connection_url(config) == "hive://localhost:10000?Key=Value"
+
+
+def test_hive_url_with_kerberos_auth():
+    config = HiveConnectionConfig(
+        scheme=HiveScheme.hive.value,
+        hostPort="localhost:10000",
+        connectionArguments={
+            "auth": "KERBEROS",
+            "kerberos_service_name": "hive",
+        },
+    )
+    assert HiveConnection.get_connection_url(config) == "hive://localhost:10000"
+
+
+def test_hive_url_with_ldap_auth():
+    config = HiveConnectionConfig(
+        scheme=HiveScheme.hive.value,
+        username="username",
+        password="password",
+        hostPort="localhost:10000",
+        connectionArguments={"auth": "LDAP"},
+    )
+    assert HiveConnection.get_connection_url(config) == "hive://username:password@localhost:10000"
+
+
+def test_hive_url_without_auth():
+    config = HiveConnectionConfig(
+        scheme=HiveScheme.hive.value,
+        username="username",
+        password="password",
+        hostPort="localhost:10000",
+        connectionArguments={"customKey": "value"},
+    )
+    assert HiveConnection.get_connection_url(config) == "hive://username:password@localhost:10000"
+
+
+def test_hive_url_without_connection_arguments():
+    config = HiveConnectionConfig(
+        scheme=HiveScheme.hive.value,
+        username="username",
+        password="password",
+        hostPort="localhost:10000",
+    )
+    assert HiveConnection.get_connection_url(config) == "hive://username:password@localhost:10000"
+
+
+def test_hive_url_without_connection_arguments_pass():
+    config = HiveConnectionConfig(
+        scheme=HiveScheme.hive.value,
+        username="username",
+        hostPort="localhost:10000",
+    )
+    assert HiveConnection.get_connection_url(config) == "hive://username@localhost:10000"

ingestion/tests/unit/source/database/impala/test_connection.py

@@ -0,0 +1,140 @@
+#  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 Impala connection handling."""
+
+from unittest.mock import patch
+
+from metadata.generated.schema.entity.services.connections.database.impalaConnection import (
+    AuthMechanism,
+    ImpalaScheme,
+)
+from metadata.generated.schema.entity.services.connections.database.impalaConnection import (
+    ImpalaConnection as ImpalaConnectionConfig,
+)
+from metadata.ingestion.connections.connection import BaseConnection
+from metadata.ingestion.source.database.impala.connection import ImpalaConnection
+
+CONNECTION_MODULE = "metadata.ingestion.source.database.impala.connection"
+
+
+def _config(**kwargs) -> ImpalaConnectionConfig:
+    base = {
+        "scheme": ImpalaScheme.impala,
+        "hostPort": "localhost:21050",
+        "databaseSchema": "default",
+    }
+    base.update(kwargs)
+    return ImpalaConnectionConfig(**base)
+
+
+def test_impala_connection_is_base_connection():
+    assert issubclass(ImpalaConnection, BaseConnection)
+
+
+def test_get_client_injects_auth_kerberos_and_ssl_args():
+    config = _config(
+        authMechanism=AuthMechanism.GSSAPI,
+        kerberosServiceName="impala",
+        useSSL=True,
+    )
+    with patch(f"{CONNECTION_MODULE}.create_generic_db_connection") as mock_connection:
+        _ = ImpalaConnection(config).client
+    built = mock_connection.call_args.kwargs["connection"]
+    root = built.connectionArguments.root
+    assert root["auth_mechanism"] == "GSSAPI"
+    assert root["kerberos_service_name"] == "impala"
+    assert root["use_ssl"] is True
+
+
+def test_get_client_leaves_connection_arguments_unset_without_auth():
+    config = _config(authMechanism=None, useSSL=None)
+    with patch(f"{CONNECTION_MODULE}.create_generic_db_connection") as mock_connection:
+        _ = ImpalaConnection(config).client
+    assert mock_connection.call_args.kwargs["connection"].connectionArguments is None
+
+
+def test_get_client_uses_the_class_url_builder():
+    config = _config()
+    with patch(f"{CONNECTION_MODULE}.create_generic_db_connection") as mock_connection:
+        _ = ImpalaConnection(config).client
+    assert mock_connection.call_args.kwargs["get_connection_url_fn"].__name__ == "get_connection_url"
+
+
+def test_impala_url():
+    config = ImpalaConnectionConfig(scheme=ImpalaScheme.impala, hostPort="localhost:21050")
+    assert ImpalaConnection.get_connection_url(config) == "impala://localhost:21050"
+
+
+def test_impala_url_custom_auth():
+    config = ImpalaConnectionConfig(
+        scheme=ImpalaScheme.impala.value,
+        username="username",
+        password="password",
+        hostPort="localhost:21050",
+        connectionArguments={"auth": "CUSTOM"},
+    )
+    assert ImpalaConnection.get_connection_url(config) == "impala://username:password@localhost:21050"
+
+    config = ImpalaConnectionConfig(
+        scheme=ImpalaScheme.impala.value,
+        username="username@444",
+        password="password@333",
+        hostPort="localhost:21050",
+        connectionArguments={"auth": "CUSTOM"},
+    )
+    assert ImpalaConnection.get_connection_url(config) == "impala://username%40444:password%40333@localhost:21050"
+
+
+def test_impala_url_conn_options_with_db():
+    config = ImpalaConnectionConfig(
+        hostPort="localhost:21050",
+        databaseSchema="test_db",
+        connectionOptions={"Key": "Value"},
+    )
+    assert ImpalaConnection.get_connection_url(config) == "impala://localhost:21050/test_db?Key=Value"
+
+
+def test_impala_url_conn_options_without_db():
+    config = ImpalaConnectionConfig(
+        hostPort="localhost:21050",
+        connectionOptions={"Key": "Value"},
+    )
+    assert ImpalaConnection.get_connection_url(config) == "impala://localhost:21050?Key=Value"
+
+
+def test_impala_url_with_ldap_auth():
+    config = ImpalaConnectionConfig(
+        scheme=ImpalaScheme.impala.value,
+        username="username",
+        password="password",
+        hostPort="localhost:21050",
+        connectionArguments={"auth_mechanism": "LDAP"},
+    )
+    assert ImpalaConnection.get_connection_url(config) == "impala://username:password@localhost:21050"
+
+
+def test_impala_url_without_connection_arguments():
+    config = ImpalaConnectionConfig(
+        scheme=ImpalaScheme.impala.value,
+        username="username",
+        password="password",
+        hostPort="localhost:21050",
+    )
+    assert ImpalaConnection.get_connection_url(config) == "impala://username:password@localhost:21050"
+
+
+def test_impala_url_without_connection_arguments_pass():
+    config = ImpalaConnectionConfig(
+        scheme=ImpalaScheme.impala.value,
+        username="username",
+        hostPort="localhost:21050",
+    )
+    assert ImpalaConnection.get_connection_url(config) == "impala://username@localhost:21050"

ingestion/tests/unit/source/database/iomete/test_connection.py

@@ -0,0 +1,37 @@
+#  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 IOMETE connection handling."""
+
+from metadata.generated.schema.entity.services.connections.database.iometeConnection import (
+    IometeConnection as IometeConnectionConfig,
+)
+from metadata.ingestion.connections.connection import BaseConnection
+from metadata.ingestion.source.database.iomete.connection import IometeConnection
+
+
+def test_iomete_connection_is_base_connection():
+    assert issubclass(IometeConnection, BaseConnection)
+
+
+def test_get_connection_url_builds_iomete_dsn():
+    connection = IometeConnectionConfig(
+        username="openmetadata_user",
+        password="openmetadata_password",
+        hostPort="localhost:443",
+        catalog="spark_catalog",
+        cluster="my_cluster",
+        dataPlane="my_dp",
+    )
+    url = IometeConnection.get_connection_url(connection)
+    assert (
+        url.render_as_string(hide_password=False) == "iomete://openmetadata_user:openmetadata_password@localhost:443"
+        "/spark_catalog?cluster=my_cluster&data_plane=my_dp"
+    )

ingestion/tests/unit/source/database/mariadb/test_connection.py

@@ -0,0 +1,34 @@
+#  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 MariaDB connection handling."""
+
+from metadata.generated.schema.entity.services.connections.database.mariaDBConnection import (
+    MariaDBConnection as MariaDBConnectionConfig,
+)
+from metadata.generated.schema.entity.services.connections.database.mariaDBConnection import (
+    MariaDBScheme,
+)
+from metadata.ingestion.source.database.mariadb.connection import MariaDBConnection
+
+
+def test_basic_auth_builds_expected_url():
+    connection = MariaDBConnectionConfig(
+        username="openmetadata_user",
+        password="openmetadata_password",
+        hostPort="localhost:3306",
+        databaseSchema="openmetadata_db",
+        scheme=MariaDBScheme.mysql_pymysql,
+    )
+    engine = MariaDBConnection(connection).client
+    assert (
+        engine.url.render_as_string(hide_password=False)
+        == "mysql+pymysql://openmetadata_user:openmetadata_password@localhost:3306/openmetadata_db"
+    )

ingestion/tests/unit/source/database/mongodb/test_connection.py

@@ -0,0 +1,68 @@
+#  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 MongoDB BaseConnection wiring.
+
+MongoDB is the first non-Engine connector: its client is a ``pymongo.MongoClient``
+rather than a SQLAlchemy ``Engine``, so ``BaseConnection``'s client type is
+``MongoClient`` and ``test_connection`` drives pymongo calls directly.
+"""
+
+from unittest.mock import patch
+
+from metadata.generated.schema.entity.services.connections.database.mongoDBConnection import (
+    MongoDBConnection as MongoDBConnectionConfig,
+)
+from metadata.generated.schema.entity.services.connections.database.mongoDBConnection import (
+    MongoDBScheme,
+)
+from metadata.ingestion.connections.connection import BaseConnection
+from metadata.ingestion.source.database.mongodb.connection import MongoDBConnection
+
+CONNECTION_MODULE = "metadata.ingestion.source.database.mongodb.connection"
+
+
+def _config(**kwargs) -> MongoDBConnectionConfig:
+    base = {
+        "scheme": MongoDBScheme.mongodb,
+        "username": "user",
+        "password": "pass",
+        "hostPort": "localhost:27017",
+    }
+    base.update(kwargs)
+    return MongoDBConnectionConfig(**base)
+
+
+def test_mongodb_connection_is_base_connection():
+    assert issubclass(MongoDBConnection, BaseConnection)
+
+
+def test_get_client_builds_a_mongo_client_from_the_url():
+    with patch(f"{CONNECTION_MODULE}.MongoClient") as mock_mongo:
+        client = MongoDBConnection(_config()).client
+    args, kwargs = mock_mongo.call_args
+    assert args[0].startswith("mongodb://")
+    assert kwargs == {}
+    assert client is mock_mongo.return_value
+
+
+def test_get_client_passes_connection_options_as_kwargs():
+    config = _config(connectionOptions={"serverSelectionTimeoutMS": "5000"})
+    with patch(f"{CONNECTION_MODULE}.MongoClient") as mock_mongo:
+        _ = MongoDBConnection(config).client
+    assert mock_mongo.call_args.kwargs["serverSelectionTimeoutMS"] == "5000"
+
+
+def test_close_releases_the_mongo_client_pool():
+    with patch(f"{CONNECTION_MODULE}.MongoClient") as mock_mongo:
+        connection = MongoDBConnection(_config())
+        _ = connection.client
+        connection.close()
+    mock_mongo.return_value.close.assert_called_once_with()

ingestion/tests/unit/source/database/mssql/test_connection.py

@@ -0,0 +1,197 @@
+#  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 MSSQL BaseConnection wiring.
+
+URL building (including the pyodbc delegation to azuresql) is covered in
+tests/unit/test_source_url.py and tests/unit/test_source_connection.py.
+"""
+
+import socket
+from unittest.mock import MagicMock, patch
+
+from metadata.core.connections.test_connection import collect_checks
+from metadata.core.connections.test_connection.checks.database import DEFAULT_SAMPLE_ROWS, DatabaseStep
+from metadata.generated.schema.entity.services.connections.database.mssqlConnection import (
+    MssqlConnection as MssqlConnectionConfig,
+)
+from metadata.generated.schema.entity.services.connections.database.mssqlConnection import (
+    MssqlScheme,
+)
+from metadata.ingestion.connections.connection import BaseConnection
+from metadata.ingestion.source.database.mssql.connection import (
+    MSSQL_ERRORS,
+    MssqlChecks,
+    MssqlConnection,
+    _databases_enumerated,
+    get_connection_url,
+)
+from metadata.ingestion.source.database.mssql.queries import (
+    MSSQL_GET_CURRENT_DATABASE,
+    MSSQL_GET_DATABASE,
+)
+
+CONNECTION_MODULE = "metadata.ingestion.source.database.mssql.connection"
+
+
+def _config(scheme: MssqlScheme = MssqlScheme.mssql_pytds) -> MssqlConnectionConfig:
+    return MssqlConnectionConfig(
+        scheme=scheme,
+        username="user",
+        password="pass",
+        hostPort="myhost:1433",
+        database="mydb",
+    )
+
+
+def test_mssql_connection_is_base_connection():
+    assert issubclass(MssqlConnection, BaseConnection)
+
+
+def test_get_client_uses_the_module_url_builder():
+    with patch(f"{CONNECTION_MODULE}.create_generic_db_connection") as mock_connection:
+        _ = MssqlConnection(_config()).client
+    assert mock_connection.call_args.kwargs["get_connection_url_fn"].__name__ == "get_connection_url"
+
+
+def test_pyodbc_scheme_delegates_to_azuresql_url_builder():
+    with patch(f"{CONNECTION_MODULE}.get_pyodbc_connection_url", return_value="delegated") as mock_pyodbc:
+        result = get_connection_url(_config(scheme=MssqlScheme.mssql_pyodbc))
+    mock_pyodbc.assert_called_once()
+    assert result == "delegated"
+
+
+def test_non_pyodbc_scheme_uses_common_url_builder():
+    url = get_connection_url(_config(scheme=MssqlScheme.mssql_pytds))
+    assert url.startswith("mssql+pytds://")
+
+
+def _checks() -> MssqlChecks:
+    return MssqlChecks(client=MagicMock(), get_databases_statement="SELECT 1")
+
+
+def test_every_definition_step_resolves_to_a_check():
+    collected = collect_checks(_checks())
+    assert set(collected) == {
+        DatabaseStep.CheckAccess,
+        DatabaseStep.GetDatabases,
+        DatabaseStep.GetSchemas,
+        DatabaseStep.GetTables,
+        DatabaseStep.GetViews,
+        DatabaseStep.GetQueries,
+    }
+
+
+def test_close_disposes_the_engine():
+    with patch(f"{CONNECTION_MODULE}.create_generic_db_connection"):
+        connection = MssqlConnection(_config())
+        engine = connection.client
+        connection.close()
+    engine.dispose.assert_called_once()
+
+
+def test_building_checks_does_not_touch_the_network():
+    with patch(f"{CONNECTION_MODULE}.create_generic_db_connection") as mock_engine:
+        provider = MssqlConnection(_config()).checks()
+    engine = mock_engine.return_value
+    engine.connect.assert_not_called()
+    engine.exec_driver_sql.assert_not_called()
+    assert isinstance(provider, MssqlChecks)
+
+
+def test_get_databases_statement_uses_current_db_when_not_ingest_all():
+    config = _config()
+    config.ingestAllDatabases = False
+    handler = MssqlConnection(config)
+    handler._client = MagicMock()
+    assert handler._get_databases_statement() == MSSQL_GET_CURRENT_DATABASE
+
+
+def test_get_databases_statement_uses_all_dbs_when_ingest_all():
+    config = _config()
+    config.ingestAllDatabases = True
+    handler = MssqlConnection(config)
+    handler._client = MagicMock()
+    assert handler._get_databases_statement() == MSSQL_GET_DATABASE
+
+
+def _pymssql_error(number: int, message: str) -> Exception:
+    """A pymssql-style driver error: numeric SQL Server code in ``args[0]``,
+    wrapped by SQLAlchemy which preserves the original on ``.orig``."""
+    orig = Exception(number, message)
+    wrapper = Exception(f"({number}) {message}")
+    wrapper.orig = orig  # type: ignore[attr-defined]
+    return wrapper
+
+
+def _pyodbc_error(sqlstate: str, message: str) -> Exception:
+    """A pyodbc-style driver error: SQLSTATE + message text, no numeric code."""
+    return Exception(f"('{sqlstate}', \"[{sqlstate}] {message}\")")
+
+
+def test_pymssql_login_failure_classifies_as_auth():
+    diagnosis = MSSQL_ERRORS.classify(_pymssql_error(18456, "Login failed for user 'x'"))
+    assert diagnosis is not None
+    assert diagnosis.title == "Authentication failed"
+
+
+def test_pyodbc_login_failure_classifies_as_auth():
+    diagnosis = MSSQL_ERRORS.classify(_pyodbc_error("28000", "Login failed for user 'x'."))
+    assert diagnosis is not None
+    assert diagnosis.title == "Authentication failed"
+
+
+def test_pymssql_database_not_found_classifies():
+    diagnosis = MSSQL_ERRORS.classify(_pymssql_error(4060, 'Cannot open database "x" requested by the login.'))
+    assert diagnosis is not None
+    assert diagnosis.title == "Database not found or not accessible"
+
+
+def test_pyodbc_cannot_open_database_classifies():
+    diagnosis = MSSQL_ERRORS.classify(_pyodbc_error("42000", 'Cannot open database "x" requested by the login.'))
+    assert diagnosis is not None
+    assert diagnosis.title == "Database not found or not accessible"
+
+
+def test_cannot_open_database_wins_over_login_failed():
+    """The SQL Server 4060 message embeds 'The login failed.', so the database
+    rule must be ordered before the login rule (regression for live-found bug)."""
+    message = "Cannot open database \"x\" requested by the login. The login failed. Login failed for user 'y'."
+    assert MSSQL_ERRORS.classify(Exception(message)).title == "Database not found or not accessible"
+
+
+def test_pymssql_permission_denied_classifies():
+    diagnosis = MSSQL_ERRORS.classify(_pymssql_error(229, "The SELECT permission was denied on the object 'x'."))
+    assert diagnosis is not None
+    assert diagnosis.title == "Insufficient privileges"
+
+
+def test_pyodbc_permission_denied_classifies():
+    diagnosis = MSSQL_ERRORS.classify(_pyodbc_error("42000", "The SELECT permission was denied on the object 'x'."))
+    assert diagnosis is not None
+    assert diagnosis.title == "Insufficient privileges"
+
+
+def test_network_pack_is_folded_in():
+    diagnosis = MSSQL_ERRORS.classify(socket.gaierror("Name or service not known"))
+    assert diagnosis is not None
+    assert diagnosis.title == "Host could not be resolved"
+
+
+def test_unknown_error_is_not_classified():
+    assert MSSQL_ERRORS.classify(Exception("something unexpected")) is None
+
+
+def test_databases_summary_reports_exact_count_below_cap():
+    assert _databases_enumerated([object()] * 3) == "3 databases enumerated"
+
+
+def test_databases_summary_reports_floor_when_capped():
+    assert _databases_enumerated([object()] * DEFAULT_SAMPLE_ROWS) == f"{DEFAULT_SAMPLE_ROWS}+ databases enumerated"

ingestion/tests/unit/source/database/mysql/test_mysql_connection.py

@@ -0,0 +1,85 @@
+#  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.
+"""MySQL test-connection: its checks match its shipped definition."""
+
+import json
+from pathlib import Path
+
+from metadata.core.connections.test_connection.check import collect_checks
+from metadata.ingestion.source.database.mysql.connection import MYSQL_ERRORS, MySQLChecks
+
+_SEED = (
+    Path(__file__).parents[6]
+    / "openmetadata-service/src/main/resources/json/data"
+    / "testConnections/database/mysql.json"
+)
+
+
+def _check_names():
+    checks = MySQLChecks(client=None, schema=None, queries_statement="")
+    return {step.value for step in collect_checks(checks)}
+
+
+def test_mysql_checks_cover_expected_steps():
+    assert _check_names() == {
+        "CheckAccess",
+        "GetSchemas",
+        "GetTables",
+        "GetViews",
+        "GetQueries",
+    }
+
+
+def test_mysql_checks_match_definition_seed():
+    definition_steps = {step["name"] for step in json.loads(_SEED.read_text())["steps"]}
+    assert _check_names() == definition_steps
+
+
+def _driver_error(errno: int, message: str) -> Exception:
+    """A PyMySQL-style error: numeric code in args[0]."""
+    error = Exception()
+    error.args = (errno, message)
+    return error
+
+
+def test_error_pack_access_denied():
+    diagnosis = MYSQL_ERRORS.classify(_driver_error(1045, "Access denied for user 'x'@'h' (using password: YES)"))
+    assert diagnosis is not None
+    assert "Authentication" in diagnosis.title
+
+
+def test_error_pack_unknown_database():
+    diagnosis = MYSQL_ERRORS.classify(_driver_error(1049, "Unknown database 'foo'"))
+    assert diagnosis is not None
+    assert "not found" in diagnosis.title.lower()
+
+
+def test_error_pack_query_log_denied():
+    diagnosis = MYSQL_ERRORS.classify(_driver_error(1142, "SELECT command denied to user 'x' for table 'general_log'"))
+    assert diagnosis is not None
+    assert "Query history" in diagnosis.title
+
+
+def test_error_pack_cannot_connect():
+    diagnosis = MYSQL_ERRORS.classify(_driver_error(2003, "Can't connect to MySQL server on 'h' ([Errno 111] refused)"))
+    assert diagnosis is not None
+    assert "reach" in diagnosis.title.lower()
+
+
+def test_error_pack_caching_sha2_has_no_errno():
+    # PyMySQL raises this with the message in args[0], no numeric code.
+    diagnosis = MYSQL_ERRORS.classify(Exception("Couldn't receive server's public key"))
+    assert diagnosis is not None
+    assert "Secure connection" in diagnosis.title
+
+
+def test_error_pack_unmatched_returns_none():
+    assert MYSQL_ERRORS.classify(_driver_error(9999, "novel error")) is None

ingestion/tests/unit/source/database/pinotdb/test_connection.py

@@ -0,0 +1,36 @@
+#  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 PinotDB connection handling."""
+
+from metadata.generated.schema.entity.services.connections.database.pinotDBConnection import (
+    PinotDBConnection as PinotDBConnectionConfig,
+)
+from metadata.generated.schema.entity.services.connections.database.pinotDBConnection import (
+    PinotDBScheme,
+)
+from metadata.ingestion.connections.connection import BaseConnection
+from metadata.ingestion.source.database.pinotdb.connection import PinotDBConnection
+
+
+def test_pinotdb_connection_is_base_connection():
+    assert issubclass(PinotDBConnection, BaseConnection)
+
+
+def test_get_connection_url_appends_controller():
+    connection = PinotDBConnectionConfig(
+        scheme=PinotDBScheme.pinot,
+        hostPort="localhost:8099",
+        pinotControllerHost="http://localhost:9000/",
+    )
+    assert (
+        PinotDBConnection.get_connection_url(connection)
+        == "pinot://localhost:8099/query/sql?controller=http://localhost:9000/"
+    )

ingestion/tests/unit/source/database/postgres/test_connection.py

@@ -0,0 +1,264 @@
+#  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 PostgreSQL connection handling (auth strategies + checks)."""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+from azure.core.credentials import AccessToken
+from azure.identity import ClientSecretCredential
+from sqlalchemy import create_engine
+from sqlalchemy.pool import StaticPool
+
+from metadata.core.connections.test_connection.check import CheckError, collect_checks
+from metadata.core.connections.test_connection.checks.database import (
+    DEFAULT_SAMPLE_ROWS,
+    DatabaseStep,
+)
+from metadata.core.connections.test_connection.network import NetworkUnreachableError
+from metadata.generated.schema.entity.services.connections.database.common.azureConfig import (
+    AzureConfigurationSource,
+)
+from metadata.generated.schema.entity.services.connections.database.common.basicAuth import (
+    BasicAuth,
+)
+from metadata.generated.schema.entity.services.connections.database.postgresConnection import (
+    PostgresConnection as PostgresConnectionConfig,
+)
+from metadata.generated.schema.security.credentials.azureCredentials import (
+    AzureCredentials,
+)
+from metadata.ingestion.source.database.postgres.connection import (
+    POSTGRES_ERRORS,
+    PostgresChecks,
+    PostgresConnection,
+)
+
+
+def _azure_connection() -> PostgresConnectionConfig:
+    return PostgresConnectionConfig(
+        username="openmetadata_user",
+        authType=AzureConfigurationSource(
+            azureConfig=AzureCredentials(
+                clientId="clientid",
+                tenantId="tenantid",
+                clientSecret="clientsecret",
+                scopes="scope1,scope2",
+            )
+        ),
+        hostPort="localhost:5432",
+        database="openmetadata_db",
+    )
+
+
+def test_basic_auth_builds_expected_url():
+    connection = PostgresConnectionConfig(
+        username="openmetadata_user",
+        authType=BasicAuth(password="openmetadata_password"),
+        hostPort="localhost:5432",
+        database="openmetadata_db",
+    )
+    engine = PostgresConnection(connection).client
+    assert (
+        engine.url.render_as_string(hide_password=False)
+        == "postgresql+psycopg2://openmetadata_user:openmetadata_password@localhost:5432/openmetadata_db"
+    )
+
+
+def test_azure_ad_uses_token_as_password():
+    connection = _azure_connection()
+    with patch.object(
+        ClientSecretCredential,
+        "get_token",
+        return_value=AccessToken(token="mocked_token", expires_on=100),
+    ):
+        engine = PostgresConnection(connection).client
+    assert (
+        engine.url.render_as_string(hide_password=False)
+        == "postgresql+psycopg2://openmetadata_user:mocked_token@localhost:5432/openmetadata_db"
+    )
+
+
+def test_azure_ad_does_not_mutate_caller_connection():
+    connection = _azure_connection()
+    with patch.object(
+        ClientSecretCredential,
+        "get_token",
+        return_value=AccessToken(token="mocked_token", expires_on=100),
+    ):
+        assert PostgresConnection(connection).client is not None
+    assert isinstance(connection.authType, AzureConfigurationSource)
+
+
+class _Psycopg2Error(Exception):
+    """Mirror a psycopg2 DBAPI error.
+
+    The message and (for query-execution errors only) the SQLSTATE on ``.pgcode``
+    reproduce what psycopg2 actually raises - connect-phase failures carry no
+    code, captured live against PostgreSQL 15.
+    """
+
+    def __init__(self, message: str = "", pgcode: str | None = None) -> None:
+        super().__init__(message)
+        self.pgcode = pgcode
+
+
+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
+
+
+# Real first-line messages captured from psycopg2 against PostgreSQL 15.
+_AUTH_FAILED_MSG = (
+    'connection to server at "localhost" (::1), port 5432 failed: '
+    'FATAL:  password authentication failed for user "postgres"'
+)
+_DB_NOT_FOUND_MSG = (
+    'connection to server at "localhost" (::1), port 5432 failed: FATAL:  database "nope_db" does not exist'
+)
+
+
+def test_auth_failure_message_is_classified():
+    error = _SqlAlchemyError(_Psycopg2Error(_AUTH_FAILED_MSG))
+    assert POSTGRES_ERRORS.classify(error).title == "Authentication failed"
+
+
+def test_database_not_found_message_is_classified():
+    error = _SqlAlchemyError(_Psycopg2Error(_DB_NOT_FOUND_MSG))
+    assert POSTGRES_ERRORS.classify(error).title == "Database not found"
+
+
+def test_unknown_role_message_is_classified_as_auth_failure():
+    # Under trust/peer auth a missing role surfaces as `role "x" does not exist`
+    # (no SQLSTATE); it must map to Authentication failed, not Database not found -
+    # both share the `does not exist` suffix.
+    error = _SqlAlchemyError(_Psycopg2Error('FATAL:  role "ghost" does not exist'))
+    assert POSTGRES_ERRORS.classify(error).title == "Authentication failed"
+
+
+def test_insufficient_privilege_sqlstate_is_classified():
+    error = _SqlAlchemyError(_Psycopg2Error("permission denied for table secret_t", pgcode="42501"))
+    assert POSTGRES_ERRORS.classify(error).title == "Insufficient privileges"
+
+
+# Real GetQueries failure when the source is absent: the error text embeds the
+# probe SQL, which references pg_catalog.pg_database - the case that must NOT be
+# misread as "Database not found". Keyed on SQLSTATE 42P01, so it holds for any
+# configured queryStatementSource, not just the default pg_stat_statements.
+_MISSING_PGSS_MSG = (
+    'relation "pg_stat_statements" does not exist\n[SQL: SELECT u.usename, d.datname '
+    "FROM pg_stat_statements s JOIN pg_catalog.pg_database d ON s.dbid = d.oid LIMIT 1]"
+)
+
+
+def test_missing_query_history_source_is_classified():
+    error = _SqlAlchemyError(_Psycopg2Error(_MISSING_PGSS_MSG, pgcode="42P01"))
+    assert POSTGRES_ERRORS.classify(error).title == "Query history source not found"
+
+
+def test_missing_query_history_source_is_not_read_as_database_not_found():
+    # The probe SQL mentions pg_database; the diagnosis must be query-history, not db.
+    error = _SqlAlchemyError(_Psycopg2Error(_MISSING_PGSS_MSG, pgcode="42P01"))
+    assert POSTGRES_ERRORS.classify(error).title != "Database not found"
+
+
+def test_missing_query_history_source_holds_for_a_custom_source_name():
+    # queryStatementSource is configurable; a custom source missing must still be
+    # diagnosed as a query-history problem, not left unclassified.
+    error = _SqlAlchemyError(_Psycopg2Error('relation "my_query_log" does not exist', pgcode="42P01"))
+    assert POSTGRES_ERRORS.classify(error).title == "Query history source not found"
+
+
+def test_pg_hba_message_is_classified():
+    error = Exception('FATAL: no pg_hba.conf entry for host "1.2.3.4", user "u", SSL off')
+    assert POSTGRES_ERRORS.classify(error).title == "Connection not permitted by pg_hba.conf"
+
+
+def test_network_errors_classify_through_including():
+    error = NetworkUnreachableError("db:5432 is not reachable")
+    error.__cause__ = ConnectionRefusedError(61, "Connection refused")
+    assert POSTGRES_ERRORS.classify(error).title == "Connection refused"
+
+
+def test_unknown_error_returns_no_diagnosis():
+    error = _SqlAlchemyError(_Psycopg2Error("something unexpected", pgcode="99999"))
+    assert POSTGRES_ERRORS.classify(error) is None
+
+
+def test_checks_cover_exactly_the_seeded_steps():
+    engine = create_engine("sqlite://", poolclass=StaticPool)
+    checks = PostgresChecks(client=engine, query_statement_source=None)
+    collected = collect_checks(checks)
+    assert set(collected.keys()) == {
+        DatabaseStep.CheckAccess,
+        DatabaseStep.GetDatabases,
+        DatabaseStep.GetSchemas,
+        DatabaseStep.GetTables,
+        DatabaseStep.GetViews,
+        DatabaseStep.GetTags,
+        DatabaseStep.GetQueries,
+        DatabaseStep.GetColumnMetadata,
+        DatabaseStep.GetTableComments,
+        DatabaseStep.GetInformationSchemaColumns,
+    }
+
+
+def test_check_access_reports_unreachable_host_as_network_failure():
+    # Prove the wiring: check_access -> ping -> _preflight raises whatever
+    # tcp_probe raises, wrapped as a CheckError whose cause classifies as a
+    # network failure. tcp_probe is stubbed so the test is deterministic and
+    # fast - no real socket, hence no ephemeral-port TOCTOU and no 20s timeout
+    # waiting on an unreachable host. tcp_probe's own socket behaviour is
+    # covered in the network module's tests.
+    client = MagicMock()
+    client.url.host = "db.invalid"
+    client.url.port = 5432
+    checks = PostgresChecks(client=client, query_statement_source=None)
+    probe_error = NetworkUnreachableError("db.invalid:5432 is not reachable")
+    probe_error.__cause__ = ConnectionRefusedError(61, "Connection refused")
+    with (
+        patch(
+            "metadata.core.connections.test_connection.checks.database.tcp_probe",
+            side_effect=probe_error,
+        ) as mock_probe,
+        pytest.raises(CheckError) as exc,
+    ):
+        checks.check_access()
+    mock_probe.assert_called_once_with("db.invalid", 5432)
+    assert exc.value.cause is probe_error
+    assert POSTGRES_ERRORS.classify(exc.value.cause).title == "Connection refused"
+
+
+def test_query_statement_is_built_lazily_not_at_construction():
+    # Regression: the version-probe query must run inside GetQueries (behind the
+    # CheckAccess gate), never at checks() construction - otherwise an unreachable
+    # host hangs on that eager connect instead of failing fast at the preflight.
+    calls = []
+    with patch(
+        "metadata.ingestion.source.database.postgres.connection.get_postgres_time_column_name",
+        side_effect=lambda engine: calls.append(1) or "total_exec_time",
+    ):
+        engine = create_engine("sqlite://", poolclass=StaticPool)
+        PostgresChecks(client=engine, query_statement_source=None)
+        assert calls == []
+
+
+def test_get_databases_summary_marks_the_sample_cap():
+    # run_sql fetches at most DEFAULT_SAMPLE_ROWS; below the cap the count is exact,
+    # at the cap it is a lower bound and must be reported as "N+".
+    assert PostgresChecks._summarize_databases([1, 2, 3]) == "3 databases enumerated"
+    assert (
+        PostgresChecks._summarize_databases(list(range(DEFAULT_SAMPLE_ROWS)))
+        == f"{DEFAULT_SAMPLE_ROWS}+ databases enumerated"
+    )

ingestion/tests/unit/source/database/presto/test_connection.py

@@ -0,0 +1,57 @@
+#  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 Presto connection handling."""
+
+from metadata.generated.schema.entity.services.connections.database.prestoConnection import (
+    PrestoConnection as PrestoConnectionConfig,
+)
+from metadata.generated.schema.entity.services.connections.database.prestoConnection import (
+    PrestoScheme,
+)
+from metadata.ingestion.connections.connection import BaseConnection
+from metadata.ingestion.source.database.presto.connection import PrestoConnection
+
+
+def test_presto_connection_is_base_connection():
+    assert issubclass(PrestoConnection, BaseConnection)
+
+
+def test_url_with_catalog_and_no_password():
+    connection = PrestoConnectionConfig(
+        username="admin",
+        hostPort="localhost:8080",
+        scheme=PrestoScheme.presto,
+        catalog="test_catalog",
+    )
+    assert PrestoConnection.get_connection_url(connection) == "presto://admin@localhost:8080/test_catalog"
+
+
+def test_url_escapes_special_characters_in_credentials():
+    connection = PrestoConnectionConfig(
+        username="admin@333",
+        password="pass@111",
+        hostPort="localhost:8080",
+        scheme=PrestoScheme.presto,
+        catalog="test_catalog",
+    )
+    assert (
+        PrestoConnection.get_connection_url(connection) == "presto://admin%40333:pass%40111@localhost:8080/test_catalog"
+    )
+
+
+def test_url_without_catalog():
+    connection = PrestoConnectionConfig(
+        scheme=PrestoScheme.presto,
+        hostPort="localhost:8080",
+        username="username",
+        password="pass",
+    )
+    assert PrestoConnection.get_connection_url(connection) == "presto://username:pass@localhost:8080"

ingestion/tests/unit/source/database/redshift/test_connection.py

@@ -0,0 +1,232 @@
+#  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 Redshift BaseConnection wiring and test-connection checks.
+
+The IAM/basic URL building and credential fetching are covered in
+tests/unit/topology/database/test_redshift_connection.py.
+"""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+from sqlalchemy import create_engine
+from sqlalchemy.pool import StaticPool
+
+from metadata.core.connections.test_connection.check import CheckError, collect_checks
+from metadata.core.connections.test_connection.checks.database import (
+    DEFAULT_SAMPLE_ROWS,
+    DatabaseStep,
+)
+from metadata.core.connections.test_connection.network import NetworkUnreachableError
+from metadata.generated.schema.entity.services.connections.database.common.basicAuth import (
+    BasicAuth,
+)
+from metadata.generated.schema.entity.services.connections.database.redshiftConnection import (
+    RedshiftConnection as RedshiftConnectionConfig,
+)
+from metadata.ingestion.connections.connection import BaseConnection
+from metadata.ingestion.connections.test_connections import SourceConnectionException
+from metadata.ingestion.source.database.redshift.connection import (
+    REDSHIFT_ERRORS,
+    RedshiftChecks,
+    RedshiftConnection,
+    _summarize_databases,
+)
+from metadata.ingestion.source.database.redshift.models import RedshiftInstanceType
+from metadata.ingestion.source.database.redshift.queries import (
+    REDSHIFT_TEST_GET_QUERIES_MAP,
+)
+
+CONNECTION_MODULE = "metadata.ingestion.source.database.redshift.connection"
+
+
+def _config() -> RedshiftConnectionConfig:
+    return RedshiftConnectionConfig(
+        hostPort="my-cluster.abc123.us-east-1.redshift.amazonaws.com:5439",
+        username="admin",
+        authType=BasicAuth(password="secret"),
+        database="mydb",
+    )
+
+
+def test_redshift_connection_is_base_connection():
+    assert issubclass(RedshiftConnection, BaseConnection)
+
+
+def test_get_client_uses_the_redshift_url_builder():
+    with patch(f"{CONNECTION_MODULE}.create_generic_db_connection") as mock_connection:
+        _ = RedshiftConnection(_config()).client
+    assert mock_connection.call_args.kwargs["get_connection_url_fn"].__name__ == "get_redshift_connection_url"
+
+
+def test_close_disposes_the_engine_pool():
+    # The old test_connection called kill_active_connections(engine); the framework
+    # path relies on close() unwinding _on_close teardowns, so _get_client must
+    # register engine.dispose or the QueuePool leaks idle connections.
+    engine = MagicMock()
+    with patch(f"{CONNECTION_MODULE}.create_generic_db_connection", return_value=engine):
+        connection = RedshiftConnection(_config())
+        assert connection.client is engine
+        connection.close()
+    engine.dispose.assert_called_once_with()
+
+
+class _Psycopg2Error(Exception):
+    """Mirror a psycopg2 DBAPI error.
+
+    The Redshift dialect rides on psycopg2: the message and (for query-execution
+    errors only) the SQLSTATE on ``.pgcode`` reproduce what the driver raises -
+    connect-phase failures carry no code.
+    """
+
+    def __init__(self, message: str = "", pgcode: str | None = None) -> None:
+        super().__init__(message)
+        self.pgcode = pgcode
+
+
+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
+
+
+_AUTH_FAILED_MSG = (
+    'connection to server at "my-cluster" port 5439 failed: FATAL:  password authentication failed for user "admin"'
+)
+_DB_NOT_FOUND_MSG = 'connection to server at "my-cluster" port 5439 failed: FATAL:  database "nope_db" does not exist'
+
+
+def test_auth_failure_message_is_classified():
+    error = _SqlAlchemyError(_Psycopg2Error(_AUTH_FAILED_MSG))
+    assert REDSHIFT_ERRORS.classify(error).title == "Authentication failed"
+
+
+def test_auth_failure_sqlstate_is_classified():
+    error = _SqlAlchemyError(_Psycopg2Error("authorization not valid", pgcode="28000"))
+    assert REDSHIFT_ERRORS.classify(error).title == "Authentication failed"
+
+
+def test_database_not_found_message_is_classified():
+    error = _SqlAlchemyError(_Psycopg2Error(_DB_NOT_FOUND_MSG))
+    assert REDSHIFT_ERRORS.classify(error).title == "Database not found"
+
+
+def test_insufficient_privilege_sqlstate_is_classified():
+    error = _SqlAlchemyError(_Psycopg2Error("permission denied for table secret_t", pgcode="42501"))
+    assert REDSHIFT_ERRORS.classify(error).title == "Insufficient privileges"
+
+
+def test_missing_query_history_source_is_classified():
+    # An absent query-history view raises undefined_table (42P01) from the probe;
+    # keyed on SQLSTATE, so it holds whether the deployment is Provisioned or Serverless.
+    error = _SqlAlchemyError(_Psycopg2Error('relation "stl_query" does not exist', pgcode="42P01"))
+    assert REDSHIFT_ERRORS.classify(error).title == "Query history source not found"
+
+
+def test_missing_query_history_source_is_not_read_as_database_not_found():
+    error = _SqlAlchemyError(_Psycopg2Error('relation "stl_query" does not exist', pgcode="42P01"))
+    assert REDSHIFT_ERRORS.classify(error).title != "Database not found"
+
+
+def test_missing_query_privilege_is_classified():
+    # has_table_privilege returns False rather than raising; the GetQueries check
+    # turns that into a SourceConnectionException whose text the pack classifies.
+    error = Exception("Missing SELECT privilege on the Redshift stl query-history views - (True, False)")
+    assert REDSHIFT_ERRORS.classify(error).title == "Query history not accessible"
+
+
+def test_network_errors_classify_through_including():
+    error = NetworkUnreachableError("cluster:5439 is not reachable")
+    error.__cause__ = ConnectionRefusedError(61, "Connection refused")
+    assert REDSHIFT_ERRORS.classify(error).title == "Connection refused"
+
+
+def test_unknown_error_returns_no_diagnosis():
+    error = _SqlAlchemyError(_Psycopg2Error("something unexpected", pgcode="99999"))
+    assert REDSHIFT_ERRORS.classify(error) is None
+
+
+def test_database_summary_reports_exact_count_under_the_sample_cap():
+    assert _summarize_databases(["dev", "analytics"]) == "2 databases reachable"
+
+
+def test_database_summary_flags_a_capped_count_with_a_plus():
+    # run_sql fetches at most DEFAULT_SAMPLE_ROWS, so a saturated sample must read
+    # as "at least N", never as an exact total.
+    capped = list(range(DEFAULT_SAMPLE_ROWS))
+    assert _summarize_databases(capped) == f"{DEFAULT_SAMPLE_ROWS}+ databases reachable"
+
+
+def test_get_queries_failure_reports_the_attempted_command_as_evidence():
+    # A missing privilege raises from run_sql's summary callback, outside its own
+    # CheckError wrapping; get_queries must re-raise with the statement so the
+    # failed step still shows the command it ran, like every other step.
+    engine = MagicMock()
+    checks = RedshiftChecks(client=engine)
+    with (
+        patch(f"{CONNECTION_MODULE}.get_redshift_instance_type", return_value=RedshiftInstanceType.PROVISIONED),
+        patch(f"{CONNECTION_MODULE}.run_sql", side_effect=SourceConnectionException("missing privilege")) as mock_run,
+        pytest.raises(CheckError) as exc,
+    ):
+        checks.get_queries()
+    assert exc.value.evidence.command == REDSHIFT_TEST_GET_QUERIES_MAP[RedshiftInstanceType.PROVISIONED]
+    assert isinstance(exc.value.cause, SourceConnectionException)
+    assert mock_run.call_count == 1
+
+
+def test_checks_cover_exactly_the_seeded_steps():
+    engine = create_engine("sqlite://", poolclass=StaticPool)
+    checks = RedshiftChecks(client=engine)
+    collected = collect_checks(checks)
+    assert set(collected.keys()) == {
+        DatabaseStep.CheckAccess,
+        DatabaseStep.GetDatabases,
+        DatabaseStep.GetSchemas,
+        DatabaseStep.GetTables,
+        DatabaseStep.GetViews,
+        DatabaseStep.GetQueries,
+    }
+
+
+def test_check_access_reports_unreachable_host_as_network_failure():
+    client = MagicMock()
+    client.url.host = "cluster.invalid"
+    client.url.port = 5439
+    checks = RedshiftChecks(client=client)
+    probe_error = NetworkUnreachableError("cluster.invalid:5439 is not reachable")
+    probe_error.__cause__ = ConnectionRefusedError(61, "Connection refused")
+    with (
+        patch(
+            "metadata.core.connections.test_connection.checks.database.tcp_probe",
+            side_effect=probe_error,
+        ) as mock_probe,
+        pytest.raises(CheckError) as exc,
+    ):
+        checks.check_access()
+    mock_probe.assert_called_once_with("cluster.invalid", 5439)
+    assert exc.value.cause is probe_error
+    assert REDSHIFT_ERRORS.classify(exc.value.cause).title == "Connection refused"
+
+
+def test_instance_type_probe_is_run_lazily_not_at_construction():
+    # Regression: the instance-type probe must run inside GetQueries (behind the
+    # CheckAccess gate), never at checks() construction - otherwise an unreachable
+    # host hangs on that eager connect instead of failing fast at the preflight.
+    calls = []
+    with patch(
+        f"{CONNECTION_MODULE}.get_redshift_instance_type",
+        side_effect=lambda engine: calls.append(1),
+    ):
+        engine = create_engine("sqlite://", poolclass=StaticPool)
+        RedshiftChecks(client=engine)
+        assert calls == []

ingestion/tests/unit/source/database/salesforce/test_connection.py

@@ -0,0 +1,22 @@
+#  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 SalesforceConnection BaseConnection wiring (non-Engine client)."""
+
+import pytest
+
+from metadata.ingestion.connections.connection import BaseConnection
+
+# The connector SDK may be an optional dependency absent from the unit env.
+connection = pytest.importorskip("metadata.ingestion.source.database.salesforce.connection")
+
+
+def test_salesforce_connection_is_base_connection():
+    assert issubclass(connection.SalesforceConnection, BaseConnection)

ingestion/tests/unit/source/database/saperp/test_connection.py

@@ -0,0 +1,22 @@
+#  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 SapErpConnection BaseConnection wiring (non-Engine client)."""
+
+import pytest
+
+from metadata.ingestion.connections.connection import BaseConnection
+
+# The connector SDK may be an optional dependency absent from the unit env.
+connection = pytest.importorskip("metadata.ingestion.source.database.saperp.connection")
+
+
+def test_saperp_connection_is_base_connection():
+    assert issubclass(connection.SapErpConnection, BaseConnection)

ingestion/tests/unit/source/database/saphana/test_connection.py

@@ -0,0 +1,124 @@
+#  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 SAP HANA connection handling.
+
+The SAP HANA dialect is an optional extra absent from the unit-test
+environment, so URL parity is asserted via each connection-mode strategy
+rather than by instantiating an engine.
+"""
+
+from unittest.mock import patch
+
+from metadata.generated.schema.entity.services.connections.database.sapHana.sapHanaHDBConnection import (
+    SapHanaHDBConnection,
+)
+from metadata.generated.schema.entity.services.connections.database.sapHana.sapHanaSQLConnection import (
+    SapHanaSQLConnection,
+)
+from metadata.generated.schema.entity.services.connections.database.sapHanaConnection import (
+    SapHanaConnection as SapHanaConnectionConfig,
+)
+from metadata.generated.schema.entity.services.connections.database.sapHanaConnection import (
+    SapHanaScheme,
+)
+from metadata.ingestion.connections.connection import BaseConnection
+from metadata.ingestion.source.database.saphana.connection import (
+    SapHanaConnection,
+    SapHanaHDBStrategy,
+    SapHanaSQLStrategy,
+)
+
+CONNECTION_MODULE = "metadata.ingestion.source.database.saphana.connection"
+
+
+def _sql_config(**sql_kwargs) -> SapHanaConnectionConfig:
+    return SapHanaConnectionConfig(
+        scheme=SapHanaScheme.hana,
+        connection=SapHanaSQLConnection(
+            username="openmetadata_user",
+            password="openmetadata_password",
+            hostPort="localhost:39041",
+            **sql_kwargs,
+        ),
+    )
+
+
+def _hdb_config() -> SapHanaConnectionConfig:
+    return SapHanaConnectionConfig(
+        scheme=SapHanaScheme.hana,
+        connection=SapHanaHDBConnection(userKey="USER1UserKey"),
+    )
+
+
+def test_saphana_connection_is_base_connection():
+    assert issubclass(SapHanaConnection, BaseConnection)
+
+
+def test_sql_url_includes_database():
+    config = _sql_config(database="HXE")
+    url = SapHanaSQLStrategy(config, config.connection).get_url()
+    assert url == "hana://openmetadata_user:openmetadata_password@localhost:39041/HXE"
+
+
+def test_sql_url_appends_options_with_database():
+    config = SapHanaConnectionConfig(
+        scheme=SapHanaScheme.hana,
+        connection=SapHanaSQLConnection(
+            username="openmetadata_user",
+            password="openmetadata_password",
+            hostPort="localhost:39041",
+            database="HXE",
+        ),
+        connectionOptions={"encrypt": "true"},
+    )
+    url = SapHanaSQLStrategy(config, config.connection).get_url()
+    assert url == "hana://openmetadata_user:openmetadata_password@localhost:39041/HXE?encrypt=true"
+
+
+def test_sql_url_without_database():
+    config = _sql_config()
+    url = SapHanaSQLStrategy(config, config.connection).get_url()
+    assert url == "hana://openmetadata_user:openmetadata_password@localhost:39041"
+
+
+def test_hdb_url_uses_user_key():
+    config = _hdb_config()
+    url = SapHanaHDBStrategy(config, config.connection).get_url()
+    assert url == "hana://userkey=USER1UserKey"
+
+
+def test_get_client_selects_sql_strategy():
+    config = _sql_config(database="HXE")
+    with patch(f"{CONNECTION_MODULE}.create_generic_db_connection") as mock_connection:
+        _ = SapHanaConnection(config).client
+    url_fn = mock_connection.call_args.kwargs["get_connection_url_fn"]
+    assert url_fn(config) == "hana://openmetadata_user:openmetadata_password@localhost:39041/HXE"
+
+
+def test_get_client_selects_hdb_strategy():
+    config = _hdb_config()
+    with patch(f"{CONNECTION_MODULE}.create_generic_db_connection") as mock_connection:
+        _ = SapHanaConnection(config).client
+    url_fn = mock_connection.call_args.kwargs["get_connection_url_fn"]
+    assert url_fn(config) == "hana://userkey=USER1UserKey"
+
+
+@patch(f"{CONNECTION_MODULE}.inspect")
+@patch(f"{CONNECTION_MODULE}.create_generic_db_connection")
+def test_hdb_test_functions_tolerate_missing_database_schema(_mock_connection, mock_inspect):
+    mock_inspect.return_value.get_schema_names.return_value = []
+    connection = SapHanaConnection(_hdb_config())
+
+    test_functions = connection._build_test_functions()
+    test_functions["GetTables"]()
+    test_functions["GetViews"]()
+
+    mock_inspect.return_value.get_schema_names.assert_called()

ingestion/tests/unit/source/database/sas/test_connection.py

@@ -0,0 +1,22 @@
+#  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 SASConnection BaseConnection wiring (non-Engine client)."""
+
+import pytest
+
+from metadata.ingestion.connections.connection import BaseConnection
+
+# The connector SDK may be an optional dependency absent from the unit env.
+connection = pytest.importorskip("metadata.ingestion.source.database.sas.connection")
+
+
+def test_sas_connection_is_base_connection():
+    assert issubclass(connection.SASConnection, BaseConnection)

ingestion/tests/unit/source/database/singlestore/test_connection.py

@@ -0,0 +1,34 @@
+#  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 SingleStore connection handling."""
+
+from metadata.generated.schema.entity.services.connections.database.singleStoreConnection import (
+    SingleStoreConnection as SingleStoreConnectionConfig,
+)
+from metadata.generated.schema.entity.services.connections.database.singleStoreConnection import (
+    SingleStoreScheme,
+)
+from metadata.ingestion.source.database.singlestore.connection import SingleStoreConnection
+
+
+def test_basic_auth_builds_expected_url():
+    connection = SingleStoreConnectionConfig(
+        username="openmetadata_user",
+        password="openmetadata_password",
+        hostPort="localhost:3306",
+        databaseSchema="openmetadata_db",
+        scheme=SingleStoreScheme.mysql_pymysql,
+    )
+    engine = SingleStoreConnection(connection).client
+    assert (
+        engine.url.render_as_string(hide_password=False)
+        == "mysql+pymysql://openmetadata_user:openmetadata_password@localhost:3306/openmetadata_db"
+    )

ingestion/tests/unit/source/database/snowflake/test_connection.py

@@ -0,0 +1,344 @@
+#  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 Snowflake 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 scenario classifies to
+the intended diagnosis).
+"""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from metadata.core.connections.test_connection import Evidence
+from metadata.core.connections.test_connection.check import CheckError, collect_checks
+from metadata.core.connections.test_connection.checks.database import (
+    DEFAULT_SAMPLE_ROWS,
+    DatabaseStep,
+)
+from metadata.core.connections.test_connection.network import NetworkUnreachableError
+from metadata.generated.schema.entity.services.connections.database.snowflakeConnection import (
+    SnowflakeConnection as SnowflakeConnectionConfig,
+)
+from metadata.ingestion.source.database.snowflake.connection import (
+    SNOWFLAKE_ERRORS,
+    SNOWFLAKE_PORT,
+    SnowflakeChecks,
+    _count_summary,
+)
+
+
+def _config(**overrides) -> SnowflakeConnectionConfig:
+    base = {"username": "user", "account": "ue18849.us-east-2.aws", "warehouse": "wh"}
+    base.update(overrides)
+    return SnowflakeConnectionConfig(**base)
+
+
+class _SnowflakeError(Exception):
+    """Mirror a snowflake.connector error: the code lives on ``.errno`` (not in
+    ``args[0]`` as for PyMySQL), the message is the first arg."""
+
+    def __init__(self, message: str = "", errno: int | None = None) -> None:
+        super().__init__(message)
+        self.errno = errno
+
+
+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_auth_failure_errno_is_classified():
+    error = _SqlAlchemyError(_SnowflakeError("250001 (08001): Failed to connect to DB", errno=250001))
+    assert SNOWFLAKE_ERRORS.classify(error).title == "Authentication failed"
+
+
+def test_auth_failure_message_is_classified():
+    error = _SqlAlchemyError(_SnowflakeError("Incorrect username or password were specified."))
+    assert SNOWFLAKE_ERRORS.classify(error).title == "Authentication failed"
+
+
+def test_bad_account_login_404_is_classified():
+    error = _SqlAlchemyError(
+        _SnowflakeError(
+            "None: 404 Not Found: post nope99999.us-east-1.snowflakecomputing.com:443/session/v1/login-request",
+            errno=290404,
+        )
+    )
+    assert SNOWFLAKE_ERRORS.classify(error).title == "Snowflake account not found"
+
+
+def test_mfa_required_beats_generic_auth():
+    # errno 250001 alone would read as "Authentication failed"; the MFA message is
+    # ordered first so the user gets an actionable, specific diagnosis.
+    error = _SqlAlchemyError(
+        _SnowflakeError(
+            "Failed to connect to DB: acc.snowflakecomputing.com:443. "
+            "Multi-factor authentication is required for this account. Log in to Snowsight to enroll.",
+            errno=250001,
+        )
+    )
+    assert SNOWFLAKE_ERRORS.classify(error).title == "Multi-factor authentication required"
+
+
+def test_missing_role_beats_generic_auth():
+    # A missing role also raises errno 250001; match the role-specific token first.
+    error = _SqlAlchemyError(
+        _SnowflakeError(
+            "Failed to connect to DB: acc.snowflakecomputing.com:443. Role 'NO_SUCH_ROLE' specified in the "
+            "connect string is not granted to this user, or is not permitted for the credentials being used.",
+            errno=250001,
+        )
+    )
+    assert SNOWFLAKE_ERRORS.classify(error).title == "Role not granted"
+
+
+def test_account_usage_denied_is_classified():
+    error = _SqlAlchemyError(
+        _SnowflakeError(
+            "Object 'SNOWFLAKE.ACCOUNT_USAGE.QUERY_HISTORY' does not exist or not authorized.",
+            errno=2003,
+        )
+    )
+    assert SNOWFLAKE_ERRORS.classify(error).title == "Account usage not accessible"
+
+
+def test_account_usage_denied_wins_over_generic_object_not_found():
+    # The same message also matches the generic "does not exist or not authorized"
+    # rule; the account_usage rule is ordered first so the sharper diagnosis wins.
+    error = _SnowflakeError(
+        "SQL compilation error: Object 'SNOWFLAKE.ACCOUNT_USAGE.ACCESS_HISTORY' does not exist or not authorized.",
+        errno=2003,
+    )
+    assert SNOWFLAKE_ERRORS.classify(error).title != "Object not found"
+
+
+def test_account_usage_denial_honors_custom_schema():
+    # The pack is built from the configured accountUsageSchema, so a denial on a
+    # custom schema is still diagnosed as an account_usage gap - not "Object not found".
+    schema = "MYDB.MY_USAGE"
+    checks = SnowflakeChecks(client=MagicMock(), service_connection=_config(accountUsageSchema=schema))
+    error = _SnowflakeError("Object 'MYDB.MY_USAGE.QUERY_HISTORY' does not exist or not authorized.", errno=2003)
+    assert checks.errors.classify(error).title == "Account usage not accessible"
+    # the default pack, keyed on snowflake.account_usage, does not recognize the custom schema
+    assert SNOWFLAKE_ERRORS.classify(error).title == "Object not found"
+
+
+def test_insufficient_privileges_is_classified():
+    error = _SnowflakeError("SQL access control error: Insufficient privileges to operate on schema 'PUBLIC'")
+    assert SNOWFLAKE_ERRORS.classify(error).title == "Insufficient privileges"
+
+
+def test_object_not_found_errno_is_classified():
+    error = _SqlAlchemyError(_SnowflakeError("Database 'NOPE' does not exist or not authorized.", errno=2003))
+    assert SNOWFLAKE_ERRORS.classify(error).title == "Object not found"
+
+
+def test_missing_database_use_is_classified():
+    # USE DATABASE on a missing DB -> errno 2043 with a different message
+    # ("Object does not exist, or operation cannot be performed") than the 2003 form.
+    error = _SqlAlchemyError(
+        _SnowflakeError(
+            "SQL compilation error:\nObject does not exist, or operation cannot be performed.",
+            errno=2043,
+        )
+    )
+    assert SNOWFLAKE_ERRORS.classify(error).title == "Object not found"
+
+
+def test_no_active_warehouse_is_classified():
+    error = _SnowflakeError("No active warehouse selected in the current session.")
+    assert SNOWFLAKE_ERRORS.classify(error).title == "No active warehouse"
+
+
+def test_unknown_error_returns_no_diagnosis():
+    error = _SqlAlchemyError(_SnowflakeError("something unexpected", errno=99999))
+    assert SNOWFLAKE_ERRORS.classify(error) is None
+
+
+def test_network_errors_classify_through_including():
+    error = NetworkUnreachableError("acc.snowflakecomputing.com:443 is not reachable")
+    error.__cause__ = TimeoutError("timed out")
+    assert SNOWFLAKE_ERRORS.classify(error).title == "Connection timed out"
+
+
+def test_count_summary_marks_empty_count_and_cap():
+    assert _count_summary([], "table") == "no tables enumerated"
+    assert _count_summary([object()] * 3, "database") == "3 databases enumerated"
+    assert _count_summary([object()] * 3, "view") == "3 views enumerated"
+    assert _count_summary([object()] * 1, "stream") == "1 stream enumerated"
+    assert _count_summary([object()] * 1, "table") == "1 table enumerated"
+    capped = _count_summary([object()] * DEFAULT_SAMPLE_ROWS, "table")
+    assert capped == f"{DEFAULT_SAMPLE_ROWS}+ tables enumerated"
+
+
+def _fake_run_sql(returned_rows):
+    def fake(client, statement, summarize, *args, **kwargs):
+        return Evidence(summary=summarize(returned_rows), command="cmd")
+
+    return fake
+
+
+def test_get_tables_warns_when_no_user_tables():
+    # Empty probe (INFORMATION_SCHEMA filtered out) -> passing step with a caveat,
+    # which the runner records as a Warning.
+    checks = SnowflakeChecks(client=MagicMock(), service_connection=_config(database="MYDB"))
+    with patch(
+        "metadata.ingestion.source.database.snowflake.connection.run_sql",
+        side_effect=_fake_run_sql([]),
+    ):
+        evidence = checks.get_tables()
+    assert evidence.summary == "no tables enumerated"
+    assert evidence.caveat is not None
+    assert evidence.caveat.title == "No tables visible in database 'MYDB'"
+    assert "privileges" in evidence.caveat.remediation.lower()
+
+
+def test_get_tables_has_no_caveat_when_tables_exist():
+    checks = SnowflakeChecks(client=MagicMock(), service_connection=_config(database="MYDB"))
+    with patch(
+        "metadata.ingestion.source.database.snowflake.connection.run_sql",
+        side_effect=_fake_run_sql([object(), object()]),
+    ):
+        evidence = checks.get_tables()
+    assert evidence.summary == "2 tables enumerated"
+    assert evidence.caveat is None
+
+
+def test_table_and_view_probes_exclude_information_schema():
+    from metadata.ingestion.source.database.snowflake.queries import (
+        SNOWFLAKE_TEST_GET_TABLES,
+        SNOWFLAKE_TEST_GET_VIEWS,
+    )
+
+    for query in (SNOWFLAKE_TEST_GET_TABLES, SNOWFLAKE_TEST_GET_VIEWS):
+        assert "INFORMATION_SCHEMA" in query
+        assert "LIMIT 100" in query
+
+
+def test_checks_cover_exactly_the_wired_steps():
+    checks = SnowflakeChecks(client=MagicMock(), service_connection=_config())
+    collected = collect_checks(checks)
+    assert set(collected.keys()) == {
+        DatabaseStep.CheckAccess,
+        DatabaseStep.GetDatabases,
+        DatabaseStep.GetSchemas,
+        DatabaseStep.GetTables,
+        DatabaseStep.GetViews,
+        DatabaseStep.GetStreams,
+        DatabaseStep.GetTags,
+        DatabaseStep.GetQueries,
+        DatabaseStep.GetAccessHistory,
+    }
+
+
+def test_get_access_history_is_wired():
+    # ACCESS_HISTORY is the default lineage source; its probe is the GetAccessHistory step.
+    checks = SnowflakeChecks(client=MagicMock(), service_connection=_config())
+    assert DatabaseStep.GetAccessHistory in collect_checks(checks)
+
+
+def test_construction_touches_no_network():
+    # Regression for gotcha #2: building the provider must not connect or resolve a
+    # database - that would run before the gate and bypass the preflight.
+    client = MagicMock()
+    SnowflakeChecks(client=client, service_connection=_config())
+    client.connect.assert_not_called()
+    client.execute.assert_not_called()
+
+
+def test_check_access_probes_account_host_and_reports_network_failure():
+    # check_access -> tcp_probe(<account>.snowflakecomputing.com, 443); a probe
+    # failure is wrapped as a CheckError whose cause classifies via the network
+    # pack. tcp_probe is stubbed so the test is deterministic and fast.
+    checks = SnowflakeChecks(client=MagicMock(), service_connection=_config(account="acc"))
+    probe_error = NetworkUnreachableError("acc.snowflakecomputing.com:443 is not reachable")
+    probe_error.__cause__ = TimeoutError("timed out")
+    with (
+        patch(
+            "metadata.ingestion.source.database.snowflake.connection.tcp_probe",
+            side_effect=probe_error,
+        ) as mock_probe,
+        pytest.raises(CheckError) as exc,
+    ):
+        checks.check_access()
+    mock_probe.assert_called_once_with("acc.snowflakecomputing.com", SNOWFLAKE_PORT)
+    assert exc.value.cause is probe_error
+    assert SNOWFLAKE_ERRORS.classify(exc.value.cause).title == "Connection timed out"
+
+
+def test_check_access_prefers_explicit_connection_argument_host():
+    # A proxy / load balancer / PrivateLink endpoint set as connectionArguments
+    # host is what the driver dials, so the gate probe must target it - not the
+    # synthesized public account host (which may be unreachable for that deployment).
+    checks = SnowflakeChecks(
+        client=MagicMock(),
+        service_connection=_config(account="acc", connectionArguments={"host": "proxy.internal"}),
+    )
+    probe_error = NetworkUnreachableError("proxy.internal:443 is not reachable")
+    with (
+        patch(
+            "metadata.ingestion.source.database.snowflake.connection.tcp_probe",
+            side_effect=probe_error,
+        ) as mock_probe,
+        pytest.raises(CheckError),
+    ):
+        checks.check_access()
+    mock_probe.assert_called_once_with("proxy.internal", SNOWFLAKE_PORT)
+
+
+def test_check_access_honors_explicit_connection_argument_port():
+    # A custom host may listen on a non-443 port; the preflight must probe that
+    # port, not the default 443, or it false-gates the whole connection test.
+    checks = SnowflakeChecks(
+        client=MagicMock(),
+        service_connection=_config(account="acc", connectionArguments={"host": "proxy.internal", "port": 8443}),
+    )
+    with (
+        patch(
+            "metadata.ingestion.source.database.snowflake.connection.tcp_probe",
+            side_effect=NetworkUnreachableError("proxy.internal:8443 is not reachable"),
+        ) as mock_probe,
+        pytest.raises(CheckError),
+    ):
+        checks.check_access()
+    mock_probe.assert_called_once_with("proxy.internal", 8443)
+
+
+def test_get_schemas_runs_show_schemas_with_command_and_count():
+    # GetSchemas goes through run_sql (like the other probes), so it reports the
+    # database-scoped SHOW SCHEMAS command and a counted summary - and a failure
+    # would surface as a CheckError carrying that command.
+    checks = SnowflakeChecks(client=MagicMock(), service_connection=_config(database="MYDB"))
+    captured = {}
+
+    def fake(client, statement, summarize, *args, **kwargs):
+        captured["statement"] = statement
+        return Evidence(summary=summarize([object(), object()]), command=statement)
+
+    with patch("metadata.ingestion.source.database.snowflake.connection.run_sql", side_effect=fake):
+        evidence = checks.get_schemas()
+    assert 'SHOW SCHEMAS IN DATABASE "MYDB"' in captured["statement"]
+    assert evidence.summary == "2 schemas enumerated"
+    assert evidence.command == captured["statement"]
+
+
+def test_account_usage_queries_built_lazily_not_at_construction():
+    # The account_usage statements must be formatted inside their checks (behind
+    # the gate), never at construction; constructing must not read the engine.
+    client = MagicMock()
+    checks = SnowflakeChecks(client=client, service_connection=_config(account="acc"))
+    assert checks._engine_wrapper.database_name is None
+    client.connect.assert_not_called()

ingestion/tests/unit/source/database/sqlite/test_connection.py

@@ -0,0 +1,34 @@
+#  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 SQLite connection handling."""
+
+from metadata.generated.schema.entity.services.connections.database.sqliteConnection import (
+    SQLiteConnection as SQLiteConnectionConfig,
+)
+from metadata.generated.schema.entity.services.connections.database.sqliteConnection import (
+    SQLiteScheme,
+)
+from metadata.ingestion.source.database.sqlite.connection import SQLiteConnection
+
+
+def test_in_memory_url_when_no_database_mode():
+    connection = SQLiteConnectionConfig(scheme=SQLiteScheme.sqlite_pysqlite)
+    engine = SQLiteConnection(connection).client
+    assert engine.url.render_as_string(hide_password=False) == "sqlite+pysqlite:///:memory:"
+
+
+def test_file_url_when_database_mode_set():
+    connection = SQLiteConnectionConfig(
+        scheme=SQLiteScheme.sqlite_pysqlite,
+        databaseMode="/tmp/openmetadata.db",
+    )
+    engine = SQLiteConnection(connection).client
+    assert engine.url.render_as_string(hide_password=False) == "sqlite+pysqlite:////tmp/openmetadata.db"

ingestion/tests/unit/source/database/starrocks/test_connection.py

@@ -0,0 +1,34 @@
+#  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 StarRocks connection handling."""
+
+from metadata.generated.schema.entity.services.connections.database.starrocksConnection import (
+    StarRocksConnection as StarRocksConnectionConfig,
+)
+from metadata.generated.schema.entity.services.connections.database.starrocksConnection import (
+    StarrocksScheme,
+)
+from metadata.ingestion.source.database.starrocks.connection import StarRocksConnection
+
+
+def test_basic_auth_builds_expected_url():
+    connection = StarRocksConnectionConfig(
+        username="openmetadata_user",
+        password="openmetadata_password",
+        hostPort="localhost:9030",
+        databaseSchema="openmetadata_db",
+        scheme=StarrocksScheme.mysql_pymysql,
+    )
+    engine = StarRocksConnection(connection).client
+    assert (
+        engine.url.render_as_string(hide_password=False)
+        == "mysql+pymysql://openmetadata_user:openmetadata_password@localhost:9030/openmetadata_db"
+    )

ingestion/tests/unit/source/database/teradata/test_connection.py

@@ -0,0 +1,37 @@
+#  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 Teradata connection handling."""
+
+from metadata.generated.schema.entity.services.connections.database.teradataConnection import (
+    TeradataConnection as TeradataConnectionConfig,
+)
+from metadata.generated.schema.entity.services.connections.database.teradataConnection import (
+    TeradataScheme,
+)
+from metadata.ingestion.connections.connection import BaseConnection
+from metadata.ingestion.source.database.teradata.connection import TeradataConnection
+
+
+def test_teradata_connection_is_base_connection():
+    assert issubclass(TeradataConnection, BaseConnection)
+
+
+def test_get_connection_url_with_credentials_and_defaults():
+    connection = TeradataConnectionConfig(
+        scheme=TeradataScheme.teradatasql,
+        username="openmetadata_user",
+        password="openmetadata_password",
+        hostPort="localhost:1025",
+    )
+    assert (
+        TeradataConnection.get_connection_url(connection) == "teradatasql://localhost:1025/?user=openmetadata_user"
+        "&password=openmetadata_password&logmech=TD2&tmode=DEFAULT"
+    )

ingestion/tests/unit/source/database/test_database_service.py

@@ -0,0 +1,25 @@
+#  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.
+"""
+Structural tests for DatabaseServiceSource — guards against re-introducing
+progress-projection responsibilities that now live in the topology runner
+and connector-specific modules.
+"""
+
+from metadata.ingestion.source.database.database_service import DatabaseServiceSource
+
+
+class TestDatabaseServiceHasNoProgressProjection:
+    def test_no_container_expected_count(self):
+        assert "container_expected_count" not in vars(DatabaseServiceSource)
+
+    def test_no_progress_snapshot(self):
+        assert "progress_snapshot" not in vars(DatabaseServiceSource)

ingestion/tests/unit/source/database/test_json_schema_extractor.py

@@ -0,0 +1,902 @@
+#  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 JSON schema extraction from sampled data.
+"""
+
+import json
+
+from metadata.generated.schema.entity.data.table import DataType
+from metadata.ingestion.source.database.json_schema_extractor import (
+    _build_column_children,
+    _build_json_schema,
+    _merge_json_structures,
+    _parse_json_values,
+    infer_json_schema_from_sample,
+)
+
+
+class TestParseJsonValues:
+    """Tests for _parse_json_values function."""
+
+    def test_parse_dict_values(self):
+        """Test parsing already-parsed dict values."""
+        values = [{"name": "John"}, {"name": "Jane"}]
+        result = _parse_json_values(values)
+        assert len(result) == 2
+        assert result[0] == {"name": "John"}
+        assert result[1] == {"name": "Jane"}
+
+    def test_parse_string_json_values(self):
+        """Test parsing JSON string values."""
+        values = ['{"name": "John"}', '{"name": "Jane"}']
+        result = _parse_json_values(values)
+        assert len(result) == 2
+        assert result[0] == {"name": "John"}
+        assert result[1] == {"name": "Jane"}
+
+    def test_parse_mixed_values(self):
+        """Test parsing mixed dict and string values."""
+        values = [{"name": "John"}, '{"name": "Jane"}']
+        result = _parse_json_values(values)
+        assert len(result) == 2
+
+    def test_filter_none_values(self):
+        """Test that None values are filtered out."""
+        values = [{"name": "John"}, None, {"name": "Jane"}]
+        result = _parse_json_values(values)
+        assert len(result) == 2
+
+    def test_filter_non_dict_values(self):
+        """Test that non-dict JSON values are filtered out."""
+        values = [{"name": "John"}, '"just a string"', "[1, 2, 3]"]
+        result = _parse_json_values(values)
+        assert len(result) == 1
+        assert result[0] == {"name": "John"}
+
+    def test_handle_invalid_json(self):
+        """Test that invalid JSON strings are filtered out."""
+        values = [{"name": "John"}, "not valid json", '{"name": "Jane"}']
+        result = _parse_json_values(values)
+        assert len(result) == 2
+
+    def test_empty_list(self):
+        """Test handling empty list."""
+        result = _parse_json_values([])
+        assert result == []
+
+
+class TestMergeJsonStructures:
+    """Tests for _merge_json_structures function."""
+
+    def test_merge_simple_objects(self):
+        """Test merging simple objects with same keys."""
+        dicts = [{"name": "John", "age": 30}, {"name": "Jane", "age": 25}]
+        result = _merge_json_structures(dicts)
+        assert "name" in result
+        assert "age" in result
+
+    def test_merge_objects_with_different_keys(self):
+        """Test merging objects with different keys."""
+        dicts = [{"name": "John"}, {"age": 30}, {"city": "NYC"}]
+        result = _merge_json_structures(dicts)
+        assert "name" in result
+        assert "age" in result
+        assert "city" in result
+
+    def test_merge_nested_objects(self):
+        """Test merging nested objects."""
+        dicts = [
+            {"user": {"name": "John"}},
+            {"user": {"age": 30}},
+        ]
+        result = _merge_json_structures(dicts)
+        assert "user" in result
+        assert isinstance(result["user"], dict)
+        assert "name" in result["user"]
+        assert "age" in result["user"]
+
+    def test_merge_with_arrays(self):
+        """Test merging objects containing arrays."""
+        dicts = [
+            {"tags": ["python", "java"]},
+            {"tags": ["javascript"]},
+        ]
+        result = _merge_json_structures(dicts)
+        assert "tags" in result
+        assert isinstance(result["tags"], list)
+
+    def test_merge_array_of_objects(self):
+        """Test merging arrays containing objects."""
+        dicts = [
+            {"items": [{"id": 1}]},
+            {"items": [{"name": "Item 2"}]},
+        ]
+        result = _merge_json_structures(dicts)
+        assert "items" in result
+        assert isinstance(result["items"], list)
+        assert len(result["items"]) == 1
+        assert "id" in result["items"][0]
+        assert "name" in result["items"][0]
+
+    def test_dict_takes_precedence_over_scalar(self):
+        """Test that dict value takes precedence over scalar."""
+        dicts = [
+            {"field": "string_value"},
+            {"field": {"nested": "value"}},
+        ]
+        result = _merge_json_structures(dicts)
+        assert isinstance(result["field"], dict)
+        assert result["field"]["nested"] == "value"
+
+    def test_handle_null_values(self):
+        """Test handling null values in objects."""
+        dicts = [
+            {"name": "John", "email": None},
+            {"name": "Jane", "email": "jane@example.com"},
+        ]
+        result = _merge_json_structures(dicts)
+        assert "email" in result
+        assert result["email"] == "jane@example.com"
+
+
+class TestBuildJsonSchema:
+    """Tests for _build_json_schema function."""
+
+    def test_build_simple_object_schema(self):
+        """Test building schema for simple object."""
+        structure = {"name": "John", "age": 30}
+        schema = _build_json_schema(structure)
+        assert schema["type"] == "object"
+        assert "properties" in schema
+        assert schema["properties"]["name"]["type"] == "string"
+        assert schema["properties"]["age"]["type"] == "integer"
+
+    def test_build_nested_object_schema(self):
+        """Test building schema for nested object."""
+        structure = {"user": {"name": "John"}}
+        schema = _build_json_schema(structure)
+        assert schema["type"] == "object"
+        assert schema["properties"]["user"]["type"] == "object"
+        assert schema["properties"]["user"]["properties"]["name"]["type"] == "string"
+
+    def test_build_array_schema(self):
+        """Test building schema for array."""
+        structure = {"tags": ["python"]}
+        schema = _build_json_schema(structure)
+        assert schema["properties"]["tags"]["type"] == "array"
+        assert schema["properties"]["tags"]["items"]["type"] == "string"
+
+    def test_build_array_of_objects_schema(self):
+        """Test building schema for array of objects."""
+        structure = {"items": [{"id": 1, "name": "test"}]}
+        schema = _build_json_schema(structure)
+        assert schema["properties"]["items"]["type"] == "array"
+        assert schema["properties"]["items"]["items"]["type"] == "object"
+        assert "id" in schema["properties"]["items"]["items"]["properties"]
+
+    def test_build_schema_with_boolean(self):
+        """Test building schema with boolean type."""
+        structure = {"active": True}
+        schema = _build_json_schema(structure)
+        assert schema["properties"]["active"]["type"] == "boolean"
+
+    def test_build_schema_with_float(self):
+        """Test building schema with float type."""
+        structure = {"price": 19.99}
+        schema = _build_json_schema(structure)
+        assert schema["properties"]["price"]["type"] == "number"
+
+    def test_build_schema_with_null(self):
+        """Test building schema with null type."""
+        structure = {"value": None}
+        schema = _build_json_schema(structure)
+        assert schema["properties"]["value"]["type"] == "null"
+
+
+class TestBuildColumnChildren:
+    """Tests for _build_column_children function."""
+
+    def test_build_simple_children(self):
+        """Test building children for simple structure."""
+        structure = {"name": "John", "age": 30}
+        children = _build_column_children(structure)
+        assert children is not None
+        assert len(children) == 2
+
+        names = {child.name.root for child in children}
+        assert "name" in names
+        assert "age" in names
+
+    def test_build_nested_children(self):
+        """Test building children with nested objects."""
+        structure = {"user": {"name": "John", "email": "john@example.com"}}
+        children = _build_column_children(structure)
+        assert children is not None
+        assert len(children) == 1
+
+        user_child = children[0]
+        assert user_child.name.root == "user"
+        assert user_child.dataType == DataType.JSON
+        assert user_child.children is not None
+        assert len(user_child.children) == 2
+
+    def test_build_array_children(self):
+        """Test building children with array type."""
+        structure = {"tags": ["python", "java"]}
+        children = _build_column_children(structure)
+        assert children is not None
+
+        tags_child = children[0]
+        assert tags_child.name.root == "tags"
+        assert tags_child.dataType == DataType.ARRAY
+        assert tags_child.arrayDataType == DataType.STRING
+
+    def test_build_array_of_objects_children(self):
+        """Test building children for array of objects."""
+        structure = {"items": [{"id": 1, "name": "test"}]}
+        children = _build_column_children(structure)
+        assert children is not None
+
+        items_child = children[0]
+        assert items_child.name.root == "items"
+        assert items_child.dataType == DataType.ARRAY
+        assert items_child.arrayDataType == DataType.JSON
+        assert items_child.children is not None
+        assert len(items_child.children) == 2
+
+    def test_returns_none_for_non_dict(self):
+        """Test that non-dict input returns None."""
+        result = _build_column_children("not a dict")
+        assert result is None
+
+
+class TestInferJsonSchemaFromSample:
+    """Tests for the main infer_json_schema_from_sample function."""
+
+    def test_infer_schema_basic(self):
+        """Test basic schema inference."""
+        json_values = [
+            {"name": "John", "age": 30},
+            {"name": "Jane", "age": 25},
+        ]
+        schema_str, children = infer_json_schema_from_sample(json_values)
+
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+        assert schema["type"] == "object"
+        assert "name" in schema["properties"]
+        assert "age" in schema["properties"]
+
+        assert children is not None
+        assert len(children) == 2
+
+    def test_infer_schema_from_json_strings(self):
+        """Test schema inference from JSON strings."""
+        json_values = [
+            '{"product": "laptop", "price": 999.99}',
+            '{"product": "mouse", "price": 29.99}',
+        ]
+        schema_str, children = infer_json_schema_from_sample(json_values)
+
+        assert schema_str is not None
+        assert children is not None
+
+        schema = json.loads(schema_str)
+        assert schema["properties"]["product"]["type"] == "string"
+        assert schema["properties"]["price"]["type"] == "number"
+
+    def test_infer_schema_nested_structure(self):
+        """Test schema inference for nested structures."""
+        json_values = [
+            {
+                "user": {"name": "John", "email": "john@example.com"},
+                "active": True,
+            }
+        ]
+        schema_str, children = infer_json_schema_from_sample(json_values)  # noqa: RUF059
+
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+        assert schema["properties"]["user"]["type"] == "object"
+        assert "name" in schema["properties"]["user"]["properties"]
+
+    def test_infer_schema_with_arrays(self):
+        """Test schema inference for arrays."""
+        json_values = [
+            {"tags": ["python", "data"], "scores": [85, 90, 95]},
+        ]
+        schema_str, children = infer_json_schema_from_sample(json_values)  # noqa: RUF059
+
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+        assert schema["properties"]["tags"]["type"] == "array"
+        assert schema["properties"]["scores"]["type"] == "array"
+
+    def test_infer_schema_empty_list(self):
+        """Test that empty list returns None."""
+        schema_str, children = infer_json_schema_from_sample([])
+        assert schema_str is None
+        assert children is None
+
+    def test_infer_schema_all_none(self):
+        """Test that list of all None values returns None."""
+        schema_str, children = infer_json_schema_from_sample([None, None])
+        assert schema_str is None
+        assert children is None
+
+    def test_infer_schema_all_invalid(self):
+        """Test that list of all invalid JSON returns None."""
+        json_values = ["not json", "also not json"]
+        schema_str, children = infer_json_schema_from_sample(json_values)
+        assert schema_str is None
+        assert children is None
+
+    def test_infer_schema_handles_mixed_valid_invalid(self):
+        """Test that valid values are processed even with some invalid ones."""
+        json_values = [
+            {"name": "John"},
+            "invalid json",
+            None,
+            {"name": "Jane", "age": 25},
+        ]
+        schema_str, children = infer_json_schema_from_sample(json_values)  # noqa: RUF059
+
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+        assert "name" in schema["properties"]
+        assert "age" in schema["properties"]
+
+    def test_infer_schema_real_world_event_data(self):
+        """Test with realistic event data structure."""
+        json_values = [
+            {
+                "ab_test_group": "group_1",
+                "campaign_id": "cmp_01",
+                "user_agent": "Mozilla/5.0",
+            },
+            {
+                "ab_test_group": "group_2",
+                "campaign_id": "cmp_02",
+                "amount_cents": 5000,
+            },
+        ]
+        schema_str, children = infer_json_schema_from_sample(json_values)  # noqa: RUF059
+
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+        assert "ab_test_group" in schema["properties"]
+        assert "campaign_id" in schema["properties"]
+        assert "user_agent" in schema["properties"]
+        assert "amount_cents" in schema["properties"]
+
+    def test_infer_schema_deeply_nested(self):
+        """Test with deeply nested structure."""
+        json_values = [{"level1": {"level2": {"level3": {"value": "deep"}}}}]
+        schema_str, children = infer_json_schema_from_sample(json_values)  # noqa: RUF059
+
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+        assert schema["properties"]["level1"]["type"] == "object"
+        level2_props = schema["properties"]["level1"]["properties"]["level2"]
+        assert level2_props["type"] == "object"
+        level3_props = level2_props["properties"]["level3"]
+        assert level3_props["type"] == "object"
+        assert level3_props["properties"]["value"]["type"] == "string"
+
+
+class TestJsonSchemaExtractionEdgeCases:
+    """Edge case tests for JSON schema extraction."""
+
+    def test_empty_object(self):
+        """Test handling empty JSON objects."""
+        json_values = [{}]
+        schema_str, children = infer_json_schema_from_sample(json_values)  # noqa: RUF059
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+        assert schema["type"] == "object"
+        assert schema["properties"] == {}
+
+    def test_special_characters_in_keys(self):
+        """Test handling special characters in JSON keys."""
+        json_values = [{"user-name": "John", "email@domain": "test@test.com"}]
+        schema_str, children = infer_json_schema_from_sample(json_values)
+
+        assert schema_str is not None
+        assert children is not None
+
+    def test_unicode_values(self):
+        """Test handling unicode values."""
+        json_values = [{"name": "日本語", "emoji": "🎉"}]
+        schema_str, children = infer_json_schema_from_sample(json_values)
+
+        assert schema_str is not None
+        assert children is not None
+
+    def test_large_numbers(self):
+        """Test handling large numbers."""
+        json_values = [{"big_int": 9999999999999999, "big_float": 1.7976931348623157e308}]
+        schema_str, children = infer_json_schema_from_sample(json_values)  # noqa: RUF059
+
+        assert schema_str is not None
+
+
+class TestJsonColumnTypeDetection:
+    """Tests for JSON schema extraction with different column types."""
+
+    def test_json_column_type_json_datatype(self):
+        """Test JSON schema extraction for DataType.JSON columns."""
+        json_values = [
+            {"user_id": 123, "preferences": {"theme": "dark"}},
+            {"user_id": 456, "preferences": {"theme": "light", "language": "en"}},
+        ]
+        schema_str, children = infer_json_schema_from_sample(json_values)
+
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+        assert schema["type"] == "object"
+        assert "user_id" in schema["properties"]
+        assert "preferences" in schema["properties"]
+        assert schema["properties"]["preferences"]["type"] == "object"
+
+        assert children is not None
+        child_names = {child.name.root for child in children}
+        assert "user_id" in child_names
+        assert "preferences" in child_names
+
+    def test_json_column_type_jsonb(self):
+        """Test JSON schema extraction for JSONB columns (PostgreSQL)."""
+        json_values = [
+            {"metadata": {"created_at": "2024-01-01", "updated_at": "2024-01-02"}},
+            {"metadata": {"created_at": "2024-02-01", "version": 2}},
+        ]
+        schema_str, children = infer_json_schema_from_sample(json_values)
+
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+        assert "metadata" in schema["properties"]
+        assert children is not None
+
+    def test_json_column_type_variant(self):
+        """Test JSON schema extraction for VARIANT columns (Snowflake)."""
+        json_values = [
+            {"event_type": "click", "payload": {"x": 100, "y": 200}},
+            {"event_type": "scroll", "payload": {"direction": "down", "amount": 50}},
+        ]
+        schema_str, children = infer_json_schema_from_sample(json_values)
+
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+        assert "event_type" in schema["properties"]
+        assert "payload" in schema["properties"]
+        assert children is not None
+
+    def test_json_column_type_object(self):
+        """Test JSON schema extraction for OBJECT columns (Snowflake/Databricks)."""
+        json_values = [
+            {"config": {"enabled": True, "max_retries": 3}},
+            {"config": {"enabled": False, "timeout": 30}},
+        ]
+        schema_str, children = infer_json_schema_from_sample(json_values)
+
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+        assert "config" in schema["properties"]
+        assert children is not None
+
+
+class TestStringColumnTypeAsJson:
+    """Tests for JSON schema extraction from STRING type columns containing JSON data."""
+
+    def test_string_column_with_json_data(self):
+        """Test JSON schema extraction from STRING columns containing JSON."""
+        json_values = [
+            '{"name": "Alice", "score": 95}',
+            '{"name": "Bob", "score": 87}',
+        ]
+        schema_str, children = infer_json_schema_from_sample(json_values)
+
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+        assert schema["type"] == "object"
+        assert "name" in schema["properties"]
+        assert "score" in schema["properties"]
+        assert schema["properties"]["name"]["type"] == "string"
+        assert schema["properties"]["score"]["type"] == "integer"
+
+        assert children is not None
+        child_names = {child.name.root for child in children}
+        assert "name" in child_names
+        assert "score" in child_names
+
+    def test_varchar_column_with_json_data(self):
+        """Test JSON schema extraction from VARCHAR columns containing JSON."""
+        json_values = [
+            '{"product": "laptop", "price": 999.99, "in_stock": true}',
+            '{"product": "mouse", "price": 29.99, "in_stock": false}',
+        ]
+        schema_str, children = infer_json_schema_from_sample(json_values)  # noqa: RUF059
+
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+        assert "product" in schema["properties"]
+        assert "price" in schema["properties"]
+        assert "in_stock" in schema["properties"]
+        assert schema["properties"]["price"]["type"] == "number"
+        assert schema["properties"]["in_stock"]["type"] == "boolean"
+
+    def test_text_column_with_json_data(self):
+        """Test JSON schema extraction from TEXT columns containing JSON."""
+        json_values = [
+            '{"log_level": "INFO", "message": "Application started", "timestamp": 1704067200}',
+            '{"log_level": "ERROR", "message": "Connection failed", "details": {"code": 500}}',
+        ]
+        schema_str, children = infer_json_schema_from_sample(json_values)
+
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+        assert "log_level" in schema["properties"]
+        assert "message" in schema["properties"]
+        assert "details" in schema["properties"]
+
+        assert children is not None
+
+    def test_string_column_with_nested_json(self):
+        """Test JSON schema extraction from STRING columns with deeply nested JSON."""
+        json_values = [
+            '{"user": {"profile": {"address": {"city": "NYC", "zip": "10001"}}}}',
+            '{"user": {"profile": {"address": {"city": "LA", "state": "CA"}}}}',
+        ]
+        schema_str, children = infer_json_schema_from_sample(json_values)
+
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+        assert "user" in schema["properties"]
+        user_props = schema["properties"]["user"]
+        assert user_props["type"] == "object"
+        assert "profile" in user_props["properties"]
+
+        assert children is not None
+        user_child = next(c for c in children if c.name.root == "user")
+        assert user_child.dataType == DataType.JSON
+        assert user_child.children is not None
+
+    def test_string_column_with_array_json(self):
+        """Test JSON schema extraction from STRING columns with JSON arrays."""
+        json_values = [
+            '{"items": [{"id": 1, "name": "Item 1"}, {"id": 2, "name": "Item 2"}]}',
+            '{"items": [{"id": 3, "name": "Item 3", "price": 10.99}]}',
+        ]
+        schema_str, children = infer_json_schema_from_sample(json_values)
+
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+        assert "items" in schema["properties"]
+        items_schema = schema["properties"]["items"]
+        assert items_schema["type"] == "array"
+        assert items_schema["items"]["type"] == "object"
+
+        assert children is not None
+        items_child = next(c for c in children if c.name.root == "items")
+        assert items_child.dataType == DataType.ARRAY
+        assert items_child.arrayDataType == DataType.JSON
+
+    def test_string_column_with_mixed_valid_invalid_json(self):
+        """Test STRING column where some values are valid JSON and some are not."""
+        json_values = [
+            '{"valid": true}',
+            "plain text not json",
+            '{"also_valid": 123}',
+            None,
+            "",
+        ]
+        schema_str, children = infer_json_schema_from_sample(json_values)  # noqa: RUF059
+
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+        assert "valid" in schema["properties"]
+        assert "also_valid" in schema["properties"]
+
+    def test_string_column_all_invalid_json(self):
+        """Test STRING column where all values are invalid JSON."""
+        json_values = [
+            "plain text",
+            "another plain text",
+            "not json at all",
+        ]
+        schema_str, children = infer_json_schema_from_sample(json_values)
+
+        assert schema_str is None
+        assert children is None
+
+    def test_string_column_with_json_primitives(self):
+        """Test STRING column with JSON primitives (not objects) - should be filtered."""
+        json_values = [
+            '"just a string"',
+            "123",
+            "true",
+            "[1, 2, 3]",
+        ]
+        schema_str, children = infer_json_schema_from_sample(json_values)
+
+        assert schema_str is None
+        assert children is None
+
+
+class TestAllJsonColumnTypes:
+    """Comprehensive tests covering all JSON_COLUMN_TYPES defined in sql_column_handler.py."""
+
+    def test_all_json_types_with_complex_structure(self):
+        """Test that all JSON column types can handle complex nested structures."""
+        complex_json = [
+            {
+                "id": 1,
+                "name": "Test",
+                "attributes": {
+                    "color": "red",
+                    "size": "large",
+                    "tags": ["tag1", "tag2"],
+                },
+                "measurements": [
+                    {"width": 10, "height": 20},
+                    {"width": 15, "height": 25},
+                ],
+                "metadata": {
+                    "created": "2024-01-01",
+                    "nested": {"level1": {"level2": {"value": 42}}},
+                },
+                "active": True,
+                "score": 98.5,
+                "nullable_field": None,
+            }
+        ]
+        schema_str, children = infer_json_schema_from_sample(complex_json)
+
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+
+        assert schema["properties"]["id"]["type"] == "integer"
+        assert schema["properties"]["name"]["type"] == "string"
+        assert schema["properties"]["attributes"]["type"] == "object"
+        assert schema["properties"]["measurements"]["type"] == "array"
+        assert schema["properties"]["metadata"]["type"] == "object"
+        assert schema["properties"]["active"]["type"] == "boolean"
+        assert schema["properties"]["score"]["type"] == "number"
+        assert schema["properties"]["nullable_field"]["type"] == "null"
+
+        assert children is not None
+        child_names = {c.name.root for c in children}
+        assert child_names == {
+            "id",
+            "name",
+            "attributes",
+            "measurements",
+            "metadata",
+            "active",
+            "score",
+            "nullable_field",
+        }
+
+    def test_json_type_with_empty_objects_and_arrays(self):
+        """Test JSON columns with empty objects and arrays."""
+        json_values = [
+            {"empty_obj": {}, "empty_arr": [], "normal": "value"},
+            {"empty_obj": {"key": "value"}, "empty_arr": [1, 2], "normal": "other"},
+        ]
+        schema_str, children = infer_json_schema_from_sample(json_values)  # noqa: RUF059
+
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+        assert "empty_obj" in schema["properties"]
+        assert "empty_arr" in schema["properties"]
+        assert "normal" in schema["properties"]
+
+    def test_jsonb_type_with_special_postgres_patterns(self):
+        """Test JSONB columns with patterns common in PostgreSQL."""
+        json_values = [
+            {
+                "audit_log": {
+                    "action": "UPDATE",
+                    "old_values": {"status": "pending"},
+                    "new_values": {"status": "completed"},
+                    "changed_by": "user_123",
+                }
+            },
+            {
+                "audit_log": {
+                    "action": "INSERT",
+                    "new_values": {"id": 456, "name": "New Record"},
+                    "changed_by": "system",
+                }
+            },
+        ]
+        schema_str, children = infer_json_schema_from_sample(json_values)  # noqa: RUF059
+
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+        audit_props = schema["properties"]["audit_log"]["properties"]
+        assert "action" in audit_props
+        assert "old_values" in audit_props
+        assert "new_values" in audit_props
+        assert "changed_by" in audit_props
+
+    def test_variant_type_with_semi_structured_data(self):
+        """Test VARIANT columns with semi-structured data (Snowflake pattern)."""
+        json_values = [
+            {
+                "raw_event": {
+                    "event_id": "evt_001",
+                    "timestamp": "2024-01-15T10:30:00Z",
+                    "properties": {"source": "web", "browser": "Chrome"},
+                }
+            },
+            {
+                "raw_event": {
+                    "event_id": "evt_002",
+                    "timestamp": "2024-01-15T10:31:00Z",
+                    "properties": {"source": "mobile", "os": "iOS", "version": "17.0"},
+                }
+            },
+        ]
+        schema_str, children = infer_json_schema_from_sample(json_values)  # noqa: RUF059
+
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+        event_props = schema["properties"]["raw_event"]["properties"]
+        assert "event_id" in event_props
+        assert "timestamp" in event_props
+        assert "properties" in event_props
+
+    def test_object_type_with_nested_structures(self):
+        """Test OBJECT columns with nested structures (Databricks/Snowflake pattern)."""
+        json_values = [
+            {
+                "customer": {
+                    "id": 1001,
+                    "details": {
+                        "first_name": "John",
+                        "last_name": "Doe",
+                        "contacts": {
+                            "email": "john@example.com",
+                            "phones": ["+1-555-0100", "+1-555-0101"],
+                        },
+                    },
+                }
+            },
+            {
+                "customer": {
+                    "id": 1002,
+                    "details": {
+                        "first_name": "Jane",
+                        "last_name": "Smith",
+                        "contacts": {"email": "jane@example.com"},
+                    },
+                }
+            },
+        ]
+        schema_str, children = infer_json_schema_from_sample(json_values)  # noqa: RUF059
+
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+        customer_schema = schema["properties"]["customer"]
+        assert customer_schema["type"] == "object"
+        details_schema = customer_schema["properties"]["details"]
+        assert "first_name" in details_schema["properties"]
+        assert "contacts" in details_schema["properties"]
+
+
+class TestAllStringColumnTypes:
+    """Comprehensive tests covering all STRING_COLUMN_TYPES defined in sql_column_handler.py."""
+
+    def test_string_datatype_value(self):
+        """Test DataType.STRING.value column with JSON data."""
+        json_values = [
+            {"field1": "value1", "field2": 100},
+            {"field1": "value2", "field2": 200, "field3": True},
+        ]
+        schema_str, children = infer_json_schema_from_sample(json_values)  # noqa: RUF059
+
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+        assert "field1" in schema["properties"]
+        assert "field2" in schema["properties"]
+        assert "field3" in schema["properties"]
+
+    def test_string_literal_type(self):
+        """Test literal 'STRING' type column with JSON data."""
+        json_values = [
+            '{"api_response": {"status": 200, "data": {"items": [1, 2, 3]}}}',
+        ]
+        schema_str, children = infer_json_schema_from_sample(json_values)  # noqa: RUF059
+
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+        assert "api_response" in schema["properties"]
+        api_response = schema["properties"]["api_response"]
+        assert "status" in api_response["properties"]
+        assert "data" in api_response["properties"]
+
+    def test_varchar_type_with_various_lengths(self):
+        """Test VARCHAR type columns with JSON of various complexities."""
+        short_json = [{"a": 1}]
+        schema_str, children = infer_json_schema_from_sample(short_json)
+        assert schema_str is not None
+
+        long_json = [
+            {
+                "very_long_key_name_that_might_be_stored_in_varchar": {
+                    "nested": {"deeply": {"structured": {"data": "value"}}}
+                }
+            }
+        ]
+        schema_str, children = infer_json_schema_from_sample(long_json)  # noqa: RUF059
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+        assert "very_long_key_name_that_might_be_stored_in_varchar" in schema["properties"]
+
+    def test_text_type_with_large_json(self):
+        """Test TEXT type columns that might store large JSON documents."""
+        large_json = [
+            {
+                "document": {
+                    "sections": [{"title": f"Section {i}", "content": f"Content for section {i}"} for i in range(10)],
+                    "metadata": {
+                        "author": "Test Author",
+                        "created": "2024-01-01",
+                        "tags": ["tag1", "tag2", "tag3"],
+                    },
+                }
+            }
+        ]
+        schema_str, children = infer_json_schema_from_sample(large_json)  # noqa: RUF059
+
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+        doc_schema = schema["properties"]["document"]
+        assert "sections" in doc_schema["properties"]
+        assert "metadata" in doc_schema["properties"]
+
+    def test_string_types_with_unicode_json(self):
+        """Test STRING types with JSON containing unicode characters."""
+        json_values = [
+            {"name": "日本語テスト", "emoji": "🎉🚀", "special": "café résumé naïve"},
+            {"name": "中文测试", "emoji": "✨💻", "special": "über straße"},
+        ]
+        schema_str, children = infer_json_schema_from_sample(json_values)
+
+        assert schema_str is not None
+        assert children is not None
+
+    def test_string_types_with_escaped_json(self):
+        """Test STRING types with JSON containing escaped characters."""
+        json_values = [
+            '{"path": "C:\\\\Users\\\\test", "quote": "\\"quoted\\"", "newline": "line1\\nline2"}',
+        ]
+        schema_str, children = infer_json_schema_from_sample(json_values)  # noqa: RUF059
+
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+        assert "path" in schema["properties"]
+        assert "quote" in schema["properties"]
+        assert "newline" in schema["properties"]
+
+    def test_string_types_with_numeric_keys(self):
+        """Test STRING types with JSON containing numeric-like keys."""
+        json_values = [
+            {"123": "numeric key", "456abc": "mixed key", "normal_key": "value"},
+        ]
+        schema_str, children = infer_json_schema_from_sample(json_values)  # noqa: RUF059
+
+        assert schema_str is not None
+        schema = json.loads(schema_str)
+        assert "123" in schema["properties"]
+        assert "456abc" in schema["properties"]
+        assert "normal_key" in schema["properties"]

ingestion/tests/unit/source/database/test_mysql_cloudsql.py

@@ -0,0 +1,206 @@
+#  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.
+
+"""
+Tests for GCP CloudSQL MySQL connection handling
+"""
+
+import inspect
+import sys
+from types import ModuleType
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from metadata.generated.schema.entity.services.connections.database.common.gcpCloudSqlConfig import (
+    GcpCloudsqlConfigurationSource,
+)
+from metadata.generated.schema.entity.services.connections.database.mysqlConnection import (
+    MysqlConnection,
+)
+from metadata.ingestion.source.database.mysql.connection import _CloudSqlStrategy
+
+_NAMESPACE_MODULES = (
+    "google",
+    "google.cloud",
+    "google.cloud.sql",
+    "google.cloud.sql.connector",
+)
+
+
+def _package_module(name: str) -> ModuleType:
+    """Build a stub namespace module that behaves like a package.
+
+    A ``__path__`` is required so that ``from google.cloud.sql.connector import
+    Connector`` resolves; without it Python treats the parent as a non-package
+    and raises "not a package" when the real google namespace isn't installed."""
+    module = ModuleType(name)
+    module.__path__ = []  # type: ignore[attr-defined]
+    return module
+
+
+@pytest.fixture(autouse=True)
+def mock_connector():
+    """Inject a fake google.cloud.sql.connector module into sys.modules so the
+    lazy import inside _CloudSqlStrategy.build works without installing
+    cloud-sql-python-connector."""
+    connector_instance = MagicMock()
+    connector_cls = MagicMock(return_value=connector_instance)
+
+    connector_mod = ModuleType("google.cloud.sql.connector")
+    connector_mod.Connector = connector_cls
+
+    originals = {k: sys.modules.get(k) for k in _NAMESPACE_MODULES}
+
+    sys.modules.setdefault("google", _package_module("google"))
+    sys.modules.setdefault("google.cloud", _package_module("google.cloud"))
+    sys.modules.setdefault("google.cloud.sql", _package_module("google.cloud.sql"))
+    sys.modules["google.cloud.sql.connector"] = connector_mod
+
+    yield connector_cls, connector_instance
+
+    for key, val in originals.items():
+        if val is None:
+            sys.modules.pop(key, None)
+        else:
+            sys.modules[key] = val
+
+
+class TestMySQLCloudSQLConnection:
+    @patch("metadata.ingestion.source.database.mysql.connection.create_generic_db_connection")
+    def test_cloudsql_password_auth(self, mock_create_conn, mock_connector):
+        mock_connector_cls, mock_connector_inst = mock_connector  # noqa: RUF059
+        mock_create_conn.return_value = MagicMock()
+
+        connection = MysqlConnection(
+            hostPort="my-project:us-central1:my-instance",
+            username="dbuser",
+            authType=GcpCloudsqlConfigurationSource(
+                password="dbpassword",
+            ),
+        )
+
+        _CloudSqlStrategy(connection).build()
+
+        mock_create_conn.assert_called_once()
+        call_kwargs = mock_create_conn.call_args
+        assert "creator" in call_kwargs.kwargs
+
+        creator_fn = call_kwargs.kwargs["creator"]
+        creator_fn()
+
+        mock_connector_inst.connect.assert_called_once()
+        connect_kwargs = mock_connector_inst.connect.call_args.kwargs
+        assert connect_kwargs["instance_connection_string"] == "my-project:us-central1:my-instance"
+        assert connect_kwargs["driver"] == "pymysql"
+        assert connect_kwargs["user"] == "dbuser"
+        assert connect_kwargs["password"] == "dbpassword"
+        assert "enable_iam_auth" not in connect_kwargs
+
+    @patch("metadata.ingestion.source.database.mysql.connection.create_generic_db_connection")
+    def test_cloudsql_iam_auth(self, mock_create_conn, mock_connector):
+        _, mock_connector_inst = mock_connector
+        mock_create_conn.return_value = MagicMock()
+
+        connection = MysqlConnection(
+            hostPort="my-project:us-central1:my-instance",
+            username="sa@my-project.iam",
+            authType=GcpCloudsqlConfigurationSource(
+                enableIamAuth=True,
+            ),
+        )
+
+        _CloudSqlStrategy(connection).build()
+
+        creator_fn = mock_create_conn.call_args.kwargs["creator"]
+        creator_fn()
+
+        connect_kwargs = mock_connector_inst.connect.call_args.kwargs
+        assert connect_kwargs["enable_iam_auth"] is True
+        assert "password" not in connect_kwargs
+
+    @patch("metadata.ingestion.source.database.mysql.connection.create_generic_db_connection")
+    def test_cloudsql_url_is_bare_scheme(self, mock_create_conn, mock_connector):
+        mock_create_conn.return_value = MagicMock()
+
+        connection = MysqlConnection(
+            hostPort="my-project:us-central1:my-instance",
+            username="dbuser",
+            authType=GcpCloudsqlConfigurationSource(password="pw"),
+        )
+
+        _CloudSqlStrategy(connection).build()
+
+        url_fn = mock_create_conn.call_args.kwargs["get_connection_url_fn"]
+        assert url_fn(connection) == "mysql+pymysql://"
+
+    @patch("metadata.ingestion.source.database.mysql.connection.set_google_credentials")
+    @patch("metadata.ingestion.source.database.mysql.connection.create_generic_db_connection")
+    def test_cloudsql_sets_gcp_credentials_when_provided(self, mock_create_conn, mock_set_creds, mock_connector):
+        mock_create_conn.return_value = MagicMock()
+
+        gcp_config = MagicMock()
+        connection = MysqlConnection(
+            hostPort="my-project:us-central1:my-instance",
+            username="dbuser",
+            authType=GcpCloudsqlConfigurationSource(password="pw"),
+        )
+        connection.authType.gcpConfig = gcp_config
+
+        _CloudSqlStrategy(connection).build()
+
+        mock_set_creds.assert_called_once_with(gcp_config)
+
+    @patch("metadata.ingestion.source.database.mysql.connection.set_google_credentials")
+    @patch("metadata.ingestion.source.database.mysql.connection.create_generic_db_connection")
+    def test_cloudsql_skips_gcp_credentials_when_not_provided(self, mock_create_conn, mock_set_creds, mock_connector):
+        mock_create_conn.return_value = MagicMock()
+
+        connection = MysqlConnection(
+            hostPort="my-project:us-central1:my-instance",
+            username="dbuser",
+            authType=GcpCloudsqlConfigurationSource(password="pw"),
+        )
+
+        _CloudSqlStrategy(connection).build()
+
+        mock_set_creds.assert_not_called()
+
+    @patch("metadata.ingestion.source.database.mysql.connection.create_generic_db_connection")
+    def test_cloudsql_passes_database_schema(self, mock_create_conn, mock_connector):
+        _, mock_connector_inst = mock_connector
+        mock_create_conn.return_value = MagicMock()
+
+        connection = MysqlConnection(
+            hostPort="my-project:us-central1:my-instance",
+            username="dbuser",
+            databaseSchema="mydb",
+            authType=GcpCloudsqlConfigurationSource(password="pw"),
+        )
+
+        _CloudSqlStrategy(connection).build()
+
+        creator_fn = mock_create_conn.call_args.kwargs["creator"]
+        creator_fn()
+
+        connect_kwargs = mock_connector_inst.connect.call_args.kwargs
+        assert connect_kwargs["db"] == "mydb"
+
+    def test_cloudsql_imports_singular_connector_module(self):
+        """Guard against the plural-module typo (google.cloud.sql.connectors).
+
+        The other tests mock the module into sys.modules, so they pass under
+        either name; this asserts the real source imports the correct singular
+        google.cloud.sql.connector module the library actually exposes."""
+        source = inspect.getsource(_CloudSqlStrategy.build)
+
+        assert "from google.cloud.sql.connector import Connector" in source
+        assert "google.cloud.sql.connectors" not in source

ingestion/tests/unit/source/database/test_mysql_iam.py

@@ -0,0 +1,322 @@
+#  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.
+"""
+Tests for MySQL RDS IAM authentication (finding B1 fix).
+
+RDS IAM tokens expire after ~15 minutes. The MySQL connector must inject a fresh
+token on every new connection (via a SQLAlchemy ``do_connect`` listener) instead
+of baking a single token into the connection URL, so that connections opened
+later in a long ingestion still authenticate successfully.
+"""
+
+import datetime
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+import metadata.ingestion.source.database.mysql.connection as connection_module
+from metadata.clients.aws_client import RdsIamAuthTokenManager
+from metadata.generated.schema.entity.services.connections.database.common.iamAuthConfig import (
+    IamAuthConfigurationSource,
+)
+from metadata.generated.schema.entity.services.connections.database.mysqlConnection import (
+    MysqlConnection,
+)
+from metadata.generated.schema.security.credentials.awsCredentials import (
+    AWSCredentials,
+)
+from metadata.ingestion.source.database.mysql.connection import _IamStrategy
+
+HOST = "myrds.abc.us-east-2.rds.amazonaws.com"
+PORT = "3306"
+USERNAME = "iam_user"
+REGION = "us-east-2"
+
+
+def _iam_connection() -> MysqlConnection:
+    return MysqlConnection(
+        username=USERNAME,
+        hostPort=f"{HOST}:{PORT}",
+        authType=IamAuthConfigurationSource(awsConfig=AWSCredentials(awsRegion=REGION)),
+    )
+
+
+def _presigned_token(issued_at: datetime.datetime, expires_seconds: int = 900) -> str:
+    amz_date = issued_at.strftime("%Y%m%dT%H%M%SZ")
+    return f"{HOST}:{PORT}/?Action=connect&DBUser={USERNAME}&X-Amz-Date={amz_date}&X-Amz-Expires={expires_seconds}"
+
+
+class TestMySQLIamEngine:
+    """The connector wires a per-connection token-refresh listener for IAM auth."""
+
+    @patch.object(connection_module, "RdsIamAuthTokenManager")
+    @patch.object(connection_module, "listen")
+    @patch.object(connection_module, "create_generic_db_connection")
+    def test_iam_token_is_not_baked_into_url(self, mock_create, mock_listen, mock_token_manager):
+        mock_token_manager.return_value.get_token.return_value = "FRESH_TOKEN"
+        connection = _iam_connection()
+
+        _IamStrategy(connection).build()
+
+        url_fn = mock_create.call_args.kwargs["get_connection_url_fn"]
+        url = url_fn(connection)
+        assert "FRESH_TOKEN" not in url
+        assert "Action=connect" not in url
+        assert url.startswith(f"mysql+pymysql://{USERNAME}:@{HOST}:{PORT}")
+
+    @patch.object(connection_module, "RdsIamAuthTokenManager")
+    @patch.object(connection_module, "listen")
+    @patch.object(connection_module, "create_generic_db_connection")
+    def test_username_is_url_encoded(self, mock_create, mock_listen, mock_token_manager):
+        connection = MysqlConnection(
+            username="iam@user/db",
+            hostPort=f"{HOST}:{PORT}",
+            authType=IamAuthConfigurationSource(awsConfig=AWSCredentials(awsRegion=REGION)),
+        )
+
+        _IamStrategy(connection).build()
+
+        url_fn = mock_create.call_args.kwargs["get_connection_url_fn"]
+        url = url_fn(connection)
+        assert "iam%40user%2Fdb" in url
+        assert "iam@user/db@" not in url
+
+    @patch.object(connection_module, "RdsIamAuthTokenManager")
+    @patch.object(connection_module, "listen")
+    @patch.object(connection_module, "create_generic_db_connection")
+    def test_url_preserves_database_schema_and_options(self, mock_create, mock_listen, mock_token_manager):
+        connection = MysqlConnection(
+            username=USERNAME,
+            hostPort=f"{HOST}:{PORT}",
+            databaseSchema="analytics",
+            connectionOptions={"charset": "utf8mb4"},
+            authType=IamAuthConfigurationSource(awsConfig=AWSCredentials(awsRegion=REGION)),
+        )
+
+        _IamStrategy(connection).build()
+
+        url_fn = mock_create.call_args.kwargs["get_connection_url_fn"]
+        url = url_fn(connection)
+        assert "/analytics" in url
+        assert "charset=utf8mb4" in url
+        assert "Action=connect" not in url
+
+    @patch.object(connection_module, "RdsIamAuthTokenManager")
+    @patch.object(connection_module, "listen")
+    @patch.object(connection_module, "create_generic_db_connection")
+    def test_do_connect_listener_is_registered(self, mock_create, mock_listen, mock_token_manager):
+        mock_token_manager.return_value.get_token.return_value = "FRESH_TOKEN"
+        connection = _iam_connection()
+
+        _IamStrategy(connection).build()
+
+        event_name = mock_listen.call_args.args[1]
+        assert event_name == "do_connect"
+
+    @patch.object(connection_module, "RdsIamAuthTokenManager")
+    @patch.object(connection_module, "listen")
+    @patch.object(connection_module, "create_generic_db_connection")
+    def test_listener_injects_fresh_token_per_connection(self, mock_create, mock_listen, mock_token_manager):
+        mock_token_manager.return_value.get_token.return_value = "FRESH_TOKEN"
+        connection = _iam_connection()
+
+        _IamStrategy(connection).build()
+        listener = mock_listen.call_args.args[2]
+
+        first_cparams = {}
+        listener(None, None, None, first_cparams)
+        second_cparams = {}
+        listener(None, None, None, second_cparams)
+
+        assert first_cparams["password"] == "FRESH_TOKEN"
+        assert second_cparams["password"] == "FRESH_TOKEN"
+        assert mock_token_manager.return_value.get_token.call_count == 2
+
+    @patch.object(connection_module, "RdsIamAuthTokenManager")
+    @patch.object(connection_module, "listen")
+    @patch.object(connection_module, "create_generic_db_connection")
+    def test_listener_enables_ssl_required_by_pymysql_for_iam(self, mock_create, mock_listen, mock_token_manager):
+        mock_token_manager.return_value.get_token.return_value = "FRESH_TOKEN"
+        connection = _iam_connection()
+
+        _IamStrategy(connection).build()
+        listener = mock_listen.call_args.args[2]
+
+        cparams = {}
+        listener(None, None, None, cparams)
+        assert "ssl" in cparams
+        # Must be truthy so PyMySQL treats TLS as required (not PREFERRED).
+        assert cparams["ssl"]
+
+    @patch.object(connection_module, "RdsIamAuthTokenManager")
+    @patch.object(connection_module, "listen")
+    @patch.object(connection_module, "create_generic_db_connection")
+    def test_injected_ssl_makes_pymysql_require_tls(self, mock_create, mock_listen, mock_token_manager):
+        """Assert real PyMySQL behavior: the injected ssl value enables required TLS.
+
+        An empty dict only yields PREFERRED mode (_ssl_required=False), which can
+        fall back to plaintext; the injected value must produce _ssl_required=True.
+        """
+        import pymysql
+
+        mock_token_manager.return_value.get_token.return_value = "FRESH_TOKEN"
+        connection = _iam_connection()
+        _IamStrategy(connection).build()
+        listener = mock_listen.call_args.args[2]
+
+        cparams = {}
+        listener(None, None, None, cparams)
+
+        with patch.object(pymysql.connections.Connection, "connect", lambda self: None):
+            conn = pymysql.connections.Connection(host="h", user="u", ssl=cparams["ssl"])
+        assert conn.ssl is True
+        assert conn._ssl_required is True
+
+    @patch.object(connection_module, "RdsIamAuthTokenManager")
+    @patch.object(connection_module, "listen")
+    @patch.object(connection_module, "create_generic_db_connection")
+    def test_listener_preserves_existing_ssl_config(self, mock_create, mock_listen, mock_token_manager):
+        mock_token_manager.return_value.get_token.return_value = "FRESH_TOKEN"
+        connection = _iam_connection()
+
+        _IamStrategy(connection).build()
+        listener = mock_listen.call_args.args[2]
+
+        existing_ssl = {"ssl_ca": "/path/to/ca.pem"}
+        cparams = {"ssl": existing_ssl}
+        listener(None, None, None, cparams)
+        assert cparams["ssl"] == existing_ssl
+
+    @patch.object(connection_module, "RdsIamAuthTokenManager")
+    @patch.object(connection_module, "listen")
+    @patch.object(connection_module, "create_generic_db_connection")
+    def test_listener_does_not_overwrite_explicit_empty_ssl(self, mock_create, mock_listen, mock_token_manager):
+        mock_token_manager.return_value.get_token.return_value = "FRESH_TOKEN"
+        connection = _iam_connection()
+
+        _IamStrategy(connection).build()
+        listener = mock_listen.call_args.args[2]
+
+        sentinel_ssl = {}
+        cparams = {"ssl": sentinel_ssl}
+        listener(None, None, None, cparams)
+        assert cparams["ssl"] is sentinel_ssl
+
+    @patch.object(connection_module, "RdsIamAuthTokenManager")
+    @patch.object(connection_module, "listen")
+    @patch.object(connection_module, "create_generic_db_connection")
+    def test_token_manager_built_with_split_host_and_port(self, mock_create, mock_listen, mock_token_manager):
+        connection = _iam_connection()
+
+        _IamStrategy(connection).build()
+
+        assert mock_token_manager.call_args.kwargs["host"] == HOST
+        assert mock_token_manager.call_args.kwargs["port"] == PORT
+        assert mock_token_manager.call_args.kwargs["username"] == USERNAME
+
+
+class TestRdsIamAuthTokenManager:
+    """The token manager caches a token and refreshes it before expiry."""
+
+    @pytest.fixture
+    def mock_rds(self):
+        with patch("metadata.clients.aws_client.AWSClient") as mock_aws_client:
+            rds_client = MagicMock()
+            mock_aws_client.return_value.get_rds_client.return_value = rds_client
+            yield rds_client
+
+    def _manager(self) -> RdsIamAuthTokenManager:
+        return RdsIamAuthTokenManager(
+            host=HOST,
+            port=PORT,
+            username=USERNAME,
+            aws_config=AWSCredentials(awsRegion=REGION),
+        )
+
+    def test_first_call_generates_token(self, mock_rds):
+        now = datetime.datetime.now(datetime.timezone.utc)
+        mock_rds.generate_db_auth_token.return_value = _presigned_token(now)
+
+        token = self._manager().get_token()
+
+        assert token == _presigned_token(now)
+        assert mock_rds.generate_db_auth_token.call_count == 1
+
+    def test_token_cached_within_ttl(self, mock_rds):
+        now = datetime.datetime.now(datetime.timezone.utc)
+        mock_rds.generate_db_auth_token.return_value = _presigned_token(now)
+
+        manager = self._manager()
+        manager.get_token()
+        manager.get_token()
+
+        assert mock_rds.generate_db_auth_token.call_count == 1
+
+    def test_token_refreshed_after_expiry(self, mock_rds):
+        now = datetime.datetime.now(datetime.timezone.utc)
+        mock_rds.generate_db_auth_token.return_value = _presigned_token(now)
+
+        manager = self._manager()
+        manager.get_token()
+        manager._expires_at = now - datetime.timedelta(seconds=1)
+        manager.get_token()
+
+        assert mock_rds.generate_db_auth_token.call_count == 2
+
+    def test_expiry_parsed_from_presigned_url(self, mock_rds):
+        now = datetime.datetime.now(datetime.timezone.utc)
+        mock_rds.generate_db_auth_token.return_value = _presigned_token(now, expires_seconds=900)
+
+        manager = self._manager()
+        manager.get_token()
+
+        seconds_until_expiry = (manager._expires_at - now).total_seconds()
+        assert 890 <= seconds_until_expiry <= 910
+
+    def test_malformed_token_falls_back_to_default_ttl(self, mock_rds):
+        mock_rds.generate_db_auth_token.return_value = "not-a-presigned-url"
+
+        manager = self._manager()
+        manager.get_token()
+
+        assert manager._expires_at is not None
+
+    def test_concurrent_get_token_refreshes_only_once(self, mock_rds):
+        """Many threads hitting a cold manager must trigger a single refresh.
+
+        Each worker thread calls engine.connect() -> the shared do_connect listener
+        -> the shared manager's get_token(). Without locking, threads racing on a
+        cold/expired token each call generate_db_auth_token. A slow token generator
+        widens the race window so the assertion is meaningful.
+        """
+        now = datetime.datetime.now(datetime.timezone.utc)
+
+        def slow_generate(**_kwargs):
+            time.sleep(0.05)
+            return _presigned_token(now)
+
+        mock_rds.generate_db_auth_token.side_effect = slow_generate
+        manager = self._manager()
+
+        barrier = threading.Barrier(10)
+
+        def call():
+            barrier.wait()
+            return manager.get_token()
+
+        with ThreadPoolExecutor(max_workers=10) as executor:
+            tokens = [future.result() for future in [executor.submit(call) for _ in range(10)]]
+
+        assert mock_rds.generate_db_auth_token.call_count == 1
+        assert all(token == _presigned_token(now) for token in tokens)

ingestion/tests/unit/source/database/test_mysql_progress_count.py

@@ -0,0 +1,52 @@
+#  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.
+"""MySQL declares the single-database ``Database`` denominator and defers the
+``DatabaseSchema`` total to runner reconciliation."""
+
+from types import SimpleNamespace
+
+import pytest
+
+from metadata.ingestion.progress.modes import ProgressMode, TotalsDeclarer
+from metadata.ingestion.progress.tracking import ProgressTracking
+from metadata.ingestion.source.database.mysql.metadata import MysqlSource
+
+
+def _source(service_connection):
+    source = object.__new__(MysqlSource)
+    source.service_connection = service_connection
+    source.__dict__["_progress_tracking"] = ProgressTracking(ProgressMode.AUTO, "Test")
+    return source
+
+
+@pytest.fixture
+def mysql_source():
+    return _source(SimpleNamespace(databaseName=None, database="mydb"))
+
+
+def test_declare_progress_totals_seeds_single_database(mysql_source):
+    mysql_source.declare_progress_totals(TotalsDeclarer(mysql_source.progress_tracking.registry))
+    counters = {t: (done, total) for t, done, total in mysql_source.progress_tracking.registry.global_counters()}
+    assert counters["Database"] == (0, 1)
+
+
+def test_declare_progress_totals_marks_schema_reconcilable(mysql_source):
+    mysql_source.declare_progress_totals(TotalsDeclarer(mysql_source.progress_tracking.registry))
+    assert mysql_source.progress_tracking.registry.is_reconcilable("DatabaseSchema") is True
+    counters = {t: (done, total) for t, done, total in mysql_source.progress_tracking.registry.global_counters()}
+    assert counters["DatabaseSchema"] == (0, None)
+
+
+def test_declare_progress_totals_defaults_database_name_when_unset(mysql_source):
+    source = _source(SimpleNamespace(databaseName=None, database=None))
+    source.declare_progress_totals(TotalsDeclarer(source.progress_tracking.registry))
+    counters = {t: (done, total) for t, done, total in source.progress_tracking.registry.global_counters()}
+    assert counters["Database"] == (0, 1)

ingestion/tests/unit/source/database/test_snowflake_progress_count.py

@@ -0,0 +1,138 @@
+#  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.
+"""Snowflake supplies cheap progress denominators by enumerating + filtering
+names, without running the heavy per-database/per-schema setup."""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from metadata.ingestion.progress.modes import TotalsDeclarer
+from metadata.ingestion.source.database.snowflake import metadata as snowflake_metadata
+
+SnowflakeSource = snowflake_metadata.SnowflakeSource
+
+
+def _source(raw_databases, configured_db=None):
+    source = object.__new__(SnowflakeSource)
+    source.config = SimpleNamespace(
+        serviceConnection=SimpleNamespace(root=SimpleNamespace(config=SimpleNamespace(database=configured_db)))
+    )
+    source.source_config = SimpleNamespace(useFqnForFiltering=False, databaseFilterPattern=None)
+    source.metadata = MagicMock()
+    source.status = MagicMock()
+    source.context = MagicMock()
+    source.context.get.return_value = SimpleNamespace(database_service="svc")
+    source.get_database_names_raw = MagicMock(return_value=iter(raw_databases))
+    return source
+
+
+class TestSnowflakeDatabaseCount:
+    def test_count_is_the_filtered_name_count(self):
+        source = _source(["A", "B", "C"])
+        with patch.object(snowflake_metadata, "filter_by_database", return_value=False):
+            assert len(source._filtered_database_names()) == 3
+
+    def test_filtered_out_databases_are_excluded_from_count(self):
+        source = _source(["KEEP", "DROP"])
+        with patch.object(snowflake_metadata, "filter_by_database", side_effect=lambda _pattern, name: name == "DROP"):
+            assert len(source._filtered_database_names()) == 1
+            source.status.filter.assert_called_once()
+
+    def test_configured_database_counts_as_one(self):
+        source = _source([], configured_db="ONLY_DB")
+        assert len(source._filtered_database_names()) == 1
+        source.get_database_names_raw.assert_not_called()
+
+    def test_filtered_names_are_cached_so_status_is_not_double_emitted(self):
+        source = _source(["A", "DROP"])
+        with patch.object(snowflake_metadata, "filter_by_database", side_effect=lambda _pattern, name: name == "DROP"):
+            source._filtered_database_names()
+            source._filtered_database_names()
+        source.get_database_names_raw.assert_called_once()
+        source.status.filter.assert_called_once()
+
+    def test_producer_setup_runs_per_kept_database_without_re_filtering(self):
+        from metadata.ingestion.progress.modes import ProgressMode
+        from metadata.ingestion.progress.tracking import ProgressTracking
+
+        source = _source(["A", "B"])
+        for setup in (
+            "set_inspector",
+            "set_session_query_tag",
+            "set_partition_details",
+            "set_schema_description_map",
+            "set_database_description_map",
+            "set_external_location_map",
+            "set_schema_tags_map",
+            "set_database_tags_map",
+        ):
+            setattr(source, setup, MagicMock())
+        source.__dict__["_progress_tracking"] = ProgressTracking(ProgressMode.AUTO, "Test")
+        with patch.object(snowflake_metadata, "filter_by_database", return_value=False):
+            source._filtered_database_names()  # warms the database names cache
+            yielded = list(source.get_database_names())
+        assert yielded == ["A", "B"]
+        assert source.set_inspector.call_count == 2  # setup ran per database, lazily
+        source.get_database_names_raw.assert_called_once()  # enumeration reused from cache
+
+
+@pytest.fixture
+def snowflake_source():
+    """Minimal SnowflakeSource stub for push-totals tests."""
+    from metadata.ingestion.progress.modes import ProgressMode
+    from metadata.ingestion.progress.tracking import ProgressTracking
+
+    source = object.__new__(SnowflakeSource)
+    source.config = SimpleNamespace(
+        serviceConnection=SimpleNamespace(root=SimpleNamespace(config=SimpleNamespace(database=None)))
+    )
+    source.source_config = SimpleNamespace(
+        useFqnForFiltering=False,
+        databaseFilterPattern=None,
+        schemaFilterPattern=None,
+    )
+    source.metadata = MagicMock()
+    source.status = MagicMock()
+    source.context = MagicMock()
+    source.context.get.return_value = SimpleNamespace(database=None, database_service="svc")
+    source.__dict__["_progress_tracking"] = ProgressTracking(ProgressMode.AUTO, "Test")
+    return source
+
+
+def test_declare_progress_totals_seeds_database_and_schema(snowflake_source):
+    snowflake_source._filtered_database_names = lambda: ["db1", "db2"]
+    snowflake_source._schema_names_by_database = lambda: {"db1": ["s1", "s2"], "db2": ["s3"]}
+    snowflake_source._is_schema_filtered = lambda db, sch: False
+    snowflake_source.declare_progress_totals(TotalsDeclarer(snowflake_source.progress_tracking.registry))
+    counters = {t: (d, total) for t, d, total in snowflake_source.progress_tracking.registry.global_counters()}
+    assert counters["Database"] == (0, 2)
+    assert counters["DatabaseSchema"] == (0, 3)
+
+
+def test_declare_progress_totals_applies_schema_filter(snowflake_source):
+    snowflake_source._filtered_database_names = lambda: ["db1"]
+    snowflake_source._schema_names_by_database = lambda: {"db1": ["keep", "drop"]}
+    snowflake_source._is_schema_filtered = lambda db, sch: sch == "drop"
+    snowflake_source.declare_progress_totals(TotalsDeclarer(snowflake_source.progress_tracking.registry))
+    counters = {t: (d, total) for t, d, total in snowflake_source.progress_tracking.registry.global_counters()}
+    assert counters["DatabaseSchema"] == (0, 1)
+
+
+def test_declare_progress_totals_falls_back_when_account_show_unavailable(snowflake_source):
+    snowflake_source._filtered_database_names = lambda: ["db1"]
+    snowflake_source._schema_names_by_database = lambda: None
+    snowflake_source.declare_progress_totals(TotalsDeclarer(snowflake_source.progress_tracking.registry))
+    assert snowflake_source.progress_tracking.registry.is_reconcilable("DatabaseSchema") is True
+    counters = {t: (d, total) for t, d, total in snowflake_source.progress_tracking.registry.global_counters()}
+    assert counters["Database"] == (0, 1)
+    assert counters["DatabaseSchema"] == (0, None)

ingestion/tests/unit/source/database/test_snowflake_progress_path.py

@@ -0,0 +1,114 @@
+#  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.
+"""Snowflake uses the inherited lazy progress path; no COUNT pass remains."""
+
+import inspect
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+from metadata.ingestion.progress.modes import ProgressMode
+from metadata.ingestion.progress.registry import ProgressRegistry
+from metadata.ingestion.source.database import database_service
+from metadata.ingestion.source.database.snowflake import metadata as snowflake_metadata
+
+
+class _ConcreteSource(database_service.DatabaseServiceSource):
+    """Minimal concrete subclass so object.__new__ works on the abstract base."""
+
+    close = create = get_database_names = get_database_schema_names = None  # type: ignore[assignment]
+    get_stored_procedures = get_tables_name_and_type = yield_database = None  # type: ignore[assignment]
+    yield_database_schema = yield_stored_procedure = yield_table = yield_tag = None  # type: ignore[assignment]
+
+
+def _path_for(database, schema, entity_type_name):
+    source = object.__new__(_ConcreteSource)
+    ctx = SimpleNamespace(database=database, database_schema=schema)
+    source.context = MagicMock()
+    source.context.get.return_value = ctx
+    return source.progress_tracker.current_path(entity_type_name)
+
+
+class TestSnowflakeProgressPath:
+    def test_database_node_opens_at_root(self):
+        assert _path_for("db", None, "Database") == []
+
+    def test_schema_node_opens_under_its_database(self):
+        assert _path_for("db", None, "DatabaseSchema") == ["db"]
+
+    def test_table_node_opens_under_database_and_schema(self):
+        assert _path_for("db", "sales", "Table") == ["db", "sales"]
+
+    def test_stored_procedure_node_opens_under_database_and_schema(self):
+        assert _path_for("db", "sales", "StoredProcedure") == ["db", "sales"]
+
+    def test_schema_node_ignores_stale_sibling_schema_context(self):
+        # database_schema still holds the previous database's last schema; the
+        # schema level must NOT inherit it, or the new database's schema-count
+        # open lands one level too deep and the database never completes.
+        assert _path_for("db2", "stale_schema_from_db1", "DatabaseSchema") == ["db2"]
+
+    def test_no_declare_totals_on_snowflake_source(self):
+        assert not hasattr(snowflake_metadata.SnowflakeSource, "declare_totals")
+
+    def test_connector_references_no_count_constants(self):
+        source = inspect.getsource(snowflake_metadata)
+        assert "SNOWFLAKE_COUNT" not in source
+        assert "TotalsRunner" not in source
+        assert "declare_totals" not in source
+
+
+class TestProgressOptIn:
+    def test_database_family_is_auto_by_default(self):
+        assert database_service.DatabaseServiceSource.progress_mode is ProgressMode.AUTO
+
+    def test_concrete_db_source_inherits_auto(self):
+        assert _ConcreteSource.progress_mode is ProgressMode.AUTO
+
+    def test_snowflake_is_auto(self):
+        assert snowflake_metadata.SnowflakeSource.progress_mode is ProgressMode.AUTO
+
+
+class TestStaleContextDoesNotStrandDatabases:
+    """End-to-end: drive the registry through the real path logic across two
+    single-threaded databases, with database_schema left stale when the second
+    database starts (exactly the runner's behaviour). Every database must reach
+    completion and prune; before the fix the second one stranded as 'live'."""
+
+    def _path(self, source, ctx, entity_type_name):
+        source.context.get.return_value = ctx
+        return source.progress_tracker.current_path(entity_type_name)
+
+    def test_second_database_completes_and_prunes(self):
+        source = object.__new__(_ConcreteSource)
+        source.context = MagicMock()
+        registry = ProgressRegistry()
+        ctx = SimpleNamespace(database=None, database_schema=None)
+
+        registry.open(self._path(source, ctx, "Database"), "Database", 2)
+
+        ctx.database, ctx.database_schema = "db1", None
+        registry.open(self._path(source, ctx, "DatabaseSchema"), "DatabaseSchema", 1)
+        ctx.database_schema = "s1"
+        registry.open(self._path(source, ctx, "Table"), "Table", 1)
+        registry.advance(self._path(source, ctx, "Table"))  # db1 fully done
+
+        ctx.database = "db2"  # database_schema is left as the stale "s1"
+        registry.open(self._path(source, ctx, "DatabaseSchema"), "DatabaseSchema", 1)
+        ctx.database_schema = "x1"
+        registry.open(self._path(source, ctx, "Table"), "Table", 1)
+        registry.advance(self._path(source, ctx, "Table"))  # db2 fully done
+
+        snapshot = registry.snapshot()
+        assert snapshot.child_type == "Database"
+        assert snapshot.processed == 2
+        assert snapshot.expected == 2
+        assert snapshot.active is False
+        assert snapshot.children == ()  # both databases completed and pruned

ingestion/tests/unit/source/database/timescale/test_connection.py

@@ -0,0 +1,85 @@
+#  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 TimescaleDB connection handling (auth strategies)."""
+
+from unittest.mock import patch
+
+from azure.core.credentials import AccessToken
+from azure.identity import ClientSecretCredential
+
+from metadata.generated.schema.entity.services.connections.database.common.azureConfig import (
+    AzureConfigurationSource,
+)
+from metadata.generated.schema.entity.services.connections.database.common.basicAuth import (
+    BasicAuth,
+)
+from metadata.generated.schema.entity.services.connections.database.timescaleConnection import (
+    TimescaleConnection as TimescaleConnectionConfig,
+)
+from metadata.generated.schema.security.credentials.azureCredentials import (
+    AzureCredentials,
+)
+from metadata.ingestion.source.database.timescale.connection import TimescaleConnection
+
+
+def _azure_connection() -> TimescaleConnectionConfig:
+    return TimescaleConnectionConfig(
+        username="openmetadata_user",
+        authType=AzureConfigurationSource(
+            azureConfig=AzureCredentials(
+                clientId="clientid",
+                tenantId="tenantid",
+                clientSecret="clientsecret",
+                scopes="scope1,scope2",
+            )
+        ),
+        hostPort="localhost:5432",
+        database="openmetadata_db",
+    )
+
+
+def test_basic_auth_builds_expected_url():
+    connection = TimescaleConnectionConfig(
+        username="openmetadata_user",
+        authType=BasicAuth(password="openmetadata_password"),
+        hostPort="localhost:5432",
+        database="openmetadata_db",
+    )
+    engine = TimescaleConnection(connection).client
+    assert (
+        engine.url.render_as_string(hide_password=False)
+        == "postgresql+psycopg2://openmetadata_user:openmetadata_password@localhost:5432/openmetadata_db"
+    )
+
+
+def test_azure_ad_uses_token_as_password():
+    connection = _azure_connection()
+    with patch.object(
+        ClientSecretCredential,
+        "get_token",
+        return_value=AccessToken(token="mocked_token", expires_on=100),
+    ):
+        engine = TimescaleConnection(connection).client
+    assert (
+        engine.url.render_as_string(hide_password=False)
+        == "postgresql+psycopg2://openmetadata_user:mocked_token@localhost:5432/openmetadata_db"
+    )
+
+
+def test_azure_ad_does_not_mutate_caller_connection():
+    connection = _azure_connection()
+    with patch.object(
+        ClientSecretCredential,
+        "get_token",
+        return_value=AccessToken(token="mocked_token", expires_on=100),
+    ):
+        assert TimescaleConnection(connection).client is not None
+    assert isinstance(connection.authType, AzureConfigurationSource)

ingestion/tests/unit/source/database/trino/test_connection.py

@@ -0,0 +1,139 @@
+#  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 Trino connection handling
+"""
+
+import pytest
+
+from metadata.generated.schema.entity.services.connections.connectionBasicType import (
+    ConnectionArguments,
+)
+from metadata.generated.schema.entity.services.connections.database.common.basicAuth import (
+    BasicAuth,
+)
+from metadata.generated.schema.entity.services.connections.database.common.jwtAuth import (
+    JwtAuth,
+)
+from metadata.generated.schema.entity.services.connections.database.common.noConfigAuthenticationTypes import (
+    NoConfigAuthenticationTypes,
+)
+from metadata.generated.schema.entity.services.connections.database.trinoConnection import (
+    TrinoConnection as TrinoConnectionConfig,
+)
+from metadata.generated.schema.entity.services.connections.database.trinoConnection import (
+    TrinoScheme,
+)
+from metadata.ingestion.connections.builders import init_empty_connection_arguments
+from metadata.ingestion.source.database.trino.connection import TrinoConnection
+
+
+class TestTrinoConnectionHttpScheme:
+    """Test http_scheme handling in Trino connection auth methods"""
+
+    @pytest.fixture
+    def basic_connection_config(self) -> TrinoConnectionConfig:
+        return TrinoConnectionConfig(
+            scheme=TrinoScheme.trino,
+            hostPort="localhost:8080",
+            username="test_user",
+            authType=BasicAuth(password="test_password"),
+        )
+
+    @pytest.fixture
+    def jwt_connection_config(self) -> TrinoConnectionConfig:
+        return TrinoConnectionConfig(
+            scheme=TrinoScheme.trino,
+            hostPort="localhost:8080",
+            username="test_user",
+            authType=JwtAuth(jwt="test_jwt_token"),
+        )
+
+    @pytest.fixture
+    def oauth2_connection_config(self) -> TrinoConnectionConfig:
+        return TrinoConnectionConfig(
+            scheme=TrinoScheme.trino,
+            hostPort="localhost:8080",
+            username="test_user",
+            authType=NoConfigAuthenticationTypes.OAuth2,
+        )
+
+    def test_basic_auth_defaults_to_https(self, basic_connection_config):
+        connection_args = init_empty_connection_arguments()
+
+        TrinoConnection.set_basic_auth(basic_connection_config, connection_args)
+
+        assert connection_args.root["http_scheme"] == "https"
+        assert "auth" in connection_args.root
+
+    def test_basic_auth_preserves_explicit_http_scheme(self, basic_connection_config):
+        connection_args = init_empty_connection_arguments()
+        connection_args.root["http_scheme"] = "http"
+
+        TrinoConnection.set_basic_auth(basic_connection_config, connection_args)
+
+        assert connection_args.root["http_scheme"] == "http"
+        assert "auth" in connection_args.root
+
+    def test_basic_auth_preserves_custom_https(self, basic_connection_config):
+        connection_args = init_empty_connection_arguments()
+        connection_args.root["http_scheme"] = "https"
+
+        TrinoConnection.set_basic_auth(basic_connection_config, connection_args)
+
+        assert connection_args.root["http_scheme"] == "https"
+
+    def test_jwt_auth_defaults_to_https(self, jwt_connection_config):
+        connection_args = init_empty_connection_arguments()
+
+        TrinoConnection.set_jwt_auth(jwt_connection_config, connection_args)
+
+        assert connection_args.root["http_scheme"] == "https"
+        assert "auth" in connection_args.root
+
+    def test_jwt_auth_preserves_explicit_http_scheme(self, jwt_connection_config):
+        connection_args = init_empty_connection_arguments()
+        connection_args.root["http_scheme"] = "http"
+
+        TrinoConnection.set_jwt_auth(jwt_connection_config, connection_args)
+
+        assert connection_args.root["http_scheme"] == "http"
+        assert "auth" in connection_args.root
+
+    def test_oauth2_auth_defaults_to_https(self, oauth2_connection_config):
+        connection_args = init_empty_connection_arguments()
+
+        TrinoConnection.set_oauth2_auth(oauth2_connection_config, connection_args)
+
+        assert connection_args.root["http_scheme"] == "https"
+        assert "auth" in connection_args.root
+
+    def test_oauth2_auth_preserves_explicit_http_scheme(self, oauth2_connection_config):
+        connection_args = init_empty_connection_arguments()
+        connection_args.root["http_scheme"] = "http"
+
+        TrinoConnection.set_oauth2_auth(oauth2_connection_config, connection_args)
+
+        assert connection_args.root["http_scheme"] == "http"
+        assert "auth" in connection_args.root
+
+    def test_build_connection_args_preserves_http_scheme(self, basic_connection_config):
+        basic_connection_config.connectionArguments = ConnectionArguments(root={"http_scheme": "http"})
+
+        result = TrinoConnection.build_connection_args(basic_connection_config)
+
+        assert result.root["http_scheme"] == "http"
+
+    def test_build_connection_args_defaults_http_scheme(self, basic_connection_config):
+        result = TrinoConnection.build_connection_args(basic_connection_config)
+
+        assert result.root["http_scheme"] == "https"

ingestion/tests/unit/source/database/trino/test_lineage.py

@@ -0,0 +1,157 @@
+#  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.
+"""Regression tests for Trino cross-database lineage (Issue #27419)."""
+
+from unittest.mock import MagicMock, patch
+
+from metadata.generated.schema.entity.data.database import Database
+from metadata.generated.schema.entity.data.table import Table
+from metadata.ingestion.api.models import Either
+from metadata.ingestion.source.database.trino.lineage import TrinoLineageSource
+
+
+class TrinoLineageSourceTestDouble(TrinoLineageSource):
+    """Minimal Trino lineage source for unit testing."""
+
+    def __init__(self, metadata):
+        self.metadata = metadata
+        self.config = MagicMock()
+        self.config.serviceName = "repro_trino"
+        self.source_config = MagicMock()
+        self.source_config.crossDatabaseServiceNames = ["repro_postgres"]
+
+
+def _mock_column(column_name):
+    column = MagicMock()
+    column.name.root = column_name
+    return column
+
+
+def test_check_same_table_is_case_insensitive_for_names_and_columns():
+    """Issue #27419: table and column comparisons should ignore case."""
+    metadata = MagicMock()
+    lineage_source = TrinoLineageSourceTestDouble(metadata)
+
+    source_table = MagicMock()
+    source_table.name.root = "CUSTOMER"
+    source_table.columns = [_mock_column("ID"), _mock_column("NAME")]
+
+    target_table = MagicMock()
+    target_table.name.root = "customer"
+    target_table.columns = [_mock_column("id"), _mock_column("name")]
+
+    assert lineage_source.check_same_table(source_table, target_table)
+
+
+def test_yield_cross_database_lineage_finds_uppercase_source_table():
+    """Issue #27419: resolve uppercase Postgres source table in cross-db lineage."""
+    metadata = MagicMock()
+
+    trino_database = MagicMock()
+    trino_database.fullyQualifiedName.root = "repro_trino.postgres"
+
+    source_database = MagicMock()
+    source_database.fullyQualifiedName.root = "repro_postgres.source_db"
+
+    source_schema = MagicMock()
+    source_schema.name.root = "SOURCE_SCHEMA"
+    source_schema.fullyQualifiedName.root = "repro_postgres.source_db.SOURCE_SCHEMA"
+
+    trino_table = MagicMock()
+    trino_table.id.root = "11111111-1111-1111-1111-111111111111"
+    trino_table.fullyQualifiedName.root = "repro_trino.postgres.source_schema.customer"
+    trino_table.name.root = "customer"
+    trino_table.databaseSchema.name.root = "source_schema"
+    trino_table.databaseSchema.fullyQualifiedName.root = "repro_trino.postgres.source_schema"
+    trino_table.columns = [_mock_column("id"), _mock_column("name")]
+
+    source_table = MagicMock()
+    source_table.id.root = "22222222-2222-2222-2222-222222222222"
+    source_table.fullyQualifiedName.root = "repro_postgres.source_db.SOURCE_SCHEMA.CUSTOMER"
+    source_table.name.root = "CUSTOMER"
+    source_table.databaseSchema.name.root = "SOURCE_SCHEMA"
+    source_table.databaseSchema.fullyQualifiedName.root = "repro_postgres.source_db.SOURCE_SCHEMA"
+    source_table.columns = [_mock_column("id"), _mock_column("name")]
+
+    def list_all_entities_side_effect(entity, params=None, **_kwargs):
+        if entity is Database and params == {"service": "repro_trino"}:
+            return [trino_database]
+        if entity is Database and params == {"service": "repro_postgres"}:
+            return [source_database]
+        if entity is Table and params == {"database": "repro_trino.postgres"}:
+            return [trino_table]
+        return []
+
+    metadata.list_all_entities.side_effect = list_all_entities_side_effect
+    metadata.get_by_name.return_value = None
+
+    lineage_source = TrinoLineageSourceTestDouble(metadata)
+
+    with (
+        patch.object(
+            TrinoLineageSource,
+            "get_cross_database_lineage",
+            return_value=Either(right="cross-database-edge"),
+        ) as mock_get_cross_database_lineage,
+        patch(
+            "metadata.ingestion.source.database.trino.lineage.fqn.search_database_schema_from_es",
+            return_value=[source_schema],
+        ) as mock_search_database_schema,
+        patch(
+            "metadata.ingestion.source.database.trino.lineage.fqn.search_table_from_es",
+            return_value=[source_table],
+        ) as mock_search_table,
+    ):
+        result = list(lineage_source.yield_cross_database_lineage())
+
+    assert len(result) == 1
+    assert result[0].right == "cross-database-edge"
+    mock_get_cross_database_lineage.assert_called_once_with(source_table, trino_table)
+    mock_search_database_schema.assert_called_once_with(
+        metadata=metadata,
+        database_name="source_db",
+        schema_name="source_schema",
+        service_name="repro_postgres",
+        fetch_multiple_entities=True,
+        fields="fullyQualifiedName,name",
+    )
+    mock_search_table.assert_called_once_with(
+        metadata=metadata,
+        database_name="source_db",
+        schema_name="SOURCE_SCHEMA",
+        service_name="repro_postgres",
+        table_name="customer",
+        fetch_multiple_entities=True,
+        fields="fullyQualifiedName,name,columns,databaseSchema",
+    )
+
+
+def test_get_cross_database_schema_fqn_parses_quoted_schema_from_fqn():
+    """Issue #27419: parse quoted schema names with dots from table FQNs."""
+    metadata = MagicMock()
+
+    trino_table = MagicMock()
+    trino_table.databaseSchema = None
+    trino_table.fullyQualifiedName.root = 'repro_trino.postgres."source.schema".customer'
+
+    lineage_source = TrinoLineageSourceTestDouble(metadata)
+
+    with patch(
+        "metadata.ingestion.source.database.trino.lineage.fqn.search_database_schema_from_es",
+        return_value=None,
+    ):
+        result = lineage_source._get_cross_database_schema_fqn(
+            "repro_postgres.source_db",
+            trino_table,
+            {},
+        )
+
+    assert result == 'repro_postgres.source_db."source.schema"'

ingestion/tests/unit/source/database/unitycatalog/test_connection.py

@@ -0,0 +1,42 @@
+#  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 Unity Catalog BaseConnection wiring (non-Engine: WorkspaceClient)."""
+
+from unittest.mock import patch
+
+from metadata.generated.schema.entity.services.connections.database.databricks.personalAccessToken import (
+    PersonalAccessToken,
+)
+from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import (
+    UnityCatalogConnection as UnityCatalogConnectionConfig,
+)
+from metadata.ingestion.connections.connection import BaseConnection
+from metadata.ingestion.source.database.unitycatalog.connection import UnityCatalogConnection
+
+CONNECTION_MODULE = "metadata.ingestion.source.database.unitycatalog.connection"
+
+
+def _config() -> UnityCatalogConnectionConfig:
+    return UnityCatalogConnectionConfig(
+        hostPort="test-host:443",
+        authType=PersonalAccessToken(token="test-token"),
+    )
+
+
+def test_unitycatalog_connection_is_base_connection():
+    assert issubclass(UnityCatalogConnection, BaseConnection)
+
+
+def test_get_client_delegates_to_the_workspace_client_builder():
+    with patch(f"{CONNECTION_MODULE}.get_connection") as mock_get_connection:
+        client = UnityCatalogConnection(_config()).client
+    mock_get_connection.assert_called_once()
+    assert client is mock_get_connection.return_value

ingestion/tests/unit/source/database/vertica/test_connection.py

@@ -0,0 +1,49 @@
+#  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 Vertica connection handling."""
+
+from metadata.generated.schema.entity.services.connections.database.verticaConnection import (
+    VerticaConnection as VerticaConnectionConfig,
+)
+from metadata.generated.schema.entity.services.connections.database.verticaConnection import (
+    VerticaScheme,
+)
+from metadata.ingestion.source.database.vertica.connection import VerticaConnection
+
+
+def test_basic_auth_builds_expected_url():
+    connection = VerticaConnectionConfig(
+        username="openmetadata_user",
+        password="openmetadata_password",
+        hostPort="localhost:5433",
+        database="openmetadata_db",
+        scheme=VerticaScheme.vertica_vertica_python,
+    )
+    engine = VerticaConnection(connection).client
+    assert (
+        engine.url.render_as_string(hide_password=False)
+        == "vertica+vertica_python://openmetadata_user:openmetadata_password@localhost:5433/openmetadata_db"
+    )
+
+
+def test_special_characters_in_credentials_are_escaped():
+    connection = VerticaConnectionConfig(
+        username="openmetadata_user@444",
+        password="openmetadata_password@123",
+        hostPort="localhost:5433",
+        database="openmetadata_db",
+        scheme=VerticaScheme.vertica_vertica_python,
+    )
+    engine = VerticaConnection(connection).client
+    assert (
+        engine.url.render_as_string(hide_password=False)
+        == "vertica+vertica_python://openmetadata_user%40444:openmetadata_password%40123@localhost:5433/openmetadata_db"
+    )