Merge pull request duckdb#11726 from motherduckdb/which_secret_tf

feat: rewrite which_secret() into a table function

.github/patches/extensions/aws/0001-update-tests.patch

@@ -0,0 +1,53 @@
+From e8e6c286376d97e0a695284fc32b3b67a77e35af Mon Sep 17 00:00:00 2001
+From: stephaniewang <qw79@cornell.edu>
+Date: Thu, 18 Apr 2024 21:41:29 -0400
+Subject: [PATCH] update tests
+
+---
+ test/sql/aws_minio_secret.test | 2 +-
+ test/sql/aws_secret_gcs.test   | 2 +-
+ test/sql/aws_secret_r2.test    | 2 +-
+ 3 files changed, 3 insertions(+), 3 deletions(-)
+
+diff --git a/test/sql/aws_minio_secret.test b/test/sql/aws_minio_secret.test
+index 2ddc29f..34c4c92 100644
+--- a/test/sql/aws_minio_secret.test
++++ b/test/sql/aws_minio_secret.test
+@@ -28,7 +28,7 @@ CREATE SECRET my_aws_secret (
+ );
+
+ query I
+-SELECT which_secret('s3://test-bucket/aws_minio_secret/secret1/test.csv', 's3')
++SELECT name FROM which_secret('s3://test-bucket/aws_minio_secret/secret1/test.csv', 's3')
+ ----
+ my_aws_secret
+
+diff --git a/test/sql/aws_secret_gcs.test b/test/sql/aws_secret_gcs.test
+index 0b1fd40..cbed048 100644
+--- a/test/sql/aws_secret_gcs.test
++++ b/test/sql/aws_secret_gcs.test
+@@ -18,7 +18,7 @@ CREATE SECRET s1 (
+ );
+
+ query I
+-SELECT which_secret('gcs://haha/hoehoe.parkoe', 'gcs')
++SELECT name FROM which_secret('gcs://haha/hoehoe.parkoe', 'gcs')
+ ----
+ s1
+
+diff --git a/test/sql/aws_secret_r2.test b/test/sql/aws_secret_r2.test
+index 01be38b..19ebd1e 100644
+--- a/test/sql/aws_secret_r2.test
++++ b/test/sql/aws_secret_r2.test
+@@ -19,7 +19,7 @@ CREATE SECRET s1 (
+ );
+
+ query I
+-SELECT which_secret('r2://haha/hoehoe.parkoe', 'r2')
++SELECT name FROM which_secret('r2://haha/hoehoe.parkoe', 'r2')
+ ----
+ s1
+
+-- 
+2.39.2 (Apple Git-143)
+

scripts/generate_functions.py

@@ -15,7 +15,6 @@
     'math',
     'operators',
     'random',
-    'secret',
     'string',
     'debug',
     'struct',

src/core_functions/function_list.cpp

