From 25ce44d06b1f013a264a09cee31223a58ee1b2b4 Mon Sep 17 00:00:00 2001 From: "Y." Date: Mon, 21 Oct 2024 11:16:19 +0200 Subject: [PATCH] Synchronise extensions before reconnect (#316) --- include/pgduckdb/pgduckdb_duckdb.hpp | 25 ++++++- sql/pg_duckdb--0.0.1.sql | 17 +++++ src/pgduckdb_duckdb.cpp | 30 ++++---- src/pgduckdb_options.cpp | 3 +- test/regression/expected/extensions.out | 96 +++++++++++++++++++++++++ test/regression/expected/secrets.out | 1 - test/regression/schedule | 1 + test/regression/sql/extensions.sql | 28 ++++++++ test/regression/sql/secrets.sql | 2 - 9 files changed, 185 insertions(+), 18 deletions(-) create mode 100644 test/regression/expected/extensions.out create mode 100644 test/regression/sql/extensions.sql diff --git a/include/pgduckdb/pgduckdb_duckdb.hpp b/include/pgduckdb/pgduckdb_duckdb.hpp index 56f6b936..8c25b506 100644 --- a/include/pgduckdb/pgduckdb_duckdb.hpp +++ b/include/pgduckdb/pgduckdb_duckdb.hpp @@ -41,14 +41,35 @@ class DuckDBManager { void Initialize(); void InitializeDatabase(); - bool CheckSecretsSeq(); + void LoadSecrets(duckdb::ClientContext &); void DropSecrets(duckdb::ClientContext &); void LoadExtensions(duckdb::ClientContext &); void LoadFunctions(duckdb::ClientContext &); + inline bool + IsSecretSeqLessThan(int64 seq) const { + return secret_table_current_seq < seq; + } + + inline bool + IsExtensionsSeqLessThan(int64 seq) const { + return extensions_table_current_seq < seq; + } + + inline void + UpdateSecretSeq(int64 seq) { + secret_table_current_seq = seq; + } + + inline void + UpdateExtensionsSeq(int64 seq) { + extensions_table_current_seq = seq; + } + int secret_table_num_rows; - int secret_table_current_seq; + int64 secret_table_current_seq; + int64 extensions_table_current_seq; /* * FIXME: Use a unique_ptr instead of a raw pointer. For now this is not * possible though, as the MotherDuck extension causes an ABORT when the diff --git a/sql/pg_duckdb--0.0.1.sql b/sql/pg_duckdb--0.0.1.sql index 4303b6e5..f660d44d 100644 --- a/sql/pg_duckdb--0.0.1.sql +++ b/sql/pg_duckdb--0.0.1.sql @@ -197,11 +197,28 @@ $$ LANGUAGE PLpgSQL; CREATE TRIGGER duckdb_secret_r2_tr BEFORE INSERT OR UPDATE ON secrets FOR EACH ROW EXECUTE PROCEDURE duckdb_secret_r2_check(); +-- Extensions + +CREATE SEQUENCE extensions_table_seq START WITH 1 INCREMENT BY 1; +SELECT setval('extensions_table_seq', 1); + CREATE TABLE extensions ( name TEXT NOT NULL PRIMARY KEY, enabled BOOL DEFAULT TRUE ); +CREATE OR REPLACE FUNCTION duckdb_update_extensions_table_seq() +RETURNS TRIGGER AS +$$ +BEGIN + PERFORM nextval('duckdb.extensions_table_seq'); + RETURN NEW; +END; +$$ LANGUAGE PLpgSQL; + +CREATE TRIGGER extensions_table_seq_tr AFTER INSERT OR UPDATE OR DELETE ON extensions +EXECUTE FUNCTION duckdb_update_extensions_table_seq(); + -- The following might seem unnecesasry, but it's needed to know if a dropped -- table was a DuckDB table or not. See the comments and code in -- duckdb_drop_table_trigger for details. diff --git a/src/pgduckdb_duckdb.cpp b/src/pgduckdb_duckdb.cpp index 7d8854c0..041a8bc5 100644 --- a/src/pgduckdb_duckdb.cpp +++ b/src/pgduckdb_duckdb.cpp @@ -124,17 +124,11 @@ DuckDBManager::LoadFunctions(duckdb::ClientContext &context) { context.transaction.Commit(); } -bool -DuckDBManager::CheckSecretsSeq() { +int64 +GetSeqLastValue(const char *seq_name) { Oid duckdb_namespace = get_namespace_oid("duckdb", false); - 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) { - secret_table_current_seq = seq; - return true; - } - return false; + Oid table_seq_oid = get_relname_relid(seq_name, duckdb_namespace); + return PostgresFunctionGuard(DirectFunctionCall1Coll, pg_sequence_last_value, InvalidOid, table_seq_oid); } void @@ -176,8 +170,9 @@ 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); - pgduckdb::DuckDBQueryOrThrow(drop_secret_cmd); + pgduckdb::DuckDBQueryOrThrow(context, drop_secret_cmd); } + secret_table_num_rows = 0; } @@ -202,6 +197,8 @@ DuckDBManager::LoadExtensions(duckdb::ClientContext &context) { if (extension.name == "httpfs") { continue; } + + DuckDBQueryOrThrow(context, "INSTALL " + extension.name); DuckDBQueryOrThrow(context, "LOAD " + extension.name); } } @@ -211,9 +208,18 @@ DuckDBManager::CreateConnection() { auto &instance = Get(); auto connection = duckdb::make_uniq(*instance.database); auto &context = *connection->context; - if (instance.CheckSecretsSeq()) { + + const auto secret_table_last_seq = GetSeqLastValue("secrets_table_seq"); + if (instance.IsSecretSeqLessThan(secret_table_last_seq)) { instance.DropSecrets(context); instance.LoadSecrets(context); + instance.UpdateSecretSeq(secret_table_last_seq); + } + + const auto extensions_table_last_seq = GetSeqLastValue("extensions_table_seq"); + if (instance.IsExtensionsSeqLessThan(extensions_table_last_seq)) { + instance.LoadExtensions(context); + instance.UpdateExtensionsSeq(extensions_table_last_seq); } auto http_file_cache_set_dir_query = diff --git a/src/pgduckdb_options.cpp b/src/pgduckdb_options.cpp index 1c7411de..0648e4e3 100644 --- a/src/pgduckdb_options.cpp +++ b/src/pgduckdb_options.cpp @@ -209,7 +209,8 @@ PG_FUNCTION_INFO_V1(pgduckdb_raw_query); Datum pgduckdb_raw_query(PG_FUNCTION_ARGS) { const char *query = text_to_cstring(PG_GETARG_TEXT_PP(0)); - auto result = pgduckdb::DuckDBQueryOrThrow(query); + typedef duckdb::unique_ptr (*DuckDBQueryOrThrow)(const std::string &); + auto result = pgduckdb::DuckDBFunctionGuard, DuckDBQueryOrThrow>(pgduckdb::DuckDBQueryOrThrow, "pgduckdb_raw_query", query); elog(NOTICE, "result: %s", result->ToString().c_str()); PG_RETURN_BOOL(true); } diff --git a/test/regression/expected/extensions.out b/test/regression/expected/extensions.out new file mode 100644 index 00000000..ea0be3fb --- /dev/null +++ b/test/regression/expected/extensions.out @@ -0,0 +1,96 @@ +SET duckdb.force_execution TO false; +SELECT * FROM duckdb.raw_query($$ SELECT extension_name, loaded, installed FROM duckdb_extensions() WHERE loaded and extension_name != 'jemalloc' $$); +NOTICE: result: extension_name loaded installed +VARCHAR BOOLEAN BOOLEAN +[ Rows: 4] +cached_httpfs true true +json true true +parquet true true +pgduckdb true false + + + raw_query +----------- + +(1 row) + +SELECT last_value FROM duckdb.extensions_table_seq; + last_value +------------ + 1 +(1 row) + +-- INSERT SHOULD TRIGGER UPDATE OF EXTENSIONS +INSERT INTO duckdb.extensions (name, enabled) VALUES ('icu', TRUE); +SELECT last_value FROM duckdb.extensions_table_seq; + last_value +------------ + 2 +(1 row) + +SELECT * FROM duckdb.raw_query($$ SELECT extension_name, loaded, installed FROM duckdb_extensions() WHERE loaded and extension_name != 'jemalloc' $$); +NOTICE: result: extension_name loaded installed +VARCHAR BOOLEAN BOOLEAN +[ Rows: 5] +cached_httpfs true true +icu true true +json true true +parquet true true +pgduckdb true false + + + raw_query +----------- + +(1 row) + +INSERT INTO duckdb.extensions (name, enabled) VALUES ('aws', TRUE); +SELECT last_value FROM duckdb.extensions_table_seq; + last_value +------------ + 3 +(1 row) + +SELECT * FROM duckdb.raw_query($$ SELECT extension_name, loaded, installed FROM duckdb_extensions() WHERE loaded and extension_name != 'jemalloc' $$); +NOTICE: result: extension_name loaded installed +VARCHAR BOOLEAN BOOLEAN +[ Rows: 6] +aws true true +cached_httpfs true true +icu true true +json true true +parquet true true +pgduckdb true false + + + raw_query +----------- + +(1 row) + +-- DELETE SHOULD TRIGGER UPDATE OF EXTENSIONS +-- But we do not unload for now (would require a restart of DuckDB) +DELETE FROM duckdb.extensions WHERE name = 'aws'; +SELECT last_value FROM duckdb.extensions_table_seq; + last_value +------------ + 4 +(1 row) + +SELECT * FROM duckdb.raw_query($$ SELECT extension_name, loaded, installed FROM duckdb_extensions() WHERE loaded and extension_name != 'jemalloc' $$); +NOTICE: result: extension_name loaded installed +VARCHAR BOOLEAN BOOLEAN +[ Rows: 6] +aws true true +cached_httpfs true true +icu true true +json true true +parquet true true +pgduckdb true false + + + raw_query +----------- + +(1 row) + diff --git a/test/regression/expected/secrets.out b/test/regression/expected/secrets.out index da28a665..479b6a77 100644 --- a/test/regression/expected/secrets.out +++ b/test/regression/expected/secrets.out @@ -78,4 +78,3 @@ pgduckb_secret_0 (1 row) -SET duckdb.force_execution TO true; diff --git a/test/regression/schedule b/test/regression/schedule index d54393b5..d711987e 100644 --- a/test/regression/schedule +++ b/test/regression/schedule @@ -3,6 +3,7 @@ test: concurrency test: correct_null_conversions test: search_path test: execution_error +test: extensions test: type_support test: array_type_support test: views diff --git a/test/regression/sql/extensions.sql b/test/regression/sql/extensions.sql new file mode 100644 index 00000000..09f6df46 --- /dev/null +++ b/test/regression/sql/extensions.sql @@ -0,0 +1,28 @@ + +SET duckdb.force_execution TO false; + +SELECT * FROM duckdb.raw_query($$ SELECT extension_name, loaded, installed FROM duckdb_extensions() WHERE loaded and extension_name != 'jemalloc' $$); + +SELECT last_value FROM duckdb.extensions_table_seq; + +-- INSERT SHOULD TRIGGER UPDATE OF EXTENSIONS + +INSERT INTO duckdb.extensions (name, enabled) VALUES ('icu', TRUE); + +SELECT last_value FROM duckdb.extensions_table_seq; + +SELECT * FROM duckdb.raw_query($$ SELECT extension_name, loaded, installed FROM duckdb_extensions() WHERE loaded and extension_name != 'jemalloc' $$); + +INSERT INTO duckdb.extensions (name, enabled) VALUES ('aws', TRUE); + +SELECT last_value FROM duckdb.extensions_table_seq; + +SELECT * FROM duckdb.raw_query($$ SELECT extension_name, loaded, installed FROM duckdb_extensions() WHERE loaded and extension_name != 'jemalloc' $$); + +-- DELETE SHOULD TRIGGER UPDATE OF EXTENSIONS +-- But we do not unload for now (would require a restart of DuckDB) +DELETE FROM duckdb.extensions WHERE name = 'aws'; + +SELECT last_value FROM duckdb.extensions_table_seq; + +SELECT * FROM duckdb.raw_query($$ SELECT extension_name, loaded, installed FROM duckdb_extensions() WHERE loaded and extension_name != 'jemalloc' $$); diff --git a/test/regression/sql/secrets.sql b/test/regression/sql/secrets.sql index 6d6c6c6a..e6fc6f82 100644 --- a/test/regression/sql/secrets.sql +++ b/test/regression/sql/secrets.sql @@ -27,5 +27,3 @@ DELETE FROM duckdb.secrets WHERE key_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.force_execution TO true;