postgres.go

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

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