Add MSK IAM authentication via SASL/OAUTHBEARER

AGENT.md

@@ -80,6 +80,17 @@ while not shutdown:
 
 ## Key Design Decisions
 
+### Credential Isolation
+
+Millpond has two independent credential paths that must not interfere:
+
+- **Kafka (MSK)**: SASL/OAUTHBEARER via IRSA. The `aws-msk-iam-sasl-signer-python` library uses the standard AWS credential chain (IRSA projected token → STS → temporary creds).
+- **S3 (DuckLake data)**: Static IAM credentials via `DUCKDB_S3_ACCESS_KEY_ID` / `DUCKDB_S3_SECRET_ACCESS_KEY`. DuckDB's aws extension does not support IRSA ([duckdb/duckdb-aws#31](https://github.com/duckdb/duckdb-aws/issues/31)).
+
+The S3 credentials use DuckDB-specific env var names, **not** the standard `AWS_ACCESS_KEY_ID` / `AWS_SECRET_ACCESS_KEY`. This is intentional — standard AWS env vars would shadow IRSA in the credential chain and break MSK IAM auth. Do not change the S3 credential env var names to standard AWS names.
+
+`aws-msk-iam-sasl-signer-python` is an optional dependency (`pip install millpond[msk-iam]`) but the Dockerfile always installs it (`--extra msk-iam`). All production deployments use IAM auth. The optional dep is for local dev where boto3/botocore (~15MB) may not be wanted.
+
 ### Language: Python
 
 The hot path is all C/C++ (librdkafka, orjson, PyArrow, DuckDB). Python is glue — it touches each record once to pass a parsed dict into a list. Performance bottleneck is S3 write latency, not Python.

Dockerfile

@@ -5,12 +5,12 @@ COPY --from=uv /uv /usr/local/bin/uv
 
 WORKDIR /app
 COPY pyproject.toml uv.lock LICENSE ./
-RUN uv sync --frozen --no-dev --no-install-project
+RUN uv sync --frozen --no-dev --no-install-project --extra msk-iam
 
 COPY millpond/ millpond/
 ARG MILLPOND_VERSION=0.0.0.dev0
 ENV SETUPTOOLS_SCM_PRETEND_VERSION=$MILLPOND_VERSION
-RUN uv sync --frozen --no-dev
+RUN uv sync --frozen --no-dev --extra msk-iam
 
 # Install DuckDB CLI (pinned to match the Python package version)
 RUN uv tool install "duckdb-cli>=1.4,<1.5"

README.md

@@ -263,6 +263,19 @@ pipelines:
 
 One `range` over `pipelines` in the StatefulSet template produces N independent StatefulSets. Adding a pipeline is adding a block to `values.yaml` and running `helm upgrade`.
 
