From 004251845d20fe66bb513b69e0fb2c70f7ea89ab Mon Sep 17 00:00:00 2001 From: Paolo Quadri Date: Thu, 20 Nov 2025 21:32:03 +0100 Subject: [PATCH 1/4] feat: implement rds auth via aws-cli --- src/include/storage/postgres_catalog.hpp | 6 +- .../storage/postgres_connection_pool.hpp | 8 +- src/postgres_extension.cpp | 6 + src/postgres_scanner.cpp | 4 +- src/postgres_storage.cpp | 3 +- src/storage/postgres_catalog.cpp | 120 +++++++++++++++++- src/storage/postgres_connection_pool.cpp | 21 +-- src/storage/postgres_transaction.cpp | 2 +- 8 files changed, 149 insertions(+), 21 deletions(-) diff --git a/src/include/storage/postgres_catalog.hpp b/src/include/storage/postgres_catalog.hpp index d6f273a9a..5fff4feea 100644 --- a/src/include/storage/postgres_catalog.hpp +++ b/src/include/storage/postgres_catalog.hpp @@ -21,11 +21,13 @@ class PostgresSchemaEntry; class PostgresCatalog : public Catalog { public: explicit PostgresCatalog(AttachedDatabase &db_p, string connection_string, string attach_path, - AccessMode access_mode, string schema_to_load, PostgresIsolationLevel isolation_level); + AccessMode access_mode, string schema_to_load, PostgresIsolationLevel isolation_level, + string secret_name = string()); ~PostgresCatalog(); string connection_string; string attach_path; + string secret_name; AccessMode access_mode; PostgresIsolationLevel isolation_level; @@ -40,6 +42,8 @@ class PostgresCatalog : public Catalog { static string GetConnectionString(ClientContext &context, const string &attach_path, string secret_name); + string GetFreshConnectionString(ClientContext &context); + optional_ptr CreateSchema(CatalogTransaction transaction, CreateSchemaInfo &info) override; void ScanSchemas(ClientContext &context, std::function callback) override; diff --git a/src/include/storage/postgres_connection_pool.hpp b/src/include/storage/postgres_connection_pool.hpp index 313598c92..6711cf254 100644 --- a/src/include/storage/postgres_connection_pool.hpp +++ b/src/include/storage/postgres_connection_pool.hpp @@ -44,10 +44,10 @@ class PostgresConnectionPool { PostgresConnectionPool(PostgresCatalog &postgres_catalog, idx_t maximum_connections = DEFAULT_MAX_CONNECTIONS); public: - bool TryGetConnection(PostgresPoolConnection &connection); - PostgresPoolConnection GetConnection(); + bool TryGetConnection(PostgresPoolConnection &connection, optional_ptr context = nullptr); + PostgresPoolConnection GetConnection(optional_ptr context = nullptr); //! Always returns a connection - even if the connection slots are exhausted - PostgresPoolConnection ForceGetConnection(); + PostgresPoolConnection ForceGetConnection(optional_ptr context = nullptr); void ReturnConnection(PostgresConnection connection); void SetMaximumConnections(idx_t new_max); @@ -61,7 +61,7 @@ class PostgresConnectionPool { vector connection_cache; private: - PostgresPoolConnection GetConnectionInternal(); + PostgresPoolConnection GetConnectionInternal(optional_ptr context = nullptr); }; } // namespace duckdb diff --git a/src/postgres_extension.cpp b/src/postgres_extension.cpp index 42465559f..ed9a89671 100644 --- a/src/postgres_extension.cpp +++ b/src/postgres_extension.cpp @@ -94,6 +94,10 @@ unique_ptr CreatePostgresSecretFunction(ClientContext &context, Crea result->secret_map["port"] = named_param.second.ToString(); } else if (lower_name == "passfile") { result->secret_map["passfile"] = named_param.second.ToString(); + } else if (lower_name == "use_rds_iam_auth") { + result->secret_map["use_rds_iam_auth"] = named_param.second.ToString(); + } else if (lower_name == "aws_region") { + result->secret_map["aws_region"] = named_param.second.ToString(); } else { throw InternalException("Unknown named parameter passed to CreatePostgresSecretFunction: " + lower_name); } @@ -112,6 +116,8 @@ void SetPostgresSecretParameters(CreateSecretFunction &function) { function.named_parameters["database"] = LogicalType::VARCHAR; // alias for dbname function.named_parameters["dbname"] = LogicalType::VARCHAR; function.named_parameters["passfile"] = LogicalType::VARCHAR; + function.named_parameters["use_rds_iam_auth"] = LogicalType::BOOLEAN; + function.named_parameters["aws_region"] = LogicalType::VARCHAR; } void SetPostgresNullByteReplacement(ClientContext &context, SetScope scope, Value ¶meter) { diff --git a/src/postgres_scanner.cpp b/src/postgres_scanner.cpp index bb4136f41..786b44916 100644 --- a/src/postgres_scanner.cpp +++ b/src/postgres_scanner.cpp @@ -391,7 +391,7 @@ bool PostgresGlobalState::TryOpenNewConnection(ClientContext &context, PostgresL } else { // we cannot use the main thread but we haven't initiated ANY scan yet // we HAVE to open a new connection - lstate.pool_connection = pg_catalog->GetConnectionPool().ForceGetConnection(); + lstate.pool_connection = pg_catalog->GetConnectionPool().ForceGetConnection(&context); lstate.connection = PostgresConnection(lstate.pool_connection.GetConnection().GetConnection()); } used_main_thread = true; @@ -400,7 +400,7 @@ bool PostgresGlobalState::TryOpenNewConnection(ClientContext &context, PostgresL } if (pg_catalog) { - if (!pg_catalog->GetConnectionPool().TryGetConnection(lstate.pool_connection)) { + if (!pg_catalog->GetConnectionPool().TryGetConnection(lstate.pool_connection, &context)) { return false; } lstate.connection = PostgresConnection(lstate.pool_connection.GetConnection().GetConnection()); diff --git a/src/postgres_storage.cpp b/src/postgres_storage.cpp index cfff7466e..cf3b9c122 100644 --- a/src/postgres_storage.cpp +++ b/src/postgres_storage.cpp @@ -45,7 +45,8 @@ static unique_ptr PostgresAttach(optional_ptr sto } auto connection_string = PostgresCatalog::GetConnectionString(context, attach_path, secret_name); return make_uniq(db, std::move(connection_string), std::move(attach_path), - attach_options.access_mode, std::move(schema_to_load), isolation_level); + attach_options.access_mode, std::move(schema_to_load), isolation_level, + std::move(secret_name)); } static unique_ptr PostgresCreateTransactionManager(optional_ptr storage_info, diff --git a/src/storage/postgres_catalog.cpp b/src/storage/postgres_catalog.cpp index f91c1f245..b7b5e1a2e 100644 --- a/src/storage/postgres_catalog.cpp +++ b/src/storage/postgres_catalog.cpp @@ -7,14 +7,17 @@ #include "duckdb/parser/parsed_data/create_schema_info.hpp" #include "duckdb/main/attached_database.hpp" #include "duckdb/main/secret/secret_manager.hpp" +#include "duckdb/common/printer.hpp" +#include namespace duckdb { PostgresCatalog::PostgresCatalog(AttachedDatabase &db_p, string connection_string_p, string attach_path_p, - AccessMode access_mode, string schema_to_load, PostgresIsolationLevel isolation_level) + AccessMode access_mode, string schema_to_load, PostgresIsolationLevel isolation_level, + string secret_name_p) : Catalog(db_p), connection_string(std::move(connection_string_p)), attach_path(std::move(attach_path_p)), - access_mode(access_mode), isolation_level(isolation_level), schemas(*this, schema_to_load), - connection_pool(*this), default_schema(schema_to_load) { + secret_name(std::move(secret_name_p)), access_mode(access_mode), isolation_level(isolation_level), + schemas(*this, schema_to_load), connection_pool(*this), default_schema(schema_to_load) { if (default_schema.empty()) { default_schema = "public"; } @@ -72,6 +75,71 @@ unique_ptr GetSecret(ClientContext &context, const string &secret_n return nullptr; } +string GenerateRdsAuthToken(const string &hostname, const string &port, const string &username, + const string &aws_region) { + + auto escape_shell_arg = [](const string &arg) -> string { + string escaped = "'"; + for (char c : arg) { + if (c == '\'') { + escaped += "'\\''"; + } else { + escaped += c; + } + } + escaped += "'"; + return escaped; + }; + + string command = "aws rds generate-db-auth-token --hostname " + escape_shell_arg(hostname) + + " --port " + escape_shell_arg(port) + " --username " + escape_shell_arg(username); + + if (!aws_region.empty()) { + command += " --region " + escape_shell_arg(aws_region); + } + + command += " 2>&1"; + + FILE *pipe = popen(command.c_str(), "r"); + if (!pipe) { + throw IOException("Failed to execute AWS CLI command to generate RDS auth token. " + "Make sure AWS CLI is installed and configured."); + } + + string token; + char buffer[128]; + while (fgets(buffer, sizeof(buffer), pipe) != nullptr) { + token += buffer; + } + + int status = pclose(pipe); + + if (!token.empty() && token.back() == '\n') { + token.pop_back(); + } + if (status != 0) { + + throw IOException("Failed to generate RDS auth token: %s. " + "Make sure AWS CLI is installed, configured, and you have the necessary permissions.", + token.empty() ? "Unknown error" : token.c_str()); + } + + if (!token.empty() && token.back() == '\n') { + token.pop_back(); + } + + + if (PostgresConnection::DebugPrintQueries()) { + string debug_msg = StringUtil::Format( + "[RDS IAM Auth] Generated auth token for hostname=%s, port=%s, username=%s, region=%s\n, token=%s", + hostname.c_str(), port.c_str(), username.c_str(), + aws_region.empty() ? "(default)" : aws_region.c_str(), token.c_str()); + Printer::Print(debug_msg); + } + + return token; +} + string PostgresCatalog::GetConnectionString(ClientContext &context, const string &attach_path, string secret_name) { // if no secret is specified we default to the unnamed postgres secret, if it exists string connection_string = attach_path; @@ -87,8 +155,48 @@ string PostgresCatalog::GetConnectionString(ClientContext &context, const string const auto &kv_secret = dynamic_cast(*secret_entry->secret); string new_connection_info; + // Check if RDS IAM authentication is enabled + Value use_rds_iam_auth_val = kv_secret.TryGetValue("use_rds_iam_auth"); + bool use_rds_iam_auth = false; + if (!use_rds_iam_auth_val.IsNull()) { + use_rds_iam_auth = BooleanValue::Get(use_rds_iam_auth_val); + } + new_connection_info += AddConnectionOption(kv_secret, "user"); - new_connection_info += AddConnectionOption(kv_secret, "password"); + + if (use_rds_iam_auth) { + Value host_val = kv_secret.TryGetValue("host"); + Value port_val = kv_secret.TryGetValue("port"); + Value user_val = kv_secret.TryGetValue("user"); + Value aws_region_val = kv_secret.TryGetValue("aws_region"); + + if (host_val.IsNull() || port_val.IsNull() || user_val.IsNull()) { + throw BinderException( + "RDS IAM authentication requires 'host', 'port', and 'user' to be set in the secret"); + } + + string hostname = host_val.ToString(); + string port = port_val.ToString(); + string username = user_val.ToString(); + string aws_region; + + + if (!aws_region_val.IsNull()) { + aws_region = aws_region_val.ToString(); + } + + try { + string rds_token = GenerateRdsAuthToken(hostname, port, username, aws_region); + new_connection_info += "password="; + new_connection_info += EscapeConnectionString(rds_token); + new_connection_info += " "; + } catch (const std::exception &e) { + throw BinderException("Failed to generate RDS auth token: %s", e.what()); + } + } else { + new_connection_info += AddConnectionOption(kv_secret, "password"); + } + new_connection_info += AddConnectionOption(kv_secret, "host"); new_connection_info += AddConnectionOption(kv_secret, "port"); new_connection_info += AddConnectionOption(kv_secret, "dbname"); @@ -102,6 +210,10 @@ string PostgresCatalog::GetConnectionString(ClientContext &context, const string return connection_string; } +string PostgresCatalog::GetFreshConnectionString(ClientContext &context) { + return GetConnectionString(context, attach_path, secret_name); +} + PostgresCatalog::~PostgresCatalog() = default; void PostgresCatalog::Initialize(bool load_builtin) { diff --git a/src/storage/postgres_connection_pool.cpp b/src/storage/postgres_connection_pool.cpp index c80a2e127..9430d1604 100644 --- a/src/storage/postgres_connection_pool.cpp +++ b/src/storage/postgres_connection_pool.cpp @@ -44,7 +44,7 @@ PostgresConnectionPool::PostgresConnectionPool(PostgresCatalog &postgres_catalog : postgres_catalog(postgres_catalog), active_connections(0), maximum_connections(maximum_connections_p) { } -PostgresPoolConnection PostgresConnectionPool::GetConnectionInternal() { +PostgresPoolConnection PostgresConnectionPool::GetConnectionInternal(optional_ptr context) { active_connections++; // check if we have any cached connections left if (!connection_cache.empty()) { @@ -54,21 +54,26 @@ PostgresPoolConnection PostgresConnectionPool::GetConnectionInternal() { } // no cached connections left but there is space to open a new one - open it + // If we have a context, generate a fresh connection string (with new RDS token if needed) + string connection_string_to_use = postgres_catalog.connection_string; + if (context) { + connection_string_to_use = postgres_catalog.GetFreshConnectionString(*context); + } return PostgresPoolConnection( - this, PostgresConnection::Open(postgres_catalog.connection_string, postgres_catalog.attach_path)); + this, PostgresConnection::Open(connection_string_to_use, postgres_catalog.attach_path)); } -PostgresPoolConnection PostgresConnectionPool::ForceGetConnection() { +PostgresPoolConnection PostgresConnectionPool::ForceGetConnection(optional_ptr context) { lock_guard l(connection_lock); - return GetConnectionInternal(); + return GetConnectionInternal(context); } -bool PostgresConnectionPool::TryGetConnection(PostgresPoolConnection &connection) { +bool PostgresConnectionPool::TryGetConnection(PostgresPoolConnection &connection, optional_ptr context) { lock_guard l(connection_lock); if (active_connections >= maximum_connections) { return false; } - connection = GetConnectionInternal(); + connection = GetConnectionInternal(context); return true; } @@ -79,9 +84,9 @@ void PostgresConnectionPool::PostgresSetConnectionCache(ClientContext &context, pg_use_connection_cache = BooleanValue::Get(parameter); } -PostgresPoolConnection PostgresConnectionPool::GetConnection() { +PostgresPoolConnection PostgresConnectionPool::GetConnection(optional_ptr context) { PostgresPoolConnection result; - if (!TryGetConnection(result)) { + if (!TryGetConnection(result, context)) { throw IOException( "Failed to get connection from PostgresConnectionPool - maximum connection count exceeded (%llu/%llu max)", active_connections, maximum_connections); diff --git a/src/storage/postgres_transaction.cpp b/src/storage/postgres_transaction.cpp index 7b198cab5..a87ee31fa 100644 --- a/src/storage/postgres_transaction.cpp +++ b/src/storage/postgres_transaction.cpp @@ -11,7 +11,7 @@ PostgresTransaction::PostgresTransaction(PostgresCatalog &postgres_catalog, Tran ClientContext &context) : Transaction(manager, context), access_mode(postgres_catalog.access_mode), isolation_level(postgres_catalog.isolation_level) { - connection = postgres_catalog.GetConnectionPool().GetConnection(); + connection = postgres_catalog.GetConnectionPool().GetConnection(&context); } PostgresTransaction::~PostgresTransaction() = default; From f0bf4895d86e59b021e0607020516e49c8a50dc1 Mon Sep 17 00:00:00 2001 From: Paolo Quadri Date: Thu, 20 Nov 2025 21:37:08 +0100 Subject: [PATCH 2/4] chore: update readme --- README.md | 72 +++++++++++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 68 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index a3d3b65dd..1743c3a8f 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,68 @@ host=localhost port=5432 dbname=mydb connect_timeout=10 | dbname | Database Name | [user] | | passfile | Name of file passwords are stored in | ~/.pgpass | +### AWS RDS IAM Authentication + +The extension supports AWS RDS IAM-based authentication, which allows you to connect to RDS PostgreSQL instances using IAM database authentication instead of static passwords. This feature automatically generates temporary authentication tokens using the AWS CLI. + +#### Requirements + +- AWS CLI installed and configured +- RDS instance with IAM database authentication enabled +- IAM user/role with `rds-db:connect` permission for the RDS instance +- AWS credentials configured (via `AWS_PROFILE`, `AWS_ACCESS_KEY_ID`/`AWS_SECRET_ACCESS_KEY`, or IAM role) + +#### Usage + +To use RDS IAM authentication, create a Postgres secret with the `USE_RDS_IAM_AUTH` parameter set to `true`: + +```sql +CREATE SECRET rds_secret ( + TYPE POSTGRES, + HOST 'my-db-instance.xxxxxx.us-west-2.rds.amazonaws.com', + PORT '5432', + USER 'my_iam_user', + DATABASE 'mydb', + USE_RDS_IAM_AUTH true, + AWS_REGION 'us-west-2' -- Optional: uses AWS CLI default if not specified +); + +ATTACH '' AS rds_db (TYPE POSTGRES, SECRET rds_secret); +``` + +#### Secret Parameters for RDS IAM Authentication + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `USE_RDS_IAM_AUTH` | BOOLEAN | Yes | Enable RDS IAM authentication | +| `HOST` | VARCHAR | Yes | RDS instance hostname | +| `PORT` | VARCHAR | Yes | RDS instance port (typically 5432) | +| `USER` | VARCHAR | Yes | IAM database username | +| `DATABASE` or `DBNAME` | VARCHAR | No | Database name | +| `AWS_REGION` | VARCHAR | No | AWS region (optional, uses AWS CLI default if not specified) | + +#### Important Notes + +- **Token Expiration**: RDS auth tokens expire after 15 minutes. The extension automatically generates fresh tokens when creating new connections, so long-running queries will continue to work. +- **AWS CLI**: The extension uses the `aws rds generate-db-auth-token` command. Make sure the AWS CLI is installed and configured with appropriate credentials. +- **Environment Variables**: The AWS CLI command inherits environment variables from the parent process, so `AWS_PROFILE`, `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, `AWS_SESSION_TOKEN`, and `AWS_REGION` will be available to the AWS CLI. + + +#### Example with AWS Profile + +```sql + +CREATE SECRET rds_secret ( + TYPE POSTGRES, + HOST 'my-db.xxxxxx.us-east-1.rds.amazonaws.com', + PORT '5432', + USER 'my_iam_user', + DATABASE 'mydb', + USE_RDS_IAM_AUTH true +); + +ATTACH '' AS rds_db (TYPE POSTGRES, SECRET rds_secret); +``` The tables in the file can be read as if they were normal DuckDB tables, but the underlying data is read directly from Postgres at query time. @@ -62,16 +124,18 @@ git pull --recurse-submodules ``` To build, type -``` +```bash make ``` To run, run the bundled `duckdb` shell: -``` - ./build/release/duckdb -unsigned # allow unsigned extensions + +```bash +./build/release/duckdb -unsigned # allow unsigned extensions ``` Then, load the Postgres extension like so: -```SQL + +```sql LOAD 'build/release/extension/postgres_scanner/postgres_scanner.duckdb_extension'; ``` From 9556f0be31226bca2e74f1b536713f3ffa2890c9 Mon Sep 17 00:00:00 2001 From: Paolo Quadri Date: Fri, 21 Nov 2025 10:09:08 +0100 Subject: [PATCH 3/4] chore: remove comment --- src/storage/postgres_catalog.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/src/storage/postgres_catalog.cpp b/src/storage/postgres_catalog.cpp index b7b5e1a2e..109fa55e7 100644 --- a/src/storage/postgres_catalog.cpp +++ b/src/storage/postgres_catalog.cpp @@ -155,7 +155,6 @@ string PostgresCatalog::GetConnectionString(ClientContext &context, const string const auto &kv_secret = dynamic_cast(*secret_entry->secret); string new_connection_info; - // Check if RDS IAM authentication is enabled Value use_rds_iam_auth_val = kv_secret.TryGetValue("use_rds_iam_auth"); bool use_rds_iam_auth = false; if (!use_rds_iam_auth_val.IsNull()) { From ef825b1124c24edf4a554d92a46b9a9bcf23856b Mon Sep 17 00:00:00 2001 From: Paolo Quadri Date: Fri, 21 Nov 2025 10:16:05 +0100 Subject: [PATCH 4/4] chore: reorganize readme to not break another section --- README.md | 75 +++++++++++++++++++++---------------------------------- 1 file changed, 29 insertions(+), 46 deletions(-) diff --git a/README.md b/README.md index 1743c3a8f..88c435358 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,31 @@ host=localhost port=5432 dbname=mydb connect_timeout=10 | dbname | Database Name | [user] | | passfile | Name of file passwords are stored in | ~/.pgpass | + +The tables in the file can be read as if they were normal DuckDB tables, but the underlying data is read directly from Postgres at query time. + +```sql +D SHOW ALL TABLES; +┌───────────────────────────────────────┐ +│ name │ +│ varchar │ +├───────────────────────────────────────┤ +│ uuids │ +└───────────────────────────────────────┘ +D SELECT * FROM postgres_db.uuids; +┌──────────────────────────────────────┐ +│ u │ +│ uuid │ +├──────────────────────────────────────┤ +│ 6d3d2541-710b-4bde-b3af-4711738636bf │ +│ NULL │ +│ 00000000-0000-0000-0000-000000000001 │ +│ ffffffff-ffff-ffff-ffff-ffffffffffff │ +└──────────────────────────────────────┘ +``` + +For more information on how to use the connector, refer to the [Postgres documentation on the website](https://duckdb.org/docs/extensions/postgres). + ### AWS RDS IAM Authentication The extension supports AWS RDS IAM-based authentication, which allows you to connect to RDS PostgreSQL instances using IAM database authentication instead of static passwords. This feature automatically generates temporary authentication tokens using the AWS CLI. @@ -74,46 +99,6 @@ ATTACH '' AS rds_db (TYPE POSTGRES, SECRET rds_secret); - **Environment Variables**: The AWS CLI command inherits environment variables from the parent process, so `AWS_PROFILE`, `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, `AWS_SESSION_TOKEN`, and `AWS_REGION` will be available to the AWS CLI. -#### Example with AWS Profile - -```sql - -CREATE SECRET rds_secret ( - TYPE POSTGRES, - HOST 'my-db.xxxxxx.us-east-1.rds.amazonaws.com', - PORT '5432', - USER 'my_iam_user', - DATABASE 'mydb', - USE_RDS_IAM_AUTH true -); - -ATTACH '' AS rds_db (TYPE POSTGRES, SECRET rds_secret); -``` - -The tables in the file can be read as if they were normal DuckDB tables, but the underlying data is read directly from Postgres at query time. - -```sql -D SHOW ALL TABLES; -┌───────────────────────────────────────┐ -│ name │ -│ varchar │ -├───────────────────────────────────────┤ -│ uuids │ -└───────────────────────────────────────┘ -D SELECT * FROM postgres_db.uuids; -┌──────────────────────────────────────┐ -│ u │ -│ uuid │ -├──────────────────────────────────────┤ -│ 6d3d2541-710b-4bde-b3af-4711738636bf │ -│ NULL │ -│ 00000000-0000-0000-0000-000000000001 │ -│ ffffffff-ffff-ffff-ffff-ffffffffffff │ -└──────────────────────────────────────┘ -``` - -For more information on how to use the connector, refer to the [Postgres documentation on the website](https://duckdb.org/docs/extensions/postgres). - ## Building & Loading the Extension The DuckDB submodule must be initialized prior to building. @@ -124,18 +109,16 @@ git pull --recurse-submodules ``` To build, type -```bash +``` make ``` To run, run the bundled `duckdb` shell: - -```bash -./build/release/duckdb -unsigned # allow unsigned extensions +``` + ./build/release/duckdb -unsigned # allow unsigned extensions ``` Then, load the Postgres extension like so: - -```sql +```SQL LOAD 'build/release/extension/postgres_scanner/postgres_scanner.duckdb_extension'; ```