Add "--generate-helper-explain-analyze-role" argument for helper SQL

This optionally generates the SQL for the pganalyze.explain_analyze helper as part of the "--generate--helper-sql" command. The argument takes the role that the helper function should be owned by as a value. Note that the role creation, as well as any GRANT statements for the role must be taken care of separately.
pganalyze · Dec 24, 2024 · 94d5ab0 · 94d5ab0
1 parent 1530416
commit 94d5ab0
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 30 deletions.
diff --git a/main.go b/main.go
@@ -50,6 +50,7 @@ func main() {
 	var testExplain bool
 	var testSection string
 	var generateHelperSql string
+	var generateHelperExplainAnalyzeRole string
 	var forceStateUpdate bool
 	var configFilename string
 	var stateFilename string
@@ -75,6 +76,7 @@ func main() {
 	flag.StringVar(&testSection, "test-section", "", "Tests a particular section of the config file, i.e. a specific server, and ignores all other config sections")
 	flag.StringVar(&generateHelperSql, "generate-stats-helper-sql", "", "Deprecated alias for --generate-helper-sql")
 	flag.StringVar(&generateHelperSql, "generate-helper-sql", "", "Generates a SQL script for the given server (name of section in the config file, or \"default\" for env variables), that can be run with \"psql -f\" for installing the collector helpers on all configured databases")
+	flag.StringVar(&generateHelperExplainAnalyzeRole, "generate-helper-explain-analyze-role", "", "Includes pganalyze.explain_analyze helper in generated helper SQL, and sets owner role of helper function")
 	flag.BoolVar(&reloadRun, "reload", false, "Reloads the collector daemon that's running on the host")
 	flag.BoolVar(&noReload, "no-reload", false, "Disables automatic config reloading during a test run")
 	flag.BoolVarP(&logger.Verbose, "verbose", "v", false, "Outputs additional debugging information, use this if you're encountering errors or other problems")
@@ -151,28 +153,29 @@ func main() {
 	}
 
 	globalCollectionOpts := state.CollectionOpts{
-		StartedAt:                time.Now(),
-		SubmitCollectedData:      !benchmark && true,
-		TestRun:                  testRun,
-		TestRunLogs:              testRunLogs || dryRunLogs,
-		TestExplain:              testExplain,
-		TestSection:              testSection,
-		GenerateHelperSql:        generateHelperSql,
-		DebugLogs:                debugLogs,
-		DiscoverLogLocation:      discoverLogLocation,
-		CollectPostgresRelations: !noPostgresRelations,
-		CollectPostgresSettings:  !noPostgresSettings,
-		CollectPostgresLocks:     !noPostgresLocks,
-		CollectPostgresFunctions: !noPostgresFunctions,
-		CollectPostgresBloat:     !noPostgresBloat,
-		CollectPostgresViews:     !noPostgresViews,
-		CollectLogs:              !noLogs,
-		CollectExplain:           !noExplain,
-		CollectSystemInformation: !noSystemInformation,
-		StateFilename:            stateFilename,
-		WriteStateUpdate:         (!dryRun && !dryRunLogs && !testRun) || forceStateUpdate,
-		ForceEmptyGrant:          dryRun || dryRunLogs || testRunLogs || benchmark,
-		OutputAsJson:             !benchmark,
+		StartedAt:                        time.Now(),
+		SubmitCollectedData:              !benchmark && true,
+		TestRun:                          testRun,
+		TestRunLogs:                      testRunLogs || dryRunLogs,
+		TestExplain:                      testExplain,
+		TestSection:                      testSection,
+		GenerateHelperSql:                generateHelperSql,
+		GenerateHelperExplainAnalyzeRole: generateHelperExplainAnalyzeRole,
+		DebugLogs:                        debugLogs,
+		DiscoverLogLocation:              discoverLogLocation,
+		CollectPostgresRelations:         !noPostgresRelations,
+		CollectPostgresSettings:          !noPostgresSettings,
+		CollectPostgresLocks:             !noPostgresLocks,
+		CollectPostgresFunctions:         !noPostgresFunctions,
+		CollectPostgresBloat:             !noPostgresBloat,
+		CollectPostgresViews:             !noPostgresViews,
+		CollectLogs:                      !noLogs,
+		CollectExplain:                   !noExplain,
+		CollectSystemInformation:         !noSystemInformation,
+		StateFilename:                    stateFilename,
+		WriteStateUpdate:                 (!dryRun && !dryRunLogs && !testRun) || forceStateUpdate,
+		ForceEmptyGrant:                  dryRun || dryRunLogs || testRunLogs || benchmark,
+		OutputAsJson:                     !benchmark,
 	}
 
 	if reloadRun && !testRun {

diff --git a/runner/generate_helper_sql.go b/runner/generate_helper_sql.go
@@ -63,6 +63,13 @@ func GenerateHelperSql(ctx context.Context, server *state.Server, globalCollecti
 		for _, helper := range statsHelpers {
 			output.WriteString(helper + "\n")
 		}
+		if globalCollectionOpts.GenerateHelperExplainAnalyzeRole != "" {
+			output.WriteString(fmt.Sprintf("GRANT CREATE ON SCHEMA pganalyze TO %s;\n", globalCollectionOpts.GenerateHelperExplainAnalyzeRole))
+			output.WriteString(fmt.Sprintf("SET ROLE %s;\n", globalCollectionOpts.GenerateHelperExplainAnalyzeRole))
+			output.WriteString(util.ExplainAnalyzeHelper + "\n")
+			output.WriteString("RESET ROLE;\n")
+			output.WriteString(fmt.Sprintf("REVOKE CREATE ON SCHEMA pganalyze FROM %s;\n", globalCollectionOpts.GenerateHelperExplainAnalyzeRole))
+		}
 		output.WriteString("\n")
 	}
 

diff --git a/state/state.go b/state/state.go
@@ -212,14 +212,15 @@ type CollectionOpts struct {
 
 	DiffStatements bool
 
-	SubmitCollectedData bool
-	TestRun             bool
-	TestRunLogs         bool
-	TestExplain         bool
-	TestSection         string
-	GenerateHelperSql   string
-	DebugLogs           bool
-	DiscoverLogLocation bool
+	SubmitCollectedData              bool
+	TestRun                          bool
+	TestRunLogs                      bool
+	TestExplain                      bool
+	TestSection                      string
+	GenerateHelperSql                string
+	GenerateHelperExplainAnalyzeRole string
+	DebugLogs                        bool
+	DiscoverLogLocation              bool
 
 	StateFilename    string
 	WriteStateUpdate bool