From 3c283e33e26a46bf88c824314dacd434f3118838 Mon Sep 17 00:00:00 2001 From: mkaruza Date: Mon, 7 Oct 2024 14:19:25 +0200 Subject: [PATCH 1/2] Sync after each insert/update/delete of `duckdb.secrets` table If duckdb.secrets table is inserted, updated or deleted in same connection we would not sync changes for next query. Introduce sequence that keeps track of each change so we can compare with and sync if required. --- include/pgduckdb/pgduckdb_duckdb.hpp | 15 ++++++-- sql/pg_duckdb--0.0.1.sql | 15 ++++++++ src/pgduckdb_duckdb.cpp | 40 +++++++++++++++++-- test/regression/expected/secrets.out | 57 ++++++++++++++++++++++++++++ test/regression/schedule | 1 + test/regression/sql/secrets.sql | 23 +++++++++++ 6 files changed, 145 insertions(+), 6 deletions(-) create mode 100644 test/regression/expected/secrets.out create mode 100644 test/regression/sql/secrets.sql diff --git a/include/pgduckdb/pgduckdb_duckdb.hpp b/include/pgduckdb/pgduckdb_duckdb.hpp index 1243678d..8620e50c 100644 --- a/include/pgduckdb/pgduckdb_duckdb.hpp +++ b/include/pgduckdb/pgduckdb_duckdb.hpp @@ -12,25 +12,34 @@ namespace pgduckdb { class DuckDBManager { public: - static inline const DuckDBManager & + static inline DuckDBManager & Get() { static DuckDBManager instance; return instance; } inline duckdb::DuckDB & - GetDatabase() const { + GetDatabase() { + if(CheckSecretsSeq()) { + auto connection = duckdb::make_uniq(*database); + auto &context = *connection->context; + DropSecrets(context); + LoadSecrets(context); + } return *database; } private: DuckDBManager(); void InitializeDatabase(); - + bool CheckSecretsSeq(); void LoadSecrets(duckdb::ClientContext &); + void DropSecrets(duckdb::ClientContext &); void LoadExtensions(duckdb::ClientContext &); void LoadFunctions(duckdb::ClientContext &); + int secret_table_num_rows; + int secret_table_current_seq; duckdb::unique_ptr database; }; diff --git a/sql/pg_duckdb--0.0.1.sql b/sql/pg_duckdb--0.0.1.sql index 202e74b7..ba250396 100644 --- a/sql/pg_duckdb--0.0.1.sql +++ b/sql/pg_duckdb--0.0.1.sql @@ -108,6 +108,9 @@ $func$; CREATE SCHEMA duckdb; SET search_path TO duckdb; +CREATE SEQUENCE secret_table_seq START WITH 1 INCREMENT BY 1; +SELECT setval('secret_table_seq', 1); + CREATE TABLE secrets ( type TEXT NOT NULL, id TEXT NOT NULL, @@ -120,6 +123,18 @@ CREATE TABLE secrets ( CONSTRAINT type_constraint CHECK (type IN ('S3', 'GCS', 'R2')) ); +CREATE OR REPLACE FUNCTION duckdb_update_secret_table_seq() +RETURNS TRIGGER AS +$$ +BEGIN + PERFORM nextval('duckdb.secret_table_seq'); + RETURN NEW; +END; +$$ LANGUAGE PLpgSQL; + +CREATE TRIGGER secrets_table_seq_tr AFTER INSERT OR UPDATE OR DELETE ON secrets +EXECUTE FUNCTION duckdb_update_secret_table_seq(); + CREATE OR REPLACE FUNCTION duckdb_secret_r2_check() RETURNS TRIGGER AS $$ diff --git a/src/pgduckdb_duckdb.cpp b/src/pgduckdb_duckdb.cpp index 9c8e6c78..19c87150 100644 --- a/src/pgduckdb_duckdb.cpp +++ b/src/pgduckdb_duckdb.cpp @@ -5,6 +5,13 @@ #include "duckdb/catalog/catalog_search_path.hpp" #include "duckdb/main/extension_install_info.hpp" +extern "C" { +#include "postgres.h" +#include "catalog/namespace.h" +#include "utils/lsyscache.h" +#include "utils/fmgrprotos.h" +} + #include "pgduckdb/pgduckdb_options.hpp" #include "pgduckdb/pgduckdb_duckdb.hpp" #include "pgduckdb/scan/postgres_scan.hpp" @@ -69,7 +76,7 @@ GetExtensionDirectory() { return duckdb_extension_directory; } -DuckDBManager::DuckDBManager() { +DuckDBManager::DuckDBManager() : secret_table_num_rows(0), secret_table_current_seq(0) { elog(DEBUG2, "(PGDuckDB/DuckDBManager) Creating DuckDB instance"); duckdb::DBConfig config; @@ -88,7 +95,6 @@ DuckDBManager::DuckDBManager() { auto &context = *connection->context; context.Query("ATTACH DATABASE 'pgduckdb' (TYPE pgduckdb)", false); LoadFunctions(context); - LoadSecrets(context); LoadExtensions(context); } @@ -105,6 +111,19 @@ DuckDBManager::LoadFunctions(duckdb::ClientContext &context) { context.transaction.Commit(); } +bool +DuckDBManager::CheckSecretsSeq() { + Oid duckdb_namespace = get_namespace_oid("duckdb", false); + Oid secret_table_seq_oid = get_relname_relid("secret_table_seq", duckdb_namespace); + int64 seq = + PostgresFunctionGuard(DirectFunctionCall1Coll, pg_sequence_last_value, InvalidOid, secret_table_seq_oid); + if (secret_table_current_seq < seq) { + secret_table_current_seq = seq; + return true; + } + return false; +} + void DuckDBManager::LoadSecrets(duckdb::ClientContext &context) { auto duckdb_secrets = ReadDuckdbSecrets(); @@ -113,7 +132,7 @@ DuckDBManager::LoadSecrets(duckdb::ClientContext &context) { for (auto &secret : duckdb_secrets) { StringInfo secret_key = makeStringInfo(); bool is_r2_cloud_secret = (secret.type.rfind("R2", 0) == 0); - appendStringInfo(secret_key, "CREATE SECRET duckdbSecret_%d ", secret_id); + appendStringInfo(secret_key, "CREATE SECRET pgduckb_secret_%d ", secret_id); appendStringInfo(secret_key, "(TYPE %s, KEY_ID '%s', SECRET '%s'", secret.type.c_str(), secret.id.c_str(), secret.secret.c_str()); if (secret.region.length() && !is_r2_cloud_secret) { @@ -140,6 +159,21 @@ DuckDBManager::LoadSecrets(duckdb::ClientContext &context) { pfree(secret_key->data); secret_id++; } + + secret_table_num_rows = secret_id; +} + +void +DuckDBManager::DropSecrets(duckdb::ClientContext &context) { + + for (auto secret_id = 0; secret_id < secret_table_num_rows; secret_id++) { + auto drop_secret_cmd = duckdb::StringUtil::Format("DROP SECRET pgduckb_secret_%d;", secret_id); + auto res = context.Query(drop_secret_cmd, false); + if (res->HasError()) { + elog(ERROR, "(PGDuckDB/DropSecrets) secret `%d` could not be dropped.", secret_id); + } + } + secret_table_num_rows = 0; } void diff --git a/test/regression/expected/secrets.out b/test/regression/expected/secrets.out new file mode 100644 index 00000000..c408ed8b --- /dev/null +++ b/test/regression/expected/secrets.out @@ -0,0 +1,57 @@ +SET duckdb.execution TO false; +SELECT * FROM duckdb.raw_query($$ SELECT name FROM duckdb_secrets() $$); +NOTICE: result: name +VARCHAR +[ Rows: 0] + + + raw_query +----------- + +(1 row) + +-- INSERT SHOULD TRIGGER UPDATE OF SECRETS +INSERT INTO duckdb.secrets (type, id, secret, session_token, region) +VALUES ('S3', 'access_key_id_1', 'secret_access_key', 'session_token', 'us-east-1'); +SELECT * FROM duckdb.raw_query($$ SELECT name FROM duckdb_secrets() $$); +NOTICE: result: name +VARCHAR +[ Rows: 1] +pgduckb_secret_0 + + + raw_query +----------- + +(1 row) + +INSERT INTO duckdb.secrets (type, id, secret, session_token, region) +VALUES ('S3', 'access_key_id_2', 'secret_access_key', 'session_token', 'us-east-1'); +SELECT * FROM duckdb.raw_query($$ SELECT name FROM duckdb_secrets() $$); +NOTICE: result: name +VARCHAR +[ Rows: 2] +pgduckb_secret_0 +pgduckb_secret_1 + + + raw_query +----------- + +(1 row) + +-- DELETE SHOULD TRIGGER UPDATE OF SECRETS +DELETE FROM duckdb.secrets WHERE id = 'access_key_id_1'; +SELECT * FROM duckdb.raw_query($$ SELECT name FROM duckdb_secrets() $$); +NOTICE: result: name +VARCHAR +[ Rows: 1] +pgduckb_secret_0 + + + raw_query +----------- + +(1 row) + +SET duckdb.execution TO true; diff --git a/test/regression/schedule b/test/regression/schedule index ffe347cd..26043aa1 100644 --- a/test/regression/schedule +++ b/test/regression/schedule @@ -15,3 +15,4 @@ test: standard_conforming_strings test: query_filter test: altered_tables test: transaction_errors +test: secrets diff --git a/test/regression/sql/secrets.sql b/test/regression/sql/secrets.sql new file mode 100644 index 00000000..94a0a563 --- /dev/null +++ b/test/regression/sql/secrets.sql @@ -0,0 +1,23 @@ + +SET duckdb.execution TO false; + +SELECT * FROM duckdb.raw_query($$ SELECT name FROM duckdb_secrets() $$); + +-- INSERT SHOULD TRIGGER UPDATE OF SECRETS + +INSERT INTO duckdb.secrets (type, id, secret, session_token, region) +VALUES ('S3', 'access_key_id_1', 'secret_access_key', 'session_token', 'us-east-1'); + +SELECT * FROM duckdb.raw_query($$ SELECT name FROM duckdb_secrets() $$); + +INSERT INTO duckdb.secrets (type, id, secret, session_token, region) +VALUES ('S3', 'access_key_id_2', 'secret_access_key', 'session_token', 'us-east-1'); + +SELECT * FROM duckdb.raw_query($$ SELECT name FROM duckdb_secrets() $$); + +-- DELETE SHOULD TRIGGER UPDATE OF SECRETS +DELETE FROM duckdb.secrets WHERE id = 'access_key_id_1'; + +SELECT * FROM duckdb.raw_query($$ SELECT name FROM duckdb_secrets() $$); + +SET duckdb.execution TO true; From 68ac98804ee3c8d0e8a86a5facc700b8b5e15c87 Mon Sep 17 00:00:00 2001 From: mkaruza Date: Mon, 7 Oct 2024 16:30:50 +0200 Subject: [PATCH 2/2] Rename to secrets_table_seq and added extra test steps --- sql/pg_duckdb--0.0.1.sql | 10 +++++----- src/pgduckdb_duckdb.cpp | 2 +- test/regression/expected/secrets.out | 24 ++++++++++++++++++++++++ test/regression/sql/secrets.sql | 8 ++++++++ 4 files changed, 38 insertions(+), 6 deletions(-) diff --git a/sql/pg_duckdb--0.0.1.sql b/sql/pg_duckdb--0.0.1.sql index ba250396..d8fda6df 100644 --- a/sql/pg_duckdb--0.0.1.sql +++ b/sql/pg_duckdb--0.0.1.sql @@ -108,8 +108,8 @@ $func$; CREATE SCHEMA duckdb; SET search_path TO duckdb; -CREATE SEQUENCE secret_table_seq START WITH 1 INCREMENT BY 1; -SELECT setval('secret_table_seq', 1); +CREATE SEQUENCE secrets_table_seq START WITH 1 INCREMENT BY 1; +SELECT setval('secrets_table_seq', 1); CREATE TABLE secrets ( type TEXT NOT NULL, @@ -123,17 +123,17 @@ CREATE TABLE secrets ( CONSTRAINT type_constraint CHECK (type IN ('S3', 'GCS', 'R2')) ); -CREATE OR REPLACE FUNCTION duckdb_update_secret_table_seq() +CREATE OR REPLACE FUNCTION duckdb_update_secrets_table_seq() RETURNS TRIGGER AS $$ BEGIN - PERFORM nextval('duckdb.secret_table_seq'); + PERFORM nextval('duckdb.secrets_table_seq'); RETURN NEW; END; $$ LANGUAGE PLpgSQL; CREATE TRIGGER secrets_table_seq_tr AFTER INSERT OR UPDATE OR DELETE ON secrets -EXECUTE FUNCTION duckdb_update_secret_table_seq(); +EXECUTE FUNCTION duckdb_update_secrets_table_seq(); CREATE OR REPLACE FUNCTION duckdb_secret_r2_check() RETURNS TRIGGER AS diff --git a/src/pgduckdb_duckdb.cpp b/src/pgduckdb_duckdb.cpp index 19c87150..a9d3fdf9 100644 --- a/src/pgduckdb_duckdb.cpp +++ b/src/pgduckdb_duckdb.cpp @@ -114,7 +114,7 @@ DuckDBManager::LoadFunctions(duckdb::ClientContext &context) { bool DuckDBManager::CheckSecretsSeq() { Oid duckdb_namespace = get_namespace_oid("duckdb", false); - Oid secret_table_seq_oid = get_relname_relid("secret_table_seq", duckdb_namespace); + Oid secret_table_seq_oid = get_relname_relid("secrets_table_seq", duckdb_namespace); int64 seq = PostgresFunctionGuard(DirectFunctionCall1Coll, pg_sequence_last_value, InvalidOid, secret_table_seq_oid); if (secret_table_current_seq < seq) { diff --git a/test/regression/expected/secrets.out b/test/regression/expected/secrets.out index c408ed8b..9ae0cc70 100644 --- a/test/regression/expected/secrets.out +++ b/test/regression/expected/secrets.out @@ -10,9 +10,21 @@ VARCHAR (1 row) +SELECT last_value FROM duckdb.secrets_table_seq; + last_value +------------ + 1 +(1 row) + -- INSERT SHOULD TRIGGER UPDATE OF SECRETS INSERT INTO duckdb.secrets (type, id, secret, session_token, region) VALUES ('S3', 'access_key_id_1', 'secret_access_key', 'session_token', 'us-east-1'); +SELECT last_value FROM duckdb.secrets_table_seq; + last_value +------------ + 2 +(1 row) + SELECT * FROM duckdb.raw_query($$ SELECT name FROM duckdb_secrets() $$); NOTICE: result: name VARCHAR @@ -27,6 +39,12 @@ pgduckb_secret_0 INSERT INTO duckdb.secrets (type, id, secret, session_token, region) VALUES ('S3', 'access_key_id_2', 'secret_access_key', 'session_token', 'us-east-1'); +SELECT last_value FROM duckdb.secrets_table_seq; + last_value +------------ + 3 +(1 row) + SELECT * FROM duckdb.raw_query($$ SELECT name FROM duckdb_secrets() $$); NOTICE: result: name VARCHAR @@ -42,6 +60,12 @@ pgduckb_secret_1 -- DELETE SHOULD TRIGGER UPDATE OF SECRETS DELETE FROM duckdb.secrets WHERE id = 'access_key_id_1'; +SELECT last_value FROM duckdb.secrets_table_seq; + last_value +------------ + 4 +(1 row) + SELECT * FROM duckdb.raw_query($$ SELECT name FROM duckdb_secrets() $$); NOTICE: result: name VARCHAR diff --git a/test/regression/sql/secrets.sql b/test/regression/sql/secrets.sql index 94a0a563..00a0ac81 100644 --- a/test/regression/sql/secrets.sql +++ b/test/regression/sql/secrets.sql @@ -3,21 +3,29 @@ SET duckdb.execution TO false; SELECT * FROM duckdb.raw_query($$ SELECT name FROM duckdb_secrets() $$); +SELECT last_value FROM duckdb.secrets_table_seq; + -- INSERT SHOULD TRIGGER UPDATE OF SECRETS INSERT INTO duckdb.secrets (type, id, secret, session_token, region) VALUES ('S3', 'access_key_id_1', 'secret_access_key', 'session_token', 'us-east-1'); +SELECT last_value FROM duckdb.secrets_table_seq; + SELECT * FROM duckdb.raw_query($$ SELECT name FROM duckdb_secrets() $$); INSERT INTO duckdb.secrets (type, id, secret, session_token, region) VALUES ('S3', 'access_key_id_2', 'secret_access_key', 'session_token', 'us-east-1'); +SELECT last_value FROM duckdb.secrets_table_seq; + SELECT * FROM duckdb.raw_query($$ SELECT name FROM duckdb_secrets() $$); -- DELETE SHOULD TRIGGER UPDATE OF SECRETS DELETE FROM duckdb.secrets WHERE id = 'access_key_id_1'; +SELECT last_value FROM duckdb.secrets_table_seq; + SELECT * FROM duckdb.raw_query($$ SELECT name FROM duckdb_secrets() $$); SET duckdb.execution TO true;