@@ -385,7 +385,6 @@ static const StaticFunctionDefinition internal_functions[] = {
 	DUCKDB_SCALAR_FUNCTION_SET(WeekFun),
 	DUCKDB_SCALAR_FUNCTION_SET(WeekDayFun),
 	DUCKDB_SCALAR_FUNCTION_SET(WeekOfYearFun),
-	DUCKDB_SCALAR_FUNCTION(WhichSecretFun),
 	DUCKDB_SCALAR_FUNCTION_SET(BitwiseXorFun),
 	DUCKDB_SCALAR_FUNCTION_SET(YearFun),
 	DUCKDB_SCALAR_FUNCTION_SET(YearWeekFun),

src/core_functions/scalar/CMakeLists.txt

@@ -9,7 +9,6 @@ add_subdirectory(map)
 add_subdirectory(math)
 add_subdirectory(operators)
 add_subdirectory(random)
-add_subdirectory(secret)
 add_subdirectory(string)
 add_subdirectory(struct)
 add_subdirectory(union)

src/function/table/system/CMakeLists.txt

@@ -13,6 +13,7 @@ add_library_unity(
   duckdb_optimizers.cpp
   duckdb_schemas.cpp
   duckdb_secrets.cpp
+  duckdb_which_secret.cpp
   duckdb_sequences.cpp
   duckdb_settings.cpp
   duckdb_tables.cpp

src/function/table/system/duckdb_which_secret.cpp

@@ -0,0 +1,75 @@
+#include "duckdb/function/table/system_functions.hpp"
+
+#include "duckdb/common/file_system.hpp"
+#include "duckdb/common/map.hpp"
+#include "duckdb/common/string_util.hpp"
+#include "duckdb/common/multi_file_reader.hpp"
+#include "duckdb/function/function_set.hpp"
+#include "duckdb/main/client_context.hpp"
+#include "duckdb/main/database.hpp"
+#include "duckdb/main/extension_helper.hpp"
+#include "duckdb/main/secret/secret_manager.hpp"
+
+namespace duckdb {
+
+struct DuckDBWhichSecretData : public GlobalTableFunctionState {
+	DuckDBWhichSecretData() : finished(false) {
+	}
+	bool finished;
+};
+
+struct DuckDBWhichSecretBindData : public TableFunctionData {
+	explicit DuckDBWhichSecretBindData(TableFunctionBindInput &tf_input) : inputs(tf_input.inputs) {};
+
+	duckdb::vector<duckdb::Value> inputs;
+};
+
+static unique_ptr<FunctionData> DuckDBWhichSecretBind(ClientContext &context, TableFunctionBindInput &input,
+                                                      vector<LogicalType> &return_types, vector<string> &names) {
+	names.emplace_back("name");
+	return_types.emplace_back(LogicalType::VARCHAR);
+
+	names.emplace_back("persistent");
+	return_types.emplace_back(LogicalType::VARCHAR);
+
+	names.emplace_back("storage");
+	return_types.emplace_back(LogicalType::VARCHAR);
+
+	return make_uniq<DuckDBWhichSecretBindData>(input);
+}
+
+unique_ptr<GlobalTableFunctionState> DuckDBWhichSecretInit(ClientContext &context, TableFunctionInitInput &input) {
+	return make_uniq<DuckDBWhichSecretData>();
+}
+
+void DuckDBWhichSecretFunction(ClientContext &context, TableFunctionInput &data_p, DataChunk &output) {
+	auto &data = data_p.global_state->Cast<DuckDBWhichSecretData>();
+	if (data.finished) {
+		// finished returning values
+		return;
+	}
+	auto &bind_data = data_p.bind_data->Cast<DuckDBWhichSecretBindData>();
+
+	auto &secret_manager = SecretManager::Get(context);
+	auto transaction = CatalogTransaction::GetSystemCatalogTransaction(context);
+
+	auto &inputs = bind_data.inputs;
+	auto path = inputs[0].ToString();
+	auto type = inputs[1].ToString();
+	auto secret_match = secret_manager.LookupSecret(transaction, path, type);
+	if (secret_match.HasMatch()) {
+		auto &secret_entry = *secret_match.secret_entry;
+		output.SetCardinality(1);
+		output.SetValue(0, 0, secret_entry.secret->GetName());
+		output.SetValue(1, 0, EnumUtil::ToString(secret_entry.persist_type));
+		output.SetValue(2, 0, secret_entry.storage_mode);
+	}
+	data.finished = true;
+}
+
+void DuckDBWhichSecretFun::RegisterFunction(BuiltinFunctions &set) {
+	set.AddFunction(TableFunction("which_secret", {duckdb::LogicalType::VARCHAR, duckdb::LogicalType::VARCHAR},
+	                              DuckDBWhichSecretFunction, DuckDBWhichSecretBind, DuckDBWhichSecretInit));
+}
+
+} // namespace duckdb

src/function/table/system_functions.cpp

@@ -30,6 +30,7 @@ void BuiltinFunctions::RegisterSQLiteFunctions() {
 	DuckDBMemoryFun::RegisterFunction(*this);
 	DuckDBOptimizersFun::RegisterFunction(*this);
 	DuckDBSecretsFun::RegisterFunction(*this);
+	DuckDBWhichSecretFun::RegisterFunction(*this);
 	DuckDBSequencesFun::RegisterFunction(*this);
 	DuckDBSettingsFun::RegisterFunction(*this);
 	DuckDBTablesFun::RegisterFunction(*this);

src/include/duckdb/function/table/system_functions.hpp

@@ -57,6 +57,10 @@ struct DuckDBSecretsFun {
 	static void RegisterFunction(BuiltinFunctions &set);
 };
 
+struct DuckDBWhichSecretFun {
+	static void RegisterFunction(BuiltinFunctions &set);
+};
+
 struct DuckDBDatabasesFun {
 	static void RegisterFunction(BuiltinFunctions &set);
 };

test/secrets/test_custom_secret_storage.cpp

@@ -91,8 +91,8 @@ TEST_CASE("Test secret lookups by secret type", "[secret][.]") {
 	REQUIRE_NO_FAIL(con.Query("CREATE SECRET s2 (TYPE secret_type_2, SCOPE '')"));
 
 	// Note that the secrets collide completely, except for their types
-	auto res1 = con.Query("SELECT which_secret('blablabla', 'secret_type_1')");
-	auto res2 = con.Query("SELECT which_secret('blablabla', 'secret_type_2')");
+	auto res1 = con.Query("SELECT name FROM which_secret('blablabla', 'secret_type_1')");
+	auto res2 = con.Query("SELECT name FROM which_secret('blablabla', 'secret_type_2')");
 
 	// Correct secret is selected
 	REQUIRE(res1->GetValue(0, 0).ToString() == "s1");
@@ -154,14 +154,14 @@ TEST_CASE("Test adding a custom secret storage", "[secret][.]") {
 	REQUIRE(secret_ptr->secret->GetName() == "s2");
 
 	// Now try resolve secret by path -> this will return s1 because its scope matches best
-	auto which_secret_result = con.Query("SELECT which_secret('s3://foo/bar.csv', 'S3');");
+	auto which_secret_result = con.Query("SELECT name FROM which_secret('s3://foo/bar.csv', 'S3');");
 	REQUIRE(which_secret_result->GetValue(0, 0).ToString() == "s1");
 
 	// Exclude the storage from lookups
 	storage_ref.include_in_lookups = false;
 
 	// Now the lookup will choose the other storage
-	which_secret_result = con.Query("SELECT which_secret('s3://foo/bar.csv', 's3');");
+	which_secret_result = con.Query("SELECT name FROM which_secret('s3://foo/bar.csv', 's3');");
 	REQUIRE(which_secret_result->GetValue(0, 0).ToString() == "s2");
 
 	// Lets drop stuff now
@@ -222,17 +222,17 @@ TEST_CASE("Test tie-break behaviour for custom secret storage", "[secret][.]") {
 	REQUIRE(result->GetValue(0, 2).ToString() == "s3");
 	REQUIRE(result->GetValue(1, 2).ToString() == "test_storage_before");
 
-	result = con.Query("SELECT which_secret('s3://', 's3');");
+	result = con.Query("SELECT name FROM which_secret('s3://', 's3');");
 	REQUIRE(result->GetValue(0, 0).ToString() == "s3");
 
 	REQUIRE_NO_FAIL(con.Query("DROP SECRET s3"));
 
-	result = con.Query("SELECT which_secret('s3://', 's3');");
+	result = con.Query("SELECT name FROM which_secret('s3://', 's3');");
 	REQUIRE(result->GetValue(0, 0).ToString() == "s1");
 
 	REQUIRE_NO_FAIL(con.Query("DROP SECRET s1"));
 
-	result = con.Query("SELECT which_secret('s3://', 's3');");
+	result = con.Query("SELECT name FROM which_secret('s3://', 's3');");
 	REQUIRE(result->GetValue(0, 0).ToString() == "s2");
 
 	REQUIRE_NO_FAIL(con.Query("DROP SECRET s2"));

test/sql/secrets/create_secret_scope_matching.test

@@ -12,6 +12,11 @@ require httpfs
 statement ok
 set secret_directory='__TEST_DIR__/create_secret_scope_matching'
 
+# No match
+query I
+SELECT name FROM which_secret('s3://', 's3')
+----
+
 statement ok
 CREATE TEMPORARY SECRET t1 ( TYPE S3 )
 
@@ -24,24 +29,29 @@ CREATE SECRET p1 IN LOCAL_FILE ( TYPE S3 )
 # This ties within the same storage: the two temporary secrets s1 and s2 both score identically. We solve this by
 # tie-breaking on secret name alphabetical ordering
 query I
-SELECT which_secret('s3://', 's3')
+SELECT name FROM which_secret('s3://', 's3')
 ----
 t1
 
+query III
+FROM which_secret('s3://', 's3')
+----
+t1	TEMPORARY	memory
+
 statement ok
 DROP SECRET t1
 
 # Temporary secrets take preference over temporary ones
 query I
-SELECT which_secret('s3://', 's3')
+SELECT name FROM which_secret('s3://', 's3')
 ----
 t2
 
 statement ok
 DROP SECRET t2
 
 query I
-SELECT which_secret('s3://', 's3')
+SELECT name FROM which_secret('s3://', 's3')
 ----
 p1