From b37dd684da3ff6f9c00719c27b89f1e39f6961a4 Mon Sep 17 00:00:00 2001 From: Michael Paquier Date: Wed, 4 Oct 2023 14:28:20 +0900 Subject: [PATCH v3 1/2] Redesign database and role handling in worker_spi worker_spi_launch gains two arguments for a database ID and a role ID. If these are InvalidOid, the dynamic worker falls back to the GUCs defined to launch the workers. Tests are refactored in consequence. bgw_extra is used to pass down these two OIDs to the main bgworker routine. Workers loaded with shared_preload_libraries just rely on the GUCs. --- .../modules/worker_spi/t/001_worker_spi.pl | 21 +++++--- .../modules/worker_spi/worker_spi--1.0.sql | 2 +- src/test/modules/worker_spi/worker_spi.c | 54 ++++++++++++++++++- 3 files changed, 67 insertions(+), 10 deletions(-) diff --git a/src/test/modules/worker_spi/t/001_worker_spi.pl b/src/test/modules/worker_spi/t/001_worker_spi.pl index 2965acd789..a026042c65 100644 --- a/src/test/modules/worker_spi/t/001_worker_spi.pl +++ b/src/test/modules/worker_spi/t/001_worker_spi.pl @@ -20,9 +20,11 @@ $node->safe_psql('postgres', 'CREATE EXTENSION worker_spi;'); # This consists in making sure that a table name "counted" is created # on a new schema whose name includes the index defined in input argument # of worker_spi_launch(). -# By default, dynamic bgworkers connect to the "postgres" database. +# By default, dynamic bgworkers connect to the "postgres" database with +# an undefined role, falling back to the GUC defaults (InvalidOid for +# worker_spi_launch). my $result = - $node->safe_psql('postgres', 'SELECT worker_spi_launch(4) IS NOT NULL;'); + $node->safe_psql('postgres', 'SELECT worker_spi_launch(4, 0, 0) IS NOT NULL;'); is($result, 't', "dynamic bgworker launched"); $node->poll_query_until( 'postgres', @@ -58,6 +60,7 @@ note "testing bgworkers loaded with shared_preload_libraries"; # Create the database first so as the workers can connect to it when # the library is loaded. $node->safe_psql('postgres', q(CREATE DATABASE mydb;)); +$node->safe_psql('postgres', q(CREATE ROLE myrole SUPERUSER LOGIN;)); $node->safe_psql('mydb', 'CREATE EXTENSION worker_spi;'); # Now load the module as a shared library. @@ -80,16 +83,18 @@ ok( $node->poll_query_until( # Ask worker_spi to launch dynamic bgworkers with the library loaded, then # check their existence. Use IDs that do not overlap with the schemas created -# by the previous workers. -my $worker1_pid = $node->safe_psql('mydb', 'SELECT worker_spi_launch(10);'); -my $worker2_pid = $node->safe_psql('mydb', 'SELECT worker_spi_launch(11);'); +# by the previous workers. These use a specific role. +my $myrole_id = $node->safe_psql('mydb', "SELECT oid FROM pg_roles where rolname = 'myrole';"); +my $mydb_id = $node->safe_psql('mydb', "SELECT oid FROM pg_database where datname = 'mydb';"); +my $worker1_pid = $node->safe_psql('mydb', "SELECT worker_spi_launch(10, $mydb_id, $myrole_id);"); +my $worker2_pid = $node->safe_psql('mydb', "SELECT worker_spi_launch(11, $mydb_id, $myrole_id);"); ok( $node->poll_query_until( 'mydb', - qq[SELECT datname, count(datname), wait_event FROM pg_stat_activity + qq[SELECT datname, usename, count(datname), wait_event FROM pg_stat_activity WHERE backend_type = 'worker_spi dynamic' AND - pid IN ($worker1_pid, $worker2_pid) GROUP BY datname, wait_event;], - 'mydb|2|worker_spi_main'), + pid IN ($worker1_pid, $worker2_pid) GROUP BY datname, usename, wait_event;], + 'mydb|myrole|2|worker_spi_main'), 'dynamic bgworkers all launched' ) or die "Timed out while waiting for dynamic bgworkers to be launched"; diff --git a/src/test/modules/worker_spi/worker_spi--1.0.sql b/src/test/modules/worker_spi/worker_spi--1.0.sql index e9d5b07373..f5b91128ff 100644 --- a/src/test/modules/worker_spi/worker_spi--1.0.sql +++ b/src/test/modules/worker_spi/worker_spi--1.0.sql @@ -3,7 +3,7 @@ -- complain if script is sourced in psql, rather than via CREATE EXTENSION \echo Use "CREATE EXTENSION worker_spi" to load this file. \quit -CREATE FUNCTION worker_spi_launch(pg_catalog.int4) +CREATE FUNCTION worker_spi_launch(pg_catalog.int4, pg_catalog.oid, pg_catalog.oid) RETURNS pg_catalog.int4 STRICT AS 'MODULE_PATHNAME' LANGUAGE C; diff --git a/src/test/modules/worker_spi/worker_spi.c b/src/test/modules/worker_spi/worker_spi.c index 98f8d4194b..f35897ddac 100644 --- a/src/test/modules/worker_spi/worker_spi.c +++ b/src/test/modules/worker_spi/worker_spi.c @@ -34,10 +34,12 @@ /* these headers are used by this particular worker's code */ #include "access/xact.h" +#include "commands/dbcommands.h" #include "executor/spi.h" #include "fmgr.h" #include "lib/stringinfo.h" #include "pgstat.h" +#include "utils/acl.h" #include "utils/builtins.h" #include "utils/snapmgr.h" #include "tcop/utility.h" @@ -52,6 +54,7 @@ PGDLLEXPORT void worker_spi_main(Datum main_arg) pg_attribute_noreturn(); static int worker_spi_naptime = 10; static int worker_spi_total_workers = 2; static char *worker_spi_database = NULL; +static char *worker_spi_role = NULL; /* value cached, fetched from shared memory */ static uint32 worker_spi_wait_event_main = 0; @@ -138,12 +141,21 @@ worker_spi_main(Datum main_arg) worktable *table; StringInfoData buf; char name[20]; + Oid dboid; + Oid roleoid; + char *p; table = palloc(sizeof(worktable)); sprintf(name, "schema%d", index); table->schema = pstrdup(name); table->name = pstrdup("counted"); + /* fetch database and role OIDs, these are set for a dynamic worker */ + p = MyBgworkerEntry->bgw_extra; + memcpy(&dboid, p, sizeof(Oid)); + p += sizeof(Oid); + memcpy(&roleoid, p, sizeof(Oid)); + /* Establish signal handlers before unblocking signals. */ pqsignal(SIGHUP, SignalHandlerForConfigReload); pqsignal(SIGTERM, die); @@ -152,7 +164,10 @@ worker_spi_main(Datum main_arg) BackgroundWorkerUnblockSignals(); /* Connect to our database */ - BackgroundWorkerInitializeConnection(worker_spi_database, NULL, 0); + if (OidIsValid(dboid)) + BackgroundWorkerInitializeConnectionByOid(dboid, roleoid, 0); + else + BackgroundWorkerInitializeConnection(worker_spi_database, worker_spi_role, 0); elog(LOG, "%s initialized with %s.%s", MyBgworkerEntry->bgw_name, table->schema, table->name); @@ -316,6 +331,15 @@ _PG_init(void) 0, NULL, NULL, NULL); + DefineCustomStringVariable("worker_spi.role", + "Role to connect with.", + NULL, + &worker_spi_role, + NULL, + PGC_SIGHUP, + 0, + NULL, NULL, NULL); + if (!process_shared_preload_libraries_in_progress) return; @@ -346,6 +370,9 @@ _PG_init(void) /* * Now fill in worker-specific data, and do the actual registrations. + * + * bgw_extra can optionally include a dabatase OID and a role OID. This + * is left empty here to fallback to the related GUCs at startup. */ for (int i = 1; i <= worker_spi_total_workers; i++) { @@ -364,10 +391,13 @@ Datum worker_spi_launch(PG_FUNCTION_ARGS) { int32 i = PG_GETARG_INT32(0); + Oid dboid = PG_GETARG_OID(1); + Oid roleoid = PG_GETARG_OID(2); BackgroundWorker worker; BackgroundWorkerHandle *handle; BgwHandleStatus status; pid_t pid; + char *p; memset(&worker, 0, sizeof(worker)); worker.bgw_flags = BGWORKER_SHMEM_ACCESS | @@ -382,6 +412,28 @@ worker_spi_launch(PG_FUNCTION_ARGS) /* set bgw_notify_pid so that we can use WaitForBackgroundWorkerStartup */ worker.bgw_notify_pid = MyProcPid; + /* + * Register database and role to use for the worker started in bgw_extra. + * If none have been provided, fallback to the GUCs. + */ + if (!OidIsValid(dboid)) + dboid = get_database_oid(worker_spi_database, false); + if (!OidIsValid(roleoid)) + { + /* + * worker_spi_role is NULL by default, so just pass down an invalid + * OID to let the main() routine do its connection work. + */ + if (worker_spi_role) + roleoid = get_role_oid(worker_spi_role, false); + else + roleoid = InvalidOid; + } + p = worker.bgw_extra; + memcpy(p, &dboid, sizeof(Oid)); + p += sizeof(Oid); + memcpy(p, &roleoid, sizeof(Oid)); + if (!RegisterDynamicBackgroundWorker(&worker, &handle)) PG_RETURN_NULL(); -- 2.42.0