postgres.go

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

     1  package sinks
     2  
     3  import (
     4  	"context"
     5  	_ "embed"
     6  	"errors"
     7  	"fmt"
     8  	"maps"
     9  	"slices"
    10  	"strings"
    11  	"sync"
    12  	"time"
    13  
    14  	jsoniter "github.com/json-iterator/go"
    15  
    16  	"github.com/cybertec-postgresql/pgwatch/v6/internal/db"
    17  	"github.com/cybertec-postgresql/pgwatch/v6/internal/log"
    18  	"github.com/cybertec-postgresql/pgwatch/v6/internal/metrics"
    19  	migrator "github.com/cybertec-postgresql/pgx-migrator"
    20  	"github.com/jackc/pgx/v5"
    21  	"github.com/jackc/pgx/v5/pgconn"
    22  	"github.com/jackc/pgx/v5/pgxpool"
    23  )
    24  
    25  var (
    26  	cacheLimit      = 256
    27  	highLoadTimeout = time.Second * 5
    28  	targetColumns   = [...]string{"time", "dbname", "data", "tag_data"}
    29  )
    30  
    31  //go:embed sql/admin_schema.sql
    32  var sqlMetricAdminSchema string
    33  
    34  //go:embed sql/admin_functions.sql
    35  var sqlMetricAdminFunctions string
    36  
    37  //go:embed sql/ensure_partition_postgres.sql
    38  var sqlMetricEnsurePartitionPostgres string
    39  
    40  //go:embed sql/ensure_partition_timescale.sql
    41  var sqlMetricEnsurePartitionTimescale string
    42  
    43  //go:embed sql/change_chunk_interval.sql
    44  var sqlMetricChangeChunkIntervalTimescale string
    45  
    46  //go:embed sql/change_compression_interval.sql
    47  var sqlMetricChangeCompressionIntervalTimescale string
    48  
    49  var (
    50  	metricSchemaSQLs = []string{
    51  		sqlMetricAdminSchema,
    52  		sqlMetricAdminFunctions,
    53  		sqlMetricEnsurePartitionPostgres,
    54  		sqlMetricEnsurePartitionTimescale,
    55  		sqlMetricChangeChunkIntervalTimescale,
    56  		sqlMetricChangeCompressionIntervalTimescale,
    57  	}
    58  )
    59  
    60  // PostgresWriter is a sink that writes metric measurements to a Postgres database.
    61  // At the moment, it supports both Postgres and TimescaleDB as a storage backend.
    62  // However, one is able to use any Postgres-compatible database as a storage backend,
    63  // e.g. PGEE, Citus, Greenplum, CockroachDB, etc.
    64  type PostgresWriter struct {
    65  	ctx                     context.Context
    66  	sinkDb                  db.PgxPoolIface
    67  	metricSchema            DbStorageSchemaType
    68  	opts                    *CmdOpts
    69  	retentionInterval       time.Duration
    70  	maintenanceInterval     time.Duration
    71  	input                   chan metrics.MeasurementEnvelope
    72  	lastError               chan error
    73  	forceRecreatePartitions bool                             // to signal override PG metrics storage cache
    74  	partitionMapMetric      map[string]ExistingPartitionInfo // metric = min/max bounds
    75  	// mu guards partitionMapMetric and serializes the DDL issued by SyncMetric.
    76  	mu sync.Mutex
    77  }
    78  
    79  // make sure *dbMetricReaderWriter implements the Migrator interface
    80  var _ db.Migrator = (*PostgresWriter)(nil)
    81  
    82  func NewPostgresWriter(ctx context.Context, connstr string, opts *CmdOpts) (pgw *PostgresWriter, err error) {
    83  	var conn db.PgxPoolIface
    84  	if conn, err = db.New(ctx, connstr); err != nil {
    85  		return
    86  	}
    87  	return NewWriterFromPostgresConn(ctx, conn, opts)
    88  }
    89  
    90  var ErrNeedsMigration = errors.New("sink database schema is outdated, please run migrations using `pgwatch config upgrade` command")
    91  
    92  func NewWriterFromPostgresConn(ctx context.Context, conn db.PgxPoolIface, opts *CmdOpts) (pgw *PostgresWriter, err error) {
    93  	l := log.GetLogger(ctx).WithField("sink", "postgres").WithField("db", conn.Config().ConnConfig.Database)
    94  	ctx = log.WithLogger(ctx, l)
    95  	pgw = &PostgresWriter{
    96  		ctx:                     ctx,
    97  		opts:                    opts,
    98  		input:                   make(chan metrics.MeasurementEnvelope, cacheLimit),
    99  		lastError:               make(chan error),
   100  		sinkDb:                  conn,
   101  		forceRecreatePartitions: false,
   102  		partitionMapMetric:      make(map[string]ExistingPartitionInfo),
   103  	}
   104  	l.Info("initialising measurements database...")
   105  	if err = pgw.init(); err != nil {
   106  		return nil, err
   107  	}
   108  	if err = pgw.ReadMetricSchemaType(); err != nil {
   109  		return nil, err
   110  	}
   111  	if err = pgw.EnsureBuiltinMetricDummies(); err != nil {
   112  		return nil, err
   113  	}
   114  	pgw.scheduleJob(pgw.maintenanceInterval, func() {
   115  		pgw.DeleteOldPartitions()
   116  		pgw.MaintainUniqueSources()
   117  	})
   118  	go pgw.poll()
   119  	l.Info(`measurements sink is activated`)
   120  	return
   121  }
   122  
   123  func (pgw *PostgresWriter) init() (err error) {
   124  	return db.Init(pgw.ctx, pgw.sinkDb, func(ctx context.Context, conn db.PgxIface) error {
   125  		var isValidPartitionInterval bool
   126  		if err = conn.QueryRow(ctx,
   127  			"SELECT extract(epoch from $1::interval), extract(epoch from $2::interval), $3::interval >= '1h'::interval",
   128  			pgw.opts.RetentionInterval, pgw.opts.MaintenanceInterval, pgw.opts.PartitionInterval,
   129  		).Scan(&pgw.retentionInterval, &pgw.maintenanceInterval, &isValidPartitionInterval); err != nil {
   130  			return err
   131  		}
   132  
   133  		// epoch returns seconds but time.Duration represents nanoseconds
   134  		pgw.retentionInterval *= time.Second
   135  		pgw.maintenanceInterval *= time.Second
   136  
   137  		if !isValidPartitionInterval {
   138  			return fmt.Errorf("--partition-interval must be at least 1 hour, got: %s", pgw.opts.PartitionInterval)
   139  		}
   140  		if pgw.maintenanceInterval < 0 {
   141  			return errors.New("--maintenance-interval must be a positive PostgreSQL interval or 0 to disable it")
   142  		}
   143  		if pgw.retentionInterval < time.Hour && pgw.retentionInterval != 0 {
   144  			return errors.New("--retention must be at least 1 hour PostgreSQL interval or 0 to disable it")
   145  		}
   146  
   147  		exists, err := db.DoesSchemaExist(ctx, conn, "admin")
   148  		if err != nil || exists {
   149  			return err
   150  		}
   151  		for _, sql := range metricSchemaSQLs {
   152  			if _, err = conn.Exec(ctx, sql); err != nil {
   153  				return err
   154  			}
   155  		}
   156  		return nil
   157  	})
   158  }
   159  
   160  type ExistingPartitionInfo struct {
   161  	StartTime time.Time
   162  	EndTime   time.Time
   163  }
   164  
   165  type MeasurementMessagePostgres struct {
   166  	Time    time.Time
   167  	DBName  string
   168  	Metric  string
   169  	Data    map[string]any
   170  	TagData map[string]string
   171  }
   172  
   173  type DbStorageSchemaType int
   174  
   175  const (
   176  	DbStorageSchemaPostgres DbStorageSchemaType = iota
   177  	DbStorageSchemaTimescale
   178  )
   179  
   180  func (pgw *PostgresWriter) scheduleJob(interval time.Duration, job func()) {
   181  	if interval > 0 {
   182  		go func() {
   183  			for {
   184  				select {
   185  				case <-pgw.ctx.Done():
   186  					return
   187  				case <-time.After(interval):
   188  					job()
   189  				}
   190  			}
   191  		}()
   192  	}
   193  }
   194  
   195  func (pgw *PostgresWriter) ReadMetricSchemaType() (err error) {
   196  	var isTs bool
   197  	pgw.metricSchema = DbStorageSchemaPostgres
   198  	sqlSchemaType := `SELECT schema_type = 'timescale' FROM admin.storage_schema_type`
   199  	if err = pgw.sinkDb.QueryRow(pgw.ctx, sqlSchemaType).Scan(&isTs); err == nil && isTs {
   200  		pgw.metricSchema = DbStorageSchemaTimescale
   201  	}
   202  	return
   203  }
   204  
   205  // SyncMetric ensures that tables exist for newly added metrics and/or sources
   206  func (pgw *PostgresWriter) SyncMetric(sourceName, metricName string, op SyncOp) error {
   207  	pgw.mu.Lock()
   208  	defer pgw.mu.Unlock()
   209  	if op == AddOp {
   210  		return errors.Join(
   211  			pgw.AddDBUniqueMetricToListingTable(sourceName, metricName),
   212  			pgw.EnsureMetricDummy(metricName), // ensure that there is at least an empty top-level table not to get ugly Grafana notifications
   213  		)
   214  	}
   215  	return nil
   216  }
   217  
   218  // EnsureBuiltinMetricDummies creates empty tables for all built-in metrics if they don't exist
   219  func (pgw *PostgresWriter) EnsureBuiltinMetricDummies() (err error) {
   220  	for _, name := range metrics.GetDefaultBuiltInMetrics() {
   221  		err = errors.Join(err, pgw.EnsureMetricDummy(name))
   222  	}
   223  	return
   224  }
   225  
   226  // EnsureMetricDummy creates an empty table for a metric measurements if it doesn't exist
   227  func (pgw *PostgresWriter) EnsureMetricDummy(metric string) (err error) {
   228  	_, err = pgw.sinkDb.Exec(pgw.ctx, "SELECT admin.ensure_dummy_metrics_table($1)", metric)
   229  	return
   230  }
   231  
   232  // Write sends the measurements to the cache channel
   233  func (pgw *PostgresWriter) Write(msg metrics.MeasurementEnvelope) error {
   234  	if pgw.ctx.Err() != nil {
   235  		return pgw.ctx.Err()
   236  	}
   237  	select {
   238  	case pgw.input <- msg:
   239  		// msgs sent
   240  	case <-time.After(highLoadTimeout):
   241  		// msgs dropped due to a huge load, check stdout or file for detailed log
   242  	}
   243  	select {
   244  	case err := <-pgw.lastError:
   245  		return err
   246  	default:
   247  		return nil
   248  	}
   249  }
   250  
   251  // poll is the main loop that reads from the input channel and flushes the data to the database
   252  func (pgw *PostgresWriter) poll() {
   253  	cache := make([]metrics.MeasurementEnvelope, 0, cacheLimit)
   254  	cacheTimeout := pgw.opts.BatchingDelay
   255  	tick := time.NewTicker(cacheTimeout)
   256  	for {
   257  		select {
   258  		case <-pgw.ctx.Done(): //check context with high priority
   259  			return
   260  		default:
   261  			select {
   262  			case entry := <-pgw.input:
   263  				cache = append(cache, entry)
   264  				if len(cache) < cacheLimit {
   265  					break
   266  				}
   267  				tick.Stop()
   268  				pgw.flush(cache)
   269  				cache = cache[:0]
   270  				tick = time.NewTicker(cacheTimeout)
   271  			case <-tick.C:
   272  				pgw.flush(cache)
   273  				cache = cache[:0]
   274  			case <-pgw.ctx.Done():
   275  				return
   276  			}
   277  		}
   278  	}
   279  }
   280  
   281  func newCopyFromMeasurements(rows []metrics.MeasurementEnvelope) *copyFromMeasurements {
   282  	return &copyFromMeasurements{envelopes: rows, envelopeIdx: -1, measurementIdx: -1}
   283  }
   284  
   285  type copyFromMeasurements struct {
   286  	envelopes      []metrics.MeasurementEnvelope
   287  	envelopeIdx    int
   288  	measurementIdx int // index of the current measurement in the envelope
   289  	metricName     string
   290  	err            error
   291  }
   292  
   293  func (c *copyFromMeasurements) NextEnvelope() bool {
   294  	c.envelopeIdx++
   295  	c.measurementIdx = -1
   296  	return c.envelopeIdx < len(c.envelopes)
   297  }
   298  
   299  func (c *copyFromMeasurements) Next() bool {
   300  	for {
   301  		// Check if we need to advance to the next envelope
   302  		if c.envelopeIdx < 0 || c.measurementIdx+1 >= len(c.envelopes[c.envelopeIdx].Data) {
   303  			// Advance to next envelope
   304  			if ok := c.NextEnvelope(); !ok {
   305  				return false // No more envelopes
   306  			}
   307  			// Set metric name from first envelope, or detect metric boundary
   308  			if c.metricName == "" {
   309  				c.metricName = c.envelopes[c.envelopeIdx].MetricName
   310  			} else if c.metricName != c.envelopes[c.envelopeIdx].MetricName {
   311  				// We've hit a different metric - we're done with current metric
   312  				// Reset position to process this envelope on next call
   313  				c.envelopeIdx--
   314  				c.measurementIdx = len(c.envelopes[c.envelopeIdx].Data) // Set to length so we've "finished" this envelope
   315  				c.metricName = ""                                       // Reset for next metric
   316  				return false
   317  			}
   318  		}
   319  
   320  		// Advance to next measurement in current envelope
   321  		c.measurementIdx++
   322  		if c.measurementIdx < len(c.envelopes[c.envelopeIdx].Data) {
   323  			return true // Found valid measurement
   324  		}
   325  		// If we reach here, we've exhausted current envelope, loop will advance to next envelope
   326  	}
   327  }
   328  
   329  func (c *copyFromMeasurements) EOF() bool {
   330  	return c.envelopeIdx >= len(c.envelopes)
   331  }
   332  
   333  func (c *copyFromMeasurements) Values() ([]any, error) {
   334  	row := maps.Clone(c.envelopes[c.envelopeIdx].Data[c.measurementIdx])
   335  	tagRow := maps.Clone(c.envelopes[c.envelopeIdx].CustomTags)
   336  	if tagRow == nil {
   337  		tagRow = make(map[string]string)
   338  	}
   339  	for k, v := range row {
   340  		if after, ok := strings.CutPrefix(k, metrics.TagPrefix); ok {
   341  			tagRow[after] = fmt.Sprintf("%v", v)
   342  			delete(row, k)
   343  		}
   344  	}
   345  	jsonTags, terr := jsoniter.ConfigFastest.MarshalToString(tagRow)
   346  	json, err := jsoniter.ConfigFastest.MarshalToString(row)
   347  	if err != nil || terr != nil {
   348  		c.err = errors.Join(err, terr)
   349  		return nil, c.err
   350  	}
   351  	return []any{time.Unix(0, c.envelopes[c.envelopeIdx].Data.GetEpoch()), c.envelopes[c.envelopeIdx].DBName, json, jsonTags}, nil
   352  }
   353  
   354  func (c *copyFromMeasurements) Err() error {
   355  	return c.err
   356  }
   357  
   358  func (c *copyFromMeasurements) MetricName() (ident pgx.Identifier) {
   359  	if c.envelopeIdx+1 < len(c.envelopes) {
   360  		// Metric name is taken from the next envelope
   361  		ident = pgx.Identifier{c.envelopes[c.envelopeIdx+1].MetricName}
   362  	}
   363  	return
   364  }
   365  
   366  // flush sends the cached measurements to the database
   367  func (pgw *PostgresWriter) flush(msgs []metrics.MeasurementEnvelope) {
   368  	if len(msgs) == 0 {
   369  		return
   370  	}
   371  	logger := log.GetLogger(pgw.ctx)
   372  	pgPartBounds := make(map[string]ExistingPartitionInfo) // metric=min/max
   373  	var err error
   374  
   375  	slices.SortFunc(msgs, func(a, b metrics.MeasurementEnvelope) int {
   376  		if a.MetricName < b.MetricName {
   377  			return -1
   378  		} else if a.MetricName > b.MetricName {
   379  			return 1
   380  		}
   381  		return 0
   382  	})
   383  
   384  	for _, msg := range msgs {
   385  		if len(msg.Data) > 0 {
   386  			epochTime := time.Unix(0, msg.Data.GetEpoch())
   387  			bounds, ok := pgPartBounds[msg.MetricName]
   388  			if !ok || (ok && epochTime.Before(bounds.StartTime)) {
   389  				bounds.StartTime = epochTime
   390  				pgPartBounds[msg.MetricName] = bounds
   391  			}
   392  			if !ok || (ok && epochTime.After(bounds.EndTime)) {
   393  				bounds.EndTime = epochTime
   394  				pgPartBounds[msg.MetricName] = bounds
   395  			}
   396  		}
   397  	}
   398  
   399  	switch pgw.metricSchema {
   400  	case DbStorageSchemaPostgres:
   401  		err = pgw.EnsureMetricTimePartsExist(pgPartBounds)
   402  	case DbStorageSchemaTimescale:
   403  		err = pgw.EnsureMetricTimescale(pgPartBounds)
   404  	default:
   405  		logger.Fatal("unknown storage schema...")
   406  	}
   407  	pgw.forceRecreatePartitions = false
   408  	if err != nil {
   409  		select {
   410  		case pgw.lastError <- err:
   411  		default:
   412  		}
   413  	}
   414  
   415  	var rowsBatched, n int64
   416  	t1 := time.Now()
   417  	cfm := newCopyFromMeasurements(msgs)
   418  	for !cfm.EOF() {
   419  		n, err = pgw.sinkDb.CopyFrom(context.Background(), cfm.MetricName(), targetColumns[:], cfm)
   420  		rowsBatched += n
   421  		if err != nil {
   422  			logger.Error(err)
   423  			if _, ok := err.(*pgconn.ConnectError); ok {
   424  				logger.Errorf("Sink DB not reachable, dropping %d cached measurements", len(msgs))
   425  				break
   426  			}
   427  			if PgError, ok := err.(*pgconn.PgError); ok {
   428  				pgw.forceRecreatePartitions = PgError.Code == "23514"
   429  			}
   430  			if pgw.forceRecreatePartitions {
   431  				logger.Warning("Some metric partitions might have been removed, halting all metric storage. Trying to re-create all needed partitions on next run")
   432  			}
   433  		}
   434  	}
   435  	diff := time.Since(t1)
   436  	if err == nil {
   437  		logger.WithField("rows", rowsBatched).WithField("elapsed", diff).Info("measurements written")
   438  		return
   439  	}
   440  	select {
   441  	case pgw.lastError <- err:
   442  	default:
   443  	}
   444  }
   445  
   446  func (pgw *PostgresWriter) EnsureMetricTimescale(pgPartBounds map[string]ExistingPartitionInfo) (err error) {
   447  	pgw.mu.Lock()
   448  	defer pgw.mu.Unlock()
   449  	logger := log.GetLogger(pgw.ctx)
   450  	sqlEnsure := `select * from admin.ensure_partition_timescale($1)`
   451  	for metric := range pgPartBounds {
   452  		if _, ok := pgw.partitionMapMetric[metric]; !ok {
   453  			if _, err = pgw.sinkDb.Exec(pgw.ctx, sqlEnsure, metric); err != nil {
   454  				logger.Errorf("Failed to create a TimescaleDB table for metric '%s': %v", metric, err)
   455  				return err
   456  			}
   457  			pgw.partitionMapMetric[metric] = ExistingPartitionInfo{}
   458  		}
   459  	}
   460  	return
   461  }
   462  
   463  func (pgw *PostgresWriter) EnsureMetricTimePartsExist(metricPartBounds map[string]ExistingPartitionInfo) error {
   464  	pgw.mu.Lock()
   465  	defer pgw.mu.Unlock()
   466  	var err error
   467  	var rows pgx.Rows
   468  	sqlEnsure := `select * from admin.ensure_partition_metric_time($1, $2, $3)`
   469  	for metric, pb := range metricPartBounds {
   470  		if pb.StartTime.IsZero() || pb.EndTime.IsZero() {
   471  			return fmt.Errorf("zero StartTime/EndTime in partitioning request: [%s:%v]", metric, pb)
   472  		}
   473  		partInfo, ok := pgw.partitionMapMetric[metric]
   474  		if !ok || pb.StartTime.Before(partInfo.StartTime) || pgw.forceRecreatePartitions {
   475  			if rows, err = pgw.sinkDb.Query(pgw.ctx, sqlEnsure, metric, pb.StartTime, pgw.opts.PartitionInterval); err != nil {
   476  				return err
   477  			}
   478  			if partInfo, err = pgx.CollectOneRow(rows, pgx.RowToStructByPos[ExistingPartitionInfo]); err != nil {
   479  				return err
   480  			}
   481  			pgw.partitionMapMetric[metric] = partInfo
   482  		}
   483  		if pb.EndTime.After(partInfo.EndTime) || pb.EndTime.Equal(partInfo.EndTime) || pgw.forceRecreatePartitions {
   484  			if rows, err = pgw.sinkDb.Query(pgw.ctx, sqlEnsure, metric, pb.EndTime, pgw.opts.PartitionInterval); err != nil {
   485  				return err
   486  			}
   487  			if partInfo, err = pgx.CollectOneRow(rows, pgx.RowToStructByPos[ExistingPartitionInfo]); err != nil {
   488  				return err
   489  			}
   490  			pgw.partitionMapMetric[metric] = partInfo
   491  		}
   492  	}
   493  	return nil
   494  }
   495  
   496  // DeleteOldPartitions is a background task that deletes old partitions from the measurements DB
   497  func (pgw *PostgresWriter) DeleteOldPartitions() {
   498  	l := log.GetLogger(pgw.ctx)
   499  	var partsDropped int
   500  	err := pgw.sinkDb.QueryRow(pgw.ctx, `SELECT admin.drop_old_time_partitions(older_than => $1::interval)`,
   501  		pgw.opts.RetentionInterval).Scan(&partsDropped)
   502  	if err != nil {
   503  		l.Error("Could not drop old time partitions:", err)
   504  	} else if partsDropped > 0 {
   505  		l.Infof("Dropped %d old time partitions", partsDropped)
   506  	}
   507  }
   508  
   509  // MaintainUniqueSources is a background task that maintains a mapping of unique sources
   510  // in each metric table in admin.all_distinct_dbname_metrics.
   511  // This is used to avoid listing the same source multiple times in Grafana dropdowns.
   512  func (pgw *PostgresWriter) MaintainUniqueSources() {
   513  	logger := log.GetLogger(pgw.ctx)
   514  	var rowsAffected int
   515  	if err := pgw.sinkDb.QueryRow(pgw.ctx, `SELECT admin.maintain_unique_sources()`).Scan(&rowsAffected); err != nil {
   516  		logger.Error("Failed to run admin.all_distinct_dbname_metrics maintenance:", err)
   517  		return
   518  	}
   519  	logger.WithField("rows", rowsAffected).Info("Successfully processed admin.all_distinct_dbname_metrics")
   520  }
   521  
   522  func (pgw *PostgresWriter) AddDBUniqueMetricToListingTable(dbUnique, metric string) error {
   523  	sql := `INSERT INTO admin.all_distinct_dbname_metrics
   524  			SELECT $1, $2
   525  			WHERE NOT EXISTS (
   526  				SELECT * FROM admin.all_distinct_dbname_metrics WHERE dbname = $1 AND metric = $2
   527  			)`
   528  	_, err := pgw.sinkDb.Exec(pgw.ctx, sql, dbUnique, metric)
   529  	return err
   530  }
   531  
   532  func NewPostgresSinkMigrator(ctx context.Context, connStr string) (db.Migrator, error) {
   533  	conn, err := pgxpool.New(ctx, connStr)
   534  	if err != nil {
   535  		return nil, err
   536  	}
   537  	pgw := &PostgresWriter{
   538  		ctx:    ctx,
   539  		sinkDb: conn,
   540  	}
   541  	exists, err := db.DoesSchemaExist(ctx, conn, "admin")
   542  	if err != nil {
   543  		return nil, err
   544  	}
   545  	if exists {
   546  		return pgw, nil
   547  	}
   548  	for _, sql := range metricSchemaSQLs {
   549  		if _, err = conn.Exec(ctx, sql); err != nil {
   550  			return nil, err
   551  		}
   552  	}
   553  	return pgw, nil
   554  }
   555  
   556  var initMigrator = func(pgw *PostgresWriter) (*migrator.Migrator, error) {
   557  	return migrator.New(
   558  		migrator.TableName("admin.migration"),
   559  		migrator.SetNotice(func(s string) {
   560  			log.GetLogger(pgw.ctx).Info(s)
   561  		}),
   562  		migrations(),
   563  	)
   564  }
   565  
   566  // Migrate upgrades database with all migrations
   567  func (pgw *PostgresWriter) Migrate() error {
   568  	m, err := initMigrator(pgw)
   569  	if err != nil {
   570  		return fmt.Errorf("cannot initialize migration: %w", err)
   571  	}
   572  	return m.Migrate(pgw.ctx, pgw.sinkDb)
   573  }
   574  
   575  // NeedsMigration checks if database needs migration
   576  func (pgw *PostgresWriter) NeedsMigration() (bool, error) {
   577  	m, err := initMigrator(pgw)
   578  	if err != nil {
   579  		return false, err
   580  	}
   581  	return m.NeedUpgrade(pgw.ctx, pgw.sinkDb)
   582  }
   583  
   584  // migrations holds function returning all upgrade migrations needed
   585  
   586  var migrations func() migrator.Option = func() migrator.Option {
   587  	return migrator.Migrations(
   588  		&migrator.Migration{
   589  			Name: "01110 Apply postgres sink schema migrations",
   590  			Func: func(context.Context, pgx.Tx) error {
   591  				// "migration" table will be created automatically
   592  				return nil
   593  			},
   594  		},
   595  
   596  		&migrator.Migration{
   597  			Name: "01180 Apply admin functions migrations for v5",
   598  			Func: func(ctx context.Context, tx pgx.Tx) error {
   599  				_, err := tx.Exec(ctx, `
   600  					DROP FUNCTION IF EXISTS admin.ensure_partition_metric_dbname_time;
   601  					DROP FUNCTION IF EXISTS admin.ensure_partition_metric_time;
   602  					DROP FUNCTION IF EXISTS admin.get_old_time_partitions(integer, text);
   603  					DROP FUNCTION IF EXISTS admin.drop_old_time_partitions(integer, boolean, text);
   604  				`)
   605  				if err != nil {
   606  					return err
   607  				}
   608  
   609  				_, err = tx.Exec(ctx, sqlMetricEnsurePartitionPostgres)
   610  				if err != nil {
   611  					return err
   612  				}
   613  				_, err = tx.Exec(ctx, sqlMetricAdminFunctions)
   614  				return err
   615  			},
   616  		},
   617  
   618  		&migrator.MigrationNoTx{
   619  			Name: "01409 Switch to time-only partitioning",
   620  			Func: func(ctx context.Context, conn migrator.PgxIface) error {
   621  				const (
   622  					sqlDropOldEnsurePartitionDbnameTime = `DROP FUNCTION IF EXISTS admin.ensure_partition_metric_dbname_time;`
   623  
   624  					sqlListMetricTables = `SELECT c.relname
   625  						FROM pg_description d
   626  						JOIN pg_class c ON c.oid = d.objoid
   627  						WHERE d.description = 'pgwatch-generated-metric-lvl'`
   628  
   629  					sqlIsTableMigrated = `SELECT EXISTS (
   630  						SELECT 1 FROM pg_partitioned_table WHERE partrelid = to_regclass($1) AND partstrat = 'r')`
   631  
   632  					sqlRenameMetricTable = `ALTER TABLE %s RENAME TO %s`
   633  
   634  					sqlMetricTableBounds = `SELECT COALESCE(MIN(time), NOW()),
   635  						COALESCE(CEIL(EXTRACT(EPOCH FROM (MAX(time) - MIN(time))::interval) / 86400) + 1, 0)
   636  						FROM %s`
   637  
   638  					sqlEnsurePartitionMetricTime = `SELECT admin.ensure_partition_metric_time($1::text, $2::timestamptz, '1 day'::interval, $3)`
   639  
   640  					sqlListSubpartitions = `SELECT relid::regclass, parentrelid::regclass 
   641  						FROM pg_partition_tree(to_regclass($1)) WHERE relid::text LIKE 'subpartitions%' 
   642  						ORDER BY isleaf DESC, relid::text`
   643  
   644  					// moves rows into the new partitioned parent, detaches and drops the old subpartition;
   645  					// multiple statements in a single Exec run in an implicit transaction
   646  					sqlMoveSubpartition = `INSERT INTO %[1]s (time, dbname, data, tag_data) SELECT time, dbname, data, tag_data FROM %[2]s;
   647  						ALTER TABLE %[3]s DETACH PARTITION %[2]s;
   648  						DROP TABLE %[2]s;`
   649  
   650  					sqlDropTableIfExists = `DROP TABLE IF EXISTS %s`
   651  				)
   652  				if _, err := conn.Exec(ctx, sqlDropOldEnsurePartitionDbnameTime+
   653  					sqlMetricAdminFunctions+
   654  					sqlMetricEnsurePartitionPostgres); err != nil {
   655  					return err
   656  				}
   657  
   658  				rows, _ := conn.Query(ctx, sqlListMetricTables)
   659  				metricTables, err := pgx.CollectRows(rows, pgx.RowTo[string])
   660  				if err != nil {
   661  					return err
   662  				}
   663  
   664  				// skip *_before_v6_migration tables to avoid double migration
   665  				// this could happen if the migration is re-run after a failed attempt
   666  				const suffix = "_before_v6_migration"
   667  				for _, metricTableRawName := range metricTables {
   668  					if strings.HasSuffix(metricTableRawName, suffix) {
   669  						continue
   670  					}
   671  
   672  					metricTableRawRename := metricTableRawName + suffix
   673  					metricTableRenamed := pgx.Identifier{metricTableRawRename}.Sanitize()
   674  					metricTable := pgx.Identifier{metricTableRawName}.Sanitize()
   675  
   676  					// check if the table is already migrated to avoid errors on re-run after a failed migration attempt
   677  					var isTableMigrated bool
   678  					if err = conn.QueryRow(ctx, sqlIsTableMigrated, metricTableRawName).Scan(&isTableMigrated); err != nil {
   679  						return err
   680  					} else if isTableMigrated {
   681  						continue
   682  					}
   683  
   684  					err = pgx.BeginFunc(ctx, conn, func(tx pgx.Tx) (ferr error) {
   685  						if _, ferr = tx.Exec(ctx, fmt.Sprintf(sqlRenameMetricTable, metricTable, metricTableRenamed)); ferr != nil {
   686  							return
   687  						}
   688  
   689  						// for an empty table MIN(time) is NULL, so COALESCE falls back to server-side
   690  						// NOW() and daysToPrecreate becomes 0, creating a single empty partition
   691  						var minTime time.Time
   692  						var daysToPrecreate int32
   693  						if ferr = tx.QueryRow(ctx, fmt.Sprintf(sqlMetricTableBounds, metricTableRenamed)).Scan(&minTime, &daysToPrecreate); ferr == nil {
   694  							_, ferr = tx.Exec(ctx, sqlEnsurePartitionMetricTime, metricTableRawName, minTime, daysToPrecreate)
   695  						}
   696  						return
   697  					})
   698  
   699  					if err != nil {
   700  						return err
   701  					}
   702  
   703  					type partitionInfo struct {
   704  						Rel       string
   705  						ParentRel string
   706  					}
   707  					rows, _ := conn.Query(ctx, sqlListSubpartitions, metricTableRawRename)
   708  					partitionsInfo, err := pgx.CollectRows(rows, pgx.RowToStructByPos[partitionInfo])
   709  					if err != nil {
   710  						return err
   711  					}
   712  
   713  					for _, partInfo := range partitionsInfo {
   714  						if _, err := conn.Exec(ctx, fmt.Sprintf(sqlMoveSubpartition, metricTable, partInfo.Rel, partInfo.ParentRel)); err != nil {
   715  							return err
   716  						}
   717  					}
   718  
   719  					if _, err := conn.Exec(ctx, fmt.Sprintf(sqlDropTableIfExists, metricTableRenamed)); err != nil {
   720  						return err
   721  					}
   722  				}
   723  
   724  				return nil
   725  			},
   726  		},
   727  
   728  		&migrator.Migration{
   729  			Name: "01474 Change drop_all_metric_tables to procedure",
   730  			Func: func(ctx context.Context, tx pgx.Tx) error {
   731  				_, err := tx.Exec(ctx, sqlMetricAdminFunctions)
   732  				return err
   733  			},
   734  		},
   735  
   736  		&migrator.Migration{
   737  			Name: "01529 Fix ensure_partition_metric_time partitioning strategy",
   738  			Func: func(ctx context.Context, tx pgx.Tx) error {
   739  				_, err := tx.Exec(ctx, sqlMetricEnsurePartitionPostgres)
   740  				return err
   741  			},
   742  		},
   743  
   744  		// adding new migration here, update "admin"."migration" in "admin_schema.sql"!
   745  
   746  		// &migrator.Migration{
   747  		// 	Name: "000XX Short description of a migration",
   748  		// 	Func: func(ctx context.Context, tx pgx.Tx) error {
   749  		// 		return executeMigrationScript(ctx, tx, "000XX.sql")
   750  		// 	},
   751  		// },
   752  	)
   753  }
   754  
   755  // registeredMigrationsCount returns the number of migrations actually registered in
   756  // migrations(). This is the single source of truth for "how many migration rows
   757  // admin.migration must contain after a full migrate"; no separate MigrationsCount
   758  // constant exists.
   759  func registeredMigrationsCount() int {
   760  	m, err := migrator.New(migrations())
   761  	if err != nil {
   762  		panic(fmt.Errorf("registeredMigrationsCount: %w", err))
   763  	}
   764  	return m.Count()
   765  }
   766