+## AWS Credential Isolation
+
+Millpond uses two separate AWS credential paths that must not interfere with each other:
+
+| Component | Auth | Credential source |
+|---|---|---|
+| Kafka (MSK) | SASL/OAUTHBEARER | IRSA (standard AWS credential chain) |
+| S3 (DuckLake data files) | Static IAM keys | `DUCKDB_S3_ACCESS_KEY_ID` / `DUCKDB_S3_SECRET_ACCESS_KEY` |
+
+DuckDB's [aws extension does not support IRSA](https://github.com/duckdb/duckdb-aws/issues/31) — it cannot perform the `AssumeRoleWithWebIdentity` token exchange that IRSA requires. Until that is resolved, S3 access requires static credentials.
+
+These static credentials are passed via DuckDB-specific env var names, not the standard `AWS_ACCESS_KEY_ID` / `AWS_SECRET_ACCESS_KEY`. This is deliberate: standard AWS env vars take precedence in the credential chain and would shadow the IRSA role used for Kafka authentication.
+
 ## Note
 This project should absolutely be called TableFowl, but that would be an [SEO](https://www.confluent.io/product/tableflow/) and linguistic palaver.
 

millpond/consumer.py

@@ -1,4 +1,5 @@
 import logging
+import os
 
 import orjson
 from confluent_kafka import Consumer, TopicPartition
@@ -10,6 +11,41 @@
 log = logging.getLogger(__name__)
 
 
+def _maybe_attach_oauth_cb(config: dict) -> dict:
+    """If SASL OAUTHBEARER is configured, attach the MSK IAM token callback."""
+    if "sasl.mechanism" in config and "sasl.mechanisms" not in config:
+        log.warning(
+            "sasl.mechanism (singular) is set but librdkafka uses sasl.mechanisms (plural). "
+            "Use KAFKA_CONSUMER_SASL_MECHANISMS instead of KAFKA_CONSUMER_SASL_MECHANISM."
+        )
+    sasl_mechanism = config.get("sasl.mechanisms")
+    if sasl_mechanism != "OAUTHBEARER":
+        return config
+    try:
+        from aws_msk_iam_auth import MSKAuthTokenProvider
+    except ImportError:
+        raise RuntimeError(
+            "aws-msk-iam-sasl-signer-python is required for OAUTHBEARER auth. "
+            "Install with: pip install millpond[msk-iam]"
+        ) from None
+
+    region = os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION")
+    if not region:
+        raise RuntimeError("AWS_REGION or AWS_DEFAULT_REGION must be set when using OAUTHBEARER auth")
+    log.info("MSK IAM auth enabled (region=%s)", region)
+
+    def oauth_cb(config_str):
+        try:
+            token, expiry_ms = MSKAuthTokenProvider.generate_auth_token(region)
+            return token, expiry_ms / 1000
+        except Exception:
+            log.exception("MSK IAM token generation failed")
+            raise
+
+    config["oauth_cb"] = oauth_cb
+    return config
+
+
 def _base_kafka_config(cfg: Config) -> dict:
     """Base config shared by all Kafka clients (AdminClient, Consumer)."""
     base = {
@@ -18,7 +54,7 @@ def _base_kafka_config(cfg: Config) -> dict:
     }
     for key, val in cfg.kafka_config_overrides:
         base[key] = val
-    return base
+    return _maybe_attach_oauth_cb(base)
 
 
 def _on_stats(stats_json: str) -> None:

pyproject.toml

@@ -22,6 +22,9 @@ dependencies = [
     "prometheus-client>=0.21",
 ]
 
+[project.optional-dependencies]
+msk-iam = ["aws-msk-iam-sasl-signer-python>=1.0.0"]
+
 [project.scripts]
 millpond = "millpond.main:main"
 

tests/unit/test_config.py

@@ -118,12 +118,12 @@ def test_partition_by_empty_string_treated_as_none(self, monkeypatch):
 
     def test_kafka_consumer_overrides(self, monkeypatch):
         monkeypatch.setenv("KAFKA_CONSUMER_SECURITY_PROTOCOL", "SASL_SSL")
-        monkeypatch.setenv("KAFKA_CONSUMER_SASL_MECHANISM", "PLAIN")
+        monkeypatch.setenv("KAFKA_CONSUMER_SASL_MECHANISMS", "PLAIN")
         monkeypatch.setenv("KAFKA_CONSUMER_SASL_USERNAME", "user1")
         cfg = load()
         overrides = dict(cfg.kafka_config_overrides)
         assert overrides["security.protocol"] == "SASL_SSL"
-        assert overrides["sasl.mechanism"] == "PLAIN"
+        assert overrides["sasl.mechanisms"] == "PLAIN"
         assert overrides["sasl.username"] == "user1"
 
     def test_kafka_consumer_overrides_default_empty(self):

tests/unit/test_consumer.py

@@ -3,7 +3,14 @@
 
 import pytest
 
-from millpond.consumer import _base_kafka_config, _on_stats, compute_assignment, create, discover_partition_count
+from millpond.consumer import (
+    _base_kafka_config,
+    _maybe_attach_oauth_cb,
+    _on_stats,
+    compute_assignment,
+    create,
+    discover_partition_count,
+)
 from millpond.config import Config
 
 
@@ -160,17 +167,70 @@ def test_includes_client_id(self):
         assert base["client.id"] == "millpond-test-topic-events-3"
 
     def test_includes_overrides(self):
-        cfg = _make_cfg(kafka_config_overrides=(("security.protocol", "SSL"), ("sasl.mechanism", "PLAIN")))
+        cfg = _make_cfg(kafka_config_overrides=(("security.protocol", "SSL"), ("sasl.mechanisms", "PLAIN")))
         base = _base_kafka_config(cfg)
         assert base["security.protocol"] == "SSL"
-        assert base["sasl.mechanism"] == "PLAIN"
+        assert base["sasl.mechanisms"] == "PLAIN"
 
     def test_empty_overrides(self):
         cfg = _make_cfg(kafka_config_overrides=())
         base = _base_kafka_config(cfg)
         assert "security.protocol" not in base
 
 
+class TestMaybeAttachOauthCb:
+    def test_noop_for_ssl_config(self):
+        config = {"bootstrap.servers": "localhost:9092", "security.protocol": "SSL"}
+        result = _maybe_attach_oauth_cb(config)
+        assert "oauth_cb" not in result
+
+    def test_noop_when_no_sasl(self):
+        config = {"bootstrap.servers": "localhost:9092"}
+        result = _maybe_attach_oauth_cb(config)
+        assert "oauth_cb" not in result
+
+    @patch.dict("os.environ", {"AWS_REGION": "us-east-1"})
+    def test_attaches_callback_for_oauthbearer(self):
+        mock_provider = MagicMock()
+        mock_provider.generate_auth_token.return_value = ("fake-token", 1700000000000)
+        with patch.dict("sys.modules", {"aws_msk_iam_auth": MagicMock(MSKAuthTokenProvider=mock_provider)}):
+            config = {"bootstrap.servers": "localhost:9092", "sasl.mechanisms": "OAUTHBEARER"}
+            result = _maybe_attach_oauth_cb(config)
+            assert "oauth_cb" in result
+            token, expiry = result["oauth_cb"]("")
+            assert token == "fake-token"
+            assert expiry == 1700000000.0
+
+    def test_singular_mechanism_ignored_with_warning(self, caplog):
+        config = {"bootstrap.servers": "localhost:9092", "sasl.mechanism": "OAUTHBEARER"}
+        result = _maybe_attach_oauth_cb(config)
+        assert "oauth_cb" not in result
+        assert "sasl.mechanisms (plural)" in caplog.text
+
+    def test_raises_when_signer_not_installed(self):
+        with patch.dict("sys.modules", {"aws_msk_iam_auth": None}):
+            config = {"bootstrap.servers": "localhost:9092", "sasl.mechanisms": "OAUTHBEARER"}
+            with pytest.raises(RuntimeError, match="aws-msk-iam-sasl-signer-python is required"):
+                _maybe_attach_oauth_cb(config)
+
+    @patch.dict("os.environ", {"AWS_REGION": "eu-central-1"})
+    def test_uses_aws_region_env(self):
+        mock_provider = MagicMock()
+        mock_provider.generate_auth_token.return_value = ("token", 1000)
+        with patch.dict("sys.modules", {"aws_msk_iam_auth": MagicMock(MSKAuthTokenProvider=mock_provider)}):
+            config = {"sasl.mechanisms": "OAUTHBEARER"}
+            result = _maybe_attach_oauth_cb(config)
+            result["oauth_cb"]("")
+            mock_provider.generate_auth_token.assert_called_with("eu-central-1")
+
+    @patch.dict("os.environ", {}, clear=True)
+    def test_raises_when_region_not_set(self):
+        with patch.dict("sys.modules", {"aws_msk_iam_auth": MagicMock()}):
+            config = {"sasl.mechanisms": "OAUTHBEARER"}
+            with pytest.raises(RuntimeError, match="AWS_REGION or AWS_DEFAULT_REGION must be set"):
+                _maybe_attach_oauth_cb(config)
+
+
 class TestDiscoverPartitionCount:
     @patch("millpond.consumer.AdminClient")
     def test_returns_partition_count(self, mock_admin_cls):