Add timeouts for botocore

Author: adammcdonagh

CHANGELOG.md

@@ -1,5 +1,10 @@
 # Changelog
 
+# v26.28.0
+
+- Add support for specifying botore timeouts and retries for s3 transfers/executions
+- Fix issue in tidy throwing an exception when the s3 client was already closed
+
 # v26.18.0
 
 - Fix a bug where the S3 client was not being properly closed and garbage collected, which could lead to resource leaks and issues with too many open connections. This was caused by the `tidy` method not setting the `s3_client` attribute to `None` after closing it, which meant that the botocore objects were not being garbage collected.

README.md

@@ -140,6 +140,20 @@ JSON configs for transfers can be defined as follows:
 ]
 ```
 
+### Default S3 client timeout behaviour
+
+S3 transfers now apply botocore defaults even when not specified in the protocol:
+
+- `botocoreReadTimeout`: `60`
+- `botocoreConnectTimeout`: `10`
+- `max_attempts`: `0`
+
+These are socket/connect and retry settings, not a total transfer deadline. A large
+upload or download that is still making progress can continue past the batch timeout
+request and still complete successfully. The defaults mainly prevent stalled S3 calls
+from sitting in botocore retries for longer than necessary. All three values can still
+be overridden per protocol when needed.
+
 ## Example S3 upload with flag files
 
 ```json

src/opentaskpy/addons/aws/remotehandlers/s3.py

@@ -7,6 +7,7 @@
 
 import boto3
 import opentaskpy.otflogging
+from botocore.config import Config
 from botocore.exceptions import ClientError
 from dateutil.tz import tzlocal
 from opentaskpy.remotehandlers.remotehandler import (
@@ -17,6 +18,32 @@
 from .creds import get_aws_client, set_aws_creds
 
 MAX_OBJECTS_PER_QUERY = 100
+DEFAULT_S3_BOTOCORE_READ_TIMEOUT = 60
+DEFAULT_S3_BOTOCORE_CONNECT_TIMEOUT = 10
+DEFAULT_S3_BOTOCORE_MAX_ATTEMPTS = 0
+
+
+def _build_botocore_config(protocol: dict) -> Config:
+    """Build a botocore client config from the task protocol settings.
+
+    Defaults are applied so S3 transfers do not inherit botocore retry behaviour that
+    can keep a batch blocked longer than expected after a timeout request.
+    """
+    config_options = {
+        "read_timeout": protocol.get(
+            "botocoreReadTimeout", DEFAULT_S3_BOTOCORE_READ_TIMEOUT
+        ),
+        "connect_timeout": protocol.get(
+            "botocoreConnectTimeout", DEFAULT_S3_BOTOCORE_CONNECT_TIMEOUT
+        ),
+        "retries": {
+            "max_attempts": protocol.get(
+                "max_attempts", DEFAULT_S3_BOTOCORE_MAX_ATTEMPTS
+            )
+        },
+    }
+
+    return Config(**config_options)
 
 
 class S3Transfer(RemoteTransferHandler):
@@ -74,12 +101,15 @@ def validate_or_refresh_creds(self) -> None:
             if self.temporary_creds:
                 self.logger.info("Renewing temporary credentials")
 
+            client_config = _build_botocore_config(self.spec["protocol"])
+
             client_result = get_aws_client(
                 "s3",
                 self.credentials,
                 token_expiry_seconds=self.token_expiry_seconds,
                 assume_role_arn=self.assume_role_arn,
                 assume_role_external_id=self.assume_role_external_id,
+                config=client_config,
             )
             self.temporary_creds = (
                 client_result["temporary_creds"]
@@ -484,7 +514,8 @@ def create_flag_files(self) -> int:
 
     def tidy(self) -> None:
         """Tidy up the S3 client."""
-        self.s3_client.close()
+        if self.s3_client and hasattr(self.s3_client, "close"):
+            self.s3_client.close()
         self.s3_client = None  # allow botocore objects to be garbage collected
 
 
@@ -593,5 +624,6 @@ def execute(self) -> bool:
 
     def tidy(self) -> None:
         """Tidy up the S3 client."""
-        self.s3_client.close()
+        if self.s3_client and hasattr(self.s3_client, "close"):
+            self.s3_client.close()
         self.s3_client = None  # allow botocore objects to be garbage collected

src/opentaskpy/addons/aws/remotehandlers/schemas/transfer/s3_destination/protocol.json

@@ -27,6 +27,15 @@
     },
     "region_name": {
       "type": "string"
+    },
+    "botocoreReadTimeout": {
+      "type": "integer"
+    },
+    "botocoreConnectTimeout": {
+      "type": "integer"
+    },
+    "max_attempts": {
+      "type": "integer"
     }
   },
   "required": ["name"],

src/opentaskpy/addons/aws/remotehandlers/schemas/transfer/s3_source/protocol.json

@@ -27,6 +27,15 @@
     },
     "region_name": {
       "type": "string"
+    },
+    "botocoreReadTimeout": {
+      "type": "integer"
+    },
+    "botocoreConnectTimeout": {
+      "type": "integer"
+    },
+    "max_attempts": {
+      "type": "integer"
     }
   },
   "required": ["name"],

tests/test_remotehandler_s3_transfer.py

@@ -12,6 +12,7 @@
 import botocore
 import freezegun
 import pytest
+from botocore.config import Config
 from opentaskpy.taskhandlers import transfer
 from pytest_shell import fs
 
@@ -477,6 +478,139 @@ def test_remote_handler():
     assert transfer_obj.dest_remote_handlers is None
 
 
