From 2d96d92925e73d6008a5ba6fccf81c31d90427be Mon Sep 17 00:00:00 2001 From: Brandon Liu Date: Mon, 18 Jul 2022 12:21:12 +0800 Subject: [PATCH] remove lambda@edge support; use python functools.lru_cache --- serverless/aws/config.py | 2 - serverless/aws/create_lambda_function.py | 15 +---- serverless/aws/lambda_function.py | 84 ++++++++++-------------- serverless/aws/test_lambda_function.py | 10 +-- 4 files changed, 35 insertions(+), 76 deletions(-) delete mode 100644 serverless/aws/config.py diff --git a/serverless/aws/config.py b/serverless/aws/config.py deleted file mode 100644 index e475043..0000000 --- a/serverless/aws/config.py +++ /dev/null @@ -1,2 +0,0 @@ -BUCKET = "example-bucket-name" -REGION = "us-east-1" diff --git a/serverless/aws/create_lambda_function.py b/serverless/aws/create_lambda_function.py index 14f06af..9aa1fdf 100644 --- a/serverless/aws/create_lambda_function.py +++ b/serverless/aws/create_lambda_function.py @@ -1,20 +1,7 @@ -import argparse -import os import zipfile -import sys - -parser = argparse.ArgumentParser( - description="Create a deployment-ready PMTiles Lambda zip." -) -parser.add_argument("region", help="AWS Region of the S3 bucket.") -parser.add_argument("bucket", help="S3 Bucket Name.") -args = parser.parse_args() with zipfile.ZipFile("lambda_function.zip", "w", zipfile.ZIP_DEFLATED) as z: z.write("lambda_function.py") z.write("../../python/pmtiles/reader.py", "pmtiles.py") - info = zipfile.ZipInfo("config.py") - info.external_attr = 0o777 << 16 - z.writestr(info, f'REGION="{args.region}"\nBUCKET="{args.bucket}"') -print(f"created lambda_function.zip with REGION {args.region} and BUCKET {args.bucket}") +print(f"created lambda_function.zip") diff --git a/serverless/aws/lambda_function.py b/serverless/aws/lambda_function.py index ad574fe..8abbd7e 100644 --- a/serverless/aws/lambda_function.py +++ b/serverless/aws/lambda_function.py @@ -1,16 +1,30 @@ -import json -import pmtiles -import re -import boto3 import base64 import collections -from config import BUCKET, REGION +from functools import lru_cache +import gzip +import json +import os +import re +import boto3 +import pmtiles Zxy = collections.namedtuple("Zxy", ["z", "x", "y"]) s3 = boto3.client("s3") -rootCache = {} + +@lru_cache +def get_object_bytes(key, offset, length): + end = offset + length - 1 + return ( + s3.get_object( + Bucket=os.environ["BUCKET"], + Key=key, + Range=f"bytes={offset}-{end}", + ) + .get("Body") + .read() + ) def parse_tile_uri(str): @@ -20,60 +34,28 @@ def parse_tile_uri(str): return (m.group(1), Zxy(int(m.group(2)), int(m.group(3)), int(m.group(4)))) -def cloudfrontResponse(status_code, body, body_b64=False, headers={}): - headers = {key: [{"value": value}] for key, value in headers.items()} - resp = {"status": status_code, "body": body, "headers": headers} - if body_b64: - resp["bodyEncoding"] = "base64" - return resp - - -def apiGatewayResponse(status_code, body, body_b64=False, headers={}): - resp = {"statusCode": status_code, "body": body, "headers": headers} - if body_b64: - resp["isBase64Encoded"] = True - return resp - - def lambda_handler(event, context): - if "Records" in event: - uri = event["Records"][0]["cf"]["request"]["uri"] # CloudFront Origin Request - lambdaResponse = cloudfrontResponse - else: - uri = event["rawPath"] # API Gateway and Lambda Function URLs - lambdaResponse = apiGatewayResponse + start = datetime.now() + uri = event["rawPath"] # API Gateway and Lambda Function URLs tileset, tile = parse_tile_uri(uri) if not tile: - return lambdaResponse(400, "Invalid tile URL") + return {"statusCode": 400, "body": "Invalid Tile URL"} def get_bytes(offset, length): - global rootCache - if offset == 0 and length == 512000 and tileset in rootCache: - return rootCache[tileset] - - end = offset + length - 1 - result = ( - s3.get_object( - Bucket=BUCKET, - Key=tileset + ".pmtiles", - Range=f"bytes={offset}-{end}", - ) - .get("Body") - .read() - ) - if offset == 0 and length == 512000: - rootCache[tileset] = result - return result + return get_object_bytes(tileset + ".pmtiles", offset, length) reader = pmtiles.Reader(get_bytes) tile_data = reader.get(tile.z, tile.x, tile.y) if not tile_data: - return lambdaResponse(404, "Tile not found") + return {"statusCode": 404, "body": "Tile not found"} - headers = { - "Content-Encoding": "gzip", - "Content-Type": "application/protobuf", - "Access-Control-Allow-Origin": "*", + if reader.header().metadata.get("compression") == "gzip": + tile_data = gzip.decompress(tile_data) + + return { + "statusCode": 200, + "body": base64.b64encode(tile_data), + "isBase64Encoded": True, + "headers": {"Content-Type": "application/protobuf"}, } - return lambdaResponse(200, base64.b64encode(tile_data), True, headers) diff --git a/serverless/aws/test_lambda_function.py b/serverless/aws/test_lambda_function.py index a3cf5c2..a48a310 100644 --- a/serverless/aws/test_lambda_function.py +++ b/serverless/aws/test_lambda_function.py @@ -1,5 +1,5 @@ import unittest -from lambda_function import parse_tile_uri, cloudfrontResponse, apiGatewayResponse +from lambda_function import parse_tile_uri class TestLambda(unittest.TestCase): @@ -13,14 +13,6 @@ class TestLambda(unittest.TestCase): tileset, tile = parse_tile_uri("abcd") self.assertEqual(tile, None) - def test_cloudfront_response(self): - resp = cloudfrontResponse(200, "ok", False, {"a": "b"}) - self.assertEqual(resp["headers"]["a"], [{"value": "b"}]) - - def test_api_gateway_response(self): - resp = apiGatewayResponse(200, "ok", False, {"a": "b"}) - self.assertEqual(resp["headers"]["a"], "b") - if __name__ == "__main__": unittest.main()