Skip to content

Commit

Permalink
Simple GUC variable for duckdb cloud secret
Browse files Browse the repository at this point in the history
* `quack.cloud_secret` provides way to store information about remote
  cloud service secret. Format of GUC should be <TYPE>#<KEY>#<SECRET>#<REGION> and variable is
  stored only in current session.
  • Loading branch information
mkaruza committed May 13, 2024
1 parent cd4417a commit 6179e2a
Show file tree
Hide file tree
Showing 4 changed files with 74 additions and 1 deletion.
1 change: 1 addition & 0 deletions include/quack/quack.h
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

// quack.c
extern int quack_max_threads_per_query;
extern char *quack_secret;
extern "C" void _PG_init(void);

// quack_hooks.c
Expand Down
24 changes: 24 additions & 0 deletions include/quack/quack_utils.hpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
#pragma once

#include <vector>
#include <string>
#include <sstream>

#include <cstdio>
#include <cstdarg>
#include <cstring>

namespace quack {

inline std::vector<std::string>
tokenizeString(char *str, const char delimiter) {
std::vector<std::string> v;
std::stringstream ss(str); // Turn the string into a stream.
std::string tok;
while (getline(ss, tok, delimiter)) {
v.push_back(tok);
}
return v;
};

} // namespace quack
39 changes: 38 additions & 1 deletion src/quack.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,12 @@ extern "C" {

#include "quack/quack.h"
#include "quack/quack_node.hpp"
#include "quack/quack_utils.hpp"

static void quack_init_guc(void);

int quack_max_threads_per_query = 1;
char *quack_secret = nullptr;

extern "C" {
PG_MODULE_MAGIC;
Expand All @@ -19,12 +21,36 @@ _PG_init(void) {
quack_init_hooks();
quack_init_node();
}

}

static bool
quack_cloud_secret_check_hooks(char **newval, void **extra, GucSource source) {

std::vector<std::string> tokens = quack::tokenizeString(*newval, '#');

if (tokens.size() == 0) {
return true;
}

if (tokens.size() != 4) {
elog(WARNING, "Incorrect quack.cloud_secret format.");
return false;
}

if (tokens[0].compare("S3")) {
elog(WARNING, "quack.cloud_secret supports only S3.");
return false;
}

return true;
}

/* clang-format off */
static void
quack_init_guc(void) {
DefineCustomIntVariable("quack.max_threads_per_query",

DefineCustomIntVariable("quack.max_threads_per_query",
gettext_noop("DuckDB max no. threads per query."),
NULL,
&quack_max_threads_per_query,
Expand All @@ -36,4 +62,15 @@ quack_init_guc(void) {
NULL,
NULL,
NULL);

DefineCustomStringVariable("quack.cloud_secret",
"Quack (duckdb) cloud secret GUC. Format is TYPE#ID#SECRET#REGION",
NULL,
&quack_secret,
"",
PGC_USERSET,
0,
&quack_cloud_secret_check_hooks,
NULL,
NULL);
}
11 changes: 11 additions & 0 deletions src/quack_planner.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ extern "C" {
#include "quack/quack_node.hpp"
#include "quack/quack_planner.hpp"
#include "quack/quack_types.hpp"
#include "quack/quack_utils.hpp"

namespace quack {

Expand Down Expand Up @@ -47,6 +48,16 @@ quack_create_plan(Query *parse, const char *query) {
catalog.CreateTableFunction(context, &heap_scan_info);
context.transaction.Commit();

if (strlen(quack_secret) != 0) {
std::vector<std::string> quackSecret = quack::tokenizeString(quack_secret, '#');
StringInfo s3SecretKey = makeStringInfo();
appendStringInfoString(s3SecretKey, "CREATE SECRET s3Secret ");
appendStringInfo(s3SecretKey, "(TYPE S3, KEY_ID '%s', SECRET '%s', REGION '%s');", quackSecret[1].c_str(),
quackSecret[2].c_str(), quackSecret[3].c_str());
context.Query(s3SecretKey->data, false);
pfree(s3SecretKey->data);
}

auto preparedQuery = context.Prepare(query);

if (preparedQuery->HasError()) {
Expand Down

0 comments on commit 6179e2a

Please sign in to comment.