+def test_s3_transfer_passes_botocore_config(monkeypatch):
+    captured = {}
+
+    class DummyClient:
+        def close(self):
+            return None
+
+    def fake_get_aws_client(
+        client_type,
+        credentials,
+        token_expiry_seconds=900,
+        assume_role_arn=None,
+        assume_role_external_id=None,
+        config=None,
+    ):
+        captured["client_type"] = client_type
+        captured["credentials"] = credentials
+        captured["token_expiry_seconds"] = token_expiry_seconds
+        captured["assume_role_arn"] = assume_role_arn
+        captured["assume_role_external_id"] = assume_role_external_id
+        captured["config"] = config
+        return {"client": DummyClient(), "temporary_creds": None}
+
+    monkeypatch.setattr(
+        "opentaskpy.addons.aws.remotehandlers.s3.get_aws_client", fake_get_aws_client
+    )
+
+    handler = S3Transfer(
+        {
+            "task_id": "s3-config-test",
+            "bucket": BUCKET_NAME,
+            "directory": "src",
+            "fileRegex": ".*\\.txt",
+            "protocol": {
+                "name": "opentaskpy.addons.aws.remotehandlers.s3.S3Transfer",
+                "botocoreReadTimeout": 120,
+                "botocoreConnectTimeout": 30,
+                "max_attempts": 2,
+            },
+        }
+    )
+
+    assert isinstance(captured["config"], Config)
+    assert captured["client_type"] == "s3"
+    assert captured["config"].read_timeout == 120
+    assert captured["config"].connect_timeout == 30
+    assert captured["config"].retries["max_attempts"] == 2
+    handler.tidy()
+
+
+def test_s3_transfer_uses_default_botocore_config(monkeypatch):
+    captured = {}
+
+    class DummyClient:
+        def close(self):
+            return None
+
+    def fake_get_aws_client(
+        client_type,
+        credentials,
+        token_expiry_seconds=900,
+        assume_role_arn=None,
+        assume_role_external_id=None,
+        config=None,
+    ):
+        captured["config"] = config
+        return {"client": DummyClient(), "temporary_creds": None}
+
+    monkeypatch.setattr(
+        "opentaskpy.addons.aws.remotehandlers.s3.get_aws_client", fake_get_aws_client
+    )
+
+    handler = S3Transfer(
+        {
+            "task_id": "s3-default-config-test",
+            "bucket": BUCKET_NAME,
+            "directory": "src",
+            "fileRegex": ".*\\.txt",
+            "protocol": {
+                "name": "opentaskpy.addons.aws.remotehandlers.s3.S3Transfer",
+            },
+        }
+    )
+
+    assert isinstance(captured["config"], Config)
+    assert captured["config"].read_timeout == 60
+    assert captured["config"].connect_timeout == 10
+    assert captured["config"].retries["max_attempts"] == 0
+    handler.tidy()
+
+
+def test_s3_transfer_tidy_is_idempotent(monkeypatch):
+    class DummyClient:
+        def __init__(self):
+            self.close_calls = 0
+
+        def close(self):
+            self.close_calls += 1
+
+    def fake_get_aws_client(
+        client_type,
+        credentials,
+        token_expiry_seconds=900,
+        assume_role_arn=None,
+        assume_role_external_id=None,
+        config=None,
+    ):
+        return {"client": DummyClient(), "temporary_creds": None}
+
+    monkeypatch.setattr(
+        "opentaskpy.addons.aws.remotehandlers.s3.get_aws_client", fake_get_aws_client
+    )
+
+    handler = S3Transfer(
+        {
+            "task_id": "s3-tidy-test",
+            "bucket": BUCKET_NAME,
+            "directory": "src",
+            "fileRegex": ".*\\.txt",
+            "protocol": {
+                "name": "opentaskpy.addons.aws.remotehandlers.s3.S3Transfer",
+            },
+        }
+    )
+
+    client = handler.s3_client
+    handler.tidy()
+    handler.tidy()
+
+    assert client.close_calls == 1
+    assert handler.s3_client is None
+
+
 def test_s3_file_watch(s3_client, setup_bucket, tmp_path):
     transfer_obj = transfer.Transfer(
         None, "s3-file-watch", s3_file_watch_task_definition

tests/test_s3_source_schema_validate.py

@@ -125,6 +125,11 @@ def test_s3_source_protocol(
     json_data["source"]["protocol"]["token_expiry_seconds"] = 10000
     assert validate_transfer_json(json_data)
 
+    json_data["source"]["protocol"]["botocoreReadTimeout"] = 120
+    json_data["source"]["protocol"]["botocoreConnectTimeout"] = 30
+    json_data["source"]["protocol"]["max_attempts"] = 2
+    assert validate_transfer_json(json_data)
+
     # Set it too low, and too high
     json_data["source"]["protocol"]["token_expiry_seconds"] = 899
     assert not validate_transfer_json(json_data)
@@ -247,6 +252,11 @@ def test_s3_destination(valid_transfer, valid_destination):
     json_data["destination"][0]["flags"] = {"fullPath": "flag.txt"}
     assert validate_transfer_json(json_data)
 
+    json_data["destination"][0]["protocol"]["botocoreReadTimeout"] = 120
+    json_data["destination"][0]["protocol"]["botocoreConnectTimeout"] = 30
+    json_data["destination"][0]["protocol"]["max_attempts"] = 2
+    assert validate_transfer_json(json_data)
+
     # Remove protocol
     del json_data["destination"][0]["protocol"]
     assert not validate_transfer_json(json_data)