Add AWS CORS

imaginate-ai · Sep 12, 2024 · a2576e5 · a2576e5
1 parent 99d4577
commit a2576e5
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 9 deletions.
diff --git a/aws/build_lambda_code.py b/aws/build_lambda_code.py
@@ -15,6 +15,7 @@
 
 
 ENV = "prod"
+FRONTEND_URL = "https://playimaginate.com"
 DIR = "aws"
 CWD = os.path.dirname(os.path.realpath(__file__))
 LAMBDA_LIBRARIES = """import os
@@ -33,18 +34,23 @@
 client = MongoClient(conn_uri)
 db = client[db_name]
 fs = GridFS(db)
+headers = {{
+    'Access-Control-Allow-Origin': '*',
+    'Access-Control-Allow-Headers': 'Content-Type,X-Amz-Date,Authorization,X-Api-Key,X-Amz-Security-Token',
+    'Access-Control-Allow-Methods': 'GET,OPTIONS'
+}}
 """
 LAMBDA_FUNC = """def handler(event, context):
     if event and 'queryStringParameters' in event and event['queryStringParameters'] and 'day' in event['queryStringParameters']:
         return images_by_date(event['queryStringParameters']['day'])
     else:
-        return {'statusCode': HTTPStatus.BAD_REQUEST, 'body': json.dumps('Invalid date')}
+        return {'statusCode': HTTPStatus.BAD_REQUEST, 'body': json.dumps('Invalid date'), 'headers': headers}
 """
 LAMBDA_SUBS = {
-  "abort\\(HTTPStatus.BAD_REQUEST": "return {'statusCode': HTTPStatus.BAD_REQUEST, 'body': json.dumps('Invalid date')}",
-  "abort\\(HTTPStatus.NOT_FOUND": "return {'statusCode': HTTPStatus.NOT_FOUND, 'body': json.dumps('Empty database')}",
+  "abort\\(HTTPStatus.BAD_REQUEST": "return {'statusCode': HTTPStatus.BAD_REQUEST, 'body': json.dumps('Invalid date'), 'headers': headers}",
+  "abort\\(HTTPStatus.NOT_FOUND": "return {'statusCode': HTTPStatus.NOT_FOUND, 'body': json.dumps('Empty database'), 'headers': headers}",
   "@bp.route": "",  # Remove function decorator entirely
-  "return jsonify": "return {'statusCode': HTTPStatus.OK, 'body': json.dumps(out)}",
+  "return jsonify": "return {'statusCode': HTTPStatus.OK, 'body': json.dumps(out), 'headers': headers}",
 }
 
 

diff --git a/aws/index.py b/aws/index.py
@@ -14,6 +14,11 @@
 client = MongoClient(conn_uri)
 db = client[db_name]
 fs = GridFS(db)
+headers = {
+  "Access-Control-Allow-Origin": "*",
+  "Access-Control-Allow-Headers": "Content-Type,X-Amz-Date,Authorization,X-Api-Key,X-Amz-Security-Token",
+  "Access-Control-Allow-Methods": "GET,OPTIONS",
+}
 
 
 class DateInfo(Enum):
@@ -63,13 +68,25 @@ def images_by_date(day):
     # This code is from GET /date/latest and is NOT internally called for aws/build_lambda_code.py
     res = next(fs.find().sort({"date": -1}).limit(1), None)  # Descending sort
     if not res:
-      return {"statusCode": HTTPStatus.NOT_FOUND, "body": json.dumps("Empty database")}
+      return {
+        "statusCode": HTTPStatus.NOT_FOUND,
+        "body": json.dumps("Empty database"),
+        "headers": headers,
+      }
 
     date = calculate_date(day, res.date)
     if not date:
-      return {"statusCode": HTTPStatus.BAD_REQUEST, "body": json.dumps("Invalid date")}
+      return {
+        "statusCode": HTTPStatus.BAD_REQUEST,
+        "body": json.dumps("Invalid date"),
+        "headers": headers,
+      }
   except ValueError:
-    return {"statusCode": HTTPStatus.BAD_REQUEST, "body": json.dumps("Invalid date")}
+    return {
+      "statusCode": HTTPStatus.BAD_REQUEST,
+      "body": json.dumps("Invalid date"),
+      "headers": headers,
+    }
 
   res = fs.find({"date": date})
   out = []
@@ -85,7 +102,7 @@ def images_by_date(day):
     encoded_data = b64encode(document.read())
     current_res["data"] = encoded_data.decode("utf-8")
     out.append(current_res)
-  return {"statusCode": HTTPStatus.OK, "body": json.dumps(out)}
+  return {"statusCode": HTTPStatus.OK, "body": json.dumps(out), "headers": headers}
 
 
 def handler(event, context):
@@ -97,4 +114,8 @@ def handler(event, context):
   ):
     return images_by_date(event["queryStringParameters"]["day"])
   else:
-    return {"statusCode": HTTPStatus.BAD_REQUEST, "body": json.dumps("Invalid date")}
+    return {
+      "statusCode": HTTPStatus.BAD_REQUEST,
+      "body": json.dumps("Invalid date"),
+      "headers": headers,
+    }