resolver.go

Documentation: github.com/cybertec-postgresql/pgwatch/v6/internal/sources

     1  package sources
     2  
     3  // This file contains the implemendation of Patroni and PostgrSQL resolvers for continuous monitoring.
     4  // Patroni resolver will return the list of databases from the Patroni cluster.
     5  // Postgres resolver will return the list of databases from the given Postgres instance.
     6  
     7  import (
     8  	"cmp"
     9  	"context"
    10  	"crypto/tls"
    11  	"crypto/x509"
    12  	"errors"
    13  	"fmt"
    14  	"net/url"
    15  	"os"
    16  	"strings"
    17  	"sync"
    18  	"time"
    19  
    20  	jsoniter "github.com/json-iterator/go"
    21  
    22  	"github.com/cybertec-postgresql/pgwatch/v6/internal/db"
    23  	"github.com/cybertec-postgresql/pgwatch/v6/internal/log"
    24  	pgx "github.com/jackc/pgx/v5"
    25  	client "go.etcd.io/etcd/client/v3"
    26  	"go.uber.org/zap"
    27  )
    28  
    29  // Resolver discovers the monitored databases behind continuous-monitoring
    30  // sources (Patroni, Postgres discovery). It owns the last-known-good fallback
    31  // caches so that a transient DCS/DB outage does not tear down monitoring of
    32  // already-known databases.
    33  //
    34  // A Resolver is safe for concurrent use. Create independent Resolvers via
    35  // NewResolver when isolated cache state is desired (e.g. in tests or when
    36  // running several unrelated resolution pipelines in one process).
    37  type Resolver struct {
    38  	mu sync.Mutex // guards the fallback caches below
    39  	// lastFoundClusterMembers is needed for cases where DCS is temporarily down;
    40  	// we don't want to immediately remove monitoring of DBs. Keyed by source name.
    41  	lastFoundClusterMembers map[string][]PatroniClusterMember
    42  	// lastFoundDatabases is keyed by the source's identity (name + conn string +
    43  	// include/exclude patterns) so a reconfigured source never inherits the
    44  	// previous target's database list.
    45  	lastFoundDatabases map[postgresDiscoveryKey]SourceConns
    46  }
    47  
    48  // NewResolver returns a Resolver with freshly initialized, empty caches.
    49  func NewResolver() *Resolver {
    50  	return &Resolver{
    51  		lastFoundClusterMembers: make(map[string][]PatroniClusterMember),
    52  		lastFoundDatabases:      make(map[postgresDiscoveryKey]SourceConns),
    53  	}
    54  }
    55  
    56  // defaultResolver backs the convenience methods on Source and Sources so that
    57  // existing callers keep a single process-wide fallback cache without having to
    58  // thread a Resolver through.
    59  var defaultResolver = NewResolver()
    60  
    61  // ResolveDatabases() updates list of monitored objects from continuous monitoring sources, e.g. patroni.
    62  // Each source is resolved concurrently so that a slow or unreachable source does not block the others.
    63  // It delegates to the package-wide defaultResolver.
    64  func (srcs Sources) ResolveDatabases(onError func(string)) (SourceConns, error) {
    65  	return defaultResolver.ResolveDatabases(srcs, onError)
    66  }
    67  
    68  // ResolveDatabases() return a slice of found databases for continuous monitoring sources, e.g. patroni.
    69  // It delegates to the package-wide defaultResolver.
    70  func (s Source) ResolveDatabases() (SourceConns, error) {
    71  	return defaultResolver.ResolveDatabase(s)
    72  }
    73  
    74  // ResolveDatabases updates the list of monitored objects from continuous monitoring
    75  // sources, e.g. patroni. Each source is resolved concurrently so that a slow or
    76  // unreachable source does not block the others.
    77  func (r *Resolver) ResolveDatabases(srcs Sources, onError func(string)) (_ SourceConns, err error) {
    78  	type result struct {
    79  		dbs SourceConns
    80  		err error
    81  	}
    82  	results := make([]result, len(srcs))
    83  	var wg sync.WaitGroup
    84  	for i, s := range srcs {
    85  		wg.Go(func() {
    86  			dbs, e := r.ResolveDatabase(s)
    87  			results[i] = result{dbs, e}
    88  		})
    89  	}
    90  	wg.Wait()
    91  	resolvedDbs := make(SourceConns, 0, len(srcs))
    92  	for i, res := range results {
    93  		if res.err != nil {
    94  			if onError != nil {
    95  				onError(srcs[i].Name)
    96  			}
    97  			logger.WithField("source", srcs[i].Name).WithError(res.err).Error("could not resolve databases from source")
    98  			err = errors.Join(err, res.err)
    99  		}
   100  		resolvedDbs = append(resolvedDbs, res.dbs...)
   101  	}
   102  	return resolvedDbs, err
   103  }
   104  
   105  // ResolveDatabase returns a slice of found databases for a single continuous
   106  // monitoring source, e.g. patroni.
   107  func (r *Resolver) ResolveDatabase(s Source) (SourceConns, error) {
   108  	switch s.Kind {
   109  	case SourcePatroniDiscovery:
   110  		return r.ResolveDatabasesFromPatroni(s)
   111  	case SourcePostgresDiscovery:
   112  		return r.ResolveDatabasesFromPostgres(s)
   113  	case SourcePrometheus:
   114  		return SourceConns{NewPromConn(s)}, nil
   115  	}
   116  	return SourceConns{NewDbConn(s)}, nil
   117  }
   118  
   119  type PatroniClusterMember struct {
   120  	Scope   string
   121  	Name    string
   122  	ConnURL string `yaml:"conn_url"`
   123  	Role    string
   124  }
   125  
   126  func (pcm PatroniClusterMember) IsPrimary() bool {
   127  	return pcm.Role == "primary" || pcm.Role == "master"
   128  }
   129  
   130  var logger log.Logger = log.FallbackLogger
   131  
   132  // postgresDiscoveryKey uniquely identifies a Postgres discovery source for caching purposes.
   133  // It MUST incorporate every field that would change the set of discovered databases; otherwise a
   134  // reconfigured source could be served the previous target's last-known-good list.
   135  type postgresDiscoveryKey struct {
   136  	name           string
   137  	connStr        string
   138  	includePattern string
   139  	excludePattern string
   140  }
   141  
   142  func postgresDiscoveryKeyOf(s Source) postgresDiscoveryKey {
   143  	return postgresDiscoveryKey{
   144  		name:           s.Name,
   145  		connStr:        s.ConnStr,
   146  		includePattern: s.IncludePattern,
   147  		excludePattern: s.ExcludePattern,
   148  	}
   149  }
   150  
   151  func getConsulClusterMembers(Source) ([]PatroniClusterMember, error) {
   152  	return nil, errors.ErrUnsupported
   153  }
   154  
   155  func getZookeeperClusterMembers(Source, HostConfig) ([]PatroniClusterMember, error) {
   156  	return nil, errors.ErrUnsupported
   157  }
   158  
   159  func jsonTextToStringMap(jsonText string) (map[string]string, error) {
   160  	retmap := make(map[string]string)
   161  	if jsonText == "" {
   162  		return retmap, nil
   163  	}
   164  	var iMap map[string]any
   165  	if err := jsoniter.ConfigFastest.Unmarshal([]byte(jsonText), &iMap); err != nil {
   166  		return nil, err
   167  	}
   168  	for k, v := range iMap {
   169  		retmap[k] = fmt.Sprintf("%v", v)
   170  	}
   171  	return retmap, nil
   172  }
   173  
   174  // getTransport builds the TLS config for the etcd client.
   175  // It returns nil when no TLS parameters are set.
   176  func getTransport(conf HostConfig) (*tls.Config, error) {
   177  	if !conf.UsesTLS() {
   178  		return nil, nil
   179  	}
   180  	tlsClientConfig := new(tls.Config)
   181  
   182  	if conf.CAFile != "" {
   183  		caCert, err := os.ReadFile(conf.CAFile)
   184  		if err != nil {
   185  			return nil, fmt.Errorf("cannot load CA file: %w", err)
   186  		}
   187  		caCertPool := x509.NewCertPool()
   188  		if !caCertPool.AppendCertsFromPEM(caCert) {
   189  			return nil, fmt.Errorf("no PEM certificates found in CA file %s", conf.CAFile)
   190  		}
   191  		tlsClientConfig.RootCAs = caCertPool
   192  	}
   193  
   194  	if conf.CertFile != "" {
   195  		cert, err := tls.LoadX509KeyPair(conf.CertFile, conf.KeyFile)
   196  		if err != nil {
   197  			return nil, fmt.Errorf("cannot load client cert or key file: %w", err)
   198  		}
   199  		tlsClientConfig.Certificates = []tls.Certificate{cert}
   200  	}
   201  
   202  	return tlsClientConfig, nil
   203  }
   204  
   205  func (r *Resolver) getEtcdClusterMembers(s Source, hc HostConfig) ([]PatroniClusterMember, error) {
   206  	var ret = make([]PatroniClusterMember, 0)
   207  	var cfg client.Config
   208  
   209  	if len(hc.DcsEndpoints) == 0 {
   210  		return ret, errors.New("missing ETCD connect info, make sure host config has a 'dcs_endpoints' key")
   211  	}
   212  
   213  	tlsConfig, err := getTransport(hc)
   214  	if err != nil {
   215  		return nil, err
   216  	}
   217  	cfg = client.Config{
   218  		Endpoints:            hc.DcsEndpoints,
   219  		TLS:                  tlsConfig,
   220  		DialKeepAliveTimeout: time.Second,
   221  		Username:             hc.Username,
   222  		Password:             hc.Password,
   223  		DialTimeout:          5 * time.Second,
   224  		Logger:               zap.NewNop(),
   225  	}
   226  
   227  	c, err := client.New(cfg)
   228  	if err != nil {
   229  		return ret, err
   230  	}
   231  	defer c.Close()
   232  
   233  	ctx, cancel := context.WithTimeoutCause(context.Background(), 5*time.Second, errors.New("etcd client timeout"))
   234  	defer cancel()
   235  
   236  	// etcd3 does not have a dir node.
   237  	// Key="/namespace/scope/leader", e.g. "/service/batman/leader"
   238  	// Key="/namespace/scope/members/node", e.g. "/service/batman/members/pg1"
   239  
   240  	resp, err := c.Get(ctx, hc.Path, client.WithPrefix())
   241  	if err != nil {
   242  		return ret, cmp.Or(context.Cause(ctx), err)
   243  	}
   244  
   245  	for _, node := range resp.Kvs {
   246  		// remove leading slash and split by "/"
   247  		parts := strings.Split(strings.TrimPrefix(string(node.Key), "/"), "/")
   248  		if len(parts) < 4 || parts[2] != "members" {
   249  			continue // skip non-member keys
   250  		}
   251  		nodeData, err := jsonTextToStringMap(string(node.Value))
   252  		if err != nil {
   253  			logger.Errorf("Could not parse ETCD node data for node \"%s\": %s", node.Key, err)
   254  			continue
   255  		}
   256  		role := nodeData["role"]
   257  		connURL := nodeData["conn_url"]
   258  		scope := parts[1]
   259  		name := parts[3]
   260  		ret = append(ret, PatroniClusterMember{Scope: scope, ConnURL: connURL, Role: role, Name: name})
   261  	}
   262  
   263  	r.mu.Lock()
   264  	r.lastFoundClusterMembers[s.Name] = ret
   265  	r.mu.Unlock()
   266  	return ret, nil
   267  }
   268  
   269  const (
   270  	dcsTypeEtcd      = "etcd"
   271  	dcsTypeZookeeper = "zookeeper"
   272  	dcsTypeConsul    = "consul"
   273  )
   274  
   275  type HostConfig struct {
   276  	DcsType      string   `yaml:"dcs_type"`
   277  	DcsEndpoints []string `yaml:"dcs_endpoints"`
   278  	Path         string
   279  	Username     string
   280  	Password     string
   281  	CAFile       string `yaml:"ca_file"`
   282  	CertFile     string `yaml:"cert_file"`
   283  	KeyFile      string `yaml:"key_file"`
   284  }
   285  
   286  // UsesTLS reports whether any TLS parameter is set.
   287  func (hc HostConfig) UsesTLS() bool {
   288  	return hc.CAFile != "" || hc.CertFile != "" || hc.KeyFile != ""
   289  }
   290  
   291  func (hc HostConfig) IsScopeSpecified() bool {
   292  	// Path is usually "/namespace/scope"
   293  	// so we check if it has at least 2 slashes
   294  	return strings.Count(hc.Path, "/") >= 2
   295  }
   296  
   297  func NewHostConfig(URI string) (hc HostConfig, err error) {
   298  	// Extract scheme
   299  	before, after, ok := strings.Cut(URI, "://")
   300  	if !ok {
   301  		return hc, fmt.Errorf("invalid URI: missing scheme")
   302  	}
   303  	scheme := before
   304  	remainder := after // skip "://"
   305  
   306  	// Find where the host portion ends (at first '/' or '?' or end of string)
   307  	hostEnd := strings.IndexAny(remainder, "/?")
   308  	var hostPart, pathAndQuery string
   309  	if hostEnd == -1 {
   310  		hostPart = remainder
   311  		pathAndQuery = ""
   312  	} else {
   313  		hostPart = remainder[:hostEnd]
   314  		pathAndQuery = remainder[hostEnd:]
   315  	}
   316  
   317  	// Check for user info (username:password@)
   318  	var userInfo string
   319  	if atIdx := strings.LastIndex(hostPart, "@"); atIdx != -1 {
   320  		userInfo = hostPart[:atIdx]
   321  		hostPart = hostPart[atIdx+1:]
   322  	}
   323  
   324  	// Split hosts by comma for multiple endpoints
   325  	hosts := strings.Split(hostPart, ",")
   326  
   327  	// Parse a clean URL with just the first host to extract other components
   328  	cleanURI := scheme + "://"
   329  	if userInfo != "" {
   330  		cleanURI += userInfo + "@"
   331  	}
   332  	cleanURI += hosts[0] + pathAndQuery
   333  
   334  	var url *url.URL
   335  	url, err = url.Parse(cleanURI)
   336  	if err != nil {
   337  		return
   338  	}
   339  
   340  	hc.Path = url.Path
   341  	hc.Username = url.User.Username()
   342  	hc.Password, _ = url.User.Password() // password is optional, so we ignore the error
   343  	hc.CAFile = url.Query().Get("ca_file")
   344  	hc.CertFile = url.Query().Get("cert_file")
   345  	hc.KeyFile = url.Query().Get("key_file")
   346  	if (hc.CertFile == "") != (hc.KeyFile == "") {
   347  		return hc, errors.New("cert_file and key_file must be set together")
   348  	}
   349  
   350  	switch url.Scheme {
   351  	case dcsTypeEtcd:
   352  		hc.DcsType = dcsTypeEtcd
   353  		// The etcd client drops the TLS config for http:// endpoints
   354  		endpointScheme := "http://"
   355  		if hc.UsesTLS() {
   356  			endpointScheme = "https://"
   357  		}
   358  		for _, h := range hosts {
   359  			hc.DcsEndpoints = append(hc.DcsEndpoints, endpointScheme+h)
   360  		}
   361  	case dcsTypeZookeeper:
   362  		hc.DcsType = dcsTypeZookeeper
   363  		hc.DcsEndpoints = hosts // Use the split hosts directly
   364  	case dcsTypeConsul:
   365  		hc.DcsType = dcsTypeConsul
   366  		hc.DcsEndpoints = hosts // Use the split hosts directly
   367  	default:
   368  		return hc, fmt.Errorf("unsupported DCS type: %s", url.Scheme)
   369  	}
   370  
   371  	return hc, nil
   372  }
   373  
   374  func (r *Resolver) ResolveDatabasesFromPatroni(source Source) (SourceConns, error) {
   375  	var mds SourceConns
   376  	var clusterMembers []PatroniClusterMember
   377  	var err error
   378  	var ok bool
   379  
   380  	hostConfig, err := NewHostConfig(source.ConnStr)
   381  	if err != nil {
   382  		return nil, err
   383  	}
   384  
   385  	switch hostConfig.DcsType {
   386  	case dcsTypeEtcd:
   387  		clusterMembers, err = r.getEtcdClusterMembers(source, hostConfig)
   388  	case dcsTypeZookeeper:
   389  		clusterMembers, err = getZookeeperClusterMembers(source, hostConfig)
   390  	case dcsTypeConsul:
   391  		clusterMembers, err = getConsulClusterMembers(source)
   392  	default:
   393  		return nil, errors.New("unknown DCS")
   394  	}
   395  	logger := logger.WithField("source", source.Name)
   396  	if err != nil {
   397  		if errors.Is(err, errors.ErrUnsupported) {
   398  			return nil, err
   399  		}
   400  		logger.Debug("failed to get info from DCS, using previous member info if any")
   401  		r.mu.Lock()
   402  		clusterMembers, ok = r.lastFoundClusterMembers[source.Name] // mask error from main loop not to remove monitored DBs due to "jitter"
   403  		r.mu.Unlock()
   404  		if ok {
   405  			err = nil
   406  		}
   407  	} else {
   408  		r.mu.Lock()
   409  		r.lastFoundClusterMembers[source.Name] = clusterMembers
   410  		r.mu.Unlock()
   411  	}
   412  	if len(clusterMembers) == 0 {
   413  		return mds, err
   414  	}
   415  
   416  	for _, patroniMember := range clusterMembers {
   417  		logger.Info("processing Patroni cluster member: ", patroniMember.Name)
   418  		if source.OnlyIfMaster && !patroniMember.IsPrimary() {
   419  			continue
   420  		}
   421  		src := *source.Clone()
   422  		src.ConnStr = patroniMember.ConnURL
   423  		if !hostConfig.IsScopeSpecified() {
   424  			src.Name += "_" + patroniMember.Scope
   425  		}
   426  		src.Name += "_" + patroniMember.Name
   427  		if dbs, err := r.ResolveDatabasesFromPostgres(src); err == nil {
   428  			mds = append(mds, dbs...)
   429  		} else {
   430  			logger.WithError(err).Error("failed to resolve databases for Patroni member: ", patroniMember.Name)
   431  		}
   432  	}
   433  	return mds, err
   434  }
   435  
   436  // ResolveDatabasesFromPostgres reads all the databases from the given cluster,
   437  // additionally matching/not matching specified regex patterns.
   438  //
   439  // On any helper error (pool create or discovery query) and the presence of a
   440  // previously-cached successful result for the same source identity, the cached
   441  // list is returned with a nil error: a transient discovery failure (DNS hiccup,
   442  // connect timeout, discovery-SQL permission error) must not tear down monitoring
   443  // of already-known databases. The accepted trade-off is that a database dropped
   444  // server-side while discovery keeps failing stays monitored from cache (surfaced
   445  // as per-cycle connect errors reported as down) until the next successful
   446  // resolution replaces the entry.
   447  func (r *Resolver) ResolveDatabasesFromPostgres(s Source) (resolvedDbs SourceConns, err error) {
   448  	resolvedDbs, err = resolveDatabasesFromPostgres(s)
   449  	key := postgresDiscoveryKeyOf(s)
   450  	r.mu.Lock()
   451  	defer r.mu.Unlock()
   452  	if err == nil {
   453  		r.lastFoundDatabases[key] = resolvedDbs
   454  		return resolvedDbs, nil
   455  	}
   456  	cached, ok := r.lastFoundDatabases[key]
   457  	if ok && len(cached) > 0 {
   458  		logger.WithField("source", s.Name).WithError(err).Warning("postgres discovery failed; serving last-known-good database list from cache")
   459  		return cached, nil
   460  	}
   461  	return nil, err
   462  }
   463  
   464  // resolveDatabasesFromPostgres is the inner helper that actually runs the discovery query.
   465  func resolveDatabasesFromPostgres(s Source) (resolvedDbs SourceConns, err error) {
   466  	var (
   467  		c      db.PgxPoolIface
   468  		dbname string
   469  		rows   pgx.Rows
   470  	)
   471  	ctx, cancel := db.WithOpTimeout(context.Background(), "resolve "+s.Name, db.ResolverTimeout)
   472  	defer cancel()
   473  
   474  	c, err = NewConn(ctx, s.ConnStr)
   475  	if err != nil {
   476  		return nil, cmp.Or(context.Cause(ctx), err)
   477  	}
   478  	defer c.Close()
   479  
   480  	sql := `select /* pgwatch_generated */
   481  	datname
   482  	from pg_database
   483  	where not datistemplate
   484  	and datallowconn
   485  	and has_database_privilege (datname, 'CONNECT')
   486  	and case when length(trim($1)) > 0 then datname ~ $1 else true end
   487  	and case when length(trim($2)) > 0 then not datname ~ $2 else true end`
   488  
   489  	if rows, err = c.Query(ctx, sql, s.IncludePattern, s.ExcludePattern); err != nil {
   490  		return nil, cmp.Or(context.Cause(ctx), err)
   491  	}
   492  	for rows.Next() {
   493  		if err = rows.Scan(&dbname); err != nil {
   494  			return nil, err
   495  		}
   496  		rdb := NewDbConn(*s.Clone())
   497  		rdb.Name += "_" + dbname
   498  		rdb.SetDatabaseName(dbname)
   499  		resolvedDbs = append(resolvedDbs, rdb)
   500  	}
   501  
   502  	if err := rows.Err(); err != nil {
   503  		return nil, err
   504  	}
   505  	return
   506  }
   507