diff --git a/.github/workflows/IntegrationTests.yml b/.github/workflows/IntegrationTests.yml index e31d747ae..627962ae8 100644 --- a/.github/workflows/IntegrationTests.yml +++ b/.github/workflows/IntegrationTests.yml @@ -114,6 +114,7 @@ jobs: PGUSER: postgres PGPASSWORD: postgres POSTGRES_TEST_DATABASE_AVAILABLE: 1 + POSTGRES_CONNECTION_IS_DIRECT: 1 LOCAL_EXTENSION_REPO: 'build/release/repository' run: | source ./create-postgres-tables.sh @@ -345,6 +346,7 @@ jobs: PGPASSWORD: postgres PGSSLMODE: require POSTGRES_TEST_DATABASE_AVAILABLE: 1 + POSTGRES_CONNECTION_IS_DIRECT: 1 SANITIZER_BUILD: 1 LOCAL_EXTENSION_REPO: 'build/relassert/repository' run: | @@ -471,6 +473,7 @@ jobs: PGPASSWORD: postgres PGSSLMODE: require POSTGRES_TEST_DATABASE_AVAILABLE: 1 + POSTGRES_CONNECTION_IS_DIRECT: 1 LOCAL_EXTENSION_REPO: 'build/release/repository' shell: bash run: | diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index f8138d09d..51653ccb5 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -10,6 +10,7 @@ add_library( postgres_binary_file_reader.cpp postgres_binary_parser.cpp postgres_binary_reader.cpp + postgres_connection_string.cpp postgres_connection.cpp postgres_copy_from.cpp postgres_copy_to.cpp diff --git a/src/include/postgres_connection_string.hpp b/src/include/postgres_connection_string.hpp new file mode 100644 index 000000000..8c8f30314 --- /dev/null +++ b/src/include/postgres_connection_string.hpp @@ -0,0 +1,29 @@ +//===----------------------------------------------------------------------===// +// DuckDB +// +// postgres_binary_writer.hpp +// +// +//===----------------------------------------------------------------------===// + +#pragma once + +#include + +#include "duckdb.hpp" + +namespace duckdb { + +enum class PostgresConnectionStringFormat { KEY_VALUE, URI }; + +struct PostgresConnectionString { + static string Escape(const string &input); + + static string ExtractOption(const KeyValueSecret &kv_secret, const string &name); + + static bool ContainsApplicationName(const std::string &conn_str); + + static void AppendDefaultApplicationName(string &conn_string, PostgresConnectionStringFormat format); +}; + +} // namespace duckdb diff --git a/src/include/postgres_utils.hpp b/src/include/postgres_utils.hpp index 8fc9f0ad2..8fd839e6b 100644 --- a/src/include/postgres_utils.hpp +++ b/src/include/postgres_utils.hpp @@ -98,8 +98,6 @@ class PostgresUtils { static PostgresVersion ExtractPostgresVersion(const string &version); - static string EscapeConnectionString(const string &input); - static string ExtractConnectionOption(const KeyValueSecret &kv_secret, const string &name); static string WriteLiteral(const string &identifier); static string WriteLiteralsCommaSeparated(const vector &literals); static string WriteIdentifier(const string &identifier); diff --git a/src/postgres_aws.cpp b/src/postgres_aws.cpp index 44d34606d..9aafa040a 100644 --- a/src/postgres_aws.cpp +++ b/src/postgres_aws.cpp @@ -9,6 +9,7 @@ #include "dbconnector/defer.hpp" +#include "postgres_connection_string.hpp" #include "postgres_secrets.hpp" #include "postgres_utils.hpp" @@ -119,12 +120,20 @@ PostgresAwsRdsTokenConfig PostgresAws::ExtractTokenConfigFromSecret(optional_ptr "specified - generated IAM token will be used instead of the password"); } + bool has_app_name = !kv.TryGetValue("application_name").IsNull(); + // Build the base connection string (without password) for (const string &opt_name : PostgresSecrets::ConnectionOptionNames()) { if (opt_name == "password" || opt_name == "passfile") { continue; } - config.base_connection_string += PostgresUtils::ExtractConnectionOption(kv, opt_name); + config.base_connection_string += PostgresConnectionString::ExtractOption(kv, opt_name); + } + + if (!has_app_name) { + PostgresConnectionString::AppendDefaultApplicationName(config.base_connection_string, + PostgresConnectionStringFormat::KEY_VALUE); + config.base_connection_string.push_back(' '); } return config; diff --git a/src/postgres_connection.cpp b/src/postgres_connection.cpp index eab1c3b11..a896df3de 100644 --- a/src/postgres_connection.cpp +++ b/src/postgres_connection.cpp @@ -1,14 +1,16 @@ +#include "postgres_connection.hpp" + #include "duckdb/storage/table_storage_info.hpp" #include "duckdb/logging/logger.hpp" #include "duckdb/main/client_context.hpp" #include "duckdb/parser/column_list.hpp" #include "duckdb/parser/parser.hpp" -#include "postgres_connection.hpp" -#include "postgres_logging.hpp" #include "duckdb/common/types/uuid.hpp" #include "duckdb/common/shared_ptr.hpp" #include "duckdb/common/helper.hpp" +#include "postgres_logging.hpp" + namespace duckdb { static bool debug_postgres_print_queries = false; diff --git a/src/postgres_connection_string.cpp b/src/postgres_connection_string.cpp new file mode 100644 index 000000000..9695a571f --- /dev/null +++ b/src/postgres_connection_string.cpp @@ -0,0 +1,85 @@ +#include "postgres_connection_string.hpp" + +#include +#include +#include +#include + +#include "duckdb.hpp" + +#include "dbconnector/query/query_writer.hpp" + +namespace duckdb { + +string PostgresConnectionString::Escape(const string &input) { + using namespace dbconnector::query; + auto config = QueryWriter::CreateConfig('\'', QuoteEscapeStyle::BACKSLASH); + return QueryWriter::WriteQuotedAndEscaped(config, input); +} + +string PostgresConnectionString::ExtractOption(const KeyValueSecret &kv_secret, const string &name) { + Value input_val = kv_secret.TryGetValue(Identifier(name)); + if (input_val.IsNull()) { + // not provided + return string(); + } + string result; + result += name; + result += "="; + result += Escape(input_val.ToString()); + result += " "; + return result; +} + +bool PostgresConnectionString::ContainsApplicationName(const std::string &conninfo_str) { + char *errmsg = nullptr; + + // Parse the connection string into an array of PQconninfoOption structs + PQconninfoOption *options = PQconninfoParse(conninfo_str.c_str(), &errmsg); + + if (!options) { + std::string err_str; + if (errmsg) { + err_str = std::string(errmsg); + PQfreemem(errmsg); + } + throw InvalidInputException("Failed to parse connection string: \"%s\"", err_str); + } + + bool found = false; + + for (PQconninfoOption *opt = options; opt->keyword != nullptr; ++opt) { + if (std::strcmp(opt->keyword, "application_name") == 0) { + // opt->val will be non-null if it was explicitly defined in the string + if (opt->val != nullptr && std::strlen(opt->val) > 0) { + found = true; + } + break; + } + } + + // cleanup + PQconninfoFree(options); + if (errmsg) { + PQfreemem(errmsg); + } + + return found; +} + +void PostgresConnectionString::AppendDefaultApplicationName(string &conn_string, + PostgresConnectionStringFormat format) { + string app_name = "application_name=DuckDB"; + if (format == PostgresConnectionStringFormat::KEY_VALUE) { + if (conn_string.length() == 0 || conn_string[conn_string.length() - 1] != ' ') { + conn_string.push_back(' '); + } + } else { + bool first = conn_string.find('?') == std::string::npos; + char separator = first ? '?' : '&'; + conn_string.push_back(separator); + } + conn_string += app_name; +} + +} // namespace duckdb diff --git a/src/postgres_utils.cpp b/src/postgres_utils.cpp index cd6e48c83..7d2ec7780 100644 --- a/src/postgres_utils.cpp +++ b/src/postgres_utils.cpp @@ -713,26 +713,6 @@ string PostgresUtils::QuotePostgresIdentifier(const string &text) { return PostgresUtils::WriteIdentifier(text); } -string PostgresUtils::EscapeConnectionString(const string &input) { - using namespace dbconnector::query; - auto config = QueryWriter::CreateConfig('\'', QuoteEscapeStyle::BACKSLASH); - return QueryWriter::WriteQuotedAndEscaped(config, input); -} - -string PostgresUtils::ExtractConnectionOption(const KeyValueSecret &kv_secret, const string &name) { - Value input_val = kv_secret.TryGetValue(Identifier(name)); - if (input_val.IsNull()) { - // not provided - return string(); - } - string result; - result += name; - result += "="; - result += EscapeConnectionString(input_val.ToString()); - result += " "; - return result; -} - string PostgresUtils::WriteLiteral(const string &literal) { using namespace dbconnector::query; auto config = QueryWriter::CreateConfig('\'', QuoteEscapeStyle::DOUBLE_QUOTE); diff --git a/src/storage/postgres_catalog.cpp b/src/storage/postgres_catalog.cpp index 5eb3e362f..3f4b49ce6 100644 --- a/src/storage/postgres_catalog.cpp +++ b/src/storage/postgres_catalog.cpp @@ -1,14 +1,8 @@ #include "storage/postgres_catalog.hpp" + #include "duckdb/main/client_context.hpp" #include "duckdb/catalog/catalog.hpp" #include "duckdb/catalog/catalog_entry/schema_catalog_entry.hpp" -#include "storage/postgres_schema_entry.hpp" -#include "storage/postgres_transaction.hpp" -#include "postgres_oauth.hpp" -#include "postgres_connection.hpp" -#include "postgres_secrets.hpp" -#include "storage/postgres_secret_storage.hpp" -#include "postgres_utils.hpp" #include "duckdb/storage/database_size.hpp" #include "duckdb/parser/parsed_data/drop_info.hpp" #include "duckdb/parser/parsed_data/create_schema_info.hpp" @@ -20,6 +14,15 @@ #include "duckdb/main/attached_database.hpp" #include "duckdb/main/secret/secret_manager.hpp" +#include "postgres_connection_string.hpp" +#include "postgres_connection.hpp" +#include "postgres_oauth.hpp" +#include "postgres_secrets.hpp" +#include "postgres_utils.hpp" +#include "storage/postgres_schema_entry.hpp" +#include "storage/postgres_transaction.hpp" +#include "storage/postgres_secret_storage.hpp" + namespace duckdb { unique_ptr PostgresCatalog::RemoteExecute(ClientContext &context, const string &sql) { @@ -127,6 +130,9 @@ unique_ptr PostgresCatalog::GetSecretEntry(ClientContext &ctx, cons string PostgresCatalog::CreateConnectionString(optional_ptr secret_entry, const string &attach_path) { string connection_string = attach_path; + bool user_specified_app_name = PostgresConnectionString::ContainsApplicationName(connection_string); + bool secret_has_app_name = false; + string secret_connection_info; if (secret_entry) { // secret found - read data const auto &kv_secret = dynamic_cast(*secret_entry->secret); @@ -146,15 +152,28 @@ string PostgresCatalog::CreateConnectionString(optional_ptr secret_ throw BinderException("ATTACH path must be empty when 'URI' option is specified, was: \"%s\"", attach_path); } - return uri_val.ToString(); + string uri = uri_val.ToString(); + bool uri_has_app_name = PostgresConnectionString::ContainsApplicationName(uri); + if (!uri_has_app_name) { + PostgresConnectionString::AppendDefaultApplicationName(uri, PostgresConnectionStringFormat::URI); + } + return uri; } - string new_connection_info; + secret_has_app_name = !kv_secret.TryGetValue("application_name").IsNull(); + for (const string &opt_name : PostgresSecrets::ConnectionOptionNames()) { - new_connection_info += PostgresUtils::ExtractConnectionOption(kv_secret, opt_name); + secret_connection_info += PostgresConnectionString::ExtractOption(kv_secret, opt_name); } - connection_string = new_connection_info + connection_string; } + + if (!(user_specified_app_name || secret_has_app_name)) { + PostgresConnectionString::AppendDefaultApplicationName(secret_connection_info, + PostgresConnectionStringFormat::KEY_VALUE); + secret_connection_info.push_back(' '); + } + connection_string = secret_connection_info + connection_string; + return connection_string; } @@ -179,7 +198,7 @@ string PostgresCatalog::GetConnectionString() { this->rds_token = PostgresAws::GenerateRdsAuthToken(GetAttached(), rds_token_config); this->rds_token_last_refreshed = now; this->connection_string = rds_token_config.base_connection_string + - "password=" + PostgresUtils::EscapeConnectionString(rds_token) + " " + attach_path; + "password=" + PostgresConnectionString::Escape(rds_token) + " " + attach_path; } return std::string(connection_string.data(), connection_string.length()); } diff --git a/test/sql/scanner/application_name.test b/test/sql/scanner/application_name.test new file mode 100644 index 000000000..339d86d7f --- /dev/null +++ b/test/sql/scanner/application_name.test @@ -0,0 +1,234 @@ +# name: test/sql/scanner/application_name.test +# description: Check default and custom application_name +# group: [scanner] + +require postgres_scanner + +require-env POSTGRES_TEST_DATABASE_AVAILABLE + +require-env POSTGRES_CONNECTION_IS_DIRECT + +# default key-value + +statement ok +ATTACH 'dbname=postgresscanner' AS s (TYPE POSTGRES) + +statement ok +SET VARIABLE conn = postgres_pin_connection('s') + +query II +FROM postgres_query('s', +' +SELECT application_name, state FROM pg_stat_activity +WHERE application_name = ''DuckDB'' +', +connection=getvariable('conn')) +---- +DuckDB active + +statement ok +SELECT postgres_close_pinned_connection('s', getvariable('conn')) + +statement ok +RESET VARIABLE conn + +statement ok +DETACH s + +# name in conn_string + +statement ok +ATTACH 'dbname=postgresscanner application_name=custom1' AS s (TYPE POSTGRES) + +statement ok +SET VARIABLE conn = postgres_pin_connection('s') + +query II +FROM postgres_query('s', +' +SELECT application_name, state FROM pg_stat_activity +WHERE application_name = ''DuckDB'' +', +connection=getvariable('conn')) +---- + +query II +FROM postgres_query('s', +' +SELECT application_name, state FROM pg_stat_activity +WHERE application_name = ''custom1'' +', +connection=getvariable('conn')) +---- +custom1 active + +statement ok +SELECT postgres_close_pinned_connection('s', getvariable('conn')) + +statement ok +RESET VARIABLE conn + +statement ok +DETACH s + +# name in secret + +statement ok +CREATE TEMPORARY SECRET pgsec1 ( + TYPE postgres, + APPLICATION_NAME custom2 +) + +statement ok +ATTACH 'dbname=postgresscanner' AS s (TYPE POSTGRES, SECRET pgsec1) + +statement ok +SET VARIABLE conn = postgres_pin_connection('s') + +query II +FROM postgres_query('s', +' +SELECT application_name, state FROM pg_stat_activity +WHERE application_name = ''DuckDB'' +', +connection=getvariable('conn')) +---- + +query II +FROM postgres_query('s', +' +SELECT application_name, state FROM pg_stat_activity +WHERE application_name = ''custom2'' +', +connection=getvariable('conn')) +---- +custom2 active + +statement ok +SELECT postgres_close_pinned_connection('s', getvariable('conn')) + +statement ok +RESET VARIABLE conn + +statement ok +DETACH s + +statement ok +DROP SECRET pgsec1 + +# default URI + +statement ok +CREATE TEMPORARY SECRET pgsec1 ( + TYPE postgres, + URI 'postgresql://127.0.0.1:/postgresscanner' +) + +statement ok +ATTACH '' AS s (TYPE POSTGRES, SECRET pgsec1) + +statement ok +SET VARIABLE conn = postgres_pin_connection('s') + +query II +FROM postgres_query('s', +' +SELECT application_name, state FROM pg_stat_activity +WHERE application_name = ''DuckDB'' +', +connection=getvariable('conn')) +---- +DuckDB active + +statement ok +SELECT postgres_close_pinned_connection('s', getvariable('conn')) + +statement ok +RESET VARIABLE conn + +statement ok +DETACH s + +statement ok +DROP SECRET pgsec1 + +# default URI 2 + +statement ok +CREATE TEMPORARY SECRET pgsec1 ( + TYPE postgres, + URI 'postgresql://127.0.0.1:/postgresscanner?sslmode=allow' +) + +statement ok +ATTACH '' AS s (TYPE POSTGRES, SECRET pgsec1) + +statement ok +SET VARIABLE conn = postgres_pin_connection('s') + +query II +FROM postgres_query('s', +' +SELECT application_name, state FROM pg_stat_activity +WHERE application_name = ''DuckDB'' +', +connection=getvariable('conn')) +---- +DuckDB active + +statement ok +SELECT postgres_close_pinned_connection('s', getvariable('conn')) + +statement ok +RESET VARIABLE conn + +statement ok +DETACH s + +statement ok +DROP SECRET pgsec1 + +# custom URI + +statement ok +CREATE TEMPORARY SECRET pgsec1 ( + TYPE postgres, + URI 'postgresql://127.0.0.1:/postgresscanner?sslmode=allow&application_name=custom2' +) + +statement ok +ATTACH '' AS s (TYPE POSTGRES, SECRET pgsec1) + +statement ok +SET VARIABLE conn = postgres_pin_connection('s') + +query II +FROM postgres_query('s', +' +SELECT application_name, state FROM pg_stat_activity +WHERE application_name = ''DuckDB'' +', +connection=getvariable('conn')) +---- + +query II +FROM postgres_query('s', +' +SELECT application_name, state FROM pg_stat_activity +WHERE application_name = ''custom2'' +', +connection=getvariable('conn')) +---- +custom2 active + +statement ok +SELECT postgres_close_pinned_connection('s', getvariable('conn')) + +statement ok +RESET VARIABLE conn + +statement ok +DETACH s + +statement ok +DROP SECRET pgsec1 diff --git a/test/sql/storage/attach_rds_iam.test b/test/sql/storage/attach_rds_iam.test index e969ee98c..4228127e5 100644 --- a/test/sql/storage/attach_rds_iam.test +++ b/test/sql/storage/attach_rds_iam.test @@ -43,6 +43,55 @@ FROM postgres_query('rds_attached', 'SELECT current_user') ---- postgres +query II +FROM postgres_query('rds_attached', +' +SELECT application_name, state FROM pg_stat_activity +WHERE application_name = ''DuckDB'' +') +---- +DuckDB active + +statement ok +DROP SECRET rds_secret + +statement ok +DETACH rds_attached + +# custom app name + +statement ok +CREATE SECRET rds_secret ( + TYPE postgres, + HOST 'database-1-instance-1.xxx.eu-west-1.rds.amazonaws.com', + PORT '5432', + USER 'postgres', + DATABASE 'postgres', + SSLMODE require, + AWS_RDS_SECRET rds_template_secret, + APPLICATION_NAME custom1 +); + +statement ok +ATTACH '' AS rds_attached (TYPE POSTGRES, SECRET rds_secret) + +query II +FROM postgres_query('rds_attached', +' +SELECT application_name, state FROM pg_stat_activity +WHERE application_name = ''DuckDB'' +') +---- + +query II +FROM postgres_query('rds_attached', +' +SELECT application_name, state FROM pg_stat_activity +WHERE application_name = ''custom1'' +') +---- +custom1 active + statement ok DROP SECRET rds_secret diff --git a/test/sql/storage/attach_secret.test b/test/sql/storage/attach_secret.test index 0d479a982..4dd83bbc7 100644 --- a/test/sql/storage/attach_secret.test +++ b/test/sql/storage/attach_secret.test @@ -152,7 +152,7 @@ CREATE OR REPLACE SECRET postgres_db ( ); statement error -ATTACH 'bar' AS secret_attach (TYPE POSTGRES, SECRET postgres_db) +ATTACH 'sslmode=require' AS secret_attach (TYPE POSTGRES, SECRET postgres_db) ---- 'URI' option