Avoid establishing separate connection to read Postgres settings

In almost all cases we have an existing database connection, so use that to retrieve the value of Postgres settings.
pganalyze · Jan 8, 2025 · 620a0e3 · 620a0e3
1 parent e07d78d
commit 620a0e3
Show file tree

Hide file tree

Showing 5 changed files with 16 additions and 13 deletions.
diff --git a/input/full.go b/input/full.go
@@ -76,7 +76,7 @@ func CollectFull(ctx context.Context, server *state.Server, connection *sql.DB,
 		logger.PrintError("Error collecting pg_stat_statements: %s", err)
 		var e *pq.Error
 		if errors.As(err, &e) && e.Code == "55000" && globalCollectionOpts.TestRun { // object_not_in_prerequisite_state
-			shared_preload_libraries, _ := postgres.GetPostgresSetting(ctx, "shared_preload_libraries", server, globalCollectionOpts, logger)
+			shared_preload_libraries, _ := postgres.GetPostgresSetting(ctx, connection, "shared_preload_libraries")
 			logger.PrintInfo("HINT - Current shared_preload_libraries setting: '%s'. Your Postgres server may need to be restarted for changes to take effect.", shared_preload_libraries)
 			server.SelfTest.HintCollectionAspect(state.CollectionAspectPgStatStatements, "Current shared_preload_libraries setting: '%s'. Your Postgres server may need to be restarted for changes to take effect.", shared_preload_libraries)
 		}

diff --git a/input/postgres/helpers.go b/input/postgres/helpers.go
@@ -9,7 +9,6 @@ import (
 
 	"github.com/guregu/null"
 	"github.com/pganalyze/collector/state"
-	"github.com/pganalyze/collector/util"
 )
 
 func unpackPostgresInt32Array(input null.String) (result []int32) {
@@ -72,16 +71,10 @@ SELECT setting
 	FROM pg_settings
  WHERE name = '%s'`
 
-func GetPostgresSetting(ctx context.Context, settingName string, server *state.Server, globalCollectionOpts state.CollectionOpts, prefixedLogger *util.Logger) (string, error) {
+func GetPostgresSetting(ctx context.Context, db *sql.DB, settingName string) (string, error) {
 	var value string
 
-	db, err := EstablishConnection(ctx, server, prefixedLogger, globalCollectionOpts, "")
-	if err != nil {
-		return "", fmt.Errorf("Could not connect to database to retrieve \"%s\": %s", settingName, err)
-	}
-
-	err = db.QueryRow(QueryMarkerSQL + fmt.Sprintf(settingValueSQL, settingName)).Scan(&value)
-	db.Close()
+	err := db.QueryRow(QueryMarkerSQL + fmt.Sprintf(settingValueSQL, settingName)).Scan(&value)
 	if err != nil {
 		return "", fmt.Errorf("Could not read \"%s\" setting: %s", settingName, err)
 	}

diff --git a/input/postgres/plans.go b/input/postgres/plans.go
@@ -35,7 +35,7 @@ func GetPlans(ctx context.Context, server *state.Server, logger *util.Logger, db
 		return nil, nil, nil
 	}
 
-	computePlanIdEnabled, err := GetPostgresSetting(ctx, "aurora_compute_plan_id", server, globalCollectionOpts, logger)
+	computePlanIdEnabled, err := GetPostgresSetting(ctx, db, "aurora_compute_plan_id")
 	if err != nil {
 		if globalCollectionOpts.TestRun {
 			logger.PrintInfo("Function aurora_stat_plans() is not supported because Aurora version is too old. Upgrade to Aurora PostgreSQL version 14.10, 15.5, or later versions to collect query plans and stats.")

diff --git a/runner/activity.go b/runner/activity.go
@@ -76,7 +76,7 @@ func processActivityForServer(ctx context.Context, server *state.Server, globalC
 		defer connection.Close()
 	}
 
-	trackActivityQuerySize, err := postgres.GetPostgresSetting(ctx, "track_activity_query_size", server, globalCollectionOpts, logger)
+	trackActivityQuerySize, err := postgres.GetPostgresSetting(ctx, connection, "track_activity_query_size")
 	if err != nil {
 		activity.TrackActivityQuerySize = -1
 	} else {

diff --git a/runner/logs.go b/runner/logs.go
@@ -376,7 +376,7 @@ func TestLogsForAllServers(ctx context.Context, servers []*state.Server, globalC
 			continue
 		}
 
-		logLinePrefix, err := postgres.GetPostgresSetting(ctx, "log_line_prefix", server, globalCollectionOpts, prefixedLogger)
+		logLinePrefix, err := getLogLinePrefix(ctx, server, globalCollectionOpts, prefixedLogger)
 		if err != nil {
 			prefixedLogger.PrintError("ERROR - Could not check log_line_prefix for server: %s", err)
 			hasFailedServers = true
@@ -426,6 +426,16 @@ func TestLogsForAllServers(ctx context.Context, servers []*state.Server, globalC
 	return
 }
 
+func getLogLinePrefix(ctx context.Context, server *state.Server, globalCollectionOpts state.CollectionOpts, prefixedLogger *util.Logger) (string, error) {
+	db, err := postgres.EstablishConnection(ctx, server, prefixedLogger, globalCollectionOpts, "")
+	if err != nil {
+		return "", fmt.Errorf("Could not connect to database to retrieve log_line_prefix: %s", err)
+	}
+	defer db.Close()
+
+	return postgres.GetPostgresSetting(ctx, db, "log_line_prefix")
+}
+
 func testLocalLogTail(ctx context.Context, wg *sync.WaitGroup, server *state.Server, globalCollectionOpts state.CollectionOpts, logger *util.Logger) bool {
 	logger.PrintInfo("Testing log collection (local)...")