diff --git a/ingestion/src/metadata/core/connections/test_connection/checks/storage.py b/ingestion/src/metadata/core/connections/test_connection/checks/storage.py new file mode 100644 index 000000000000..97172f1cfade --- /dev/null +++ b/ingestion/src/metadata/core/connections/test_connection/checks/storage.py @@ -0,0 +1,127 @@ +# 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. +""" +Storage service step identity and shared check helpers. + +Object stores have no SQL: the reported ``command`` is the API operation the +check exercised (e.g. ``s3:ListBuckets``). Connectors reuse these helpers from +their ``@check`` methods. On failure the helpers raise ``CheckError`` carrying +the operation they attempted, so a failed step still reports its ``Evidence`` +to the backend. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from metadata.core.connections.test_connection.check import CheckError, StepName +from metadata.core.connections.test_connection.records import Diagnosis, Evidence + +if TYPE_CHECKING: + from collections.abc import Sequence + + from botocore.client import BaseClient + + +# A check only needs to prove a listing is reachable, not enumerate the whole +# account, so every listing probe stops at this many assets and reports whether +# more exist beyond the cap. +DEFAULT_LIST_LIMIT = 100 + + +class StorageStep(StepName): + """The steps a storage connector can be asked to verify.""" + + ListBuckets = "ListBuckets" + GetMetrics = "GetMetrics" + + +def _count(n: int, noun: str) -> str: + """``3 buckets`` / ``1 bucket`` - pluralize the noun to match the count.""" + return f"{n} {noun if n == 1 else noun + 's'}" + + +def _more_suffix(shown: int, more: bool) -> str: + """Mark a summary as capped when the listing has more assets beyond ``shown``.""" + return f" (showing first {shown}; more exist)" if more else "" + + +def list_buckets(client: BaseClient, limit: int = DEFAULT_LIST_LIMIT) -> Evidence: + """Enumerate the buckets the identity can see, reporting at most ``limit``. + + The cap is applied to the reported count locally, not pushed to the API: the + request stays a parameter-less ``list_buckets()`` so it behaves identically + on AWS and on S3-compatible stores (MinIO, Ceph, ...) reached via an + ``endPointURL`` that may not understand AWS's newer ``MaxBuckets`` parameter. + A bucket list is inherently small (account-bounded), so this never pulls an + unbounded payload; the summary flags when more than ``limit`` exist. + + An empty listing never raises - the account may simply hold no buckets - + so 'none visible' surfaces as a non-blocking caveat for the user to judge. + A missing list permission is a different case: it raises ``AccessDenied``, + which fails the step rather than producing a caveat. + """ + command = "s3:ListBuckets" + try: + response = client.list_buckets() # pyright: ignore[reportAttributeAccessIssue] + except Exception as cause: + raise CheckError(cause, Evidence(command=command)) from cause + buckets = response.get("Buckets", []) + caveat = None + if not buckets: + caveat = Diagnosis( + title="No buckets visible", + remediation="Verify the identity can list buckets, or configure bucketNames explicitly.", + ) + shown = min(len(buckets), limit) + summary = f"{_count(shown, 'bucket')} enumerated" + _more_suffix(shown, len(buckets) > limit) + return Evidence(summary=summary, command=command, caveat=caveat) + + +def probe_buckets(client: BaseClient, buckets: Sequence[str]) -> Evidence: + """Prove each configured bucket exists and its objects can be listed. + + This is an access probe over an explicit, config-bounded set - not an open + enumeration - so it caps each bucket's listing at a single key: proving + access never needs the contents, and a large bucket must not turn a + connection test into an expensive full listing. + """ + for bucket in buckets: + try: + client.list_objects(Bucket=bucket, MaxKeys=1) # pyright: ignore[reportAttributeAccessIssue] + except Exception as cause: + raise CheckError(cause, Evidence(command=f"s3:ListBucket ({bucket})")) from cause + return Evidence( + summary=f"{_count(len(buckets), 'configured bucket')} accessible", + command="s3:ListBucket", + ) + + +def list_metrics(client: BaseClient, namespace: str, limit: int = DEFAULT_LIST_LIMIT) -> Evidence: + """Enumerate the CloudWatch metrics for ``namespace``, up to ``limit``. + + ``list_metrics`` has no page-size parameter (AWS returns up to 500 per page), + so the count is capped to ``limit`` here; a ``NextToken`` (or a page already + past the cap) means more exist and the summary says so. + + Metrics feed object-count and size extraction; without them ingestion still + works, so an empty listing only reports what was (not) found. + """ + command = f"cloudwatch:ListMetrics (Namespace={namespace})" + try: + response = client.list_metrics(Namespace=namespace) # pyright: ignore[reportAttributeAccessIssue] + except Exception as cause: + raise CheckError(cause, Evidence(command=command)) from cause + metrics = response.get("Metrics", []) + more = bool(response.get("NextToken")) or len(metrics) > limit + shown = min(len(metrics), limit) + summary = f"{_count(shown, 'metric')} visible in namespace '{namespace}'" + _more_suffix(shown, more) + return Evidence(summary=summary, command=command) diff --git a/ingestion/src/metadata/ingestion/source/storage/s3/connection.py b/ingestion/src/metadata/ingestion/source/storage/s3/connection.py index 6454498723f3..b9f3e188d7a7 100644 --- a/ingestion/src/metadata/ingestion/source/storage/s3/connection.py +++ b/ingestion/src/metadata/ingestion/source/storage/s3/connection.py @@ -15,26 +15,82 @@ the cloudwatch:GetMetricData permissions """ +from __future__ import annotations + from dataclasses import dataclass -from functools import partial -from typing import Any, Optional +from typing import TYPE_CHECKING, Any -from botocore.client import BaseClient +from botocore.exceptions import EndpointConnectionError, NoCredentialsError from metadata.clients.aws_client import AWSClient -from metadata.generated.schema.entity.automations.workflow import ( - Workflow as AutomationWorkflow, +from metadata.core.connections.test_connection import ErrorPack, Matchers, check, when +from metadata.core.connections.test_connection.checks.storage import ( + StorageStep, + list_buckets, + list_metrics, + probe_buckets, ) +from metadata.core.connections.test_connection.network import NETWORK_ERRORS from metadata.generated.schema.entity.services.connections.storage.s3Connection import ( S3Connection as S3ConnectionConfig, ) -from metadata.generated.schema.entity.services.connections.testConnectionResult import ( - TestConnectionResult, -) from metadata.ingestion.connections.connection import BaseConnection -from metadata.ingestion.connections.test_connections import test_connection_steps -from metadata.ingestion.ometa.ometa_api import OpenMetadata -from metadata.utils.constants import THREE_MIN + +if TYPE_CHECKING: + from collections.abc import Callable + + from botocore.client import BaseClient + + from metadata.core.connections.test_connection import ChecksProvider + from metadata.core.connections.test_connection.records import Evidence + + +# botocore raises every AWS-side rejection as a ClientError whose message embeds +# the service error code ("An error occurred (InvalidAccessKeyId) when calling +# ..."), so the pack matches on those stable codes. Client-side failures +# (missing credentials, unreachable endpoint) surface as dedicated botocore +# exception types and are matched by type. +S3_ERRORS = ErrorPack( + when(Matchers.contains("InvalidAccessKeyId")).diagnose( + "Invalid AWS access key", + fix="The awsAccessKeyId does not exist in AWS; check the configured credentials.", + ), + when(Matchers.contains("SignatureDoesNotMatch")).diagnose( + "AWS secret key does not match", + fix="The awsSecretAccessKey is wrong for this awsAccessKeyId; re-enter the credential pair.", + ), + when(Matchers.contains("UnrecognizedClientException")).diagnose( + "AWS credentials not recognized", + fix="The security token or access key is invalid; check the configured credentials.", + ), + when(Matchers.contains("InvalidClientTokenId")).diagnose( + "AWS security token is invalid", + fix="The awsSessionToken (or access key) is invalid for this region; refresh the credentials.", + ), + when(Matchers.contains("ExpiredToken")).diagnose( + "AWS session token expired", + fix="Temporary credentials have expired; refresh the awsSessionToken.", + ), + when(Matchers.contains("NoSuchBucket")).diagnose( + "Bucket not found", + fix="Verify the configured bucketNames exist in this AWS account and region.", + ), + when(Matchers.contains("AccessDenied")).diagnose( + "Not authorized", + fix="Grant s3:ListAllMyBuckets (or s3:ListBucket on the configured buckets) " + "and cloudwatch:ListMetrics to the identity used.", + ), + when(Matchers.exception(NoCredentialsError)).diagnose( + "No AWS credentials found", + fix="No credentials were configured or resolvable; set awsAccessKeyId/awsSecretAccessKey " + "or make an IAM role available where ingestion runs.", + ), + when(Matchers.exception(EndpointConnectionError)).diagnose( + "Cannot reach the AWS endpoint", + fix="Check awsRegion (and endPointURL for S3-compatible services), and that the " + "network allows access to it from where ingestion runs.", + ), +).including(NETWORK_ERRORS) @dataclass @@ -59,39 +115,40 @@ def get_connection(connection: S3ConnectionConfig) -> S3ObjectStoreClient: ) +class S3Checks: + """Test-connection checks for S3. + + The client is built lazily inside the checks: an assume-role configuration + calls STS while the boto3 session is created, so building it while the + provider is constructed would touch the network before the runner's gate + (and outside its per-step timeout). ``connect`` is ``BaseConnection.client`` + underneath, so both steps share the one cached client. + """ + + errors = S3_ERRORS + + def __init__(self, connect: Callable[[], S3ObjectStoreClient], bucket_names: list[str] | None) -> None: + self._connect = connect + self.bucket_names = bucket_names + + @check(StorageStep.ListBuckets) + def check_buckets(self) -> Evidence: + client = self._connect() + if self.bucket_names: + return probe_buckets(client.s3_client, self.bucket_names) + return list_buckets(client.s3_client) + + @check(StorageStep.GetMetrics) + def get_metrics(self) -> Evidence: + return list_metrics(self._connect().cloudwatch_client, "AWS/S3") + + class S3Connection(BaseConnection[S3ConnectionConfig, S3ObjectStoreClient]): def _get_client(self) -> S3ObjectStoreClient: return get_connection(self.service_connection) - def test_connection( - self, - metadata: OpenMetadata, - automation_workflow: Optional[AutomationWorkflow] = None, # noqa: UP045 - timeout_seconds: Optional[int] = THREE_MIN, # noqa: UP045 - ) -> TestConnectionResult: - """ - Test connection. This can be executed either as part - of a metadata workflow or during an Automation Workflow - """ - client = self.client - service_connection = self.service_connection - - def test_buckets(connection: S3ConnectionConfig, client: S3ObjectStoreClient): - if connection.bucketNames: - for bucket_name in connection.bucketNames: - client.s3_client.list_objects(Bucket=bucket_name) # pyright: ignore[reportAttributeAccessIssue] - return - client.s3_client.list_buckets() # pyright: ignore[reportAttributeAccessIssue] - - test_fn = { - "ListBuckets": partial(test_buckets, client=client, connection=service_connection), - "GetMetrics": partial(client.cloudwatch_client.list_metrics, Namespace="AWS/S3"), # pyright: ignore[reportAttributeAccessIssue] - } - - return test_connection_steps( - metadata=metadata, - test_fn=test_fn, - service_type=service_connection.type.value, # pyright: ignore[reportOptionalMemberAccess] - automation_workflow=automation_workflow, - timeout_seconds=timeout_seconds, + def checks(self) -> ChecksProvider: + return S3Checks( + connect=lambda: self.client, + bucket_names=self.service_connection.bucketNames, ) diff --git a/ingestion/tests/unit/source/storage/s3/test_connection.py b/ingestion/tests/unit/source/storage/s3/test_connection.py index 71004b6b5d93..1f8f9af75132 100644 --- a/ingestion/tests/unit/source/storage/s3/test_connection.py +++ b/ingestion/tests/unit/source/storage/s3/test_connection.py @@ -8,15 +8,38 @@ # 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 S3 object store connection handling.""" +"""S3 test-connection: its checks match its shipped definition.""" +import json +from pathlib import Path from unittest.mock import MagicMock, patch +import pytest +from botocore.exceptions import ClientError, EndpointConnectionError, NoCredentialsError + +from metadata.core.connections.test_connection.check import CheckError, collect_checks +from metadata.core.connections.test_connection.checks.storage import list_buckets from metadata.ingestion.connections.connection import BaseConnection -from metadata.ingestion.source.storage.s3.connection import S3Connection +from metadata.ingestion.source.storage.s3.connection import ( + S3_ERRORS, + S3Checks, + S3Connection, +) CONNECTION_MODULE = "metadata.ingestion.source.storage.s3.connection" +_SEED = ( + Path(__file__).parents[6] / "openmetadata-service/src/main/resources/json/data" / "testConnections/storage/s3.json" +) + + +def _checks(client=None, bucket_names=None): + return S3Checks(connect=lambda: client, bucket_names=bucket_names) + + +def _check_names(): + return {step.value for step in collect_checks(_checks())} + def test_s3_connection_is_base_connection(): assert issubclass(S3Connection, BaseConnection) @@ -31,10 +54,170 @@ def test_get_client_delegates_to_get_connection(): mock_get.assert_called_once_with(conn.service_connection) -def test_test_connection_runs_steps(): - conn = S3Connection(MagicMock()) - conn._client = MagicMock() - with patch(f"{CONNECTION_MODULE}.test_connection_steps") as mock_steps: - result = conn.test_connection(metadata=MagicMock()) +def test_s3_checks_cover_expected_steps(): + assert _check_names() == {"ListBuckets", "GetMetrics"} + + +def test_s3_checks_match_definition_seed(): + definition_steps = {step["name"] for step in json.loads(_SEED.read_text())["steps"]} + assert _check_names() == definition_steps + + +def test_building_checks_does_not_build_the_client(): + with patch.object(S3Connection, "_get_client") as mock_build: + S3Connection(MagicMock()).checks() + + mock_build.assert_not_called() + + +def test_check_buckets_lists_all_buckets_when_none_configured(): + client = MagicMock() + client.s3_client.list_buckets.return_value = {"Buckets": [{"Name": "a"}, {"Name": "b"}]} + + evidence = _checks(client).check_buckets() + + assert evidence.summary == "2 buckets enumerated" + assert evidence.command == "s3:ListBuckets" + assert evidence.caveat is None + client.s3_client.list_buckets.assert_called_once_with() + client.s3_client.list_objects.assert_not_called() + + +def test_list_buckets_caps_reported_count_and_flags_more(): + client = MagicMock() + client.list_buckets.return_value = {"Buckets": [{"Name": "a"}, {"Name": "b"}, {"Name": "c"}]} + + evidence = list_buckets(client, limit=2) + + assert evidence.summary == "2 buckets enumerated (showing first 2; more exist)" + client.list_buckets.assert_called_once_with() + + +def test_check_buckets_warns_when_no_buckets_visible(): + client = MagicMock() + client.s3_client.list_buckets.return_value = {"Buckets": []} + + evidence = _checks(client).check_buckets() + + assert evidence.caveat is not None + assert "No buckets visible" in evidence.caveat.title + + +def test_check_buckets_probes_each_configured_bucket(): + client = MagicMock() + + evidence = _checks(client, bucket_names=["one", "two"]).check_buckets() + + assert evidence.summary == "2 configured buckets accessible" + client.s3_client.list_objects.assert_any_call(Bucket="one", MaxKeys=1) + client.s3_client.list_objects.assert_any_call(Bucket="two", MaxKeys=1) + client.s3_client.list_buckets.assert_not_called() + + +def test_check_buckets_failure_reports_the_bucket_it_probed(): + client = MagicMock() + cause = _client_error("NoSuchBucket", "ListObjects") + client.s3_client.list_objects.side_effect = cause + + with pytest.raises(CheckError) as failure: + _checks(client, bucket_names=["missing"]).check_buckets() + + assert failure.value.cause is cause + assert failure.value.evidence.command == "s3:ListBucket (missing)" + + +def test_get_metrics_lists_the_s3_namespace(): + client = MagicMock() + client.cloudwatch_client.list_metrics.return_value = {"Metrics": [{}]} + + evidence = _checks(client).get_metrics() + + assert evidence.summary == "1 metric visible in namespace 'AWS/S3'" + client.cloudwatch_client.list_metrics.assert_called_once_with(Namespace="AWS/S3") + + +def test_get_metrics_flags_truncated_results(): + client = MagicMock() + client.cloudwatch_client.list_metrics.return_value = {"Metrics": [{}], "NextToken": "more"} + + evidence = _checks(client).get_metrics() + + assert "more exist" in evidence.summary + + +def _client_error(code: str, operation: str) -> ClientError: + return ClientError({"Error": {"Code": code, "Message": "denied"}}, operation) + + +def test_error_pack_invalid_access_key(): + diagnosis = S3_ERRORS.classify(_client_error("InvalidAccessKeyId", "ListBuckets")) + assert diagnosis is not None + assert "access key" in diagnosis.title.lower() + + +def test_error_pack_signature_mismatch(): + diagnosis = S3_ERRORS.classify(_client_error("SignatureDoesNotMatch", "ListBuckets")) + assert diagnosis is not None + assert "secret key" in diagnosis.title.lower() + + +def test_error_pack_unrecognized_client(): + diagnosis = S3_ERRORS.classify(_client_error("UnrecognizedClientException", "ListMetrics")) + assert diagnosis is not None + assert "not recognized" in diagnosis.title.lower() + + +def test_error_pack_invalid_client_token(): + diagnosis = S3_ERRORS.classify(_client_error("InvalidClientTokenId", "ListBuckets")) + assert diagnosis is not None + assert "security token" in diagnosis.title.lower() + + +def test_error_pack_expired_token(): + diagnosis = S3_ERRORS.classify(_client_error("ExpiredToken", "ListBuckets")) + assert diagnosis is not None + assert "expired" in diagnosis.title.lower() + + +def test_error_pack_bucket_not_found(): + diagnosis = S3_ERRORS.classify(_client_error("NoSuchBucket", "ListObjects")) + assert diagnosis is not None + assert "not found" in diagnosis.title.lower() + + +def test_error_pack_access_denied(): + diagnosis = S3_ERRORS.classify( + ClientError( + { + "Error": { + "Code": "AccessDenied", + "Message": "User is not authorized to perform: s3:ListAllMyBuckets", + } + }, + "ListBuckets", + ) + ) + assert diagnosis is not None + assert "authorized" in diagnosis.title.lower() + + +def test_error_pack_access_denied_exception_variant(): + diagnosis = S3_ERRORS.classify(_client_error("AccessDeniedException", "ListMetrics")) + assert diagnosis is not None + assert "authorized" in diagnosis.title.lower() + + +def test_error_pack_no_credentials(): + diagnosis = S3_ERRORS.classify(NoCredentialsError()) + assert diagnosis is not None + assert "credentials" in diagnosis.title.lower() + + +def test_error_pack_endpoint_unreachable(): + diagnosis = S3_ERRORS.classify(EndpointConnectionError(endpoint_url="https://s3.bad-region.amazonaws.com")) + assert diagnosis is not None + assert "endpoint" in diagnosis.title.lower() + - assert result is mock_steps.return_value +def test_error_pack_unmatched_returns_none(): + assert S3_ERRORS.classify(Exception("novel error")) is None diff --git a/openmetadata-service/src/main/resources/json/data/testConnections/storage/s3.json b/openmetadata-service/src/main/resources/json/data/testConnections/storage/s3.json index 8015f320eaae..3fa09d2d549b 100644 --- a/openmetadata-service/src/main/resources/json/data/testConnections/storage/s3.json +++ b/openmetadata-service/src/main/resources/json/data/testConnections/storage/s3.json @@ -8,6 +8,7 @@ "description": "List all the buckets available to the user.", "errorMessage": "Failed to fetch buckets, please validate the credentials if the user has access to list buckets", "shortCircuit": true, + "category": "ConnectionGate", "mandatory": true }, {