Merge branch 'main' into gdubya/main

duckdb · Aug 30, 2024 · 690dc5a · 690dc5a
2 parents f2c97d8 + 181966a
commit 690dc5a
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 25 deletions.
diff --git a/duckdb b/duckdb
diff --git a/extension-ci-tools b/extension-ci-tools
diff --git a/src/azure_filesystem.cpp b/src/azure_filesystem.cpp
@@ -170,21 +170,11 @@ shared_ptr<AzureContextState> AzureStorageFileSystem::GetOrCreateStorageContext(
 		auto context_key = GetContextPrefix() + parsed_url.storage_account_name;
 
 		auto &registered_state = client_context->registered_state;
-		auto storage_account_it = registered_state.find(context_key);
-		if (storage_account_it == registered_state.end()) {
+
+		result = registered_state->Get<AzureContextState>(context_key);
+		if (!result || !result->IsValid()) {
 			result = CreateStorageContext(opener, path, parsed_url);
-			registered_state.insert(std::make_pair(context_key, result));
-		} else {
-			auto *azure_context_state = static_cast<AzureContextState *>(storage_account_it->second.get());
-			// We keep the context valid until the QueryEnd (cf: AzureBlobContextState#QueryEnd())
-			// we do so because between queries the user can change the secret/variable that has been set
-			// the side effect of that is that we will reconnect (potentially retrieve a new token) on each request
-			if (!azure_context_state->IsValid()) {
-				result = CreateStorageContext(opener, path, parsed_url);
-				registered_state[context_key] = result;
-			} else {
-				result = shared_ptr<AzureContextState>(storage_account_it->second, azure_context_state);
-			}
+			registered_state->Insert(context_key, result);
 		}
 	} else {
 		result = CreateStorageContext(opener, path, parsed_url);

diff --git a/src/azure_http_state.cpp b/src/azure_http_state.cpp
@@ -14,15 +14,7 @@ void AzureHTTPState::Reset() {
 }
 
 shared_ptr<AzureHTTPState> AzureHTTPState::TryGetState(ClientContext &context) {
-	auto lookup = context.registered_state.find("azure_http_state");
-
-	if (lookup != context.registered_state.end()) {
-		return shared_ptr_cast<ClientContextState, AzureHTTPState>(lookup->second);
-	}
-
-	auto http_state = make_shared_ptr<AzureHTTPState>();
-	context.registered_state["azure_http_state"] = http_state;
-	return http_state;
+	return context.registered_state->GetOrCreate<AzureHTTPState>("azure_http_state");
 }
 
 shared_ptr<AzureHTTPState> AzureHTTPState::TryGetState(optional_ptr<FileOpener> opener) {