Fixes #29768: refactor(ingestion): migrate S3 test-connection to declarative framework

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)

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,
         )