Merge pull request #154 from MITLibraries/TIMX-515-static-duckdb-file-prep

TIMX 515 - static DuckDB file prep

Author: ghukill

Makefile

@@ -1,3 +1,5 @@
+-include .env
+
 SHELL=/bin/bash
 DATETIME:=$(shell date -u +%Y%m%dT%H%M%SZ)
 
@@ -53,4 +55,17 @@ black-apply: # Apply changes with 'black'
 	pipenv run black .
 
 ruff-apply: # Resolve 'fixable errors' with 'ruff'
-	pipenv run ruff check --fix .
+	pipenv run ruff check --fix .
+
+
+######################
+# Minio S3 Instance
+######################
+minio-start:
+	docker run \
+   -p 9000:9000 \
+   -p 9001:9001 \
+   -v $(MINIO_DATA):/data \
+   -e "MINIO_ROOT_USER=$(MINIO_USERNAME)" \
+   -e "MINIO_ROOT_PASSWORD=$(MINIO_PASSWORD)" \
+   quay.io/minio/minio server /data --console-address ":9001"

README.md

@@ -45,8 +45,31 @@ None at this time.
 ```shell
 TDA_LOG_LEVEL=# log level for timdex-dataset-api, accepts [DEBUG, INFO, WARNING, ERROR], default INFO
 WARNING_ONLY_LOGGERS=# Comma-seperated list of logger names to set as WARNING only, e.g. 'botocore,charset_normalizer,smart_open'
+
+MINIO_S3_ENDPOINT_URL=# If set, informs the library to use this Minio S3 instance.  Requires the http(s):// protocol. 
+MINIO_USERNAME=# Username / AWS Key for Minio; required when MINIO_S3_ENDPOINT_URL is set
+MINIO_PASSWORD=# Pasword / AWS Secret for Minio; required when MINIO_S3_ENDPOINT_URL is set
+MINIO_DATA=# Path to persist MinIO data if started via Makefile command 
+```
+
+## Local S3 via MinIO
+
+Set env vars:
+```shell
+MINIO_S3_ENDPOINT_URL=http://localhost:9000
+MINIO_USERNAME="admin"
+MINIO_PASSWORD="password"
+MINIO_DATA=<path to persist MinIO data, e.g. /tmp/minio>
 ```
 
+Use a `Makefile` command that will start a MinIO instance:
+
+```shell
+make minio-start
+```
+
+With the env var `MINIO_S3_ENDPOINT_URL` set, this library will configure `pyarrow` and DuckDB connections to point at this local MinIO S3 instance.
+
 ## Usage
 
 Currently, the most common use cases are:

tests/conftest.py

@@ -21,6 +21,7 @@ def _test_env(monkeypatch):
     monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "fake_secret_key")
     monkeypatch.setenv("AWS_SESSION_TOKEN", "fake_session_token")
     monkeypatch.setenv("AWS_DEFAULT_REGION", "us-east-1")
+    monkeypatch.delenv("MINIO_S3_ENDPOINT_URL", raising=False)
 
 
 @pytest.fixture

timdex_dataset_api/dataset.py

