Postgres: Switch the datasource plugin from lib/pq to pgx (#103961)

* Create libpqToPGX feature toggle * Refactor PostgreSQL datasource to support PGX with feature toggle - Updated `ProvideService` to accept feature toggles for enabling PGX. - Modified integration tests to use the new PGX connection method. - Introduced new functions for handling PGX connections and queries. - Enhanced TLS configuration handling for PostgreSQL connections. - Updated existing tests to ensure compatibility with PGX and new connection methods. * Update PostgreSQL datasource to enhance connection pooling and error handling - Increased `MaxOpenConns` to 10 in integration tests for improved connection management. - Refactored connection handling in `newPostgresPGX` to return a connection pool instead of a single connection. - Updated health check error handling to utilize context and feature toggles for better error reporting. - Adjusted `DisposePGX` method to close the connection pool properly. - Enhanced query execution to acquire connections from the pool, ensuring efficient resource usage. * Cleanup * Revert postgres_test unnecessary changes * Rename feature toggle from `libpqToPGX` to `postgresDSUsePGX` * Add null check to dispose method * Fix lint issues * Refactor connection string generation * Address comment in health check file * Rename p to pool * Refactor executeQueryPGX and split into multiple functions * Fix lint issues * The returning error message from PGX is enough no need to separate the error code. * Move TLS handling to newPostgresPGX function * Disable ssl for integration tests * Use MaxIdleConns option * Remove old feature toggle * Rename`generateConnectionConfigPGX` to `generateConnectionStringPGX` * Add back part of the error messages * Don't show max idle connections option when PGX enabled * Address comments from Sriram * Add back Sriram's changes * PostgreSQL: Rework tls manager to use temporary files instead (#105330) * Rework tls manager to use temporary files instead * Lint and test fixes * Update pkg/tsdb/grafana-postgresql-datasource/postgres.go Co-authored-by: Ivana Huckova <30407135+ivanahuckova@users.noreply.github.com> * Update betterer --------- Co-authored-by: Ivana Huckova <30407135+ivanahuckova@users.noreply.github.com> --------- Co-authored-by: Ivana Huckova <30407135+ivanahuckova@users.noreply.github.com>
2025-07-31 05:32:15 +08:00 · 2025-05-26 08:54:18 +02:00
parent 745eda2848
commit 1e383b0c1e
25 changed files with 1233 additions and 778 deletions
--- a/pkg/tsdb/grafana-postgresql-datasource/postgres.go
+++ b/pkg/tsdb/grafana-postgresql-datasource/postgres.go
@ -15,27 +15,30 @@ import (
 	"github.com/grafana/grafana-plugin-sdk-go/backend/instancemgmt"
 	"github.com/grafana/grafana-plugin-sdk-go/data"
 	"github.com/grafana/grafana-plugin-sdk-go/data/sqlutil"
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
+	"github.com/jackc/pgx/v5/pgxpool"
 	"github.com/lib/pq"

 	"github.com/grafana/grafana-plugin-sdk-go/backend/log"
-	"github.com/grafana/grafana/pkg/setting"
 	"github.com/grafana/grafana/pkg/tsdb/grafana-postgresql-datasource/sqleng"
 )

-func ProvideService(cfg *setting.Cfg) *Service {
+func ProvideService(features featuremgmt.FeatureToggles) *Service {
 	logger := backend.NewLoggerWith("logger", "tsdb.postgres")
 	s := &Service{
-		tlsManager: newTLSManager(logger, cfg.DataPath),
+		tlsManager: newTLSManager(logger),
 		logger:     logger,
+		features:   features,
 	}
 	s.im = datasource.NewInstanceManager(s.newInstanceSettings())
 	return s
 }

 type Service struct {
-	tlsManager tlsSettingsProvider
+	tlsManager *tlsManager
 	im         instancemgmt.InstanceManager
 	logger     log.Logger
+	features   featuremgmt.FeatureToggles
 }

 func (s *Service) getDSInfo(ctx context.Context, pluginCtx backend.PluginContext) (*sqleng.DataSourceHandler, error) {
@ -52,6 +55,11 @@ func (s *Service) QueryData(ctx context.Context, req *backend.QueryDataRequest)
 	if err != nil {
 		return nil, err
 	}
+
+	if s.features.IsEnabled(ctx, featuremgmt.FlagPostgresDSUsePGX) {
+		return dsInfo.QueryDataPGX(ctx, req)
+	}
+
 	return dsInfo.QueryData(ctx, req)
 }

@ -93,6 +101,13 @@ func newPostgres(ctx context.Context, userFacingDefaultError string, rowLimit in
 	db.SetMaxIdleConns(config.DSInfo.JsonData.MaxIdleConns)
 	db.SetConnMaxLifetime(time.Duration(config.DSInfo.JsonData.ConnMaxLifetime) * time.Second)

+	// We need to ping the database to ensure that the connection is valid and the temporary files are not deleted
+	// before the connection is used.
+	if err := db.Ping(); err != nil {
+		logger.Error("Failed to ping Postgres database", "error", err)
+		return nil, nil, backend.DownstreamError(fmt.Errorf("failed to ping Postgres database: %w", err))
+	}
+
 	handler, err := sqleng.NewQueryDataHandler(userFacingDefaultError, db, config, &queryResultTransformer, newPostgresMacroEngine(dsInfo.JsonData.Timescaledb),
 		logger)
 	if err != nil {
@ -104,6 +119,62 @@ func newPostgres(ctx context.Context, userFacingDefaultError string, rowLimit in
 	return db, handler, nil
 }

+func newPostgresPGX(ctx context.Context, userFacingDefaultError string, rowLimit int64, dsInfo sqleng.DataSourceInfo, cnnstr string, logger log.Logger, settings backend.DataSourceInstanceSettings) (*pgxpool.Pool, *sqleng.DataSourceHandler, error) {
+	pgxConf, err := pgxpool.ParseConfig(cnnstr)
+	if err != nil {
+		logger.Error("postgres config creation failed", "error", err)
+		return nil, nil, fmt.Errorf("postgres config creation failed")
+	}
+
+	proxyClient, err := settings.ProxyClient(ctx)
+	if err != nil {
+		logger.Error("postgres proxy creation failed", "error", err)
+		return nil, nil, fmt.Errorf("postgres proxy creation failed")
+	}
+
+	if proxyClient.SecureSocksProxyEnabled() {
+		dialer, err := proxyClient.NewSecureSocksProxyContextDialer()
+		if err != nil {
+			logger.Error("postgres proxy creation failed", "error", err)
+			return nil, nil, fmt.Errorf("postgres proxy creation failed")
+		}
+
+		pgxConf.ConnConfig.DialFunc = newPgxDialFunc(dialer)
+	}
+
+	// by default pgx resolves hostnames to ip addresses. we must avoid this.
+	// (certain socks-proxy related functionality relies on the hostname being preserved)
+	pgxConf.ConnConfig.LookupFunc = func(_ context.Context, host string) ([]string, error) {
+		return []string{host}, nil
+	}
+
+	config := sqleng.DataPluginConfiguration{
+		DSInfo:            dsInfo,
+		MetricColumnTypes: []string{"unknown", "text", "varchar", "char", "bpchar"},
+		RowLimit:          rowLimit,
+	}
+
+	queryResultTransformer := postgresQueryResultTransformer{}
+	pgxConf.MaxConnLifetime = time.Duration(config.DSInfo.JsonData.ConnMaxLifetime) * time.Second
+	pgxConf.MaxConns = int32(config.DSInfo.JsonData.MaxOpenConns)
+
+	p, err := pgxpool.NewWithConfig(ctx, pgxConf)
+	if err != nil {
+		logger.Error("Failed connecting to Postgres", "err", err)
+		return nil, nil, err
+	}
+
+	handler, err := sqleng.NewQueryDataHandlerPGX(userFacingDefaultError, p, config, &queryResultTransformer, newPostgresMacroEngine(dsInfo.JsonData.Timescaledb),
+		logger)
+	if err != nil {
+		logger.Error("Failed connecting to Postgres", "err", err)
+		return nil, nil, err
+	}
+
+	logger.Debug("Successfully connected to Postgres")
+	return p, handler, nil
+}
+
 func (s *Service) newInstanceSettings() datasource.InstanceFactoryFunc {
 	logger := s.logger
 	return func(ctx context.Context, settings backend.DataSourceInstanceSettings) (instancemgmt.Instance, error) {
@ -143,7 +214,16 @@ func (s *Service) newInstanceSettings() datasource.InstanceFactoryFunc {
 			DecryptedSecureJSONData: settings.DecryptedSecureJSONData,
 		}

-		cnnstr, err := s.generateConnectionString(dsInfo)
+		tlsSettings, err := s.tlsManager.getTLSSettings(dsInfo)
+		if err != nil {
+			return "", err
+		}
+
+		// Ensure cleanupCertFiles is called after the connection is opened
+		defer s.tlsManager.cleanupCertFiles(tlsSettings)
+
+		isPGX := s.features.IsEnabled(ctx, featuremgmt.FlagPostgresDSUsePGX)
+		cnnstr, err := s.generateConnectionString(dsInfo, tlsSettings, isPGX)
 		if err != nil {
 			return nil, err
 		}
@ -153,7 +233,12 @@ func (s *Service) newInstanceSettings() datasource.InstanceFactoryFunc {
 			return nil, err
 		}

-		_, handler, err := newPostgres(ctx, userFacingDefaultError, sqlCfg.RowLimit, dsInfo, cnnstr, logger, settings)
+		var handler instancemgmt.Instance
+		if isPGX {
+			_, handler, err = newPostgresPGX(ctx, userFacingDefaultError, sqlCfg.RowLimit, dsInfo, cnnstr, logger, settings)
+		} else {
+			_, handler, err = newPostgres(ctx, userFacingDefaultError, sqlCfg.RowLimit, dsInfo, cnnstr, logger, settings)
+		}

 		if err != nil {
 			logger.Error("Failed connecting to Postgres", "err", err)
@ -170,65 +255,100 @@ func escape(input string) string {
 	return strings.ReplaceAll(strings.ReplaceAll(input, `\`, `\\`), "'", `\'`)
 }

-func (s *Service) generateConnectionString(dsInfo sqleng.DataSourceInfo) (string, error) {
-	logger := s.logger
-	var host string
-	var port int
+type connectionParams struct {
+	host     string
+	port     int
+	user     string
+	password string
+	database string
+}
+
+func parseConnectionParams(dsInfo sqleng.DataSourceInfo, logger log.Logger) (connectionParams, error) {
+	var params connectionParams
+	var err error
+
 	if strings.HasPrefix(dsInfo.URL, "/") {
-		host = dsInfo.URL
+		params.host = dsInfo.URL
 		logger.Debug("Generating connection string with Unix socket specifier", "address", dsInfo.URL)
 	} else {
-		index := strings.LastIndex(dsInfo.URL, ":")
-		v6Index := strings.Index(dsInfo.URL, "]")
-		sp := strings.SplitN(dsInfo.URL, ":", 2)
-		host = sp[0]
-		if v6Index == -1 {
-			if len(sp) > 1 {
-				var err error
-				port, err = strconv.Atoi(sp[1])
-				if err != nil {
-					logger.Debug("Error parsing the IPv4 address", "address", dsInfo.URL)
-					return "", sqleng.ErrParsingPostgresURL
-				}
-				logger.Debug("Generating IPv4 connection string with network host/port pair", "host", host, "port", port, "address", dsInfo.URL)
-			} else {
-				logger.Debug("Generating IPv4 connection string with network host", "host", host, "address", dsInfo.URL)
-			}
-		} else {
-			if index == v6Index+1 {
-				host = dsInfo.URL[1 : index-1]
-				var err error
-				port, err = strconv.Atoi(dsInfo.URL[index+1:])
-				if err != nil {
-					logger.Debug("Error parsing the IPv6 address", "address", dsInfo.URL)
-					return "", sqleng.ErrParsingPostgresURL
-				}
-				logger.Debug("Generating IPv6 connection string with network host/port pair", "host", host, "port", port, "address", dsInfo.URL)
-			} else {
-				host = dsInfo.URL[1 : len(dsInfo.URL)-1]
-				logger.Debug("Generating IPv6 connection string with network host", "host", host, "address", dsInfo.URL)
-			}
+		params.host, params.port, err = parseNetworkAddress(dsInfo.URL, logger)
+		if err != nil {
+			return connectionParams{}, err
 		}
 	}

-	connStr := fmt.Sprintf("user='%s' password='%s' host='%s' dbname='%s'",
-		escape(dsInfo.User), escape(dsInfo.DecryptedSecureJSONData["password"]), escape(host), escape(dsInfo.Database))
-	if port > 0 {
-		connStr += fmt.Sprintf(" port=%d", port)
+	params.user = dsInfo.User
+	params.password = dsInfo.DecryptedSecureJSONData["password"]
+	params.database = dsInfo.Database
+
+	return params, nil
+}
+
+func parseNetworkAddress(url string, logger log.Logger) (string, int, error) {
+	index := strings.LastIndex(url, ":")
+	v6Index := strings.Index(url, "]")
+	sp := strings.SplitN(url, ":", 2)
+	host := sp[0]
+	port := 0
+
+	if v6Index == -1 {
+		if len(sp) > 1 {
+			var err error
+			port, err = strconv.Atoi(sp[1])
+			if err != nil {
+				logger.Debug("Error parsing the IPv4 address", "address", url)
+				return "", 0, sqleng.ErrParsingPostgresURL
+			}
+			logger.Debug("Generating IPv4 connection string with network host/port pair", "host", host, "port", port, "address", url)
+		} else {
+			logger.Debug("Generating IPv4 connection string with network host", "host", host, "address", url)
+		}
+	} else {
+		if index == v6Index+1 {
+			host = url[1 : index-1]
+			var err error
+			port, err = strconv.Atoi(url[index+1:])
+			if err != nil {
+				logger.Debug("Error parsing the IPv6 address", "address", url)
+				return "", 0, sqleng.ErrParsingPostgresURL
+			}
+			logger.Debug("Generating IPv6 connection string with network host/port pair", "host", host, "port", port, "address", url)
+		} else {
+			host = url[1 : len(url)-1]
+			logger.Debug("Generating IPv6 connection string with network host", "host", host, "address", url)
+		}
 	}

-	tlsSettings, err := s.tlsManager.getTLSSettings(dsInfo)
+	return host, port, nil
+}
+
+func buildBaseConnectionString(params connectionParams) string {
+	connStr := fmt.Sprintf("user='%s' password='%s' host='%s' dbname='%s'",
+		escape(params.user), escape(params.password), escape(params.host), escape(params.database))
+	if params.port > 0 {
+		connStr += fmt.Sprintf(" port=%d", params.port)
+	}
+	return connStr
+}
+
+func (s *Service) generateConnectionString(dsInfo sqleng.DataSourceInfo, tlsSettings tlsSettings, isPGX bool) (string, error) {
+	logger := s.logger
+
+	params, err := parseConnectionParams(dsInfo, logger)
 	if err != nil {
 		return "", err
 	}

+	connStr := buildBaseConnectionString(params)
+
 	connStr += fmt.Sprintf(" sslmode='%s'", escape(tlsSettings.Mode))

 	// there is an issue with the lib/pq module, the `verify-ca` tls mode
 	// does not work correctly. ( see https://github.com/lib/pq/issues/1106 )
 	// to workaround the problem, if the `verify-ca` mode is chosen,
 	// we disable sslsni.
-	if tlsSettings.Mode == "verify-ca" {
+	if tlsSettings.Mode == "verify-ca" && !isPGX {
+		logger.Debug("Disabling sslsni for verify-ca mode")
 		connStr += " sslsni=0"
 	}

@ -262,7 +382,7 @@ func (s *Service) CheckHealth(ctx context.Context, req *backend.CheckHealthReque
 	if err != nil {
 		return sqleng.ErrToHealthCheckResult(err)
 	}
-	return dsHandler.CheckHealth(ctx, req)
+	return dsHandler.CheckHealth(ctx, req, s.features)
 }

 func (t *postgresQueryResultTransformer) GetConverterList() []sqlutil.StringConverter {