Merge branch 'master' into delete-dropped-tables-from-listing-table

Author: 0xgouda

internal/sinks/cmdopts.go

@@ -7,7 +7,7 @@ type CmdOpts struct {
 	Sinks                 []string      `long:"sink" mapstructure:"sink" description:"URI where metrics will be stored, can be used multiple times" env:"PW_SINK"`
 	BatchingDelay         time.Duration `long:"batching-delay" mapstructure:"batching-delay" description:"Sink-specific batching flush delay; may be ignored by some sinks" default:"950ms" env:"PW_BATCHING_DELAY"`
 	PartitionInterval     string        `long:"partition-interval" mapstructure:"partition-interval" description:"Time range for PostgreSQL sink time partitions. Must be a valid PostgreSQL interval." default:"1 week" env:"PW_PARTITION_INTERVAL"`
-	Retention             string        `long:"retention" mapstructure:"retention" description:"Delete metrics older than this. Set to zero to disable. Must be a valid PostgreSQL interval." default:"14 days" env:"PW_RETENTION"`
+	RetentionInterval     string        `long:"retention" mapstructure:"retention" description:"Delete metrics older than this. Set to zero to disable. Must be a valid PostgreSQL interval." default:"14 days" env:"PW_RETENTION"`
 	MaintenanceInterval   string        `long:"maintenance-interval" mapstructure:"maintenance-interval" description:"Run pgwatch maintenance tasks on sinks with this interval; Set to zero to disable. Must be a valid PostgreSQL interval." default:"12 hours" env:"PW_MAINTENANCE_INTERVAL"`
 	RealDbnameField       string        `long:"real-dbname-field" mapstructure:"real-dbname-field" description:"Tag key for real database name" env:"PW_REAL_DBNAME_FIELD" default:"real_dbname"`
 	SystemIdentifierField string        `long:"system-identifier-field" mapstructure:"system-identifier-field" description:"Tag key for system identifier value" env:"PW_SYSTEM_IDENTIFIER_FIELD" default:"sys_id"`

internal/sinks/postgres.go

@@ -59,18 +59,17 @@ func NewWriterFromPostgresConn(ctx context.Context, conn db.PgxPoolIface, opts *
 		sinkDb:    conn,
 	}
 	if err = db.Init(ctx, pgw.sinkDb, func(ctx context.Context, conn db.PgxIface) error {
-		var isValidPartitionInterval bool 
-		if err = conn.QueryRow(ctx, 
-			"SELECT extract(epoch from $1::interval), extract(epoch from $2::interval), $3::interval >= '1h'::interval", 
-			opts.Retention, opts.MaintenanceInterval, opts.PartitionInterval,
+		var isValidPartitionInterval bool
+		if err = conn.QueryRow(ctx,
+			"SELECT extract(epoch from $1::interval), extract(epoch from $2::interval), $3::interval >= '1h'::interval",
+			opts.RetentionInterval, opts.MaintenanceInterval, opts.PartitionInterval,
 		).Scan(&pgw.retentionInterval, &pgw.maintenanceInterval, &isValidPartitionInterval); err != nil {
 			return err
 		}
 
-		// multiply by (10 ^ 9) because epoch returns seconds
-		// but time.Duration works in terms of nanoseconds
-		pgw.retentionInterval *= 1_000_000_000
-		pgw.maintenanceInterval *= 1_000_000_000
+		// epoch returns seconds but time.Duration represents nanoseconds
+		pgw.retentionInterval *= time.Second
+		pgw.maintenanceInterval *= time.Second
 
 		if !isValidPartitionInterval {
 			return fmt.Errorf("--partition-interval must be at least 1 hour, got: %s", opts.PartitionInterval)
@@ -103,31 +102,8 @@ func NewWriterFromPostgresConn(ctx context.Context, conn db.PgxPoolIface, opts *
 		return
 	}
 
-	if pgw.maintenanceInterval > 0 {
-		go func() {
-			for {
-				select {
-				case <-pgw.ctx.Done():
-					return
-				case <-time.After(pgw.maintenanceInterval):
-					pgw.MaintainUniqueSources()
-				}
-			}
-		}()
-	}
-
-	if pgw.retentionInterval > 0 {
-		go func() {
-			for {
-				select {
-				case <-pgw.ctx.Done():
-					return
-				case <-time.After(pgw.retentionInterval):
-					pgw.DeleteOldPartitions()
-				}
-			}
-		}()
-	}
+	pgw.scheduleJob(pgw.maintenanceInterval, pgw.MaintainUniqueSources)
+	pgw.scheduleJob(pgw.retentionInterval, pgw.DeleteOldPartitions)
 
 	go pgw.poll()
 	l.Info(`measurements sink is activated`)
@@ -183,6 +159,21 @@ const (
 	DbStorageSchemaTimescale
 )
 
+func (pgw *PostgresWriter) scheduleJob(interval time.Duration, job func()) {
+	if interval > 0 {
+		go func() {
+			for {
+				select {
+				case <-pgw.ctx.Done():
+					return
+				case <-time.After(interval):
+					job()
+				}
+			}
+		}()
+	}
+}
+
 func (pgw *PostgresWriter) ReadMetricSchemaType() (err error) {
 	var isTs bool
 	pgw.metricSchema = DbStorageSchemaPostgres
@@ -504,15 +495,15 @@ func (pgw *PostgresWriter) DeleteOldPartitions() {
 	l := log.GetLogger(pgw.ctx)
 	var partsDropped int
 	err := pgw.sinkDb.QueryRow(pgw.ctx, `SELECT admin.drop_old_time_partitions(older_than => $1::interval)`,
-		pgw.opts.Retention).Scan(&partsDropped)
+		pgw.opts.RetentionInterval).Scan(&partsDropped)
 	if err != nil {
 		l.Error("Could not drop old time partitions:", err)
 	} else if partsDropped > 0 {
 		l.Infof("Dropped %d old time partitions", partsDropped)
 	}
 }
 
-// MaintainUniqueSources is a background task that maintains a mapping of unique sources 
+// MaintainUniqueSources is a background task that maintains a mapping of unique sources
 // in each metric table in admin.all_distinct_dbname_metrics.
 // This is used to avoid listing the same source multiple times in Grafana dropdowns.
 func (pgw *PostgresWriter) MaintainUniqueSources() {

internal/sinks/postgres_test.go

@@ -45,7 +45,7 @@ func TestNewWriterFromPostgresConn(t *testing.T) {
 
 	conn.ExpectPing()
 	conn.ExpectQuery("SELECT extract").WithArgs("1 day", "1 day", "1 hour").WillReturnRows(
-		pgxmock.NewRows([]string{"col1", "col2", "col3"}).AddRow(24 * time.Hour / 1_000_000_000, 24 * time.Hour / 1_000_000_000, true),
+		pgxmock.NewRows([]string{"col1", "col2", "col3"}).AddRow((24 * time.Hour).Seconds(), (24 * time.Hour).Seconds(), true),
 	)
 	conn.ExpectQuery("SELECT EXISTS").WithArgs("admin").WillReturnRows(pgxmock.NewRows([]string{"schema_type"}).AddRow(true))
 	conn.ExpectQuery("SELECT schema_type").WillReturnRows(pgxmock.NewRows([]string{"schema_type"}).AddRow(true))
@@ -54,10 +54,10 @@ func TestNewWriterFromPostgresConn(t *testing.T) {
 	}
 
 	opts := &CmdOpts{
-		BatchingDelay: time.Hour, 
-		Retention: "1 day",
+		BatchingDelay:       time.Hour,
+		RetentionInterval:   "1 day",
 		MaintenanceInterval: "1 day",
-		PartitionInterval: "1 hour",
+		PartitionInterval:   "1 hour",
 	}
 	pgw, err := NewWriterFromPostgresConn(ctx, conn, opts)
 	assert.NoError(t, err)
@@ -508,7 +508,7 @@ func TestCopyFromMeasurements_CopyFail(t *testing.T) {
 
 }
 
-// tests interval string validation for all 
+// tests interval string validation for all
 // cli flags that expect a PostgreSQL interval string
 func TestIntervalValidation(t *testing.T) {
 	a := assert.New(t)
@@ -529,10 +529,10 @@ func TestIntervalValidation(t *testing.T) {
 	connStr, _ := pgContainer.ConnectionString(ctx, "sslmode=disable")
 
 	opts := &CmdOpts{
-		PartitionInterval: "1 minute",
+		PartitionInterval:   "1 minute",
 		MaintenanceInterval: "-1 hours",
-		Retention: "-1 hours",
-		BatchingDelay: time.Second,
+		RetentionInterval:   "-1 hours",
+		BatchingDelay:       time.Second,
 	}
 
 	_, err = NewPostgresWriter(ctx, connStr, opts)
@@ -546,7 +546,7 @@ func TestIntervalValidation(t *testing.T) {
 	_, err = NewPostgresWriter(ctx, connStr, opts)
 	a.EqualError(err, "--retention must be a positive PostgreSQL interval or 0 to disable it")
 
-	invalidIntervals := []string {
+	invalidIntervals := []string{
 		"not an interval", "3 dayss",
 		"four hours",
 	}
@@ -562,10 +562,10 @@ func TestIntervalValidation(t *testing.T) {
 		a.Error(err)
 		opts.MaintenanceInterval = "1 hour"
 
-		opts.Retention = interval
+		opts.RetentionInterval = interval
 		_, err = NewPostgresWriter(ctx, connStr, opts)
 		a.Error(err)
-		opts.Retention = "1 hour"
+		opts.RetentionInterval = "1 hour"
 	}
 
 	validIntervals := []string{
@@ -577,7 +577,7 @@ func TestIntervalValidation(t *testing.T) {
 	for _, interval := range validIntervals {
 		opts.PartitionInterval = interval
 		opts.MaintenanceInterval = interval
-		opts.Retention = interval
+		opts.RetentionInterval = interval
 
 		_, err = NewPostgresWriter(ctx, connStr, opts)
 		a.NoError(err)
@@ -603,10 +603,10 @@ func TestPartitionInterval(t *testing.T) {
 	connStr, _ := pgContainer.ConnectionString(ctx, "sslmode=disable")
 
 	opts := &CmdOpts{
-		PartitionInterval: "3 weeks",
-		Retention: "14 days",
+		PartitionInterval:   "3 weeks",
+		RetentionInterval:   "14 days",
 		MaintenanceInterval: "12 hours",
-		BatchingDelay: time.Second,
+		BatchingDelay:       time.Second,
 	}
 
 	pgw, err := NewPostgresWriter(ctx, connStr, opts)
@@ -625,7 +625,7 @@ func TestPartitionInterval(t *testing.T) {
 	err = pgw.EnsureMetricDbnameTime(m, false)
 	r.NoError(err)
 
-	var partitionsNum int;
+	var partitionsNum int
 	err = conn.QueryRow(ctx, "SELECT COUNT(*) FROM pg_partition_tree('test_metric');").Scan(&partitionsNum)
 	a.NoError(err)
 	// 1 the metric table itself + 1 dbname partition
@@ -634,7 +634,7 @@ func TestPartitionInterval(t *testing.T) {
 
 	part := partitionMapMetricDbname["test_metric"]["test_db"]
 	// partition bounds should have a difference of 3 weeks
-	a.Equal(part.StartTime.Add(3 * 7 * 24 * time.Hour), part.EndTime)
+	a.Equal(part.StartTime.Add(3*7*24*time.Hour), part.EndTime)
 }
 
 func Test_MaintainUniqueSources_DeleteOldPartitions(t *testing.T) {
@@ -659,10 +659,10 @@ func Test_MaintainUniqueSources_DeleteOldPartitions(t *testing.T) {
 	r.NoError(err)
 
 	opts := &CmdOpts{
-		PartitionInterval: "1 hour",
-		Retention: "1 second",
+		PartitionInterval:   "1 hour",
+		RetentionInterval:   "1 second",
 		MaintenanceInterval: "0 days",
-		BatchingDelay: time.Hour,
+		BatchingDelay:       time.Hour,
 	}
 
 	pgw, err := NewPostgresWriter(ctx, connStr, opts)
@@ -675,15 +675,15 @@ func Test_MaintainUniqueSources_DeleteOldPartitions(t *testing.T) {
 
 		var numOfEntries int
 		err = conn.QueryRow(ctx, "SELECT count(*) FROM admin.all_distinct_dbname_metrics;").Scan(&numOfEntries)
-		a.NoError(err)	
+		a.NoError(err)
 		a.Equal(1, numOfEntries)
 
 		// manually call the maintenance routine
 		pgw.MaintainUniqueSources()
 
 		// entry should have been deleted, because it has no corresponding entries in `test_metric_1` table.
 		err = conn.QueryRow(ctx, "SELECT count(*) FROM admin.all_distinct_dbname_metrics;").Scan(&numOfEntries)
-		a.NoError(err)	
+		a.NoError(err)
 		a.Equal(0, numOfEntries)
 
 		message := []metrics.MeasurementEnvelope{
@@ -694,15 +694,15 @@ func Test_MaintainUniqueSources_DeleteOldPartitions(t *testing.T) {
 				},
 				DBName: "test_db",
 			},
-		} 
+		}
 		pgw.flush(message)
 
 		// manually call the maintenance routine
 		pgw.MaintainUniqueSources()
 
 		// entry should have been added, because there is a corresponding entry in `test_metric_1` table just written.
 		err = conn.QueryRow(ctx, "SELECT count(*) FROM admin.all_distinct_dbname_metrics;").Scan(&numOfEntries)
-		a.NoError(err)	
+		a.NoError(err)
 		a.Equal(1, numOfEntries)
 
 		_, err = conn.Exec(ctx, "DROP TABLE test_metric_1;")
@@ -726,31 +726,31 @@ func Test_MaintainUniqueSources_DeleteOldPartitions(t *testing.T) {
 		boundStart := time.Now().Add(-1 * 2 * 24 * time.Hour).Format("2006-01-02")
 		boundEnd := time.Now().Add(-1 * 24 * time.Hour).Format("2006-01-02")
 
-		// create the 3rd level time partition with end bound yesterday 
-		_, err = conn.Exec(ctx, 
+		// create the 3rd level time partition with end bound yesterday
+		_, err = conn.Exec(ctx,
 			fmt.Sprintf(
-			`CREATE TABLE subpartitions.test_metric_2_dbname_time 
+				`CREATE TABLE subpartitions.test_metric_2_dbname_time 
 			PARTITION OF subpartitions.test_metric_2_dbname 
-			FOR VALUES FROM ('%s') TO ('%s')`, 
-			boundStart, boundEnd),
+			FOR VALUES FROM ('%s') TO ('%s')`,
+				boundStart, boundEnd),
 		)
 		a.NoError(err)
 		_, err = conn.Exec(ctx, "COMMENT ON TABLE subpartitions.test_metric_2_dbname_time IS $$pgwatch-generated-metric-dbname-time-lvl$$")
 		a.NoError(err)
 
-		var partitionsNum int;
+		var partitionsNum int
 		err = conn.QueryRow(ctx, "SELECT COUNT(*) FROM pg_partition_tree('test_metric_2');").Scan(&partitionsNum)
 		a.NoError(err)
 		a.Equal(3, partitionsNum)
 
-		pgw.opts.Retention = "2 days"
+		pgw.opts.RetentionInterval = "2 days"
 		pgw.DeleteOldPartitions() // 1 day < 2 days, shouldn't delete anything
 
 		err = conn.QueryRow(ctx, "SELECT COUNT(*) FROM pg_partition_tree('test_metric_2');").Scan(&partitionsNum)
 		a.NoError(err)
 		a.Equal(3, partitionsNum)
 
-		pgw.opts.Retention = "1 hour"
+		pgw.opts.RetentionInterval = "1 hour"
 		pgw.DeleteOldPartitions() // 1 day > 1 hour, should delete the partition
 
 		err = conn.QueryRow(ctx, "SELECT COUNT(*) FROM pg_partition_tree('test_metric_2');").Scan(&partitionsNum)
@@ -760,23 +760,23 @@ func Test_MaintainUniqueSources_DeleteOldPartitions(t *testing.T) {
 
 	t.Run("Epoch to Duration Conversion", func(_ *testing.T) {
 		table := map[string]time.Duration{
-			"1 hour": time.Hour, 
-			"2 hours": 2 * time.Hour,
-			"4 days": 4 * 24 * time.Hour, 
-			"1 day": 24 * time.Hour,
-			"1 year": 365.25 * 24 * time.Hour,
-			"1 week": 7 * 24 * time.Hour,
-			"3 weeks": 3 * 7 * 24 * time.Hour,
+			"1 hour":   time.Hour,
+			"2 hours":  2 * time.Hour,
+			"4 days":   4 * 24 * time.Hour,
+			"1 day":    24 * time.Hour,
+			"1 year":   365.25 * 24 * time.Hour,
+			"1 week":   7 * 24 * time.Hour,
+			"3 weeks":  3 * 7 * 24 * time.Hour,
 			"2 months": 2 * 30 * 24 * time.Hour,
-			"1 month": 30 * 24 * time.Hour,
+			"1 month":  30 * 24 * time.Hour,
 		}
 
 		for k, v := range table {
 			opts := &CmdOpts{
-				PartitionInterval: "1 hour",
-				Retention: k,
+				PartitionInterval:   "1 hour",
+				RetentionInterval:   k,
 				MaintenanceInterval: k,
-				BatchingDelay: time.Hour,
+				BatchingDelay:       time.Hour,
 			}
 
 			pgw, err := NewPostgresWriter(ctx, connStr, opts)