diff --git a/Makefile b/Makefile index a6f6bcc..9450197 100644 --- a/Makefile +++ b/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 . \ No newline at end of file + 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" \ No newline at end of file diff --git a/README.md b/README.md index a8087fb..0348e5f 100644 --- a/README.md +++ b/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= ``` +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: diff --git a/tests/conftest.py b/tests/conftest.py index 8b90f80..43ac867 100644 --- a/tests/conftest.py +++ b/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 diff --git a/timdex_dataset_api/dataset.py b/timdex_dataset_api/dataset.py index 9307e35..f835074 100644 --- a/timdex_dataset_api/dataset.py +++ b/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: diff --git a/timdex_dataset_api/metadata.py b/timdex_dataset_api/metadata.py index 0b2eff2..def7533 100644 --- a/timdex_dataset_api/metadata.py +++ b/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.