From bb57ab2fa72e5ea96ddb2c3c5f86867465640766 Mon Sep 17 00:00:00 2001
From: Ru Chern Chong <hello@ruchern.xyz>
Date: Fri, 24 May 2024 13:24:57 +0800
Subject: [PATCH] Add create of python package for lambda layer

---
 .github/workflows/pulumi.preview.yml |  2 +-
 Pulumi.yaml                          |  2 +-
 __main__.py                          | 20 +++++++++++++------
 utils/create_package_zip.py          | 29 ++++++++++++++++++++++++++++
 4 files changed, 45 insertions(+), 8 deletions(-)
 create mode 100644 utils/create_package_zip.py

diff --git a/.github/workflows/pulumi.preview.yml b/.github/workflows/pulumi.preview.yml
index 7fa7e81..cb0eca8 100644
--- a/.github/workflows/pulumi.preview.yml
+++ b/.github/workflows/pulumi.preview.yml
@@ -49,7 +49,7 @@ jobs:
           python -m venv python
           source python/bin/activate
           pip install -r requirements.txt
-          zip -r package.zip python/lib/python3.12/site-packages -x "**/pulumi*/*" "**/pip*/*" "**/pytest*/*"
+          python utils/create_package_zip.py
       - uses: pulumi/actions@v5
         with:
           command: up
diff --git a/Pulumi.yaml b/Pulumi.yaml
index 18d50af..5454e79 100644
--- a/Pulumi.yaml
+++ b/Pulumi.yaml
@@ -2,7 +2,7 @@ name: lta-datasets-updater
 runtime:
   name: python
   options:
-    virtualenv: python
+    virtualenv: venv
 config:
   pulumi:tags:
     value:
diff --git a/__main__.py b/__main__.py
index 7956b13..3969a79 100644
--- a/__main__.py
+++ b/__main__.py
@@ -1,9 +1,11 @@
 import json
 import os
+import tempfile
 
 import pulumi
 import pulumi_aws as aws
 from dotenv import load_dotenv
+from utils import create_package_zip
 
 load_dotenv()
 
@@ -40,12 +42,16 @@
     policy_arn=aws.iam.ManagedPolicy.AWS_LAMBDA_BASIC_EXECUTION_ROLE,
 )
 
-package_layer = aws.lambda_.LayerVersion(
-    f"{PROJECT_NAME}-layer",
-    layer_name=f"{PROJECT_NAME}-layer",
-    code=pulumi.FileArchive("package.zip"),
-    compatible_runtimes=[RUNTIME],
-)
+temp_dir = tempfile.TemporaryDirectory().name
+zip_file = 'python.zip'
+
+create_package_zip.main(temp_dir, zip_file)
+
+package_layer = aws.lambda_.LayerVersion(f'{PROJECT_NAME}-layer',
+                                         layer_name=f"{PROJECT_NAME}-layer",
+                                         code=pulumi.AssetArchive({".": pulumi.FileArchive(f"{temp_dir}/{zip_file}")}),
+                                         compatible_runtimes=[aws.lambda_.Runtime.PYTHON3D12]
+                                         )
 
 
 def create_lambda_function(name, handler, code):
@@ -63,7 +69,9 @@ def create_lambda_function(name, handler, code):
 
 
 def create_asset_archive(update_file):
+    # TODO: Need a more recursive way to handle these file dependencies
     files = {"utils": pulumi.FileArchive("utils"),
+             "db.py": pulumi.FileAsset("db.py"),
              "download_file.py": pulumi.FileAsset("download_file.py"),
              "updater.py": pulumi.FileAsset("updater.py"),
              f"{update_file}.py": pulumi.FileAsset(f"{update_file}.py")}
diff --git a/utils/create_package_zip.py b/utils/create_package_zip.py
new file mode 100644
index 0000000..d32ca7a
--- /dev/null
+++ b/utils/create_package_zip.py
@@ -0,0 +1,29 @@
+import os
+import subprocess
+import tempfile
+
+
+def main(temp_dir, zip_file):
+    root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    package_dir = os.path.join(root_dir, 'venv')
+    excluded_patterns = ["**/__pycache__/*", "**/bin/*", "**/black*/", "**/pip*/*", "**/pulumi*/*"]
+
+    # Ensure the 'python' directory exists
+    os.makedirs(package_dir, exist_ok=True)
+    os.makedirs(temp_dir, exist_ok=True)
+
+    print(f"Temporary directory created at: {temp_dir}")
+    zip_path = os.path.join(temp_dir, zip_file)
+
+    # Construct the exclusion part of the zip command
+    exclusion_args = sum([["-x", pattern] for pattern in excluded_patterns], [])
+
+    # Zip the 'python' directory into the temporary directory
+    subprocess.run(["zip", "-r", zip_path, package_dir, *exclusion_args], stdout=subprocess.DEVNULL)
+
+    print(f"Zipped content is available at {zip_path}")
+    print(f"Zipped content size {os.path.getsize(zip_path)}")
+
+
+if __name__ == "__main__":
+    main(tempfile.TemporaryDirectory().name, 'python.zip')