sws-api-client 3.2.0__tar.gz → 3.3.0b1__tar.gz

{sws_api_client-3.2.0 → sws_api_client-3.3.0b1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: sws_api_client
-Version: 3.2.0
+Version: 3.3.0-beta.1
 Author-email: Matteo <matteo.terrinoni@fao.org>, Daniele <daniele.mansillo@fao.org>, Luca <luca.simi@fao.org>
 License-File: LICENSE
 Requires-Python: <4,>=3.10

{sws_api_client-3.2.0 → sws_api_client-3.3.0b1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "sws_api_client"
-version = "3.2.0"
+version = "3.3.0-beta.1"
 description = ""
 authors = [
     { name = "Matteo", email = "matteo.terrinoni@fao.org" },

{sws_api_client-3.2.0 → sws_api_client-3.3.0b1}/sws_api_client/__init__.py

@@ -58,4 +58,4 @@ from .tasks import PluginPayload, TaskDataset, TaskInfo, TaskManager, TaskRespon
 from .users import User, Users
 from .harvester import Harvester
 
-__version__ = "3.2.0"
+__version__ = "3.3.0-beta.1"

{sws_api_client-3.2.0 → sws_api_client-3.3.0b1}/sws_api_client/sws_api_client.py

@@ -8,6 +8,7 @@ import argparse
 import json
 import logging
 import os
+from collections.abc import Callable
 
 import boto3
 
@@ -45,6 +46,7 @@ class SwsApiClient:
         current_task_id: str | None,
         current_execution_id: str | None,
         object_cache: bool = False,
+        credentials_provider: Callable[[], AuthClientCredentials] | None = None,
     ) -> None:
         """Initialize the SWS API client.
 
@@ -55,17 +57,18 @@ class SwsApiClient:
             current_task_id: Optional current task identifier
             current_execution_id: Optional current execution identifier
             object_cache: Whether to enable client-side object caching (e.g. for codelists, dataset info)
+            credentials_provider: Optional callable returning fresh credentials, used to
+                retry a failed token request (e.g. after AWS rotates the secret)
         """
 
         logger.info("Initializing SwsApiClient")
 
         self.sws_endpoint = sws_endpoint
         self.sws_token = sws_token
-        self.authclient = authclient
         self.current_task_id = current_task_id
         self.current_execution_id = current_execution_id
         self.object_cache = object_cache
-        self.token = Token(authclient)
+        self.token = Token(authclient, credentials_provider)
         self.is_debug = self.check_debug()
         self.discoverable = Discover(
             sws_endpoint=sws_endpoint, sws_token=sws_token, token=self.token
@@ -74,6 +77,11 @@ class SwsApiClient:
             f"SwsApiClient initialized with endpoint {sws_endpoint} and token {sws_token}"
         )
 
+    @property
+    def authclient(self) -> AuthClientCredentials:
+        """Current auth credentials, updated when the secret is refreshed."""
+        return self.token.authclient
+
     @classmethod
     def __get_authclient_credentials_from_secret(
         cls, env_name="SWS_USER_CREDENTIALS_SECRET_NAME"
@@ -95,13 +103,12 @@ class SwsApiClient:
             raise ValueError(
                 f"Secret name not found in environment variable: {env_name}"
             )
-        cls.session = getattr(cls, "session", boto3.Session())
-        cls.secret_manager = getattr(
-            cls, "secret_manager", cls.session.client("secretsmanager")
-        )
+        if not hasattr(cls, "secret_manager"):
+            cls.session = getattr(cls, "session", None) or boto3.Session()
+            cls.secret_manager = cls.session.client("secretsmanager")
         response = cls.secret_manager.get_secret_value(SecretId=secret_name)
         secretContent = json.loads(response["SecretString"])
-        logger.debug(f"Secret fetched successfully: {secretContent}")
+        logger.debug(f"Secret {secret_name} fetched successfully")
         return AuthClientCredentials(
             clientId=secretContent["ID"],
             clientSecret=secretContent["SECRET"],
@@ -131,6 +138,7 @@ class SwsApiClient:
         logger.debug(
             f"Creating SwsApiClient from environment variables: {sws_endpoint_env}, {authclient_secret_name}"
         )
+        credentials_provider = None
         if (
             os.getenv("SWS_AUTH_CLIENTID")
             and os.getenv("SWS_AUTH_CLIENTSECRET")
@@ -148,11 +156,14 @@ class SwsApiClient:
                 raise ValueError(
                     f"You need ({authclient_secret_name}) or (SWS_AUTH_CLIENTID, SWS_AUTH_CLIENTSECRET, SWS_AUTH_SCOPE, SWS_AUTH_TOKENENDPOINT) environment variables to be set"
                 )
-            authclient: AuthClientCredentials = (
-                cls.__get_authclient_credentials_from_secret(
+
+            # Re-read on token failures, so a secret rotated by AWS is picked up.
+            def credentials_provider() -> AuthClientCredentials:
+                return cls.__get_authclient_credentials_from_secret(
                     env_name=authclient_secret_name
                 )
-            )
+
+            authclient = credentials_provider()
 
         sws_token = os.getenv("SWS_TOKEN")
         if not sws_token:
@@ -168,6 +179,7 @@ class SwsApiClient:
             current_task_id=os.getenv("TASK_ID"),
             current_execution_id=os.getenv("EXECUTION_ID"),
             authclient=authclient,
+            credentials_provider=credentials_provider,
         )
 
     @classmethod

{sws_api_client-3.2.0 → sws_api_client-3.3.0b1}/sws_api_client/tasks.py

@@ -335,6 +335,7 @@ class TaskManager:
         public: bool = False,
         retry: bool = False,
         emailNotification: bool = False,
+        max_duration: int | None = None,
     ) -> TaskCreateResponse:
         """Create a new plugin execution task.
 
@@ -351,6 +352,8 @@ class TaskManager:
             public: Whether task is public
             retry: Whether task can be retried
             emailNotification: Whether to send email notifications
+            max_duration: Max duration of the task in seconds; when None it is not sent
+                and the task manager's default applies
 
         Returns:
             TaskCreateResponse: Response containing the created task ID
@@ -398,6 +401,9 @@ class TaskManager:
             "parentTaskId": parentTaskId,
         }
 
+        if max_duration is not None:
+            data["config"]["maxDuration"] = max_duration
+
         if group:
             data["group"] = group
 

sws_api_client-3.3.0b1/sws_api_client/token.py

@@ -0,0 +1,93 @@
+import datetime
+import logging
+import time
+from collections.abc import Callable
+
+import requests
+
+from sws_api_client.auth import AuthClientCredentials
+from sws_api_client.case_agnostic_model import CaseAgnosticModel
+
+logger = logging.getLogger(__name__)
+
+# Wait before each retry. The credentials are refreshed before every retry, so a
+# secret rotated by AWS while a long task is running gets picked up.
+RETRY_DELAYS_SECONDS = (30, 60)
+
+
+class TokenModel(CaseAgnosticModel):
+    access_token: str
+    expires_on: datetime.datetime
+
+
+class Token:
+    def __init__(
+        self,
+        authclient: AuthClientCredentials,
+        credentials_provider: Callable[[], AuthClientCredentials] | None = None,
+    ) -> None:
+        """
+        Args:
+            authclient: Credentials used to request access tokens.
+            credentials_provider: Optional callable returning fresh credentials. When set,
+                a failed token request is retried with credentials it returns.
+        """
+        self.authclient = authclient
+        self.credentials_provider = credentials_provider
+        self.token = self.get_token()
+
+    def get_token(self) -> TokenModel:
+        if (
+            hasattr(self, "token")
+            and self.token.expires_on is not None
+            and self.token.expires_on >= datetime.datetime.now()
+        ):
+            return self.token
+
+        try:
+            self.token = self._request_token()
+        except requests.RequestException:
+            if self.credentials_provider is None:
+                raise
+            self.token = self._retry_with_fresh_credentials()
+        logger.debug(f"New token, expires on {self.token.expires_on}")
+        return self.token
+
+    def _retry_with_fresh_credentials(self) -> TokenModel:
+        retries = len(RETRY_DELAYS_SECONDS)
+        for attempt, delay in enumerate(RETRY_DELAYS_SECONDS, start=1):
+            logger.warning(
+                "Token request failed, refreshing credentials and retrying in %ss (%s/%s)",
+                delay,
+                attempt,
+                retries,
+            )
+            time.sleep(delay)
+            self.authclient = self.credentials_provider()
+            try:
+                return self._request_token()
+            except requests.RequestException:
+                if attempt == retries:
+                    logger.error(
+                        "Token request still failing after %s retries, giving up",
+                        retries,
+                    )
+                    raise
+
+    def _request_token(self) -> TokenModel:
+        response = requests.post(
+            self.authclient.token_endpoint,
+            data={
+                "grant_type": "client_credentials",
+                "client_id": self.authclient.client_id,
+                "client_secret": self.authclient.client_secret,
+                "scope": self.authclient.scope,
+            },
+        )
+        response.raise_for_status()
+        result = response.json()
+        return TokenModel(
+            access_token=result["access_token"],
+            expires_on=datetime.datetime.now()
+            + datetime.timedelta(seconds=result["expires_in"]),
+        )

sws_api_client-3.2.0/sws_api_client/token.py

@@ -1,47 +0,0 @@
-import datetime
-import logging
-import os
-
-import requests
-
-from sws_api_client.case_agnostic_model import CaseAgnosticModel
-
-logger = logging.getLogger(__name__)
-from sws_api_client.auth import AuthClientCredentials
-
-
-class TokenModel(CaseAgnosticModel):
-    access_token: str
-    expires_on: datetime.datetime
-
-
-class Token:
-    def __init__(self, authclient: AuthClientCredentials) -> None:
-        self.authclient = authclient
-        self.token = self.get_token()
-
-    def get_token(self) -> TokenModel:
-        if (
-            not hasattr(self, "token")
-            or self.token.expires_on is None
-            or self.token.expires_on < datetime.datetime.now()
-        ):
-            url = f"{self.authclient.tokenEndpoint}"
-            data = {
-                "grant_type": "client_credentials",
-                "client_id": self.authclient.clientId,
-                "client_secret": self.authclient.clientSecret,
-                "scope": self.authclient.scope,
-            }
-            response = requests.post(url, data=data)
-            response.raise_for_status()
-            result = response.json()
-            self.token = TokenModel(
-                access_token=result["access_token"],
-                expires_on=datetime.datetime.now()
-                + datetime.timedelta(seconds=result["expires_in"]),
-            )
-            logger.debug(f"Token: {self.token}")
-            return self.token
-        else:
-            return self.token