feat(storage): profile-based S3 auth and hive-partitioned datasets

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`.

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 = [

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",
     }
 )
 

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/

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",
 ]

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)