Merge pull request #423 from jbbarth/feat-boto3

boto2 -> boto3

Author: jbbarth

pyproject.toml

@@ -31,7 +31,7 @@ dynamic = ["version"]
 
 dependencies = [
     "attrs",
-    "boto>=2.49.0",
+    "boto3>=1.28.20",
     "dill>=0.3.6",
     "diskcache>=4.1.0",
     "Jinja2",
@@ -49,6 +49,7 @@ dependencies = [
 [project.optional-dependencies]
 dev = [
     "black",
+    "boto3-stubs[s3,swf]",
     "flaky",
     "invoke",
     "moto<3.0.0",

simpleflow/command.py

@@ -9,7 +9,6 @@
 from typing import TYPE_CHECKING
 from uuid import uuid4
 
-import boto.connection
 import click
 import multiprocess
 
@@ -34,24 +33,6 @@
     from simpleflow.swf.mapper.models.workflow import WorkflowType
 
 
-def disable_boto_connection_pooling():
-    # boto connection pooling doesn't work very well with multiprocessing, it
-    # provokes some errors like this:
-    #
-    #    [Errno 1] _ssl.c:1429: error:1408F119:SSL routines:SSL3_GET_RECORD:decryption failed or bad record mac
-    # when polling on analysis-testjbb-repair-a61ff96e854344748e308fefc9ddff61
-    #
-    # It's because when forking, file handles are copied and sockets are shared.
-    # Even sockets that handle SSL conections to AWS services, but SSL
-    # connections are stateful! So with multiple workers, it collides.
-    #
-    # To disable boto's connection pooling (which in practice makes boto open a
-    # *NEW* connection for each call), we make make boto believe we run on
-    # Google App Engine, where it disables connection pooling. There's no
-    # "direct" setting, so that's a hack but that works.
-    boto.connection.ON_APP_ENGINE = True
-
-
 def comma_separated_list(value):
     """
     Transforms a comma-separated list into a list of strings.
@@ -577,8 +558,6 @@ def standalone(
     with a single main process.
 
     """
-    disable_boto_connection_pooling()
-
     if force_activities and not repair:
         raise ValueError("You should only use --force-activities with --repair.")
 

simpleflow/metrology.py

@@ -163,16 +163,16 @@ def push_metrology(self, history):
         """
         Fetch workflow history and merge it with metrology
         """
-        activity_keys = [obj for obj in storage.list_keys(settings.METROLOGY_BUCKET, self.metrology_path)]
+        activity_keys = storage.list_keys(settings.METROLOGY_BUCKET, self.metrology_path)
         history_dumped = dump_history_to_json(history)
         history = json.loads(history_dumped)
 
         for key in activity_keys:
             if not key.key.startswith(os.path.join(self.metrology_path, "activity.")):
                 continue
-            contents = key.get_contents_as_string(encoding="utf-8")
+            contents = key.get()["Body"].read().decode("utf-8")
             result = json.loads(contents)
-            search = ACTIVITY_KEY_RE.search(key.name)
+            search = ACTIVITY_KEY_RE.search(key.key)
             name = search.group(1)
             for h in history:
                 if h[0] == name:

simpleflow/storage.py

@@ -1,30 +1,34 @@
 from __future__ import annotations
 
+import io
 from typing import TYPE_CHECKING
 
-from boto.exception import S3ResponseError
-from boto.s3 import connect_to_region, connection
-from boto.s3.key import Key
+import boto3
+from botocore.exceptions import ClientError
 
 from . import logger, settings
+from .swf.mapper.exceptions import extract_error_code
 
 if TYPE_CHECKING:
     from typing import Optional, Tuple  # NOQA
 
-    from boto.s3.bucket import Bucket  # NOQA
-    from boto.s3.bucketlistresultset import BucketListResultSet  # NOQA
+    from mypy_boto3_s3.service_resource import Bucket, ObjectSummary  # NOQA
 
 BUCKET_CACHE = {}
 BUCKET_LOCATIONS_CACHE = {}
 
 
-def get_connection(host_or_region: str) -> connection.S3Connection:
+def get_client() -> boto3.session.Session.client:
+    return boto3.session.Session().client("s3")
+
+
+def get_resource(host_or_region: str) -> boto3.session.Session.resource:
     # first case: we got a valid DNS (host)
     if "." in host_or_region:
-        return connection.S3Connection(host=host_or_region)
+        return boto3.resource("s3", endpoint_url=f"https://{host_or_region}")
 
     # second case: we got a region
-    return connect_to_region(host_or_region)
+    return boto3.resource("s3", region_name=host_or_region)
 
 
 def sanitize_bucket_and_host(bucket: str) -> tuple[str, str]:
@@ -48,20 +52,18 @@ def sanitize_bucket_and_host(bucket: str) -> tuple[str, str]:
 
     # second case: we got a bucket name, we need to figure out which region it's in
     try:
-        conn0 = connection.S3Connection()
-        bucket_obj = conn0.get_bucket(bucket, validate=False)
-
-        # get_location() returns a region or an empty string for us-east-1,
-        # historically named "US Standard" in some places. Maybe other S3
-        # calls support an empty string as region, but I prefer to be
-        # explicit here.
-        location = bucket_obj.get_location() or "us-east-1"
+        # get_bucket_location() returns a region or an empty string for us-east-1,
+        # historically named "US Standard" in some places. Maybe other S3 calls
+        # support an empty string as region, but I prefer to be explicit here.
+        location = get_client().get_bucket_location(Bucket=bucket)["LocationConstraint"] or "us-east-1"
 
         # save location for later use
         BUCKET_LOCATIONS_CACHE[bucket] = location
-    except S3ResponseError as e:
-        if e.error_code == "AccessDenied":
+    except ClientError as e:
+        error_code = extract_error_code(e)
+        if error_code == "AccessDenied":
             # probably not allowed to perform GetBucketLocation on this bucket
+            # TODO: consider raising instead? who forbids GetBucketLocation anyway?
             logger.warning(f"Access denied while trying to get location of bucket {bucket}")
             location = ""
         else:
@@ -74,45 +76,47 @@ def sanitize_bucket_and_host(bucket: str) -> tuple[str, str]:
     return bucket, location
 
 
-def get_bucket(bucket_name: str) -> Bucket:
+def get_bucket(bucket_name: str) -> "Bucket":
     bucket_name, location = sanitize_bucket_and_host(bucket_name)
-    conn = get_connection(location)
+    s3 = get_resource(location)
     if bucket_name not in BUCKET_CACHE:
-        bucket = conn.get_bucket(bucket_name, validate=False)
+        bucket = s3.Bucket(bucket_name)
         BUCKET_CACHE[bucket_name] = bucket
     return BUCKET_CACHE[bucket_name]
 
 
 def pull(bucket: str, path: str, dest_file: str) -> None:
-    bucket = get_bucket(bucket)
-    key = bucket.get_key(path)
-    key.get_contents_to_filename(dest_file)
+    bucket_resource = get_bucket(bucket)
+    bucket_resource.download_file(path, dest_file)
 
 
 def pull_content(bucket: str, path: str) -> str:
-    bucket = get_bucket(bucket)
-    key = bucket.get_key(path)
-    return key.get_contents_as_string(encoding="utf-8")
+    bucket_resource = get_bucket(bucket)
+    bytes_buffer = io.BytesIO()
+    bucket_resource.download_fileobj(path, bytes_buffer)
+    return bytes_buffer.getvalue().decode()
 
 
 def push(bucket: str, path: str, src_file: str, content_type: str | None = None) -> None:
-    bucket = get_bucket(bucket)
-    key = Key(bucket, path)
-    headers = {}
+    bucket_resource = get_bucket(bucket)
+    extra_args = {}
     if content_type:
-        headers["content_type"] = content_type
-    key.set_contents_from_filename(src_file, headers=headers, encrypt_key=settings.SIMPLEFLOW_S3_SSE)
+        extra_args["ContentType"] = content_type
+    if settings.SIMPLEFLOW_S3_SSE:
+        extra_args["ServerSideEncryption"] = "AES256"
+    bucket_resource.upload_file(src_file, path, ExtraArgs=extra_args)
 
 
 def push_content(bucket: str, path: str, content: str, content_type: str | None = None) -> None:
-    bucket = get_bucket(bucket)
-    key = Key(bucket, path)
-    headers = {}
+    bucket_resource = get_bucket(bucket)
+    extra_args = {}
     if content_type:
-        headers["content_type"] = content_type
-    key.set_contents_from_string(content, headers=headers, encrypt_key=settings.SIMPLEFLOW_S3_SSE)
+        extra_args["ContentType"] = content_type
+    if settings.SIMPLEFLOW_S3_SSE:
+        extra_args["ServerSideEncryption"] = "AES256"
+    bucket_resource.upload_fileobj(io.BytesIO(content.encode()), path, ExtraArgs=extra_args)
 
 
-def list_keys(bucket: str, path: str = None) -> BucketListResultSet:
-    bucket = get_bucket(bucket)
-    return bucket.list(path)
+def list_keys(bucket: str, path: str = None) -> list["ObjectSummary"]:
+    bucket_resource = get_bucket(bucket)
+    return [obj for obj in bucket_resource.objects.filter(Prefix=path or "").all()]

simpleflow/swf/executor.py

@@ -48,13 +48,13 @@
 # doesn't get the scheduled task while it should...
 @retry.with_delay(nb_times=3, delay=retry.exponential, on_exceptions=KeyError)
 def run_fake_activity_task(domain, task_list, result):
-    conn = ConnectedSWFObject().connection
-    resp = conn.poll_for_activity_task(
+    obj = ConnectedSWFObject()
+    resp = obj.poll_for_activity_task(
         domain,
         task_list,
         identity=swf_identity(),
     )
-    conn.respond_activity_task_completed(
+    obj.respond_activity_task_completed(
         resp["taskToken"],
         result,
     )
@@ -63,13 +63,13 @@ def run_fake_activity_task(domain, task_list, result):
 # Same retry condition as run_fake_activity_task
 @retry.with_delay(nb_times=3, delay=retry.exponential, on_exceptions=KeyError)
 def run_fake_child_workflow_task(domain, task_list, result=None):
-    conn = ConnectedSWFObject().connection
-    resp = conn.poll_for_decision_task(
+    obj = ConnectedSWFObject()
+    resp = obj.poll_for_decision_task(
         domain,
         task_list,
         identity=swf_identity(),
     )
-    conn.respond_decision_task_completed(
+    obj.respond_decision_task_completed(
         resp["taskToken"],
         decisions=[
             {

simpleflow/swf/mapper/actors/core.py

@@ -1,13 +1,9 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
 
 from simpleflow.swf.mapper.core import ConnectedSWFObject
 from simpleflow.swf.mapper.models.domain import Domain
 
-if TYPE_CHECKING:
-    from boto.exception import SWFResponseError
-
 
 class Actor(ConnectedSWFObject):
     """SWF Actor base class.
@@ -50,11 +46,3 @@ def stop(self):
         shutting down.
         """
         raise NotImplementedError
-
-    def get_error_message(self, e: SWFResponseError) -> str:
-        message = e.error_message
-        if not message:
-            if e.body:
-                # Expected 'message', got 'Message' ¯\_(ツ)_/¯
-                message = e.body.get("Message")
-        return message

simpleflow/swf/mapper/actors/decider.py

@@ -2,12 +2,18 @@
 
 from typing import TYPE_CHECKING, Any
 
-import boto.exception
+from botocore.exceptions import ClientError
 
 from simpleflow import format, logging_context
 from simpleflow.utils import json_dumps
 from simpleflow.swf.mapper.actors.core import Actor
-from simpleflow.swf.mapper.exceptions import DoesNotExistError, PollTimeout, ResponseError
+from simpleflow.swf.mapper.exceptions import (
+    DoesNotExistError,
+    PollTimeout,
+    ResponseError,
+    extract_error_code,
+    extract_message,
+)
 from simpleflow.swf.mapper.models.decision.base import Decision
 from simpleflow.swf.mapper.models.history.base import History
 from simpleflow.swf.mapper.models.workflow import WorkflowExecution, WorkflowType
@@ -43,14 +49,15 @@ def complete(
         if execution_context is not None and not isinstance(execution_context, str):
             execution_context = json_dumps(execution_context)
         try:
-            self.connection.respond_decision_task_completed(
+            self.respond_decision_task_completed(
                 task_token,
                 decisions,
-                format.execution_context(execution_context),
+                execution_context=format.execution_context(execution_context),
             )
-        except boto.exception.SWFResponseError as e:
-            message = self.get_error_message(e)
-            if e.error_code == "UnknownResourceFault":
+        except ClientError as e:
+            error_code = extract_error_code(e)
+            message = extract_message(e)
+            if error_code == "UnknownResourceFault":
                 raise DoesNotExistError(
                     f"Unable to complete decision task with token={task_token}",
                     message,
@@ -79,7 +86,7 @@ def poll(self, task_list=None, identity=None, **kwargs):
         logging_context.reset()
         task_list = task_list or self.task_list
 
-        task = self.connection.poll_for_decision_task(
+        task = self.poll_for_decision_task(
             self.domain.name,
             task_list=task_list,
             identity=format.identity(identity),
@@ -97,16 +104,17 @@ def poll(self, task_list=None, identity=None, **kwargs):
         next_page = task.get("nextPageToken")
         while next_page:
             try:
-                task = self.connection.poll_for_decision_task(
+                task = self.poll_for_decision_task(
                     self.domain.name,
                     task_list=task_list,
                     identity=format.identity(identity),
                     next_page_token=next_page,
                     **kwargs,
                 )
-            except boto.exception.SWFResponseError as e:
-                message = self.get_error_message(e)
-                if e.error_code == "UnknownResourceFault":
+            except ClientError as e:
+                error_code = extract_error_code(e)
+                message = extract_message(e)
+                if error_code == "UnknownResourceFault":
                     raise DoesNotExistError(
                         "Unable to poll decision task",
                         message,