Merge pull request #165 from MITLibraries/TIMX-541-extension-installation-lambda-context

TIMX 541 - handle missing HOME env in AWS Lambda

Author: ghukill

tests/test_metadata.py

@@ -1,3 +1,5 @@
+# ruff: noqa: S105, S108
+
 import glob
 import os
 from pathlib import Path
@@ -262,3 +264,72 @@ def test_tdm_current_records_most_recent_version(timdex_metadata_with_deltas):
             == most_recent.iloc[0]["run_timestamp"]
         )
         assert current_version.iloc[0]["run_id"] == most_recent.iloc[0]["run_id"]
+
+
+def test_tdm_prepare_duckdb_secret_and_extensions_home_env_var_set_and_valid(
+    monkeypatch, tmp_path_factory, timdex_dataset_with_runs
+):
+    preset_home = tmp_path_factory.mktemp("my-account")
+    monkeypatch.setenv("HOME", str(preset_home))
+
+    tdm = TIMDEXDatasetMetadata(timdex_dataset_with_runs.location)
+    df = (
+        tdm.conn.query(
+            """
+        select
+            current_setting('secret_directory') as secret_directory,
+            current_setting('extension_directory') as extension_directory
+        ;
+        """
+        )
+        .to_df()
+        .iloc[0]
+    )
+    assert "my-account" in df.secret_directory
+    assert df.extension_directory == ""  # expected and okay when HOME set
+
+
+def test_tdm_prepare_duckdb_secret_and_extensions_home_env_var_unset(
+    monkeypatch, timdex_dataset_with_runs
+):
+    monkeypatch.delenv("HOME", raising=False)
+
+    tdm = TIMDEXDatasetMetadata(timdex_dataset_with_runs.location)
+
+    df = (
+        tdm.conn.query(
+            """
+        select
+            current_setting('secret_directory') as secret_directory,
+            current_setting('extension_directory') as extension_directory
+        ;
+        """
+        )
+        .to_df()
+        .iloc[0]
+    )
+    assert df.secret_directory == "/tmp/.duckdb/secrets"
+    assert df.extension_directory == "/tmp/.duckdb/extensions"
+
+
+def test_tdm_prepare_duckdb_secret_and_extensions_home_env_var_set_but_empty(
+    monkeypatch, timdex_dataset_with_runs
+):
+    monkeypatch.setenv("HOME", "")  # simulate AWS Lambda environment
+
+    tdm = TIMDEXDatasetMetadata(timdex_dataset_with_runs.location)
+
+    df = (
+        tdm.conn.query(
+            """
+        select
+            current_setting('secret_directory') as secret_directory,
+            current_setting('extension_directory') as extension_directory
+        ;
+        """
+        )
+        .to_df()
+        .iloc[0]
+    )
+    assert df.secret_directory == "/tmp/.duckdb/secrets"
+    assert df.extension_directory == "/tmp/.duckdb/extensions"

timdex_dataset_api/dataset.py

@@ -226,8 +226,8 @@ def get_s3_filesystem() -> fs.FileSystem:
     def setup_duckdb_context(self) -> DuckDBPyConnection:
         """Create a DuckDB connection that metadata and data query and retrieval.
 
-        This relies on TIMDEXDatasetMetadata.setup_duckdb_context() to produce a DuckDB
-        connection that has all metadata already created.
+        This method extends TIMDEXDatasetMetadata's pre-existing DuckDB connection, adding
+        a 'data' schema and any other configurations needed.
         """
         start_time = time.perf_counter()
 

timdex_dataset_api/metadata.py

@@ -148,9 +148,35 @@ def configure_duckdb_connection(self, conn: DuckDBPyConnection) -> None:
 
         These configurations include things like memory settings, AWS authentication, etc.
         """
+        self._install_duckdb_extensions(conn)
         self._configure_duckdb_s3_secret(conn)
         self._configure_duckdb_memory_profile(conn)
 
+    def _install_duckdb_extensions(self, conn: DuckDBPyConnection) -> None:
+        """Ensure DuckDB capable of installing extensions and install any required."""
+        # ensure secrets and extensions paths are accessible
+        home_env = os.getenv("HOME")
+        use_fallback_home = not home_env or not Path(home_env).is_dir()
+
+        if use_fallback_home:
+            duckdb_home = Path("/tmp/.duckdb")  # noqa: S108
+            secrets_dir = duckdb_home / "secrets"
+            extensions_dir = duckdb_home / "extensions"
+
+            secrets_dir.mkdir(parents=True, exist_ok=True)
+            extensions_dir.mkdir(parents=True, exist_ok=True)
+
+            conn.execute(f"set secret_directory='{secrets_dir.as_posix()}';")
+            conn.execute(f"set extension_directory='{extensions_dir.as_posix()}';")
+
+        # install HTTPFS extension
+        conn.execute(
+            """
+            install httpfs;
+            load httpfs;
+            """
+        )
+
     def _configure_duckdb_s3_secret(
         self,
         conn: DuckDBPyConnection,
@@ -161,14 +187,6 @@ def _configure_duckdb_s3_secret(
         If a scope is provided, e.g. an S3 URI prefix like 's3://timdex', set a scope
         parameter in the config.  Else, leave it blank.
         """
-        # install httpfs extension
-        conn.execute(
-            """
-            install httpfs;
-            load httpfs;
-            """
-        )
-
         # establish scope string
         scope_str = f", scope '{scope}'" if scope else ""