fix: cache boto3 swf clients

simpleflow/boto3_utils.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+import hashlib
+from threading import local
+from typing import Any
+
+import boto3
+
+from simpleflow.utils import json_dumps
+
+
+def get_or_create_boto3_client(*, region_name: str | None, service_name: str, **kwargs: Any):
+    d = {
+        "region_name": region_name,
+        "service_name": service_name,
+    }
+    d.update(kwargs)
+    key = hashlib.sha1(json_dumps(d).encode()).hexdigest()
+    local_data = local()
+    client = getattr(local_data, key, None)
+    if client is None:
+        session = boto3.session.Session(region_name=region_name)
+        client = session.client(service_name, **kwargs)
+        setattr(local_data, key, client)
+    return client

simpleflow/storage.py

@@ -7,19 +7,18 @@
 from botocore.exceptions import ClientError
 
 from . import logger, settings
+from .boto3_utils import get_or_create_boto3_client
 from .swf.mapper.exceptions import extract_error_code
 
 if TYPE_CHECKING:
-    from typing import Optional, Tuple  # NOQA
-
-    from mypy_boto3_s3.service_resource import Bucket, ObjectSummary  # NOQA
+    from mypy_boto3_s3.service_resource import Bucket, ObjectSummary
 
 BUCKET_CACHE = {}
 BUCKET_LOCATIONS_CACHE = {}
 
 
 def get_client() -> boto3.session.Session.client:
-    return boto3.session.Session().client("s3")
+    return get_or_create_boto3_client(region_name=None, service_name="s3")
 
 
 def get_resource(host_or_region: str) -> boto3.session.Session.resource:

simpleflow/swf/mapper/core.py

@@ -15,10 +15,10 @@
 # config hosted in simpleflow. This wouldn't be the case with a standard
 # "logging.getLogger(__name__)" which would write logs under the "swf" namespace
 from simpleflow import logger
+from simpleflow.boto3_utils import get_or_create_boto3_client
+from simpleflow.swf.mapper import settings
 from simpleflow.utils import remove_none, retry
 
-from . import settings
-
 SETTINGS = settings.get()
 RETRIES = int(os.environ.get("SWF_CONNECTION_RETRIES", "5"))
 DEFAULT_AWS_REGION = "us-east-1"
@@ -47,9 +47,8 @@ def __init__(self, *args, **kwargs):
 
         self.boto3_client = kwargs.pop("boto3_client", None)
         if not self.boto3_client:
-            session = boto3.session.Session(region_name=self.region)
             # raises EndpointConnectionError if region is wrong
-            self.boto3_client = session.client("swf", **creds_)
+            self.boto3_client = get_or_create_boto3_client(region_name=self.region, service_name="swf", **creds_)
 
         logger.debug(f"initiated connection to region={self.region}")