@@ -128,7 +128,7 @@ def __init__(
 
         # reading
         self._current_records: bool = False
-        self.timdex_dataset_metadata: TIMDEXDatasetMetadata = None  # type: ignore[assignment]
+        self.metadata: TIMDEXDatasetMetadata = None  # type: ignore[assignment]
 
     @property
     def row_count(self) -> int:
@@ -173,8 +173,8 @@ def load(
         # read dataset metadata if only current records are requested
         self._current_records = current_records
         if current_records:
-            self.timdex_dataset_metadata = TIMDEXDatasetMetadata(timdex_dataset=self)
-            self.paths = self.timdex_dataset_metadata.get_current_parquet_files(**filters)
+            self.metadata = TIMDEXDatasetMetadata(timdex_dataset=self)
+            self.paths = self.metadata.get_current_parquet_files(**filters)
 
         # perform initial load of full dataset
         self.dataset = self._load_pyarrow_dataset()
@@ -285,11 +285,23 @@ def _parse_date_filters(self, run_date: str | date | None) -> DatasetFilters:
 
     @staticmethod
     def get_s3_filesystem() -> fs.FileSystem:
-        """Instantiate a pyarrow S3 Filesystem for dataset loading."""
+        """Instantiate a pyarrow S3 Filesystem for dataset loading.
+
+        If the env var 'MINIO_S3_ENDPOINT_URL' is present, assume a local MinIO S3
+        instance and configure accordingly, otherwise assume normal AWS S3.
+        """
         session = boto3.session.Session()
         credentials = session.get_credentials()
         if not credentials:
             raise RuntimeError("Could not locate AWS credentials")
+
+        if os.getenv("MINIO_S3_ENDPOINT_URL"):
+            return fs.S3FileSystem(
+                access_key=os.environ["MINIO_USERNAME"],
+                secret_key=os.environ["MINIO_PASSWORD"],
+                endpoint_override=os.environ["MINIO_S3_ENDPOINT_URL"],
+            )
+
         return fs.S3FileSystem(
             secret_key=credentials.secret_key,
             access_key=credentials.access_key,
@@ -509,9 +521,7 @@ def _yield_current_record_batches(
             - filters: pairs of column:value to filter the dataset metadata required
         """
         # get map of timdex_record_id to run_id for current version of that record
-        record_to_run_map = self.timdex_dataset_metadata.get_current_record_to_run_map(
-            **filters
-        )
+        record_to_run_map = self.metadata.get_current_record_to_run_map(**filters)
 
         # loop through batches, yielding only current records
         for batch in batches:

timdex_dataset_api/metadata.py

@@ -1,7 +1,9 @@
 """timdex_dataset_api/metadata.py"""
 
+import os
 import time
 from typing import TYPE_CHECKING, Unpack
+from urllib.parse import urlparse
 
 import duckdb
 
@@ -87,8 +89,8 @@ def _setup_database(self) -> None:
         # bump threads for high parallelization of lightweight data calls for metadata
         self.set_database_thread_usage(64)
 
-        # setup AWS credentials chain
-        self._create_aws_credential_chain()
+        # configure s3 connection
+        self._configure_s3_connection()
 
         # create a table of metadata about all rows in dataset
         self._create_full_dataset_table()
@@ -101,21 +103,41 @@ def _setup_database(self) -> None:
             f"path: '{self.db_path}'"
         )
 
-    def _create_aws_credential_chain(self) -> None:
-        """Setup AWS credentials chain in database connection.
+    def _configure_s3_connection(self) -> None:
+        """Configure S3 connection for DuckDB access.
 
-        https://duckdb.org/docs/stable/core_extensions/aws.html
+        If the env var 'MINIO_S3_ENDPOINT_URL' is present, assume a local MinIO S3
+        instance and configure accordingly, otherwise assume normal AWS S3 and setup a
+        credentials chain in DuckDB.
         """
-        logger.info("setting up AWS credentials chain")
-        query = """
-        create or replace secret secret (
-            type s3,
-            provider credential_chain,
-            chain 'sso;env;config',
-            refresh true
-        );
-        """
-        self.conn.execute(query)
+        logger.info("configuring S3 connection")
+
+        if os.getenv("MINIO_S3_ENDPOINT_URL"):
+            self.conn.execute(
+                f"""
+                create or replace secret minio_s3_secret (
+                    type s3,
+                    endpoint '{urlparse(os.environ["MINIO_S3_ENDPOINT_URL"]).netloc}',
+                    key_id '{os.environ["MINIO_USERNAME"]}',
+                    secret '{os.environ["MINIO_PASSWORD"]}',
+                    region 'us-east-1',
+                    url_style 'path',
+                    use_ssl false
+                );
+                """
+            )
+
+        else:
+            self.conn.execute(
+                """
+                create or replace secret aws_s3_secret (
+                    type s3,
+                    provider credential_chain,
+                    chain 'sso;env;config',
+                    refresh true
+                );
+                """
+            )
 
     def _create_full_dataset_table(self) -> None:
         """Create a table of metadata about all records in the parquet dataset.