Synchronise extensions before reconnect (#316)

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

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.

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<int64>(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<int64>(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<duckdb::Connection>(*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 =

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<duckdb::QueryResult> (*DuckDBQueryOrThrow)(const std::string &);
+	auto result = pgduckdb::DuckDBFunctionGuard<duckdb::unique_ptr<duckdb::QueryResult>, DuckDBQueryOrThrow>(pgduckdb::DuckDBQueryOrThrow, "pgduckdb_raw_query", query);
 	elog(NOTICE, "result: %s", result->ToString().c_str());
 	PG_RETURN_BOOL(true);
 }

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)
+

test/regression/expected/secrets.out

@@ -78,4 +78,3 @@ pgduckb_secret_0
 
 (1 row)
 
-SET duckdb.force_execution TO true;

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

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' $$);

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;