From 23fc41a7cc38804c56a63bd4c1d6c21f89151344 Mon Sep 17 00:00:00 2001 From: Gregory Mermoud Date: Wed, 22 Jul 2026 08:52:03 +0000 Subject: [PATCH] feat(storage): profile-based S3 auth and hive-partitioned datasets Read/write data to S3-compatible storage (SwitchCloud) using ~/.aws/credentials profiles, plus a polars-based layer for hive-partitioned Parquet datasets that scale to large tables. - credentials.py: shared resolve_s3_credentials() used by both the boto3 client and the PyArrow filesystem. Precedence: named profile from ~/.aws/credentials, then explicit keys, then S3_* env vars. - dataset.py: get_s3_filesystem(), write_dataset(), read_dataset() for hive-partitioned Parquet with polars; filters prune partitions. - s3.py: get_s3_client() gains a `profile` argument. - storage/README.md: usage guide incl. local-CSV -> partitioned Parquet. - Add polars + pyarrow to the `storage` extra; tests for both. --- README.md | 24 ++- pyproject.toml | 5 +- src/simlab_tools/__init__.py | 3 + src/simlab_tools/storage/README.md | 123 ++++++++++++++++ src/simlab_tools/storage/__init__.py | 8 + src/simlab_tools/storage/credentials.py | 75 ++++++++++ src/simlab_tools/storage/dataset.py | 187 ++++++++++++++++++++++++ src/simlab_tools/storage/s3.py | 41 +++--- tests/test_dataset.py | 149 +++++++++++++++++++ tests/test_s3.py | 4 +- 10 files changed, 595 insertions(+), 24 deletions(-) create mode 100644 src/simlab_tools/storage/README.md create mode 100644 src/simlab_tools/storage/credentials.py create mode 100644 src/simlab_tools/storage/dataset.py create mode 100644 tests/test_dataset.py diff --git a/README.md b/README.md index fc7af4f..92d6ebf 100644 --- a/README.md +++ b/README.md @@ -27,18 +27,32 @@ Available extras: `storage`, `geo`, `smoothing`, `weather`, `all`, and `dev` ### `simlab_tools.storage` — object storage (extra: `storage`) -Transfer files to and from S3-compatible buckets (e.g. SwitchCloud) with -progress bars and multipart transfers. +Read and write data to S3-compatible buckets (e.g. SwitchCloud): transfer files +with progress bars and multipart transfers, or read/write Hive-partitioned +Parquet datasets with polars for scalability. Credentials come from your +`~/.aws/credentials` profiles. ```python from simlab_tools.storage import get_s3_client, download_files_from_bucket -client = get_s3_client("https://zhw-a.s3.cloud.switch.ch") # reads S3_ACCESS_KEY_ID / S3_SECRET_ACCESS_KEY +client = get_s3_client("https://zhw-a.s3.cloud.switch.ch", profile="switch") download_files_from_bucket(client, "ofen", "dejection_cones_dem", "data/", file_extensions=[".tif"]) ``` -`get_s3_client` accepts explicit `key_id` / `key_secret` arguments, or falls back -to the `S3_ACCESS_KEY_ID` and `S3_SECRET_ACCESS_KEY` environment variables. +```python +import pyarrow.dataset as ds +from simlab_tools.storage import get_s3_filesystem, write_dataset, read_dataset + +fs = get_s3_filesystem("https://zhw-a.s3.cloud.switch.ch", profile="switch") +write_dataset(df, fs, "research-data", "measurements", partition_cols=["station", "year"]) +recent = read_dataset(fs, "research-data", "measurements", filters=ds.field("year") >= 2024) +``` + +Credentials resolve from a named `profile` in `~/.aws/credentials`, then explicit +`key_id` / `key_secret` arguments, then the `S3_ACCESS_KEY_ID` / +`S3_SECRET_ACCESS_KEY` environment variables. See +[`storage/README.md`](src/simlab_tools/storage/README.md) for the full guide, +including a local-CSV-to-partitioned-Parquet walkthrough. For backwards compatibility, the storage functions are also importable from the top-level package: `from simlab_tools import get_s3_client`. diff --git a/pyproject.toml b/pyproject.toml index c264e0c..af2f480 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,11 +7,14 @@ requires-python = ">=3.11" dependencies = [] [project.optional-dependencies] -# Transfer files to/from S3-compatible object storage. +# Transfer files to/from S3-compatible object storage, including reading and +# writing hive-partitioned Parquet datasets. storage = [ "boto3>=1.35.0", "mypy-boto3-s3>=1.35.0", "tqdm>=4.66.0", + "polars>=1.0", + "pyarrow>=15.0", ] # Raster handling, terrain analysis and geodata API clients. geo = [ diff --git a/src/simlab_tools/__init__.py b/src/simlab_tools/__init__.py index 7631a22..ca4b0b6 100644 --- a/src/simlab_tools/__init__.py +++ b/src/simlab_tools/__init__.py @@ -26,6 +26,9 @@ "download_file", "download_files_from_bucket", "download_specific_files_from_bucket", + "get_s3_filesystem", + "read_dataset", + "write_dataset", } ) diff --git a/src/simlab_tools/storage/README.md b/src/simlab_tools/storage/README.md new file mode 100644 index 0000000..6a1af14 --- /dev/null +++ b/src/simlab_tools/storage/README.md @@ -0,0 +1,123 @@ +# `simlab_tools.storage` + +Read and write data to S3-compatible object storage (e.g. **SwitchCloud**), +with credentials taken from your `~/.aws/credentials` profiles. Two layers: + +- **File transfer** (`s3.py`) — upload/download individual files or whole + prefixes, with progress bars and multipart transfers. +- **Datasets** (`dataset.py`) — read/write [Hive-partitioned][hive] Parquet + datasets with [polars], so large tables stay scalable to query. + +Install the extra: + +```bash +uv pip install "simlab-tools[storage] @ git+https://github.com/simlab-vs/simlab-tools.git" +``` + +## Credentials + +All entry points resolve credentials the same way (see `credentials.py`), in +this order: + +1. **A named `profile`** from `~/.aws/credentials` — the recommended path. +2. Explicit `key_id` / `key_secret` arguments. +3. The `S3_ACCESS_KEY_ID` / `S3_SECRET_ACCESS_KEY` environment variables. + +A typical `~/.aws/credentials` looks like: + +```ini +[switch] +aws_access_key_id = AKIA................ +aws_secret_access_key = .................................... +``` + +You then only ever refer to the profile name (`"switch"`), never the keys +themselves. + +## File transfer + +```python +from pathlib import Path +from simlab_tools.storage import get_s3_client, upload_file, download_files_from_bucket + +client = get_s3_client("https://zhw-a.s3.cloud.switch.ch", profile="switch") + +# Upload one file. +upload_file(client, Path("dem.tif"), bucket="ofen", key="dems/dem.tif") + +# Download every .tif under a prefix, preserving the folder structure. +download_files_from_bucket( + client, "ofen", "dejection_cones_dem", Path("data/"), file_extensions=[".tif"] +) +``` + +Other helpers: `upload_files_to_bucket`, `download_file`, +`download_specific_files_from_bucket`, `list_bucket_contents`. + +## Hive-partitioned datasets + +A dataset is a directory of Parquet files laid out as `col=value/` folders: + +``` +s3://research-data/events/year=2024/month=07/part-0.parquet + year=2024/month=08/part-0.parquet +``` + +Readers prune whole folders before downloading anything, so a query filtered on +`year` never touches the other partitions — that is what keeps reads fast as the +dataset grows. + +### Example: local CSV → partitioned Parquet on S3 + +```python +import polars as pl +import pyarrow.dataset as ds +from simlab_tools.storage import get_s3_filesystem, write_dataset, read_dataset + +# 1. Point at the bucket using an ~/.aws/credentials profile. +fs = get_s3_filesystem("https://zhw-a.s3.cloud.switch.ch", profile="switch") + +# 2. Load a local CSV. +df = pl.read_csv("measurements.csv") # columns: station, year, month, value, ... + +# 3. Write it out as a Hive-partitioned Parquet dataset, partitioned by +# station and year. Each (station, year) pair becomes its own sub-folder. +write_dataset( + df, + fs, + bucket="research-data", + path="measurements", + partition_cols=["station", "year"], +) + +# 4. Read back only what you need. The filter on `year` (a partition column) +# prunes partitions before any file is fetched; `columns` limits the load. +recent = read_dataset( + fs, + bucket="research-data", + path="measurements", + columns=["station", "year", "value"], + filters=ds.field("year") >= 2024, +) +print(recent.head()) +``` + +`write_dataset` accepts either a polars `DataFrame` or a PyArrow `Table`, and +`read_dataset` returns a polars `DataFrame`. Filters are PyArrow dataset +expressions built with `pyarrow.dataset.field(...)`. + +### Local round-trips + +Both functions take any `pyarrow.fs.FileSystem`, so the same code works against +the local disk — handy for tests or prototyping: + +```python +from pyarrow.fs import LocalFileSystem + +fs = LocalFileSystem() +write_dataset(df, fs, bucket="/data/lake", path="measurements", + partition_cols=["year"]) +``` + +[hive]: https://arrow.apache.org/docs/python/dataset.html#partitioning +[polars]: https://pola.rs/ diff --git a/src/simlab_tools/storage/__init__.py b/src/simlab_tools/storage/__init__.py index 3f72c83..dab9d16 100644 --- a/src/simlab_tools/storage/__init__.py +++ b/src/simlab_tools/storage/__init__.py @@ -1,5 +1,10 @@ """Storage utilities: transferring files to and from S3-compatible buckets.""" +from simlab_tools.storage.dataset import ( + get_s3_filesystem, + read_dataset, + write_dataset, +) from simlab_tools.storage.s3 import ( download_file, download_files_from_bucket, @@ -18,4 +23,7 @@ "download_file", "download_files_from_bucket", "download_specific_files_from_bucket", + "get_s3_filesystem", + "read_dataset", + "write_dataset", ] diff --git a/src/simlab_tools/storage/credentials.py b/src/simlab_tools/storage/credentials.py new file mode 100644 index 0000000..2d10d3c --- /dev/null +++ b/src/simlab_tools/storage/credentials.py @@ -0,0 +1,75 @@ +"""Shared credential resolution for S3-compatible storage access. + +Both the boto3 client (:mod:`simlab_tools.storage.s3`) and the PyArrow +filesystem (:mod:`simlab_tools.storage.dataset`) resolve credentials the same +way, so the logic lives here once. +""" + +import os +from typing import NamedTuple + +import boto3 + + +class S3Credentials(NamedTuple): + """Resolved S3 credentials, ready to hand to boto3 or PyArrow.""" + + access_key: str | None + secret_key: str | None + session_token: str | None = None + + +def resolve_s3_credentials( + profile: str | None = None, + key_id: str | None = None, + key_secret: str | None = None, +) -> S3Credentials: + """Resolve S3 credentials from a profile, explicit keys or the environment. + + Credentials are resolved in this order: + + 1. A named ``profile`` from ``~/.aws/credentials`` (recommended). + 2. Explicit ``key_id`` / ``key_secret`` arguments. + 3. The ``S3_ACCESS_KEY_ID`` / ``S3_SECRET_ACCESS_KEY`` environment variables. + + Parameters + ---------- + profile: str, optional + Name of an AWS shared-credentials profile (a ``[section]`` in + ``~/.aws/credentials``). When given, the other sources are ignored. + key_id: str, optional + Access key id. Falls back to ``S3_ACCESS_KEY_ID`` if omitted. + key_secret: str, optional + Secret key. Falls back to ``S3_SECRET_ACCESS_KEY`` if omitted. + + Returns + ------- + S3Credentials + The resolved ``(access_key, secret_key, session_token)`` triple. The + session token is only set when it comes from a profile. + + Raises + ------ + KeyError + If neither a profile nor explicit credentials are provided and the + corresponding environment variables are not set. + ValueError + If the named profile exists but resolves to no credentials. + + """ + if profile is not None: + session = boto3.Session(profile_name=profile) + creds = session.get_credentials() + if creds is None: + raise ValueError(f"No credentials found for AWS profile {profile!r}") + frozen = creds.get_frozen_credentials() + return S3Credentials(frozen.access_key, frozen.secret_key, frozen.token) + + if key_id is not None and key_secret is not None: + return S3Credentials(key_id, key_secret) + + if key_id is None: + key_id = os.environ["S3_ACCESS_KEY_ID"] + if key_secret is None: + key_secret = os.environ["S3_SECRET_ACCESS_KEY"] + return S3Credentials(key_id, key_secret) diff --git a/src/simlab_tools/storage/dataset.py b/src/simlab_tools/storage/dataset.py new file mode 100644 index 0000000..58f9ff5 --- /dev/null +++ b/src/simlab_tools/storage/dataset.py @@ -0,0 +1,187 @@ +"""Read and write hive-partitioned tabular datasets on S3-compatible storage. + +A "dataset" here is a directory of Parquet files laid out with Hive-style +partitioning, i.e. one nested folder per partition-column value:: + + s3://bucket/prefix/year=2024/month=07/part-0.parquet + year=2024/month=08/part-0.parquet + +Partitioning lets readers prune whole folders before touching any file, so a +query filtered on ``year`` only downloads the matching partitions. This keeps +reads fast as the dataset grows, which is why it scales better than a single +monolithic file. +""" + +from typing import Sequence +from urllib.parse import urlparse + +import polars as pl +import pyarrow as pa +import pyarrow.dataset as ds +from pyarrow.fs import FileSystem, S3FileSystem + +from simlab_tools.storage.credentials import resolve_s3_credentials + +# A dataset can be written from either a polars DataFrame or an Arrow table. +TableLike = pl.DataFrame | pa.Table + + +def get_s3_filesystem( + endpoint_url: str, + profile: str | None = None, + key_id: str | None = None, + key_secret: str | None = None, + region: str = "us-east-1", +) -> S3FileSystem: + """Build a PyArrow S3 filesystem for an S3-compatible endpoint. + + This is the filesystem counterpart of + :func:`simlab_tools.storage.get_s3_client`: pass the returned object to + :func:`read_dataset` / :func:`write_dataset`. Credentials are resolved the + same way (``profile`` from ``~/.aws/credentials``, then explicit keys, then + ``S3_*`` environment variables). + + Parameters + ---------- + endpoint_url: str + Url to the S3 endpoint (e.g. ``https://zhw-a.s3.cloud.switch.ch``). + profile: str, optional + Name of an AWS shared-credentials profile in ``~/.aws/credentials``. + key_id: str, optional + Access key id. Falls back to ``S3_ACCESS_KEY_ID`` if omitted. + key_secret: str, optional + Secret key. Falls back to ``S3_SECRET_ACCESS_KEY`` if omitted. + region: str + Region reported to the endpoint. SwitchCloud ignores it, but the S3 + client requires a value; the default is usually fine. + + Returns + ------- + S3FileSystem + Filesystem bound to the given endpoint and credentials. + + """ + parsed = urlparse(endpoint_url) + # Accept both "https://host" and a bare "host". + scheme = parsed.scheme or "https" + host = parsed.netloc or parsed.path + + creds = resolve_s3_credentials(profile, key_id, key_secret) + + return S3FileSystem( + access_key=creds.access_key, + secret_key=creds.secret_key, + session_token=creds.session_token, + endpoint_override=host, + scheme=scheme, + region=region, + ) + + +def _base_dir(bucket: str, path: str) -> str: + """Join a bucket and prefix into a filesystem path (``bucket/prefix``). + + Only the trailing slash of ``bucket`` is trimmed, so absolute local paths + (which start with ``/``) survive intact for non-S3 filesystems. + """ + base = bucket.rstrip("/") + prefix = path.strip("/") + return f"{base}/{prefix}" if prefix else base + + +def write_dataset( + data: TableLike, + filesystem: FileSystem, + bucket: str, + path: str, + partition_cols: Sequence[str] | None = None, + existing_data_behavior: str = "overwrite_or_ignore", + basename_template: str | None = None, +) -> None: + """Write a table as a Hive-partitioned Parquet dataset. + + Parameters + ---------- + data: polars.DataFrame or pyarrow.Table + The table to write. A polars DataFrame is converted to Arrow first. + filesystem: pyarrow.fs.FileSystem + Target filesystem, typically from :func:`get_s3_filesystem`. + bucket: str + Name of the destination bucket. + path: str + Key prefix under which the dataset directory is created. + partition_cols: sequence of str, optional + Columns to partition by, in nesting order. Each distinct combination of + values becomes a ``col=value`` sub-directory. If ``None``, the whole + table is written unpartitioned. + existing_data_behavior: str + Passed through to :func:`pyarrow.dataset.write_dataset`. The default + ``"overwrite_or_ignore"`` leaves untouched partitions in place and + overwrites files it re-writes; use ``"delete_matching"`` to replace a + partition's contents wholesale. + basename_template: str, optional + Template for the written file names, e.g. ``"part-{i}.parquet"``. + + """ + table = data if isinstance(data, pa.Table) else data.to_arrow() + + partitioning = None + if partition_cols: + partitioning = ds.partitioning( + pa.schema([table.schema.field(col) for col in partition_cols]), + flavor="hive", + ) + + ds.write_dataset( + table, + base_dir=_base_dir(bucket, path), + filesystem=filesystem, + format="parquet", + partitioning=partitioning, + existing_data_behavior=existing_data_behavior, + basename_template=basename_template, + ) + + +def read_dataset( + filesystem: FileSystem, + bucket: str, + path: str, + columns: Sequence[str] | None = None, + filters: ds.Expression | None = None, +) -> pl.DataFrame: + """Read a Hive-partitioned Parquet dataset into a DataFrame. + + Partition columns are restored as ordinary columns, and ``filters`` on them + prune whole partitions before any file is read. + + Parameters + ---------- + filesystem: pyarrow.fs.FileSystem + Source filesystem, typically from :func:`get_s3_filesystem`. + bucket: str + Name of the bucket to read from. + path: str + Key prefix of the dataset directory. + columns: sequence of str, optional + Subset of columns to load. ``None`` loads all columns. + filters: pyarrow.dataset.Expression, optional + Row/partition filter, e.g. ``pyarrow.dataset.field("year") == 2024``. + + Returns + ------- + polars.DataFrame + The (optionally filtered and projected) dataset. + + """ + dataset = ds.dataset( + _base_dir(bucket, path), + filesystem=filesystem, + format="parquet", + partitioning="hive", + ) + table = dataset.to_table( + columns=list(columns) if columns is not None else None, + filter=filters, + ) + return pl.from_arrow(table) diff --git a/src/simlab_tools/storage/s3.py b/src/simlab_tools/storage/s3.py index b575f82..5e2703f 100644 --- a/src/simlab_tools/storage/s3.py +++ b/src/simlab_tools/storage/s3.py @@ -1,14 +1,14 @@ """Utilities to transfer files to and from S3-compatible buckets.""" -import os from pathlib import Path -from typing import Optional import boto3 from boto3.s3.transfer import TransferConfig from mypy_boto3_s3 import S3Client from tqdm import tqdm +from simlab_tools.storage.credentials import resolve_s3_credentials + # Default multipart transfer configuration shared by up- and downloads. _TRANSFER_CONFIG = TransferConfig( multipart_threshold=1024 * 25, # 25 MB threshold @@ -19,15 +19,26 @@ def get_s3_client( endpoint_url: str, - key_id: Optional[str] = None, - key_secret: Optional[str] = None, + profile: str | None = None, + key_id: str | None = None, + key_secret: str | None = None, ) -> S3Client: """Initialize an S3 client for an S3-compatible endpoint. + Credentials are resolved in this order: + + 1. A named ``profile`` from ``~/.aws/credentials`` (recommended). + 2. Explicit ``key_id`` / ``key_secret`` arguments. + 3. The ``S3_ACCESS_KEY_ID`` / ``S3_SECRET_ACCESS_KEY`` environment variables. + Parameters ---------- endpoint_url: str Url to the S3 endpoint (e.g. ``https://zhw-a.s3.cloud.switch.ch``). + profile: str, optional + Name of an AWS shared-credentials profile (a ``[section]`` in + ``~/.aws/credentials``). When given, ``key_id`` / ``key_secret`` and the + environment variables are ignored. key_id: str, optional Access key id for S3 access. If omitted, falls back to the ``S3_ACCESS_KEY_ID`` environment variable. @@ -43,24 +54,20 @@ def get_s3_client( Raises ------ KeyError - If credentials are not provided and the corresponding environment - variables are not set. + If neither a profile nor explicit credentials are provided and the + corresponding environment variables are not set. """ - if key_id is None: - key_id = os.environ["S3_ACCESS_KEY_ID"] - if key_secret is None: - key_secret = os.environ["S3_SECRET_ACCESS_KEY"] + creds = resolve_s3_credentials(profile, key_id, key_secret) - s3_client = boto3.client( + return boto3.client( "s3", - aws_access_key_id=key_id, - aws_secret_access_key=key_secret, + aws_access_key_id=creds.access_key, + aws_secret_access_key=creds.secret_key, + aws_session_token=creds.session_token, endpoint_url=endpoint_url, ) - return s3_client - def upload_file(client: S3Client, local_file: Path, bucket: str, key: str) -> None: """Upload a single file to S3. @@ -100,7 +107,7 @@ def upload_files_to_bucket( bucket: str, prefix: str, files: list[Path], - base_folder: Optional[Path] = None, + base_folder: Path | None = None, ) -> None: """Upload files to a given S3 bucket. @@ -117,7 +124,7 @@ def upload_files_to_bucket( Prefix to add to the S3 key of each uploaded file. files: list[Path] List of pathlib.Path objects for each of the files to upload. - base_folder: Optional[Path] + base_folder: Path | None If the files to upload are all situated in subfolder hierarchies of `base_folder`, this can be provided to specify that the subfolder hierarchy of `base_folder` should be re-created in the S3 bucket. diff --git a/tests/test_dataset.py b/tests/test_dataset.py new file mode 100644 index 0000000..1bb43bf --- /dev/null +++ b/tests/test_dataset.py @@ -0,0 +1,149 @@ +"""Tests for hive-partitioned dataset I/O and shared credential resolution.""" + +import polars as pl +import pyarrow.dataset as ds +import pytest +from polars.testing import assert_frame_equal +from pyarrow.fs import LocalFileSystem, S3FileSystem + +from simlab_tools.storage import get_s3_filesystem, read_dataset, write_dataset +from simlab_tools.storage.credentials import S3Credentials, resolve_s3_credentials + + +@pytest.fixture +def sample_df(): + """A small frame spanning two years and two months per year.""" + return pl.DataFrame( + { + "year": [2023, 2023, 2024, 2024], + "month": [1, 2, 1, 2], + "value": [10, 20, 30, 40], + } + ) + + +class TestDatasetRoundTrip: + """Round-trip writes/reads on a local filesystem (no network).""" + + def test_partitioned_round_trip(self, sample_df, tmp_path): + """A hive-partitioned dataset reads back identically.""" + fs = LocalFileSystem() + + write_dataset( + sample_df, fs, str(tmp_path), "events", partition_cols=["year", "month"] + ) + + result = read_dataset(fs, str(tmp_path), "events") + + # Order/column-order are not guaranteed across partitions, so sort. + assert_frame_equal( + result.select(sample_df.columns).sort("value"), + sample_df.sort("value"), + check_dtypes=False, + ) + + def test_partition_directories_are_hive_style(self, sample_df, tmp_path): + """Partition columns become ``col=value`` directories on disk.""" + write_dataset( + sample_df, + LocalFileSystem(), + str(tmp_path), + "events", + partition_cols=["year"], + ) + + years = {p.name for p in (tmp_path / "events").iterdir() if p.is_dir()} + assert years == {"year=2023", "year=2024"} + + def test_filter_prunes_partitions(self, sample_df, tmp_path): + """A filter on a partition column returns only matching rows.""" + fs = LocalFileSystem() + write_dataset(sample_df, fs, str(tmp_path), "events", partition_cols=["year"]) + + result = read_dataset( + fs, str(tmp_path), "events", filters=ds.field("year") == 2024 + ) + + assert result["value"].sort().to_list() == [30, 40] + + def test_column_projection(self, sample_df, tmp_path): + """Reading a column subset returns only those columns.""" + fs = LocalFileSystem() + write_dataset(sample_df, fs, str(tmp_path), "events", partition_cols=["year"]) + + result = read_dataset(fs, str(tmp_path), "events", columns=["value"]) + + assert result.columns == ["value"] + + def test_unpartitioned_round_trip(self, sample_df, tmp_path): + """Datasets can be written without any partition columns.""" + fs = LocalFileSystem() + write_dataset(sample_df, fs, str(tmp_path), "flat") + + result = read_dataset(fs, str(tmp_path), "flat") + + assert_frame_equal( + result.select(sample_df.columns).sort("value"), + sample_df.sort("value"), + check_dtypes=False, + ) + + +class TestResolveCredentials: + """Shared credential precedence used by both s3 and dataset helpers.""" + + def test_explicit_keys(self): + """Explicit keys are returned as-is with no session token.""" + creds = resolve_s3_credentials(key_id="id", key_secret="secret") + assert creds == S3Credentials("id", "secret", None) + + def test_env_fallback(self, monkeypatch): + """Credentials fall back to the S3_* environment variables.""" + monkeypatch.setenv("S3_ACCESS_KEY_ID", "env-id") + monkeypatch.setenv("S3_SECRET_ACCESS_KEY", "env-secret") + + creds = resolve_s3_credentials() + + assert creds.access_key == "env-id" + assert creds.secret_key == "env-secret" + + def test_missing_env_raises(self, monkeypatch): + """A missing credential env var raises KeyError.""" + monkeypatch.delenv("S3_ACCESS_KEY_ID", raising=False) + monkeypatch.delenv("S3_SECRET_ACCESS_KEY", raising=False) + + with pytest.raises(KeyError): + resolve_s3_credentials() + + def test_profile_from_credentials_file(self, monkeypatch, tmp_path): + """A named profile is read from the shared credentials file.""" + cred_file = tmp_path / "credentials" + cred_file.write_text( + "[switch]\n" + "aws_access_key_id = profile-id\n" + "aws_secret_access_key = profile-secret\n" + ) + monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", str(cred_file)) + + creds = resolve_s3_credentials(profile="switch") + + assert creds.access_key == "profile-id" + assert creds.secret_key == "profile-secret" + + +class TestGetS3Filesystem: + """Construction of the PyArrow S3 filesystem (no network calls).""" + + def test_builds_filesystem_from_endpoint(self): + """A filesystem is built with explicit credentials and endpoint.""" + fs = get_s3_filesystem( + "https://zhw-a.s3.cloud.switch.ch", key_id="id", key_secret="secret" + ) + assert isinstance(fs, S3FileSystem) + + def test_accepts_bare_host(self): + """An endpoint without a scheme is accepted (defaults to https).""" + fs = get_s3_filesystem( + "zhw-a.s3.cloud.switch.ch", key_id="id", key_secret="secret" + ) + assert isinstance(fs, S3FileSystem) diff --git a/tests/test_s3.py b/tests/test_s3.py index b49e8be..52a5d97 100644 --- a/tests/test_s3.py +++ b/tests/test_s3.py @@ -292,7 +292,9 @@ class TestGetS3Client: def test_explicit_credentials(self): """Explicit credentials produce a client bound to the given endpoint.""" - client = get_s3_client("https://example.invalid", "id", "secret") + client = get_s3_client( + "https://example.invalid", key_id="id", key_secret="secret" + ) assert client.meta.endpoint_url == "https://example.invalid" creds = client._request_signer._credentials