Merge pull request #2 from archiewood:secret-manager

Secret manager

.gitignore

@@ -6,3 +6,5 @@ duckdb_unittest_tempdir/
 testext
 test/python/__pycache__/
 .Rhistory
+vcpkg_installed
+token.txt

src/gsheets_auth.cpp

@@ -1,5 +1,7 @@
 #include "gsheets_auth.hpp"
 #include "duckdb/common/exception.hpp"
+#include "duckdb/main/secret/secret.hpp"
+#include "duckdb/main/extension_util.hpp"
 #include <fstream>
 
 namespace duckdb {
@@ -14,4 +16,58 @@ std::string read_token_from_file(const std::string& file_path) {
     return token;
 }
 
+// This code is copied, with minor modifications from https://github.com/duckdb/duckdb_azure/blob/main/src/azure_secret.cpp
+static void CopySecret(const std::string &key, const CreateSecretInput &input, KeyValueSecret &result) {
+	auto val = input.options.find(key);
+
+	if (val != input.options.end()) {
+		result.secret_map[key] = val->second;
+	}
+}
+
+static void RegisterCommonSecretParameters(CreateSecretFunction &function) {
+	// Register google sheets common parameters
+	function.named_parameters["token"] = LogicalType::VARCHAR;
+}
+
+static void RedactCommonKeys(KeyValueSecret &result) {
+	result.redact_keys.insert("proxy_password");
+}
+
+
+// TODO: Maybe this should be a KeyValueSecret
+static unique_ptr<BaseSecret> CreateGsheetSecretFromAccessToken(ClientContext &context, CreateSecretInput &input) {
+	auto scope = input.scope;
+
+	auto result = make_uniq<KeyValueSecret>(scope, input.type, input.provider, input.name);
+
+	// Manage specific secret option
+	CopySecret("token", input, *result);
+
+	// Redact sensible keys
+	RedactCommonKeys(*result);
+	result->redact_keys.insert("token");
+
+	return std::move(result);
+}
+
+void CreateGsheetSecretFunctions::Register(DatabaseInstance &instance) {
+	string type = "gsheet";
+
+	// Register the new type
+	SecretType secret_type;
+	secret_type.name = type;
+	secret_type.deserializer = KeyValueSecret::Deserialize<KeyValueSecret>;
+	secret_type.default_provider = "access_token";
+	ExtensionUtil::RegisterSecretType(instance, secret_type);
+
+    // Register the access_token secret provider
+	CreateSecretFunction access_token_function = {type, "access_token", CreateGsheetSecretFromAccessToken};
+	access_token_function.named_parameters["access_token"] = LogicalType::VARCHAR;
+	RegisterCommonSecretParameters(access_token_function);
+	ExtensionUtil::RegisterFunction(instance, access_token_function);
+
+}
+
+
 }  // namespace duckdb

src/gsheets_extension.cpp

@@ -32,6 +32,10 @@ using json = nlohmann::json;
 
 #include <fstream>
 
