S3 message spillover (#45)

* Use middleware to spill large responses over to S3

* Fix tests

* Up requirements

* Upgrade again

* Bump middleware version

.github/workflows/ci.yml

@@ -40,4 +40,5 @@ jobs:
         DB_NAME: mydatabase
         DB_USER: myuser
         DB_PASSWORD: mypassword
-        DB_TABLE_NAME: space2stats
+        DB_TABLE_NAME: space2stats
+        S3_BUCKET_NAME: test-bucket

space2stats_api/cdk/aws_stack.py

@@ -3,6 +3,7 @@
 from aws_cdk import aws_apigatewayv2_integrations as integrations
 from aws_cdk import aws_certificatemanager as acm
 from aws_cdk import aws_lambda as _lambda
+from aws_cdk import aws_s3 as s3
 from aws_cdk.aws_lambda_python_alpha import PythonFunction
 from constructs import Construct
 from settings import AppSettings, DeploymentSettings
@@ -15,6 +16,12 @@ def __init__(self, scope: Construct, id: str, **kwargs) -> None:
         app_settings = AppSettings(_env_file="./aws_app.env")
         deployment_settings = DeploymentSettings(_env_file="./aws_deployment.env")
 
+        bucket = s3.Bucket(
+            self,
+            "LargeResponseBucket",
+            lifecycle_rules=[s3.LifecycleRule(expiration=Duration.days(1))],
+        )
+
         lambda_function = PythonFunction(
             self,
             "Space2StatsFunction",
@@ -23,10 +30,15 @@ def __init__(self, scope: Construct, id: str, **kwargs) -> None:
             index="space2stats/handler.py",
             timeout=Duration.seconds(120),
             handler="handler",
-            environment=app_settings.model_dump(),
+            environment={
+                "S3_BUCKET_NAME": bucket.bucket_name,
+                **app_settings.model_dump(),
+            },
             memory_size=1024,
         )
 
+        bucket.grant_read_write(lambda_function)
+
         certificate = acm.Certificate.from_certificate_arn(
             self, "Certificate", deployment_settings.CDK_CERTIFICATE_ARN
         )

space2stats_api/src/pyproject.toml

@@ -19,6 +19,8 @@ pydantic-settings = ">=2.0.0"
 typing_extensions = "*"
 starlette-cramjam = ">=0.3,<0.4"
 mangum = "*"
+asgi-s3-response-middleware = "^0.0.4"
+boto3 = "^1.35.11"
 
 [tool.poetry.group.lambda.dependencies]
 mangum = "*"
@@ -31,6 +33,7 @@ pre-commit = "*"
 pytest = "*" 
 pytest-cov = "*" 
 pytest-postgresql = "*"
+moto = "^5.0.13"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]

space2stats_api/src/space2stats/app.py

@@ -1,14 +1,23 @@
 from contextlib import asynccontextmanager
 from typing import Any, Dict, List
 
+import boto3
+from asgi_s3_response_middleware import S3ResponseMiddleware
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import ORJSONResponse
 from starlette.requests import Request
 from starlette_cramjam.middleware import CompressionMiddleware
 
 from .db import close_db_connection, connect_to_db
-from .main import SummaryRequest, get_available_fields, get_summaries_from_geom
+from .main import (
+    SummaryRequest,
+    get_available_fields,
+    get_summaries_from_geom,
+    settings,
+)
+
+s3_client = boto3.client("s3")
 
 
 @asynccontextmanager
@@ -34,6 +43,11 @@ async def lifespan(app: FastAPI):
     allow_headers=["*"],
 )
 app.add_middleware(CompressionMiddleware)
+app.add_middleware(
+    S3ResponseMiddleware,
+    s3_bucket_name=settings.S3_BUCKET_NAME,
+    s3_client=s3_client,
+)
 
 
 @app.post("/summary", response_model=List[Dict[str, Any]])

space2stats_api/src/space2stats/settings.py

@@ -9,6 +9,9 @@ class Settings(BaseSettings):
     DB_PASSWORD: str
     DB_TABLE_NAME: str
 
+    # Bucket for large responses
+    S3_BUCKET_NAME: str
+
     # see https://www.psycopg.org/psycopg3/docs/api/pool.html#the-connectionpool-class for options
     DB_MIN_CONN_SIZE: int = 1
     DB_MAX_CONN_SIZE: int = 10

space2stats_api/src/tests/conftest.py

@@ -0,0 +1,37 @@
+import os
+
+import boto3
+import pytest
+from moto import mock_aws
+
+
+@pytest.fixture()
+def aws_credentials():
+    """Mocked AWS Credentials for moto."""
+    os.environ["AWS_ACCESS_KEY_ID"] = "testing"
+    os.environ["AWS_SECRET_ACCESS_KEY"] = "testing"
+    os.environ["AWS_SECURITY_TOKEN"] = "testing"
+    os.environ["AWS_SESSION_TOKEN"] = "testing"
+    os.environ["AWS_DEFAULT_REGION"] = "us-east-1"
+
+
+@pytest.fixture()
+def s3_client(aws_credentials):
+    """
+    Return a mocked S3 client
+    """
+    with mock_aws():
+        yield boto3.client("s3", region_name="us-east-1")
+
+
+@pytest.fixture
+def test_bucket(s3_client) -> str:
+    bucket_name = "test-bucket"
+    s3_client.create_bucket(Bucket=bucket_name)
+
+    return bucket_name
+
+
+@pytest.fixture(autouse=True)
+def set_bucket_name(monkeypatch, test_bucket):
+    monkeypatch.setenv("S3_BUCKET_NAME", test_bucket)

space2stats_api/src/tests/test_api.py

@@ -37,7 +37,7 @@ def database(postgresql_proc):
 
 
 @pytest.fixture(autouse=True)
-def client(monkeypatch, database):
+def client(monkeypatch, database, test_bucket):
     monkeypatch.setenv("DB_HOST", database.host)
     monkeypatch.setenv("DB_PORT", str(database.port))
     monkeypatch.setenv("DB_NAME", database.dbname)