Pass DuckDB label as a default application_name

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

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

src/include/postgres_connection_string.hpp

@@ -0,0 +1,29 @@
+//===----------------------------------------------------------------------===//
+//                         DuckDB
+//
+// postgres_binary_writer.hpp
+//
+//
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include <string>
+
+#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

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<string> &literals);
 	static string WriteIdentifier(const string &identifier);

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;

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;

src/postgres_connection_string.cpp

@@ -0,0 +1,85 @@
+#include "postgres_connection_string.hpp"
+
+#include <cstring>
+#include <string>
+#include <memory>
+#include <libpq-fe.h>
+
+#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

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);

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<TableRef> PostgresCatalog::RemoteExecute(ClientContext &context, const string &sql) {
@@ -127,6 +130,9 @@ unique_ptr<SecretEntry> PostgresCatalog::GetSecretEntry(ClientContext &ctx, cons
 
 string PostgresCatalog::CreateConnectionString(optional_ptr<SecretEntry> 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<const KeyValueSecret &>(*secret_entry->secret);
@@ -146,15 +152,28 @@ string PostgresCatalog::CreateConnectionString(optional_ptr<SecretEntry> 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());
 }