Add s3 retry config and refactor s3_client creation as util

config/global_config.json

@@ -2,7 +2,11 @@
   "s3_bucket": {
     "uri": "dataclinic-gtfs-rt",
     "public_key": "",
-    "secret_key": ""
+    "secret_key": "",
+    "retries": {
+      "max_attempts": 3,
+      "mode": "standard"
+   }
   },
   "mobilitydatabase": {
     "url": "https://api.mobilitydatabase.org/v1",

src/normalize/normalize.py

@@ -9,11 +9,13 @@
 import pyarrow as pa
 import s3fs
 import structlog
+from botocore.config import Config
 from dateutil import parser
 
 from gtfs_realtime_pb2 import FeedMessage
 from src.normalize.parquet_utils import write_data, add_time_columns
 from src.normalize.protobuf_utils import protobuf_objects_to_pyarrow_table
+from src.util.s3_client import create_s3_client
 
 structlog.configure(
     processors=[
@@ -33,6 +35,10 @@ def check_config(config: dict):
     assert config["s3_bucket"]["uri"]
     assert config["s3_bucket"]["public_key"]
     assert config["s3_bucket"]["secret_key"]
+    retries_config = config["s3_bucket"].get("retries")
+    if retries_config:
+        assert retries_config["mode"], ("mode must be specified for enabling "
+                                        "retry")
 
 
 def load_config(path: str):
@@ -282,11 +288,8 @@ def main(
     for key, val in config.get("normalize_argv_override", {}).items():
         LOGGER.info(f"Overriding argv {key}={val}")
         exec(key + f"={val}")
-    s3 = boto3.client(
-        "s3",
-        aws_access_key_id=config["s3_bucket"]["public_key"],
-        aws_secret_access_key=config["s3_bucket"]["secret_key"],
-    )
+
+    s3 = create_s3_client(config["s3_bucket"])
     s3_fs = s3fs.S3FileSystem(
         key=config["s3_bucket"]["public_key"],
         secret=config["s3_bucket"]["secret_key"]

src/scraper/scrape.py

@@ -6,10 +6,12 @@
 from datetime import datetime
 
 import boto3
+from botocore.config import Config
 import click
 import structlog
 
 from src.scraper.mobilitydatabase import get_access_token, get_feed_json
+from src.util.s3_client import create_s3_client
 
 structlog.configure(
     processors=[
@@ -29,6 +31,10 @@ def check_config(config: dict):
     assert config["s3_bucket"]["uri"]
     assert config["s3_bucket"]["public_key"]
     assert config["s3_bucket"]["secret_key"]
+    retries_config = config["s3_bucket"].get("retries")
+    if retries_config:
+        assert retries_config["mode"], ("mode must be specified for enabling "
+                                        "retry")
 
     assert config["mobilitydatabase"]["url"]
     assert config["mobilitydatabase"]["token"]
@@ -41,15 +47,6 @@ def load_config(path: str):
         return config
 
 
-def create_s3_client(s3_config: dict):
-    session = boto3.Session(
-        aws_access_key_id=s3_config["public_key"],
-        aws_secret_access_key=s3_config["secret_key"],
-    )
-    s3 = session.resource("s3").Bucket(s3_config["uri"])
-    return s3
-
-
 def scrape_loop(s3_client, feed_id: str):
     module = importlib.import_module(f"src.scraper.feeds.{feed_id}")
     feed_class = getattr(module, feed_id.upper().replace("-", "_"))
@@ -107,7 +104,7 @@ def main(feed_id, config_path):
         config["mobilitydatabase"]["url"],
         config["mobilitydatabase"]["token"],
     )
-    s3 = create_s3_client(config["s3_bucket"])
+    s3 = create_s3_client(config["s3_bucket"]).resource("s3").Bucket(config["s3_bucket"]["uri"])
     scrape_loop(s3, feed_id)
 
 

src/util/s3_client.py

@@ -0,0 +1,22 @@
+import boto3
+from botocore.config import Config
+
+
+def create_s3_client(s3_config: dict):
+    retries_config = s3_config.get("retries")
+    s3 = boto3.client(
+        "s3",
+        aws_access_key_id=s3_config["public_key"],
+        aws_secret_access_key=s3_config["secret_key"],
+    )
+    if retries_config:
+        config = Config(
+            retries=retries_config
+        )
+        s3 = boto3.client(
+            "s3",
+            aws_access_key_id=s3_config["public_key"],
+            aws_secret_access_key=s3_config["secret_key"],
+            config=config
+        )
+    return s3