Add credentials refresh in more locations#165
Conversation
|
@J-Meyers: Thanks for the PR, just an heads up that it might take a bit for this to get properly through review process, but we'll get to it. Testing is indeed not trivial, unsure what's the way forward, but I'll think of something |
… expiry DuckDB's httpfs extension snapshots STS credentials into a secret at creation time and never re-queries the provider chain, even with `refresh: auto` configured. When using `CREDENTIAL_CHAIN` with `chain: sts` (AssumeRole), DuckDB requests the AWS minimum session duration of 900 seconds (15 minutes). Any SQLMesh plan that takes longer than 15 minutes to complete — such as large incremental backfills — fails with HTTP 400 (Bad Request) from S3 once the token expires. This was confirmed through controlled experiments (binary search narrowed the threshold to exactly ~15 minutes, and `duckdb_secrets()` showed the same `key_id` throughout a run — no refresh ever occurs). The upstream fix (duckdb/duckdb-httpfs#165) adds credential refresh at the httpfs layer but has not been merged yet. This patch adds a timer-based secret refresh mechanism to the DuckDB engine adapter as a workaround: - Before each SQL execution, checks if 12 minutes (80% of the 900s TTL) have elapsed since the last secret creation - If so, queries `duckdb_secrets()` for existing S3 secret names, drops them, and recreates from the original config — forcing a fresh STS AssumeRole call - Uses double-check locking to prevent concurrent refresh when `concurrent_tasks > 1` - Zero overhead for configs without S3 secrets (early return on null check) and minimal overhead otherwise (monotonic clock comparison on hot path) Changes: - sqlmesh/core/engine_adapter/duckdb.py: Add __init__, _execute override, and secret refresh methods to DuckDBEngineAdapter - sqlmesh/core/config/connection.py: Add _extra_engine_config to DuckDBConnectionConfig to pass secrets config to the adapter This workaround can be removed once the upstream duckdb-httpfs fix lands and we upgrade DuckDB. Refs: duckdb/duckdb-httpfs#165 Refs: duckdb/duckdb-aws#26
|
When can we expect this to be merged? How can I help to test this? @carlopi |
|
This also bites DuckLake's ability to auto-refresh. The chain:
Why this matters: today's only auto-refresh hook lives in ReproductionINSTALL postgres; LOAD postgres;
INSTALL ducklake; LOAD ducklake;
FORCE INSTALL aws FROM core_nightly; LOAD aws;
INSTALL httpfs; LOAD httpfs;
CREATE SECRET aws_s3 (
TYPE S3, PROVIDER credential_chain, CHAIN 'web_identity',
ASSUME_ROLE_ARN '<role>',
WEB_IDENTITY_TOKEN_FILE '<jwt-file>',
REGION 'eu-west-1'
);
ATTACH 'ducklake:postgres:host=<pg> port=5432 dbname=<db>' AS x;
EDIT: same pattern on duckdb-iceberg: e.g. https://github.com/duckdb/duckdb-iceberg/blob/b0d93913629a08b91af90da9ae7e89d846cd22ec/src/planning/iceberg_multi_file_list.cpp#L1127-L1132. |
|
This is a small shell-driver integration test that exercises the How the test works: Two service accounts under one parent user. Phase A reads with svcacct A (populates How to run# 1. minio -- same image pinned in `scripts/minio_s3.yml`
# (what `scripts/run_s3_test_server.sh` boots in CI via docker compose).
docker run -d --name duckdb-minio -p 9000:9000 \
-e MINIO_ROOT_USER=duckdb_minio_admin \
-e MINIO_ROOT_PASSWORD=duckdb_minio_admin_password \
minio/minio:RELEASE.2021-11-03T03-36-36Z server /data
until curl -sf http://localhost:9000/minio/health/live >/dev/null; do sleep 1; done
MC="docker run --rm --network=host \
-e MC_HOST_local=http://duckdb_minio_admin:duckdb_minio_admin_password@localhost:9000 \
minio/mc:RELEASE.2021-11-05T10-05-06Z"
$MC admin user add local minio_duckdb_user minio_duckdb_user_password
$MC admin policy set local readwrite user=minio_duckdb_user
$MC mb local/test-bucket
# 2. build + run the test
make release
ENDPOINT=localhost:9000 URL_STYLE=path \
DUCKDB=./build/release/duckdb \
./scripts/refresh_get_range_test.shI hope that helps moving this PR forward. |
PR duckdb#165 introduced TryRefreshCredentials, which mutates six string fields on a shared S3AuthParams field-by-field with no synchronisation. Concurrent signing code in CreateS3Header reads those fields without synchronisation either. Two distinct races result: * Reader-vs-refresher: signers can observe a tuple stitched from two different refresh generations -> SignatureDoesNotMatch under genuine STS rotation, and a data race / UB under the C++ memory model regardless of measured failure rate. * Refresher-vs-refresher: N parallel workers each independently perform an STS round-trip and each rewrite the shared fields; the secret manager's internal lock keeps the catalog entry consistent but the auth_params writes themselves still race. Fix: S3AuthParams becomes a plain immutable data tuple. The slot that holds the currently-active tuple lives on S3HTTPInput as a duckdb shared_ptr<S3AuthParams>, mutated only by atomic_store. Readers obtain a consistent snapshot via S3HTTPInput::Snapshot() (one atomic_load) and read fields off the const tuple without further synchronisation. S3HTTPInput::ExecuteWithRefresh wraps the request lambdas: it hands them a `const S3AuthParams &` snapshot, catches HTTP/IO failures, attempts a single-flight refresh (double-checked-locking on refresh_mutex), and retries the lambda once against the fresh snapshot. Multiple threads encountering 401/403 simultaneously now produce at most one STS round-trip per S3HTTPInput instance instead of N. Cross-S3HTTPInput dedup (parallel readers each opening their own handle) is not addressed here; that belongs at the SecretManager layer.
| ErrorData error(ex); | ||
| // Only attempt refresh for HTTP or IO exceptions (same logic as S3FileHandle::Initialize) | ||
| if (error.Type() == ExceptionType::IO || error.Type() == ExceptionType::HTTP) { | ||
| if (auth_params.TryRefreshCredentials()) { |
There was a problem hiding this comment.
thread safety: n threads expiring simultaneously will probably each issue their own STS call. This might require a mutex around the Credentials Refresh call and some way to detect that a refresh already happened.
| this->access_key_id = refreshed_params.access_key_id; | ||
| this->secret_access_key = refreshed_params.secret_access_key; | ||
| this->session_token = refreshed_params.session_token; | ||
| this->region = refreshed_params.region; | ||
| this->endpoint = refreshed_params.endpoint; | ||
| this->oauth2_bearer_token = refreshed_params.oauth2_bearer_token; |
There was a problem hiding this comment.
this might cause a data race -- a reading thread accessing e.g. the updated access_key_id before the secret_access_key is being written by the refreshing thread.
Might not be a big functional problem if there's the mutex around the refresh: either the reader runs with an expired token into the refresh mutex, or it does with a invalid, partially updated set of credentials and runs into the refresh mutex (assuming the exception handling for expired tokens == invalid signature).
Not sure if this is a problem, but could be solved by swapping all credential attributes atomically.
| -> decltype(request_func(auth_params)) { | ||
| try { | ||
| return request_func(auth_params); | ||
| } catch (std::exception &ex) { |
There was a problem hiding this comment.
It looks like HTTPFileSystem::PostRequest doesn't throw on 4xx [1] (and neither does Put/Delete/Head [2] [3]), so this also needs a check on 4xx. Only GetRequest does [4], and that's why the test I posted succeded...
[1] https://github.com/duckdb/duckdb-httpfs/blob/main/src/httpfs.cpp#L164
[2] https://github.com/duckdb/duckdb-httpfs/blob/main/src/httpfs.cpp#L167
[3] https://github.com/duckdb/duckdb-httpfs/blob/main/src/httpfs.cpp#L193
[4] https://github.com/duckdb/duckdb-httpfs/blob/main/src/httpfs.cpp#L242
This is very difficult to test since you must wait for the credentials to timeout, however there were a variety of locations and scenarios that previously did not automatically refresh the credentials.
After the first
S3FileHandlewasInitialized then many operations like GLOBs would never have an opportunity to refresh the handle.This adds in autorefresh around all calls, the core of this is the
ExecuteWithRefreshand keeping around theFileOpener.ExecuteWithRefreshis used to just wrap the old code exactly as it was with a mechanism to automatically refresh the credentials and then run it again.The previous refreshing in
S3FileHandle::Initializeis now removed since it is now managed at each of the lower level function calls.Should close:
duckdb/duckdb-aws#26
duckdb/duckdb-aws#64
duckdb/duckdb-aws#97
If you have any advice on how to create an automated test for this I'd appreciate it