+// Secrets
+#include "duckdb/main/secret/secret.hpp"
+#include "duckdb/main/secret/secret_manager.hpp"
+
 
 namespace duckdb {
 
@@ -143,8 +147,15 @@ SheetData parseJson(const std::string& json_str) {
             result.range = j["range"].get<std::string>();
             result.majorDimension = j["majorDimension"].get<std::string>();
             result.values = j["values"].get<std::vector<std::vector<std::string>>>();
+        } else if (j.contains("error")) {
+            string error = j["error"].get<std::string>();
+            string message = j["error"]["message"].get<std::string>();
+            string code = j["error"]["code"].get<std::string>();
+            throw std::runtime_error("Google Sheets API error: " + code + " - " + message);
         } else {
-            throw std::runtime_error("JSON does not contain expected fields");
+            std::cerr << "JSON does not contain expected fields" << std::endl;
+            std::cerr << "Raw JSON string: " << json_str << std::endl;
+            throw;
         }
     } catch (const json::exception& e) {
         std::cerr << "JSON parsing error: " << e.what() << std::endl;
@@ -211,14 +222,36 @@ static std::string extract_sheet_id(const std::string& input) {
 static unique_ptr<FunctionData> ReadSheetBind(ClientContext &context, TableFunctionBindInput &input,
                                               vector<LogicalType> &return_types, vector<string> &names) {
     auto sheet_input = input.inputs[0].GetValue<string>();
-    auto token_file_path = input.inputs[1].GetValue<string>();
-    bool header = input.inputs.size() > 2 ? input.inputs[2].GetValue<bool>() : true;
+    bool header = input.inputs.size() > 1 ? input.inputs[1].GetValue<bool>() : true;
 
     // Extract the sheet ID from the input (URL or ID)
     std::string sheet_id = extract_sheet_id(sheet_input);
 
-    // Use the read_token_from_file function from gsheets_auth.hpp
-    std::string token = read_token_from_file(token_file_path);
+    // Use the SecretManager to get the token
+    auto &secret_manager = SecretManager::Get(context);
+    auto transaction = CatalogTransaction::GetSystemCatalogTransaction(context);
+    auto secret_match = secret_manager.LookupSecret(transaction, "gsheet", "gsheet");
+
+    if (!secret_match.HasMatch()) {
+        throw InvalidInputException("No 'gsheet' secret found. Please create a secret with 'CREATE SECRET' first.");
+    }
+
+    auto &secret = secret_match.GetSecret();
+    if (secret.GetType() != "gsheet") {
+        throw InvalidInputException("Invalid secret type. Expected 'gsheet', got '%s'", secret.GetType());
+    }
+
+    const auto *kv_secret = dynamic_cast<const KeyValueSecret*>(&secret);
+    if (!kv_secret) {
+        throw InvalidInputException("Invalid secret format for 'gsheet' secret");
+    }
+
+    Value token_value;
+    if (!kv_secret->TryGetValue("token", token_value)) {
+        throw InvalidInputException("'token' not found in 'gsheet' secret");
+    }
+
+    std::string token = token_value.ToString();
 
     auto bind_data = make_uniq<ReadSheetBindData>(sheet_id, token, header);
 
@@ -249,9 +282,13 @@ static void LoadInternal(DatabaseInstance &instance) {
     OpenSSL_add_all_algorithms();
 
     // Register read_gsheet table function
-    TableFunction read_gsheet_function("read_gsheet", {LogicalType::VARCHAR, LogicalType::VARCHAR}, ReadSheetFunction, ReadSheetBind);
+    TableFunction read_gsheet_function("read_gsheet", {LogicalType::VARCHAR}, ReadSheetFunction, ReadSheetBind);
     read_gsheet_function.named_parameters["header"] = LogicalType::BOOLEAN;
     ExtensionUtil::RegisterFunction(instance, read_gsheet_function);
+
+    // Register Secret functions
+	CreateGsheetSecretFunctions::Register(instance);
+
 }
 
 void GsheetsExtension::Load(DuckDB &db) {

src/include/gsheets_auth.hpp

@@ -1,9 +1,16 @@
 #pragma once
 
 #include <string>
+#include "duckdb/main/database.hpp"
 
 namespace duckdb {
 
 std::string read_token_from_file(const std::string& file_path);
 
+struct CreateGsheetSecretFunctions {
+public:
+	//! Register all CreateSecretFunctions
+	static void Register(DatabaseInstance &instance);
+};
+
 }  // namespace duckdb

test/sql/gsheets.test

@@ -4,16 +4,20 @@
 
 # Before we load the extension, this will fail
 statement error
-FROM read_gsheet('11QdEasMWbETbFVxry-SsD8jVcdYIT1zBQszcF84MdE8', 'token.txt', header=true);
+FROM read_gsheet('11QdEasMWbETbFVxry-SsD8jVcdYIT1zBQszcF84MdE8', header=true);
 ----
 Catalog Error: Table Function with name read_gsheet does not exist!
 
 # Require statement will ensure this test is run with this extension loaded
 require gsheets
 
+# Create a secret NB must substitute a token, do not commit!
+statement ok
+create secret test_secret (type gsheet,token 'xxx');
+
 # Confirm the extension works
 query III
-FROM read_gsheet('11QdEasMWbETbFVxry-SsD8jVcdYIT1zBQszcF84MdE8', 'token.txt', header=true);
+FROM read_gsheet('11QdEasMWbETbFVxry-SsD8jVcdYIT1zBQszcF84MdE8', header=true);
 ----
 Alice	30	Toronto
 Bob	25	New York
@@ -24,11 +28,15 @@ Archie	99	NULL
 
 # Test the full URL
 query III
-FROM read_gsheet('https://docs.google.com/spreadsheets/d/11QdEasMWbETbFVxry-SsD8jVcdYIT1zBQszcF84MdE8/edit#gid=0', 'token.txt', header=true);
+FROM read_gsheet('https://docs.google.com/spreadsheets/d/11QdEasMWbETbFVxry-SsD8jVcdYIT1zBQszcF84MdE8/edit#gid=0', header=true);
 ----
 Alice	30	Toronto
 Bob	25	New York
 Charlie	45	Chicago
 Drake	NULL	NULL
 NULL	NULL	NULL
-Archie	99	NULL
+Archie	99	NULL
+
+# Drop the secret
+statement ok
+drop secret test_secret;