Skip to content

Commit

Permalink
Avoid modifying the path that is shown in duckdb_databases with the s…
Browse files Browse the repository at this point in the history 
…ecret contents
  • Loading branch information
@Mytherin
Mytherin committed Jan 29, 2025
1 parent ef6c95e commit 929bf00
Show file tree
Hide file tree
Showing 5 changed files with 92 additions and 78 deletions.
7 changes: 5 additions & 2 deletions src/include/storage/postgres_catalog.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,10 +20,11 @@ class PostgresSchemaEntry;

class PostgresCatalog : public Catalog {
public:
explicit PostgresCatalog(AttachedDatabase &db_p, const string &path, AccessMode access_mode, string schema_to_load);
explicit PostgresCatalog(AttachedDatabase &db_p, string connection_string, string attach_path, AccessMode access_mode, string schema_to_load);
~PostgresCatalog();

string path;
string connection_string;
string attach_path;
AccessMode access_mode;

public:
Expand All @@ -32,6 +33,8 @@ class PostgresCatalog : public Catalog {
return "postgres";
}

static string GetConnectionString(ClientContext &context, const string &attach_path, string secret_name);

optional_ptr<CatalogEntry> CreateSchema(CatalogTransaction transaction, CreateSchemaInfo &info) override;

void ScanSchemas(ClientContext &context, std::function<void(SchemaCatalogEntry &)> callback) override;
Expand Down
75 changes: 3 additions & 72 deletions src/postgres_storage.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,58 +4,13 @@
#include "storage/postgres_catalog.hpp"
#include "duckdb/parser/parsed_data/attach_info.hpp"
#include "storage/postgres_transaction_manager.hpp"
#include "duckdb/main/secret/secret_manager.hpp"

namespace duckdb {

string EscapeConnectionString(const string &input) {
string result = "'";
for (auto c : input) {
if (c == '\\') {
result += "\\\\";
} else if (c == '\'') {
result += "\\'";
} else {
result += c;
}
}
result += "'";
return result;
}

string AddConnectionOption(const KeyValueSecret &kv_secret, const string &name) {
Value input_val = kv_secret.TryGetValue(name);
if (input_val.IsNull()) {
// not provided
return string();
}
string result;
result += name;
result += "=";
result += EscapeConnectionString(input_val.ToString());
result += " ";
return result;
}

unique_ptr<SecretEntry> GetSecret(ClientContext &context, const string &secret_name) {
auto &secret_manager = SecretManager::Get(context);
auto transaction = CatalogTransaction::GetSystemCatalogTransaction(context);
// FIXME: this should be adjusted once the `GetSecretByName` API supports this use case
auto secret_entry = secret_manager.GetSecretByName(transaction, secret_name, "memory");
if (secret_entry) {
return secret_entry;
}
secret_entry = secret_manager.GetSecretByName(transaction, secret_name, "local_file");
if (secret_entry) {
return secret_entry;
}
return nullptr;
}

static unique_ptr<Catalog> PostgresAttach(StorageExtensionInfo *storage_info, ClientContext &context,
AttachedDatabase &db, const string &name, AttachInfo &info,
AccessMode access_mode) {
string connection_string = info.path;
string attach_path = info.path;

string secret_name;
string schema_to_load;
Expand All @@ -71,32 +26,8 @@ static unique_ptr<Catalog> PostgresAttach(StorageExtensionInfo *storage_info, Cl
throw BinderException("Unrecognized option for Postgres attach: %s", entry.first);
}
}

// if no secret is specified we default to the unnamed postgres secret, if it exists
bool explicit_secret = !secret_name.empty();
if (!explicit_secret) {
// look up settings from the default unnamed postgres secret if none is provided
secret_name = "__default_postgres";
}

auto secret_entry = GetSecret(context, secret_name);
if (secret_entry) {
// secret found - read data
const auto &kv_secret = dynamic_cast<const KeyValueSecret &>(*secret_entry->secret);
string new_connection_info;

new_connection_info += AddConnectionOption(kv_secret, "user");
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");

connection_string = new_connection_info + connection_string;
} else if (explicit_secret) {
// secret not found and one was explicitly provided - throw an error
throw BinderException("Secret with name \"%s\" not found", secret_name);
}
return make_uniq<PostgresCatalog>(db, connection_string, access_mode, std::move(schema_to_load));
auto connection_string = PostgresCatalog::GetConnectionString(context, attach_path, secret_name);
return make_uniq<PostgresCatalog>(db, std::move(connection_string), std::move(attach_path), access_mode, std::move(schema_to_load));
}

static unique_ptr<TransactionManager> PostgresCreateTransactionManager(StorageExtensionInfo *storage_info,
Expand Down
81 changes: 78 additions & 3 deletions src/storage/postgres_catalog.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,12 +6,13 @@
#include "duckdb/parser/parsed_data/drop_info.hpp"
#include "duckdb/parser/parsed_data/create_schema_info.hpp"
#include "duckdb/main/attached_database.hpp"
#include "duckdb/main/secret/secret_manager.hpp"

namespace duckdb {

PostgresCatalog::PostgresCatalog(AttachedDatabase &db_p, const string &path, AccessMode access_mode,
PostgresCatalog::PostgresCatalog(AttachedDatabase &db_p, string connection_string_p, string attach_path_p, AccessMode access_mode,
string schema_to_load)
: Catalog(db_p), path(path), access_mode(access_mode), schemas(*this, schema_to_load), connection_pool(*this),
: Catalog(db_p), connection_string(std::move(connection_string_p)), attach_path(std::move(attach_path_p)), access_mode(access_mode), schemas(*this, schema_to_load), connection_pool(*this),
default_schema(schema_to_load) {
if (default_schema.empty()) {
default_schema = "public";
Expand All @@ -26,6 +27,80 @@ PostgresCatalog::PostgresCatalog(AttachedDatabase &db_p, const string &path, Acc
this->version = connection.GetConnection().GetPostgresVersion();
}

string EscapeConnectionString(const string &input) {
string result = "'";
for (auto c : input) {
if (c == '\\') {
result += "\\\\";
} else if (c == '\'') {
result += "\\'";
} else {
result += c;
}
}
result += "'";
return result;
}

string AddConnectionOption(const KeyValueSecret &kv_secret, const string &name) {
Value input_val = kv_secret.TryGetValue(name);
if (input_val.IsNull()) {
// not provided
return string();
}
string result;
result += name;
result += "=";
result += EscapeConnectionString(input_val.ToString());
result += " ";
return result;
}

unique_ptr<SecretEntry> GetSecret(ClientContext &context, const string &secret_name) {
auto &secret_manager = SecretManager::Get(context);
auto transaction = CatalogTransaction::GetSystemCatalogTransaction(context);
// FIXME: this should be adjusted once the `GetSecretByName` API supports this use case
auto secret_entry = secret_manager.GetSecretByName(transaction, secret_name, "memory");
if (secret_entry) {
return secret_entry;
}
secret_entry = secret_manager.GetSecretByName(transaction, secret_name, "local_file");
if (secret_entry) {
return secret_entry;
}
return nullptr;
}

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;
bool explicit_secret = !secret_name.empty();
if (!explicit_secret) {
// look up settings from the default unnamed postgres secret if none is provided
secret_name = "__default_postgres";
}

auto secret_entry = GetSecret(context, secret_name);
if (secret_entry) {
// secret found - read data
const auto &kv_secret = dynamic_cast<const KeyValueSecret &>(*secret_entry->secret);
string new_connection_info;

new_connection_info += AddConnectionOption(kv_secret, "user");
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");

connection_string = new_connection_info + connection_string;
} else if (explicit_secret) {
// secret not found and one was explicitly provided - throw an error
throw BinderException("Secret with name \"%s\" not found", secret_name);
}
return connection_string;
}


PostgresCatalog::~PostgresCatalog() = default;

void PostgresCatalog::Initialize(bool load_builtin) {
Expand Down Expand Up @@ -85,7 +160,7 @@ bool PostgresCatalog::InMemory() {
}

string PostgresCatalog::GetDBPath() {
return path;
return attach_path;
}

DatabaseSize PostgresCatalog::GetDatabaseSize(ClientContext &context) {
Expand Down
2 changes: 1 addition & 1 deletion src/storage/postgres_connection_pool.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ PostgresPoolConnection PostgresConnectionPool::GetConnectionInternal() {
}

// no cached connections left but there is space to open a new one - open it
return PostgresPoolConnection(this, PostgresConnection::Open(postgres_catalog.path));
return PostgresPoolConnection(this, PostgresConnection::Open(postgres_catalog.connection_string));
}

PostgresPoolConnection PostgresConnectionPool::ForceGetConnection() {
Expand Down
5 changes: 5 additions & 0 deletions test/sql/storage/attach_secret.test
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,11 @@ CREATE SECRET postgres_db (
statement ok
ATTACH '' AS secret_attach (TYPE POSTGRES, SECRET postgres_db)

query I
SELECT path FROM duckdb_databases() WHERE database_name='secret_attach'
----
(empty)

statement ok
DETACH secret_attach

Expand Down

0 comments on commit 929bf00

Please sign in to comment.