Add DatabaseFileOpener, and rework S3 configuration to get secrets fr…

…om DatabaseInstance instead of ClientContext
duckdb · Mytherin · Mar 26, 2024 · Mar 21, 2024 · Mar 21, 2024 · Mar 21, 2024
commit e2c16a5d211a9e9c7867207475f746287f9773f1
diff --git a/extension/httpfs/s3fs.cpp b/extension/httpfs/s3fs.cpp
@@ -156,25 +156,25 @@ void AWSEnvironmentCredentialsProvider::SetExtensionOptionValue(string key, cons
 }
 
 void AWSEnvironmentCredentialsProvider::SetAll() {
-	this->SetExtensionOptionValue("s3_region", this->DEFAULT_REGION_ENV_VAR);
-	this->SetExtensionOptionValue("s3_region", this->REGION_ENV_VAR);
-	this->SetExtensionOptionValue("s3_access_key_id", this->ACCESS_KEY_ENV_VAR);
-	this->SetExtensionOptionValue("s3_secret_access_key", this->SECRET_KEY_ENV_VAR);
-	this->SetExtensionOptionValue("s3_session_token", this->SESSION_TOKEN_ENV_VAR);
-	this->SetExtensionOptionValue("s3_endpoint", this->DUCKDB_ENDPOINT_ENV_VAR);
-	this->SetExtensionOptionValue("s3_use_ssl", this->DUCKDB_USE_SSL_ENV_VAR);
+	this->SetExtensionOptionValue("s3_region", DEFAULT_REGION_ENV_VAR);
+	this->SetExtensionOptionValue("s3_region", REGION_ENV_VAR);
+	this->SetExtensionOptionValue("s3_access_key_id", ACCESS_KEY_ENV_VAR);
+	this->SetExtensionOptionValue("s3_secret_access_key", SECRET_KEY_ENV_VAR);
+	this->SetExtensionOptionValue("s3_session_token", SESSION_TOKEN_ENV_VAR);
+	this->SetExtensionOptionValue("s3_endpoint", DUCKDB_ENDPOINT_ENV_VAR);
+	this->SetExtensionOptionValue("s3_use_ssl", DUCKDB_USE_SSL_ENV_VAR);
 }
 
 S3AuthParams AWSEnvironmentCredentialsProvider::CreateParams() {
 	S3AuthParams params;
 
-	params.region = this->DEFAULT_REGION_ENV_VAR;
-	params.region = this->REGION_ENV_VAR;
-	params.access_key_id = this->ACCESS_KEY_ENV_VAR;
-	params.secret_access_key = this->SECRET_KEY_ENV_VAR;
-	params.session_token = this->SESSION_TOKEN_ENV_VAR;
-	params.endpoint = this->DUCKDB_ENDPOINT_ENV_VAR;
-	params.use_ssl = this->DUCKDB_USE_SSL_ENV_VAR;
+	params.region = DEFAULT_REGION_ENV_VAR;
+	params.region = REGION_ENV_VAR;
+	params.access_key_id = ACCESS_KEY_ENV_VAR;
+	params.secret_access_key = SECRET_KEY_ENV_VAR;
+	params.session_token = SESSION_TOKEN_ENV_VAR;
+	params.endpoint = DUCKDB_ENDPOINT_ENV_VAR;
+	params.use_ssl = DUCKDB_USE_SSL_ENV_VAR;
 
 	return params;
 }
