Add timeout options

Author: adammcdonagh

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)