From fbc7b021230e678931b8f94861c523e3d7b15851 Mon Sep 17 00:00:00 2001
From: Brandon Liu <bdon@bdon.org>
Date: Tue, 5 Jul 2022 21:26:13 +0800
Subject: [PATCH] improvements to lambda wrapper [#13]

---
 lambda/README.md          | 22 +++++++++++++++++-
 lambda/lambda_function.py | 20 ----------------
 lambda/tippecanoe.py      | 48 +++++++++++++++++++++++++++++++++++++++
 lambda/tippecanoe_test.py | 26 +++++++++++++++++++++
 lambda/tmp/.gitkeep       |  0
 5 files changed, 95 insertions(+), 21 deletions(-)
 delete mode 100644 lambda/lambda_function.py
 create mode 100644 lambda/tippecanoe.py
 create mode 100644 lambda/tippecanoe_test.py
 create mode 100644 lambda/tmp/.gitkeep

diff --git a/lambda/README.md b/lambda/README.md
index 270905c9..7838a98b 100644
--- a/lambda/README.md
+++ b/lambda/README.md
@@ -7,4 +7,24 @@
 
 2. Create a ZIP archive with the `tippecanoe` binary at `bin/tippecanoe` and upload as a Lambda layer
 3. Copy [lambda_function.py](lambda_function.py) (python 3.9 runtime)
-4. set the environment variable `OUTPUT_BUCKET` to your output bucket; give lambda IAM role access to input/output buckets
\ No newline at end of file
+4. set the environment variable `OUTPUT_BUCKET` to your output bucket; give lambda IAM role access to input/output buckets
+
+```json
+{
+    "Version": "2012-10-17",
+    "Statement": [
+        {
+            "Sid": "AllowGet",
+            "Effect": "Allow",
+            "Action": "s3:GetObject",
+            "Resource": "arn:aws:s3:::tippecanoe-lambda-input/*"
+        },
+        {
+            "Sid": "AllowPut",
+            "Effect": "Allow",
+            "Action": "s3:PutObject",
+            "Resource": "arn:aws:s3:::tippecanoe-lambda-output/*"
+        }
+    ]
+}
+```
\ No newline at end of file
diff --git a/lambda/lambda_function.py b/lambda/lambda_function.py
deleted file mode 100644
index a58a09db..00000000
--- a/lambda/lambda_function.py
+++ /dev/null
@@ -1,20 +0,0 @@
-import os
-import json
-import subprocess
-import boto3
-# if the input bucket is publicly readable
-# from botocore import UNSIGNED
-# from botocore.config import Config
-
-def lambda_handler(event, context):
-    bucket = event['Records'][0]['s3']['bucket']['name']
-    key = event['Records'][0]['s3']['object']['key']
-    #s3 = boto3.client('s3',config=Config(signature_version=UNSIGNED))
-    s3 = boto3.client('s3')
-    s3.download_file(bucket,key,'/tmp/input.geojsonseq')
-    p = subprocess.check_call(['tippecanoe','-o','/tmp/output.mbtiles','/tmp/input.geojsonseq','--force'])
-    s3.upload_file('/tmp/output.mbtiles',os.environ['OUTPUT_BUCKET'],key + ".mbtiles")
-    return {
-        'statusCode': 200,
-        'body': json.dumps('Done!')
-    }
diff --git a/lambda/tippecanoe.py b/lambda/tippecanoe.py
new file mode 100644
index 00000000..4fd38675
--- /dev/null
+++ b/lambda/tippecanoe.py
@@ -0,0 +1,48 @@
+import os
+import json
+import subprocess
+import sys
+import boto3
+# if the input bucket is publicly readable
+# from botocore import UNSIGNED
+# from botocore.config import Config
+#s3 = boto3.client('s3',config=Config(signature_version=UNSIGNED))
+
+def lambda_handler(event, context):
+    bucket_name = event['Records'][0]['s3']['bucket']['name']
+    input_key = event['Records'][0]['s3']['object']['key']
+    tmpdir = '/tmp'
+    run_tippecanoe('tippecanoe',tmpdir,bucket_name, input_key)
+    return {
+        'statusCode': 200,
+        'body': json.dumps('Done!')
+    }
+
+# input formats: "geojsonseq", "fgb", "geobuf", "geojson", "csv"
+def split_key(key):
+    if key.endswith(".gz"):
+        root, ext = os.path.splitext(key[0:-3])
+        return root, ext + ".gz"
+    else:
+        return os.path.splitext(key)
+
+def run_tippecanoe(executable, tmpdir, bucket_name, input_key):
+    s3 = boto3.client('s3')
+    root, ext = split_key(input_key)
+    input_path = os.path.join(tmpdir,'input' + ext)
+    output_path = os.path.join(tmpdir,'output.mbtiles')
+    s3.download_file(bucket_name,input_key,input_path)
+    p = subprocess.Popen([executable,'-o',output_path,input_path,'--force','-u','-U','5'],stderr=subprocess.PIPE)
+    for line in p.stderr:
+        try:
+            j = json.loads(line)
+            s3.put_object(Body=json.dumps(j),Bucket=os.environ['OUTPUT_BUCKET'],Key=root + ".progress")
+        except json.decoder.JSONDecodeError:
+            pass
+    s3.upload_file(output_path,os.environ['OUTPUT_BUCKET'],root + ".mbtiles")
+
+if __name__ == "__main__":
+    tmpdir = 'tmp'
+    bucket_name = sys.argv[1]
+    input_key = sys.argv[2]
+    run_tippecanoe('../tippecanoe', tmpdir, bucket_name, input_key)
diff --git a/lambda/tippecanoe_test.py b/lambda/tippecanoe_test.py
new file mode 100644
index 00000000..a8ce663b
--- /dev/null
+++ b/lambda/tippecanoe_test.py
@@ -0,0 +1,26 @@
+from tippecanoe import split_key
+import unittest
+
+class TestTippecanoe(unittest.TestCase):
+    def test_basic(self):
+        root, ext = split_key("file.json")
+        self.assertEqual(root,"file")
+        self.assertEqual(ext,".json")
+
+    def test_basic_gz(self):
+        root, ext = split_key("file.csv.gz")
+        self.assertEqual(root,"file")
+        self.assertEqual(ext,".csv.gz")
+
+    def test_nested_key(self):
+        root, ext = split_key("dir/file.fgb")
+        self.assertEqual(root,"dir/file")
+        self.assertEqual(ext,".fgb")
+
+    def test_nested_key_gz(self):
+        root, ext = split_key("dir/file.geojsonseq.gz")
+        self.assertEqual(root,"dir/file")
+        self.assertEqual(ext,".geojsonseq.gz")
+
+if __name__ == '__main__':
+    unittest.main()
\ No newline at end of file
diff --git a/lambda/tmp/.gitkeep b/lambda/tmp/.gitkeep
new file mode 100644
index 00000000..e69de29b