diff --git a/linkshortener/notifications.py b/linkshortener/notifications.py index 74cfd97..cf5b0e7 100644 --- a/linkshortener/notifications.py +++ b/linkshortener/notifications.py @@ -7,8 +7,8 @@ from linkshortener.shortener import connect -def generate(event): - db = connect(event) +def generate(): + db = connect() page = ( jinja2.Environment( loader=jinja2.FileSystemLoader(searchpath="./linkshortener/templates/"), @@ -37,10 +37,10 @@ def view(event, context): def summary(event, context): - if event["httpMethod"] is None: + if event.get("httpMethod") is None: db = connect() new = False - for i in db.scan(): + for i in db.scan()["Items"]: if i["uses"]["recent"] != 0: new = True if not new: @@ -50,7 +50,7 @@ def summary(event, context): Destination={"ToAddresses": [environ.get("ADMIN_CONTACT")]}, Message={ "Body": { - "Html": {"Charset": "UTF-8", "Data": generate(event)}, + "Html": {"Charset": "UTF-8", "Data": generate()}, "Text": { "Charset": "UTF-8", "Data": "This email must be viewed with HTML", diff --git a/linkshortener/shortener.py b/linkshortener/shortener.py index 505d6c4..2498e1a 100644 --- a/linkshortener/shortener.py +++ b/linkshortener/shortener.py @@ -8,11 +8,11 @@ from botocore.config import Config -def connect(event): +def connect(): return boto3.resource( "dynamodb", endpoint_url="http://localhost:8000" - if event["requestContext"]["stage"] == "dev" + if os.environ.get("STAGE") == "dev" else None, region_name="eu-west-2", config=Config(connect_timeout=5, retries={"max_attempts": 3}), @@ -40,7 +40,7 @@ def redirect(url): def shortener(event, context): - db = connect(event) + db = connect() try: url = db.get_item(Key={"code": event["pathParameters"]["id"]})["Item"]["url"] db.update_item( @@ -57,7 +57,7 @@ def shortener(event, context): def create(event, context): - db = connect(event) + db = connect() try: db.put_item( Item={ @@ -74,7 +74,7 @@ def create(event, context): def view(event, context): - db = connect(event) + db = connect() if event["queryStringParameters"] is not None and event[ "queryStringParameters" ].get("code") not in (None, ""): @@ -90,7 +90,7 @@ def view(event, context): def delete(event, context): - db = connect(event) + db = connect() try: db.delete_item( Key={"code": sanitize(json.loads(event["body"])["code"])}, @@ -103,7 +103,7 @@ def delete(event, context): def fallback(event, context): - return redirect(os.environ["FALLBACK_URL"]) + return redirect(os.environ.get("FALLBACK_URL")) def robots(event, context): diff --git a/serverless.yml b/serverless.yml index 20d96c8..6eb11f3 100644 --- a/serverless.yml +++ b/serverless.yml @@ -7,7 +7,7 @@ provider: runtime: python3.8 region: eu-west-2 memorySize: 128 - timeout: 20 + timeout: 6 apiKeys: - ${opt:stage}-adminkey environment: @@ -16,6 +16,7 @@ provider: DOMAIN: ${env:DOMAIN} SES_REGION: ${env:SES_REGION} FALLBACK_URL: ${env:FALLBACK_URL} + STAGE: ${opt:stage, self:provider.stage} iamRoleStatements: - Effect: Allow Action: