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