feat: duckdb extensions in lambda (#75)

@ceholden I took your good work in
#60
and reshaped it a bit to use the targeted extension fetching from
stac-utils/rustac-py#81. Lambda sizes:

- Before (current **main**): 29.2 MB
- With the **uv** and other tweaks, but no DuckDB extensions: 26.2 MB
- With DuckDB extensions: 93.6 MB

Things are broken, because I haven't wired up the server to actually use
those extensions, but I _think_ this is a proof out of using pre-fetched
DuckDB extensions. Lambda is here:
https://us-west-2.console.aws.amazon.com/lambda/home?region=us-west-2#/functions/stac-fastapi-geoparquet-labs-375-de-lambda8B5974B5-W1CWbXEHRA1Y?tab=code

> [!WARNING]
> When I tried to use the wheels directly, the size got too big. For
now, I'm sticking with a build-in-the-container (which is slow, ~10
minutes on my machine), but I need to dig in to why the build is smaller
than the wheel.

Author: gadomski

.github/workflows/ci.yml

@@ -9,13 +9,25 @@ on:
   release:
     types: [published]
 
+env:
+  duckdb-version: "1.2.0"
+
 jobs:
   lint:
     name: Lint
     runs-on: ubuntu-latest
+    env:
+      DUCKDB_LIB_DIR: ${{ github.workspace }}/opt/duckdb
+      LD_LIBRARY_PATH: ${{ github.workspace }}/opt/duckdb
+      DYLD_LIBRARY_PATH: ${{ github.workspace }}/opt/duckdb
     steps:
       - uses: actions/checkout@v4
       - uses: astral-sh/setup-uv@v5
+      - name: Install libduckdb
+        run: |
+          wget https://github.com/duckdb/duckdb/releases/download/v${{ env.duckdb-version }}/libduckdb-linux-amd64.zip
+          mkdir -p ${{ github.workspace }}/opt/duckdb
+          unzip libduckdb-linux-amd64.zip -d ${{ github.workspace }}/opt/duckdb
       - name: Sync uv
         run: uv sync --all-extras --all-groups
       - name: Install yarn deps

README.md

@@ -33,8 +33,8 @@ Then:
 
 ```shell
 source .venv/bin/activate
-cp .env.local .env
 cd infrastructure/aws
+cp .env.local .env
 # Make sure you're using the eoAPI sub-account
 aws sso login --profile eoapi && eval "$(aws configure export-credentials --profile eoapi --format env)" # or however you configure your AWS sessions
 cdk diff # to show any differences

infrastructure/aws/app.py

@@ -174,6 +174,8 @@ def __init__(
             ),
             environment={
                 "STAC_FASTAPI_GEOPARQUET_HREF": f"s3://{bucket.bucket_name}/{config.geoparquet_key}",
+                # find pre-fetched extensions
+                "STAC_FASTAPI_DUCKDB_EXTENSION_DIRECTORY": "/tmp/duckdb-extensions",
                 "HOME": "/tmp",  # for duckdb's home_directory
             },
         )

infrastructure/aws/lambda/Dockerfile

@@ -1,29 +1,35 @@
 ARG PYTHON_VERSION=3.12
+ARG BUILDPLATFORM=x86_64
 
-FROM public.ecr.aws/lambda/python:${PYTHON_VERSION}
-COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
+FROM --platform=${BUILDPLATFORM} ghcr.io/astral-sh/uv:0.6.6 AS uv
+FROM --platform=${BUILDPLATFORM} public.ecr.aws/lambda/python:${PYTHON_VERSION} AS builder
 
-# Install required utilities
-RUN dnf install -y findutils binutils git && \
+ENV UV_COMPILE_BYTECODE=1
+ENV UV_NO_INSTALLER_METADATA=1
+ENV UV_LINK_MODE=copy
+ENV PATH="/root/.cargo/bin:${PATH}"
+ENV MATURIN_PEP517_ARGS="--features=duckdb-bundled"
+
+RUN dnf install -y findutils binutils git gcc g++ && \
   dnf clean all && \
   rm -rf /var/cache/dnf
+RUN curl https://sh.rustup.rs -sSf | bash -s -- -y 
 
-WORKDIR /tmp
-COPY uv.lock uv.lock
-COPY pyproject.toml pyproject.toml
-COPY README.md README.md
-COPY src/ src/
-
-RUN uv pip install --compile-bytecode .[lambda] --target /asset
-
-# Reduce package size and remove useless files
+RUN --mount=from=uv,source=/uv,target=/bin/uv \
+  --mount=type=cache,target=/root/.cache/uv \
+  --mount=type=bind,source=uv.lock,target=uv.lock \
+  --mount=type=bind,source=pyproject.toml,target=pyproject.toml \
+  uv export --frozen --no-emit-workspace --no-dev --no-editable --extra lambda -o requirements.txt && \
+  uv pip install -r requirements.txt --target /asset
 WORKDIR /asset
+RUN python -c "from stacrs import DuckdbClient; DuckdbClient(use_s3_credential_chain=True, use_azure_credential_chain=False, install_extensions=True, extension_directory='/asset/duckdb-extensions')"
 RUN find . -type f -name '*.pyc' | while read f; do n=$(echo $f | sed 's/__pycache__\///' | sed 's/.cpython-[0-9]*//'); cp $f $n; done;
 RUN find . -type d -a -name '__pycache__' -print0 | xargs -0 rm -rf
 RUN find . -type f -a -name '*.py' -print0 | xargs -0 rm -f
 RUN find . -type d -a -name 'tests' -print0 | xargs -0 rm -rf
-
-# Strip debug symbols from compiled C/C++ code
 RUN find . -type f -name '*.so*' -exec strip --strip-unneeded {} \;
 
+FROM --platform=${BUILDPLATFORM} public.ecr.aws/lambda/python:${PYTHON_VERSION}
+WORKDIR /asset
+COPY --from=builder /asset /asset
 COPY infrastructure/aws/lambda/handler.py /asset/handler.py

infrastructure/aws/lambda/handler.py

@@ -1,10 +1,15 @@
 """AWS Lambda handler."""
 
 import logging
+import os
+import shutil
 
 from mangum import Mangum
 from stac_fastapi.geoparquet.main import app
 
+if not os.path.exists("/tmp/duckdb-extensions") and os.path.isdir("/duckdb-extensions"):
+    shutil.copytree("/duckdb-extensions", "/tmp/duckdb-extensions")
+
 logging.getLogger("mangum.lifespan").setLevel(logging.ERROR)
 logging.getLogger("mangum.http").setLevel(logging.ERROR)
 

pyproject.toml

@@ -4,9 +4,7 @@ version = "0.0.0"
 description = "Test the performance of stac-fastapi-geoparquet"
 readme = "README.md"
 requires-python = ">=3.12"
-dependencies = [
-    "stac-fastapi-geoparquet @ git+https://github.com/stac-utils/stac-fastapi-geoparquet",
-]
+dependencies = ["stac-fastapi-geoparquet", "stacrs"]
 
 [project.optional-dependencies]
 lambda = ["mangum==0.19.0"]
@@ -40,6 +38,10 @@ ignore_missing_imports = true
 [tool.ruff.lint]
 select = ["E", "F", "I"]
 
+[tool.uv.sources]
+stac-fastapi-geoparquet = { git = "https://github.com/stac-utils/stac-fastapi-geoparquet", branch = "main" }
+stacrs = { git = "https://github.com/stac-utils/stacrs", branch = "main" }
+
 [build-system]
 requires = ["setuptools"]
 build-backend = "setuptools.build_meta"