Merge pull request #73 from samansmink/update-duckdb-for-1.1

Update duckdb for 1.1
duckdb · Aug 30, 2024 · b3e9e08 · b3e9e08
2 parents 1907766 + f87351b
commit b3e9e08
Show file tree

Hide file tree

Showing 6 changed files with 13 additions and 30 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) {

diff --git a/test/sql/cloud/hierarchical_namespace.test b/test/sql/cloud/hierarchical_namespace.test
@@ -76,9 +76,10 @@ SET azure_http_stats = true;
 query II
 EXPLAIN ANALYZE SELECT count(*) FROM 'abfss://testing-private/partitioned/l_receipmonth=*7/l_shipmode=TRUCK/*.csv';
 ----
-analyzed_plan	<REGEX>:.*HTTP Stats\:.*in\: 322\.0 KiB.*\#HEAD\: 1.*GET\: 4.*PUT\: 0.*\#POST\: 0.*
+analyzed_plan	<REGEX>:.*HTTP Stats.*in\: 322\.0 KiB.*\#HEAD\: 1.*GET\: 4.*PUT\: 0.*\#POST\: 0.*
+
 
 query II
 EXPLAIN ANALYZE SELECT count(*) FROM 'azure://testing-private/partitioned/l_receipmonth=*7/l_shipmode=TRUCK/*.csv';
 ----
-analyzed_plan	<REGEX>:.*HTTP Stats\:.*in\: 330\.1 KiB.*\#HEAD\: 1.*GET\: 2.*PUT\: 0.*\#POST\: 0.*
+analyzed_plan	<REGEX>:.*HTTP Stats.*in\: 330\.1 KiB.*\#HEAD\: 1.*GET\: 2.*PUT\: 0.*\#POST\: 0.*
diff --git a/test/sql/cloud/network_test.test b/test/sql/cloud/network_test.test
@@ -33,16 +33,16 @@ SET azure_http_stats = true;
 query II
 EXPLAIN ANALYZE SELECT sum(l_orderkey) FROM 'az://testing-private/l.parquet';
 ----
-analyzed_plan	<REGEX>:.*HTTP Stats\:.*in\: 4\.9 MiB.*\#HEAD\: 2.*GET\: 3.*PUT\: 0.*\#POST\: 0.*
+analyzed_plan	<REGEX>:.*HTTP Stats.*in\: 4\.9 MiB.*\#HEAD\: 2.*GET\: 3.*PUT\: 0.*\#POST\: 0.*
 
 # Redoing query should still result in same request count
 query II
 EXPLAIN ANALYZE SELECT sum(l_orderkey) FROM 'az://testing-private/l.parquet';
 ----
-analyzed_plan	<REGEX>:.*HTTP Stats\:.*in\: 4\.9 MiB.*\#HEAD\: 2.*GET\: 3.*PUT\: 0.*\#POST\: 0.*
+analyzed_plan	<REGEX>:.*HTTP Stats.*in\: 4\.9 MiB.*\#HEAD\: 2.*GET\: 3.*PUT\: 0.*\#POST\: 0.*
 
 # Testing public blobs
 query II
 EXPLAIN ANALYZE SELECT COUNT(*) FROM "azure://testing-public/l.parquet";
 ----
-analyzed_plan	<REGEX>:.*HTTP Stats\:.*in\: 4\.8 MiB.*\#HEAD\: 2.*GET\: 2.*PUT\: 0.*\#POST\: 0.*
+analyzed_plan	<REGEX>:.*HTTP Stats.*in\: 4\.8 MiB.*\#HEAD\: 2.*GET\: 2.*PUT\: 0.*\#POST\: 0.*