Control extension state assigning to GUC variable

Author: kovmir

pg_show_plans--1.1--2.0.sql

@@ -2,3 +2,5 @@ DROP FUNCTION gsp_format_text;
 DROP FUNCTION gsp_format_json;
 DROP FUNCTION gsp_format_yaml;
 DROP FUNCTION gsp_format_xml;
+DROP FUNCTION pg_show_plans_enable;
+DROP FUNCTION pg_show_plans_disable;

pg_show_plans--2.0.sql

@@ -3,17 +3,6 @@
 -- complain if script is sourced in psql, rather than via CREATE EXTENSION
 \echo Use "CREATE EXTENSION pg_show_plans" to load this file. \quit
 
--- Register functions.
-CREATE FUNCTION pg_show_plans_enable()
-RETURNS void
-AS 'MODULE_PATHNAME'
-LANGUAGE C;
-
-CREATE FUNCTION pg_show_plans_disable()
-RETURNS void
-AS 'MODULE_PATHNAME'
-LANGUAGE C;
-
 CREATE FUNCTION pg_show_plans(
 	OUT pid int,
 	OUT level int,
@@ -30,7 +19,3 @@ CREATE VIEW pg_show_plans AS
 	SELECT * FROM pg_show_plans();
 
 GRANT SELECT ON pg_show_plans TO PUBLIC;
-
--- Don't want this to be available to non-superusers.
-REVOKE ALL ON FUNCTION pg_show_plans_enable() FROM PUBLIC;
-REVOKE ALL ON FUNCTION pg_show_plans_disable() FROM PUBLIC;

pg_show_plans.c

@@ -85,7 +85,7 @@ static void append_query_plan(ExplainState *es);
 /* on_shmem_exit() callback to delete hash entry on client disconnect. */
 static void cleanup(int code, Datum arg);
 /* Set extension state, either enable or disable. */
-static void set_state(const bool state);
+static void set_state(bool state, void *extra);
 static const char *show_state(void);
 /* Set query plan output format: text, json, ... */
 static void set_format(const int format);
@@ -108,14 +108,9 @@ static void pgsp_ExecutorStart(QueryDesc *queryDesc, int eflags);
 static void pgsp_ExecutorRun(QueryDesc *queryDesc, ScanDirection direction,
                              uint64 count, bool execute_once);
 
-/* Enables/Disables the extension. */
-Datum pg_show_plans_enable(PG_FUNCTION_ARGS);
-Datum pg_show_plans_disable(PG_FUNCTION_ARGS);
 /* Show query plans of all the currently running statements. */
 Datum pg_show_plans(PG_FUNCTION_ARGS);
 
-PG_FUNCTION_INFO_V1(pg_show_plans_enable);
-PG_FUNCTION_INFO_V1(pg_show_plans_disable);
 PG_FUNCTION_INFO_V1(pg_show_plans);
 
 /* Global Variables */
@@ -164,9 +159,9 @@ _PG_init(void)
 	                         NULL,
 	                         &start_enabled,
 	                         true,
-	                         PGC_POSTMASTER,
+	                         PGC_USERSET,
 	                         0,
-	                         NULL, NULL, show_state);
+	                         NULL, set_state, show_state);
 	DefineCustomIntVariable("pg_show_plans.max_plan_length",
 	                        gettext_noop("Set the maximum plan length. "
 	                                     "Note that this module allocates (max_plan_length*max_connections) "
@@ -318,10 +313,14 @@ cleanup(int code, Datum arg)
 }
 
 static void
-set_state(const bool state)
+set_state(bool state, void *extra)
 {
-	shmem_safety_check();
-	if (is_allowed_role())
+	/* Shared memory may not be fully available at server start, so we do not
+	 * check for pgsp_hash availability here. That is why the following line is
+	 * commented out. */
+	/* shmem_safety_check(); */
+
+	if (pgsp != NULL && is_allowed_role())
 		pgsp->is_enabled = state;
 }
 
@@ -510,20 +509,6 @@ pgsp_ExecutorRun(QueryDesc *queryDesc, ScanDirection direction,
 	PG_END_TRY();
 }
 
-Datum
-pg_show_plans_enable(PG_FUNCTION_ARGS)
-{
-	set_state(true);
-	PG_RETURN_VOID();
-}
-
-Datum
-pg_show_plans_disable(PG_FUNCTION_ARGS)
-{
-	set_state(false);
-	PG_RETURN_VOID();
-}
-
 Datum
 pg_show_plans(PG_FUNCTION_ARGS)
 {