From 4bfbad72d81f4c009894ce8c890b9dbc3fded127 Mon Sep 17 00:00:00 2001 From: Peter Van Bouwel Date: Mon, 8 Sep 2025 15:49:52 +0200 Subject: [PATCH 1/4] feature: handle expired token gracefully If a token is expired we can refresh it by doing another OpenEO interaction. By allowing 3 attempts for the STS service and by doing a very wide except clause we also make ourselves more robust against intermittent errors. --- openeo/extra/artifacts/_s3sts/sts.py | 41 ++++++++++++++++++++++------ 1 file changed, 33 insertions(+), 8 deletions(-) diff --git a/openeo/extra/artifacts/_s3sts/sts.py b/openeo/extra/artifacts/_s3sts/sts.py index 7b26232ec..a323fe7d1 100644 --- a/openeo/extra/artifacts/_s3sts/sts.py +++ b/openeo/extra/artifacts/_s3sts/sts.py @@ -1,5 +1,8 @@ from __future__ import annotations +import logging +from random import randint +from time import sleep from typing import TYPE_CHECKING if TYPE_CHECKING: @@ -12,12 +15,16 @@ from openeo.rest.connection import Connection from openeo.util import Rfc3339 +_log = logging.getLogger(__name__) + class OpenEOSTSClient: + _MAX_STS_ATTEMPTS = 3 + def __init__(self, config: S3STSConfig): self.config = config - def assume_from_openeo_connection(self, connection: Connection) -> AWSSTSCredentials: + def assume_from_openeo_connection(self, connection: Connection, attempt: int = 0) -> AWSSTSCredentials: """ Takes an OpenEO connection object and returns temporary credentials to interact with S3 """ @@ -27,14 +34,32 @@ def assume_from_openeo_connection(self, connection: Connection) -> AWSSTSCredent raise ProviderSpecificException("Only connections that have BearerAuth can be used.") auth_token = auth.bearer.split("/") - return AWSSTSCredentials.from_assume_role_response( - self._get_sts_client().assume_role_with_web_identity( - RoleArn=self._get_aws_access_role(), - RoleSessionName=f"artifact-helper-{Rfc3339().now_utc()}", - WebIdentityToken=auth_token[2], - DurationSeconds=43200, + try: + return AWSSTSCredentials.from_assume_role_response( + self._get_sts_client().assume_role_with_web_identity( + RoleArn=self._get_aws_access_role(), + RoleSessionName=f"artifact-helper-{Rfc3339().now_utc()}", + WebIdentityToken=auth_token[2], + DurationSeconds=43200, + ) ) - ) + except Exception as e: + _log.warning("Failed to get credentials for STS access") + + if attempt < 3: + # backoff with jitter + max_sleep_ms = 500 * (2**attempt) + sleep_ms = randint(0, max_sleep_ms) + _log.info(f"Retrying STS access in {sleep_ms} ms") + sleep(sleep_ms / 1000.0) + attempt += 1 + # Do an API call with OpenEO to trigger a refresh of our token. + connection.describe_account() + _log.info("Retrying to get credentials for STS access") + return self.assume_from_openeo_connection(connection, attempt) + else: + _log.fatal("Maximum attempts performed for STS access", exc_info=e) + raise RuntimeError("Could not get credentials from STS") from e def _get_sts_client(self) -> STSClient: return self.config.build_client("sts") From 213d00619736557e080625b2963b57b7e4147d94 Mon Sep 17 00:00:00 2001 From: Peter Van Bouwel Date: Tue, 9 Sep 2025 10:00:51 +0200 Subject: [PATCH 2/4] pr-feedback: address comments --- openeo/extra/artifacts/_s3sts/sts.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/openeo/extra/artifacts/_s3sts/sts.py b/openeo/extra/artifacts/_s3sts/sts.py index a323fe7d1..3dee2c805 100644 --- a/openeo/extra/artifacts/_s3sts/sts.py +++ b/openeo/extra/artifacts/_s3sts/sts.py @@ -1,8 +1,8 @@ from __future__ import annotations import logging +import time from random import randint -from time import sleep from typing import TYPE_CHECKING if TYPE_CHECKING: @@ -46,19 +46,18 @@ def assume_from_openeo_connection(self, connection: Connection, attempt: int = 0 except Exception as e: _log.warning("Failed to get credentials for STS access") - if attempt < 3: + if attempt < self._MAX_STS_ATTEMPTS: # backoff with jitter max_sleep_ms = 500 * (2**attempt) sleep_ms = randint(0, max_sleep_ms) _log.info(f"Retrying STS access in {sleep_ms} ms") - sleep(sleep_ms / 1000.0) + time.sleep(sleep_ms / 1000.0) attempt += 1 # Do an API call with OpenEO to trigger a refresh of our token. connection.describe_account() - _log.info("Retrying to get credentials for STS access") + _log.info(f"Retrying to get credentials for STS access {attempt}/{self._MAX_STS_ATTEMPTS}") return self.assume_from_openeo_connection(connection, attempt) else: - _log.fatal("Maximum attempts performed for STS access", exc_info=e) raise RuntimeError("Could not get credentials from STS") from e def _get_sts_client(self) -> STSClient: From 39b5b44df7757f496618f8b67cef6a8a9f787c85 Mon Sep 17 00:00:00 2001 From: Peter Van Bouwel Date: Tue, 9 Sep 2025 10:06:04 +0200 Subject: [PATCH 3/4] pr-feedback: make sure we have a valid token in any case. --- openeo/extra/artifacts/_s3sts/sts.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/openeo/extra/artifacts/_s3sts/sts.py b/openeo/extra/artifacts/_s3sts/sts.py index 3dee2c805..d474dd25a 100644 --- a/openeo/extra/artifacts/_s3sts/sts.py +++ b/openeo/extra/artifacts/_s3sts/sts.py @@ -35,6 +35,8 @@ def assume_from_openeo_connection(self, connection: Connection, attempt: int = 0 auth_token = auth.bearer.split("/") try: + # Do an API call with OpenEO to trigger a refresh of our token if it were stale. + connection.describe_account() return AWSSTSCredentials.from_assume_role_response( self._get_sts_client().assume_role_with_web_identity( RoleArn=self._get_aws_access_role(), @@ -53,8 +55,6 @@ def assume_from_openeo_connection(self, connection: Connection, attempt: int = 0 _log.info(f"Retrying STS access in {sleep_ms} ms") time.sleep(sleep_ms / 1000.0) attempt += 1 - # Do an API call with OpenEO to trigger a refresh of our token. - connection.describe_account() _log.info(f"Retrying to get credentials for STS access {attempt}/{self._MAX_STS_ATTEMPTS}") return self.assume_from_openeo_connection(connection, attempt) else: From 920de29ea5b99329cf2647c61dcc77037c516b37 Mon Sep 17 00:00:00 2001 From: Peter Van Bouwel Date: Tue, 9 Sep 2025 11:04:04 +0200 Subject: [PATCH 4/4] tests: mock /me endpoint --- tests/extra/artifacts/_s3sts/test_s3sts.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/extra/artifacts/_s3sts/test_s3sts.py b/tests/extra/artifacts/_s3sts/test_s3sts.py index 1d5bf5d99..84ae9637c 100644 --- a/tests/extra/artifacts/_s3sts/test_s3sts.py +++ b/tests/extra/artifacts/_s3sts/test_s3sts.py @@ -92,6 +92,7 @@ def conn_with_s3sts_capabilities( requests_mock, extra_api_capabilities, advertised_s3sts_config ) -> Iterator[Connection]: requests_mock.get(API_URL, json={"api_version": "1.0.0", **extra_api_capabilities}) + requests_mock.get(f"{API_URL}me", json={}) conn = Connection(API_URL) conn.auth = BearerAuth("oidc/fake/token") yield conn