@@ -183,12 +183,14 @@ unique_ptr<S3AuthParams> S3AuthParams::ReadFromStoredCredentials(optional_ptr<Fi
 	if (!opener) {
 		return nullptr;
 	}
-	auto context = opener->TryGetClientContext();
-	if (!context) {
+	auto db = opener->TryGetDatabase();
+	if (!db) {
 		return nullptr;
 	}
-	auto &secret_manager = context->db->GetSecretManager();
-	auto transaction = CatalogTransaction::GetSystemCatalogTransaction(*context);
+	auto &secret_manager = db->GetSecretManager();
+	auto context = opener->TryGetClientContext();
+	auto transaction = context ? CatalogTransaction::GetSystemCatalogTransaction(*context)
+	                           : CatalogTransaction::GetSystemTransaction(*db);
 
 	auto secret_match = secret_manager.LookupSecret(transaction, path, "s3");
 	if (!secret_match.HasMatch()) {

diff --git a/src/include/duckdb/common/file_opener.hpp b/src/include/duckdb/common/file_opener.hpp
@@ -31,8 +31,10 @@ class FileOpener {
 	virtual SettingLookupResult TryGetCurrentSetting(const string &key, Value &result, FileOpenerInfo &info);
 	virtual SettingLookupResult TryGetCurrentSetting(const string &key, Value &result) = 0;
 	virtual optional_ptr<ClientContext> TryGetClientContext() = 0;
+	virtual optional_ptr<DatabaseInstance> TryGetDatabase() = 0;
 
 	DUCKDB_API static optional_ptr<ClientContext> TryGetClientContext(optional_ptr<FileOpener> opener);
+	DUCKDB_API static optional_ptr<DatabaseInstance> TryGetDatabase(optional_ptr<FileOpener> opener);
 	DUCKDB_API static SettingLookupResult TryGetCurrentSetting(optional_ptr<FileOpener> opener, const string &key,
 	                                                           Value &result);
 	DUCKDB_API static SettingLookupResult TryGetCurrentSetting(optional_ptr<FileOpener> opener, const string &key,

diff --git a/src/include/duckdb/common/opener_file_system.hpp b/src/include/duckdb/common/opener_file_system.hpp
@@ -116,6 +116,14 @@ class OpenerFileSystem : public FileSystem {
 	std::string GetName() const override {
 		return "OpenerFileSystem - " + GetFileSystem().GetName();
 	}
+
+	void RegisterSubSystem(unique_ptr<FileSystem> sub_fs) override {
+		GetFileSystem().RegisterSubSystem(std::move(sub_fs));
+	}
+
+	void RegisterSubSystem(FileCompressionType compression_type, unique_ptr<FileSystem> fs) override {
+		GetFileSystem().RegisterSubSystem(compression_type, std::move(fs));
+	}
 };
 
 } // namespace duckdb
diff --git a/src/include/duckdb/main/client_context_file_opener.hpp b/src/include/duckdb/main/client_context_file_opener.hpp
@@ -26,7 +26,8 @@ class ClientContextFileOpener : public FileOpener {
 
 	optional_ptr<ClientContext> TryGetClientContext() override {
 		return &context;
-	};
+	}
+	optional_ptr<DatabaseInstance> TryGetDatabase() override;
 
 private:
 	ClientContext &context;

diff --git a/src/include/duckdb/main/database.hpp b/src/include/duckdb/main/database.hpp
@@ -25,6 +25,7 @@ class FileSystem;
 class TaskScheduler;
 class ObjectCache;
 struct AttachInfo;
+class DatabaseFileSystem;
 
 class DatabaseInstance : public std::enable_shared_from_this<DatabaseInstance> {
 	friend class DuckDB;
@@ -73,6 +74,7 @@ class DatabaseInstance : public std::enable_shared_from_this<DatabaseInstance> {
 	unique_ptr<ConnectionManager> connection_manager;
 	unordered_set<std::string> loaded_extensions;
 	ValidChecker db_validity;
+	unique_ptr<DatabaseFileSystem> db_file_system;
 };
 
 //! The database object. This object holds the catalog and all the

diff --git a/src/main/client_context_file_opener.cpp b/src/main/client_context_file_opener.cpp
@@ -14,13 +14,24 @@ SettingLookupResult ClientContextFileOpener::TryGetCurrentSetting(const string &
 	return context.TryGetCurrentSetting(key, result);
 }
 
+optional_ptr<DatabaseInstance> ClientContextFileOpener::TryGetDatabase() {
+	return context.db.get();
+}
+
 optional_ptr<ClientContext> FileOpener::TryGetClientContext(optional_ptr<FileOpener> opener) {
 	if (!opener) {
 		return nullptr;
 	}
 	return opener->TryGetClientContext();
 }
 
+optional_ptr<DatabaseInstance> FileOpener::TryGetDatabase(optional_ptr<FileOpener> opener) {
+	if (!opener) {
+		return nullptr;
+	}
+	return opener->TryGetDatabase();
+}
+
 SettingLookupResult FileOpener::TryGetCurrentSetting(optional_ptr<FileOpener> opener, const string &key,
                                                      Value &result) {
 	if (!opener) {

diff --git a/src/main/client_data.cpp b/src/main/client_data.cpp
@@ -24,6 +24,7 @@ class ClientFileSystem : public OpenerFileSystem {
 		auto &config = DBConfig::GetConfig(context);
 		return *config.file_system;
 	}
+
 	optional_ptr<FileOpener> GetOpener() const override {
 		return ClientData::Get(context).file_opener.get();
 	}

diff --git a/src/main/database.cpp b/src/main/database.cpp
@@ -22,6 +22,7 @@
 #include "duckdb/storage/storage_manager.hpp"
 #include "duckdb/transaction/transaction_manager.hpp"
 #include "duckdb/execution/index/index_type_set.hpp"
+#include "duckdb/main/database_file_opener.hpp"
 
 #ifndef DUCKDB_NO_THREADS
 #include "duckdb/common/thread.hpp"
@@ -221,6 +222,7 @@ void DatabaseInstance::Initialize(const char *database_path, DBConfig *user_conf
 		config.options.temporary_directory = string();
 	}
 
+	db_file_system = make_uniq<DatabaseFileSystem>(*this);
 	db_manager = make_uniq<DatabaseManager>(*this);
 	buffer_manager = make_uniq<StandardBufferManager>(*this, config.options.temporary_directory);
 	scheduler = make_uniq<TaskScheduler>(*this);
@@ -303,7 +305,7 @@ ObjectCache &DatabaseInstance::GetObjectCache() {
 }
 
 FileSystem &DatabaseInstance::GetFileSystem() {
-	return *config.file_system;
+	return *db_file_system;
 }
 
 ConnectionManager &DatabaseInstance::GetConnectionManager() {