[+] support variable partition intervals in postgres sink (#1012)

* [+] allow to specify interval for Postgres partitioning

`admin.ensure_partition_metric_dbname_time()` acccepts now
partition_period interval. The naming format is 'YYYYMMDD' for intervals
bigger than 1 day and 'YYYYMMDD_HH24' for hour-based intervals.
The function finds the last available partition bound and creates
new partition based on it. So it's possible to use mixed length
partitions, possibly with scheduler based on monitored source and/or
metrics.

* Add new `PartitionInterval` cli parameter.

* Use `PartitionInterval` parameter in `EnsureMetricDbnameTime()`

* Mention `--partition-interval` in docs.

* Add validation logic for `PartitionInterval`

* Add test for `PartitionInterval` validation

* Test actual partitions creation

* Add example usage of `--partition-interval` in docs.

---------

Co-authored-by: Pavlo Golub <pavlo.golub@gmail.com>

Author: 0xgouda

docs/reference/cli_env.md

@@ -97,6 +97,13 @@ It reads the configuration from the specified sources and metrics, then begins c
     Sink-specific batching flush delay; may be ignored by some sinks (default: 950ms).  
     ENV: `$PW_BATCHING_DELAY`
 
+- `--partition-interval=`
+    Time range for PostgreSQL sink table partitions. Must be a valid PostgreSQL interval expression. (default: 1 week)
+    ENV: `$PW_PARTITION_INTERVAL`
+
+    Example:
+    `--partition-inteval="3 weeks 4 days"`,
+
 - `--retention=`
 
     If set, metrics older than that will be deleted (default: "14 days").  

internal/sinks/cmdopts.go

@@ -6,6 +6,7 @@ import "time"
 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 table partitions. Must be a valid PostgreSQL interval expression." default:"1 week" env:"PW_PARTITION_INTERVAL"`
 	Retention             string        `long:"retention" mapstructure:"retention" description:"If set, metrics older than that will be deleted" default:"14 days" env:"PW_RETENTION"`
 	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

@@ -48,6 +48,15 @@ func NewWriterFromPostgresConn(ctx context.Context, conn db.PgxPoolIface, opts *
 		if err = conn.QueryRow(ctx, "SELECT $1::interval > '0'::interval", opts.Retention).Scan(&runDeleteOldPartitions); err != nil {
 			return err
 		}
+		var isValidInterval bool 
+		err = conn.QueryRow(ctx, "SELECT $1::interval >= '1h'::interval", opts.PartitionInterval).Scan(&isValidInterval)
+		if err != nil {
+			return err
+		}
+		if !isValidInterval {
+			return fmt.Errorf("partition interval must be at least 1 hour, got: %s", opts.PartitionInterval)
+		}
+
 		l.Info("initialising measurements database...")
 		exists, err := db.DoesSchemaExist(ctx, conn, "admin")
 		if err != nil || exists {
@@ -318,7 +327,6 @@ func (pgw *PostgresWriter) flush(msgs []metrics.MeasurementEnvelope) {
 		return
 	}
 	logger := log.GetLogger(pgw.ctx)
-	// metricsToStorePerMetric := make(map[string][]MeasurementMessagePostgres)
 	pgPartBounds := make(map[string]ExistingPartitionInfo)                  // metric=min/max
 	pgPartBoundsDbName := make(map[string]map[string]ExistingPartitionInfo) // metric=[dbname=min/max]
 	var err error
@@ -421,7 +429,7 @@ func (pgw *PostgresWriter) EnsureMetricTimescale(pgPartBounds map[string]Existin
 
 func (pgw *PostgresWriter) EnsureMetricDbnameTime(metricDbnamePartBounds map[string]map[string]ExistingPartitionInfo, force bool) (err error) {
 	var rows pgx.Rows
-	sqlEnsure := `select * from admin.ensure_partition_metric_dbname_time($1, $2, $3)`
+	sqlEnsure := `select * from admin.ensure_partition_metric_dbname_time($1, $2, $3, $4)`
 	for metric, dbnameTimestampMap := range metricDbnamePartBounds {
 		_, ok := partitionMapMetricDbname[metric]
 		if !ok {
@@ -434,7 +442,7 @@ func (pgw *PostgresWriter) EnsureMetricDbnameTime(metricDbnamePartBounds map[str
 			}
 			partInfo, ok := partitionMapMetricDbname[metric][dbname]
 			if !ok || (ok && (pb.StartTime.Before(partInfo.StartTime))) || force {
-				if rows, err = pgw.sinkDb.Query(pgw.ctx, sqlEnsure, metric, dbname, pb.StartTime); err != nil {
+				if rows, err = pgw.sinkDb.Query(pgw.ctx, sqlEnsure, metric, dbname, pb.StartTime, pgw.opts.PartitionInterval); err != nil {
 					return
 				}
 				if partInfo, err = pgx.CollectOneRow(rows, pgx.RowToStructByPos[ExistingPartitionInfo]); err != nil {
@@ -443,7 +451,7 @@ func (pgw *PostgresWriter) EnsureMetricDbnameTime(metricDbnamePartBounds map[str
 				partitionMapMetricDbname[metric][dbname] = partInfo
 			}
 			if pb.EndTime.After(partInfo.EndTime) || pb.EndTime.Equal(partInfo.EndTime) || force {
-				if rows, err = pgw.sinkDb.Query(pgw.ctx, sqlEnsure, metric, dbname, pb.StartTime); err != nil {
+				if rows, err = pgw.sinkDb.Query(pgw.ctx, sqlEnsure, metric, dbname, pb.StartTime, pgw.opts.PartitionInterval); err != nil {
 					return
 				}
 				if partInfo, err = pgx.CollectOneRow(rows, pgx.RowToStructByPos[ExistingPartitionInfo]); err != nil {

internal/sinks/postgres_test.go

@@ -44,13 +44,18 @@ func TestNewWriterFromPostgresConn(t *testing.T) {
 
 	conn.ExpectPing()
 	conn.ExpectQuery("SELECT \\$1::interval").WithArgs("1 year").WillReturnRows(pgxmock.NewRows([]string{"col"}).AddRow(true))
+	conn.ExpectQuery("SELECT \\$1::interval").WithArgs("1 hour").WillReturnRows(pgxmock.NewRows([]string{"col"}).AddRow(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))
 	for _, m := range metrics.GetDefaultBuiltInMetrics() {
 		conn.ExpectExec("SELECT admin.ensure_dummy_metrics_table").WithArgs(m).WillReturnResult(pgxmock.NewResult("EXECUTE", 1))
 	}
 
-	opts := &CmdOpts{BatchingDelay: time.Hour, Retention: "1 year"}
+	opts := &CmdOpts{
+		BatchingDelay: time.Hour, 
+		Retention: "1 year",
+		PartitionInterval: "1 hour",
+	}
 	pgw, err := NewWriterFromPostgresConn(ctx, conn, opts)
 	assert.NoError(t, err)
 	assert.NotNil(t, pgw)
@@ -499,3 +504,79 @@ func TestCopyFromMeasurements_CopyFail(t *testing.T) {
 	}
 
 }
+
+func TestPartitionInterval(t *testing.T) {
+	a := assert.New(t)
+	r := require.New(t)
+
+	const ImageName = "docker.io/postgres:17-alpine"
+	pgContainer, err := postgres.Run(ctx,
+		ImageName,
+		postgres.WithDatabase("mydatabase"),
+		testcontainers.WithWaitStrategy(
+			wait.ForLog("database system is ready to accept connections").
+				WithOccurrence(2).
+				WithStartupTimeout(5*time.Second)),
+	)
+	r.NoError(err)
+	defer func() { a.NoError(pgContainer.Terminate(ctx)) }()
+
+	connStr, _ := pgContainer.ConnectionString(ctx, "sslmode=disable")
+
+	opts := &CmdOpts{
+		PartitionInterval: "1 minute", 
+		Retention: "14 days",
+		BatchingDelay: time.Second,
+	}
+
+	t.Run("Interval Validation", func(_ *testing.T) {
+		_, err = NewPostgresWriter(ctx, connStr, opts)
+		a.EqualError(err, "partition interval must be at least 1 hour, got: 1 minute")
+
+		opts.PartitionInterval = "not an interval"
+		_, err = NewPostgresWriter(ctx, connStr, opts)
+		a.Error(err)
+
+		validIntervals := []string{
+			"3 days 4 hours", "1 year",
+			"P3D", "PT3H", "0-02", "1 00:00:00",
+			"P0-02", "P1", "2 weeks",
+		}
+
+		for _, interval := range validIntervals {
+			opts.PartitionInterval = interval
+			_, err = NewPostgresWriter(ctx, connStr, opts)
+			a.NoError(err)
+		}
+	})
+
+	t.Run("Partitions Creation", func(_ *testing.T) {
+		opts.PartitionInterval = "3 weeks"
+		pgw, err := NewPostgresWriter(ctx, connStr, opts)
+		r.NoError(err)
+
+		conn, err := pgx.Connect(ctx, connStr)
+		r.NoError(err)
+
+		m := map[string]map[string]ExistingPartitionInfo{
+			"test_metric": {
+				"test_db": {
+					time.Now(), time.Now().Add(time.Hour),
+				},
+			},
+		}
+		err = pgw.EnsureMetricDbnameTime(m, false)
+		r.NoError(err)
+
+		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
+		// + 4 time partitions (1 we asked for + 3 precreated)
+		a.Equal(6, partitionsNum)
+
+		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)
+	})
+}

internal/sinks/sql/ensure_partition_postgres.sql

@@ -1,10 +1,18 @@
--- DROP FUNCTION admin.ensure_partition_metric_dbname_time(text,text,timestamp with time zone,integer);
--- select * from admin.ensure_partition_metric_dbname_time('wal', 'kala', now());
-
+/*
+  ensure_partition_metric_dbname_time creates partitioned metric tables if not already existing
+    metric - name of the metric (top level table)
+    dbname - name of the database (2nd level partition)
+    metric_timestamp - timestamp of the metric (used to determine the time partition)
+    partition_period - interval for partitioning (e.g., '1 week', '1 day', '1 month')
+    partitions_to_precreate - how many future time partitions to create (default 3)
+    part_available_from - output parameter, start time of the time partition where the given metric_timestamp fits in
+    part_available_to - output parameter, end time of the time partition where the given metric_timestamp fits in
+*/
 CREATE OR REPLACE FUNCTION admin.ensure_partition_metric_dbname_time(
     metric text,
     dbname text,
     metric_timestamp timestamptz,
+    partition_period interval default '1 week'::interval,
     partitions_to_precreate int default 3,
     OUT part_available_from timestamptz,
     OUT part_available_to timestamptz)
@@ -15,92 +23,145 @@ RETURNS record AS
 */
 $SQL$
 DECLARE
-  l_year int;
-  l_week int;
-  l_doy  int;
   l_part_name_2nd text;
   l_part_name_3rd text;
-  l_part_start date;
-  l_part_end date;
-  l_sql text;
+  l_part_start timestamptz;
+  l_part_end timestamptz;
   ideal_length int;
   l_template_table text := 'admin.metrics_template';
   MAX_IDENT_LEN CONSTANT integer := current_setting('max_identifier_length')::int;
+  l_partition_format text;
+  l_time_suffix text;
+  l_existing_upper_bound timestamptz;
 BEGIN
+  -- Validate partition period
+  IF partition_period < interval '1 hour' THEN
+    RAISE EXCEPTION 'Partition period must be at least 1 hour, got: %', partition_period;
+  END IF;
+
+  -- Determine partition naming format based on period
+  CASE
+    WHEN partition_period >= interval '1 day' THEN
+      l_partition_format := 'YYYYMMDD';
+    ELSE
+      -- For hourly partitions (>= 1 hour, < 1 day)
+      l_partition_format := 'YYYYMMDD_HH24';
+  END CASE;
 
   PERFORM pg_advisory_xact_lock(regexp_replace( md5(metric) , E'\\D', '', 'g')::varchar(10)::int8);
 
+
   -- 1. level
-  IF NOT EXISTS (SELECT 1
-                   FROM pg_tables
-                  WHERE tablename = metric
-                    AND schemaname = 'public')
+  IF to_regclass('public.' || quote_ident(metric)) IS NULL
   THEN
-    -- RAISE NOTICE 'creating partition % ...', metric;
-    EXECUTE format('CREATE TABLE IF NOT EXISTS public.%I (LIKE admin.metrics_template INCLUDING INDEXES) PARTITION BY LIST (dbname)', metric);
+    EXECUTE format('CREATE TABLE public.%I (LIKE admin.metrics_template INCLUDING INDEXES) PARTITION BY LIST (dbname)', metric);
     EXECUTE format('COMMENT ON TABLE public.%I IS $$pgwatch-generated-metric-lvl$$', metric);
   END IF;
 
   -- 2. level
-
   l_part_name_2nd := metric || '_' || dbname;
-
   IF char_length(l_part_name_2nd) > MAX_IDENT_LEN     -- use "dbname" hash instead of name for overly long ones
   THEN
     ideal_length = MAX_IDENT_LEN - char_length(format('%s_', metric));
     l_part_name_2nd := metric || '_' || substring(md5(dbname) from 1 for ideal_length);
   END IF;
 
-  IF NOT EXISTS (SELECT 1
-                   FROM pg_tables
-                  WHERE tablename = l_part_name_2nd
-                    AND schemaname = 'subpartitions')
+  IF to_regclass('subpartitions.' || quote_ident(l_part_name_2nd)) IS NULL
   THEN
-    --RAISE NOTICE 'creating partition % ...', l_part_name_2nd; 
-    EXECUTE format('CREATE TABLE IF NOT EXISTS subpartitions.%I PARTITION OF public.%I FOR VALUES IN (%L) PARTITION BY RANGE (time)',
+    EXECUTE format('CREATE TABLE subpartitions.%I PARTITION OF public.%I FOR VALUES IN (%L) PARTITION BY RANGE (time)',
                     l_part_name_2nd, metric, dbname);
     EXECUTE format('COMMENT ON TABLE subpartitions.%s IS $$pgwatch-generated-metric-dbname-lvl$$', l_part_name_2nd);
   END IF;
 
-  -- 3. level
-  FOR i IN 0..partitions_to_precreate LOOP
+  -- 3. level 
+  
+  -- Get existing partition upper bound
+  SELECT max(substring(pg_catalog.pg_get_expr(c.relpartbound, c.oid, true) from 'TO \(''([^'']+)''')::timestamptz)
+  INTO l_existing_upper_bound
+  FROM pg_catalog.pg_class c
+  JOIN pg_catalog.pg_inherits i ON i.inhrelid = c.oid
+  JOIN pg_catalog.pg_class parent ON parent.oid = i.inhparent
+  WHERE c.relispartition
+    AND c.relnamespace = 'subpartitions'::regnamespace
+    AND parent.relname = l_part_name_2nd;
 
-      l_year := extract(isoyear from (metric_timestamp + '1week'::interval * i));
-      l_week := extract(week from (metric_timestamp + '1week'::interval * i));
+  -- Determine starting point for new partitions
+  IF l_existing_upper_bound IS NOT NULL THEN
+    -- Start from the existing upper bound to maintain continuity
+    l_part_start := l_existing_upper_bound;
+  ELSE
+    -- No existing partitions, align to clean boundaries based on period size
+    CASE
+      WHEN partition_period >= interval '1 week' THEN
+        l_part_start := date_trunc('week', metric_timestamp);
+      WHEN partition_period >= interval '1 day' THEN
+        l_part_start := date_trunc('day', metric_timestamp);
+      ELSE
+        -- For hourly periods (>= 1 hour, < 1 day)
+        l_part_start := date_trunc('hour', metric_timestamp);
+    END CASE;
+    
+    -- For the first partition, set the available range
+    part_available_from := l_part_start;
+    part_available_to := l_part_start + partition_period;
+  END IF;
 
-      IF i = 0 THEN
-          l_part_start := to_date(l_year::text || l_week::text, 'iyyyiw');
-          l_part_end := l_part_start + '1week'::interval;
+  -- Create partitions
+  FOR i IN 0..partitions_to_precreate LOOP
+      l_part_end := l_part_start + partition_period;
+      
+      -- Update the available range for the first partition only if we started from metric_timestamp
+      IF i = 0 AND l_existing_upper_bound IS NULL THEN
           part_available_from := l_part_start;
           part_available_to := l_part_end;
-      ELSE
-          l_part_start := l_part_start + '1week'::interval;
-          l_part_end := l_part_start + '1week'::interval;
-          part_available_to := l_part_end;
+      ELSIF i = 0 AND l_existing_upper_bound IS NOT NULL THEN
+          -- For existing partitions, we need to find which partition contains the metric_timestamp
+          IF metric_timestamp >= l_part_start AND metric_timestamp < l_part_end THEN
+              part_available_from := l_part_start;
+              part_available_to := l_part_end;
+          END IF;
       END IF;
 
-      l_part_name_3rd := format('%s_%s_y%sw%s', metric, dbname, l_year, to_char(l_week, 'fm00' ));
+      l_time_suffix := to_char(l_part_start, l_partition_format);
+      l_part_name_3rd := format('%s_%s_%s', metric, dbname, l_time_suffix);
 
       IF char_length(l_part_name_3rd) > MAX_IDENT_LEN     -- use "dbname" hash instead of name for overly long ones
       THEN
-          ideal_length = MAX_IDENT_LEN - char_length(format('%s__y%sw%s', metric, l_year, to_char(l_week, 'fm00')));
-          l_part_name_3rd := format('%s_%s_y%sw%s', metric, substring(md5(dbname) from 1 for ideal_length), l_year, to_char(l_week, 'fm00' ));
+          ideal_length = MAX_IDENT_LEN - char_length(format('%s__%s', metric, l_time_suffix));
+          l_part_name_3rd := format('%s_%s_%s', metric, substring(md5(dbname) from 1 for ideal_length), l_time_suffix);
       END IF;
 
-      IF NOT EXISTS (SELECT 1
-                      FROM pg_tables
-                      WHERE tablename = l_part_name_3rd
-                        AND schemaname = 'subpartitions')
+      IF to_regclass('subpartitions.' || quote_ident(l_part_name_3rd)) IS NULL
       THEN
-        l_sql := format($$CREATE TABLE IF NOT EXISTS subpartitions.%I PARTITION OF subpartitions.%I FOR VALUES FROM ('%s') TO ('%s')$$,
+        EXECUTE format('CREATE TABLE subpartitions.%I PARTITION OF subpartitions.%I FOR VALUES FROM ($$%s$$) TO ($$%s$$)',
                         l_part_name_3rd, l_part_name_2nd, l_part_start, l_part_end);
-        EXECUTE l_sql;
         EXECUTE format('COMMENT ON TABLE subpartitions.%I IS $$pgwatch-generated-metric-dbname-time-lvl$$', l_part_name_3rd);
       END IF;
 
+      l_part_start := l_part_end;
   END LOOP;
+
+  -- If we still don't have part_available_from/to set, find the partition containing metric_timestamp
+  IF part_available_from IS NULL THEN
+    SELECT lower_text::timestamptz, upper_text::timestamptz
+    INTO part_available_from, part_available_to
+    FROM (
+      SELECT substring(pg_catalog.pg_get_expr(c.relpartbound, c.oid, true) from 'FOR VALUES FROM \(''([^'']+)''') AS lower_text,
+             substring(pg_catalog.pg_get_expr(c.relpartbound, c.oid, true) from 'TO \(''([^'']+)''') AS upper_text
+      FROM pg_catalog.pg_class c
+      JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
+      JOIN pg_catalog.pg_inherits i ON i.inhrelid = c.oid
+      JOIN pg_catalog.pg_class parent ON parent.oid = i.inhparent
+      WHERE c.relispartition
+        AND n.nspname = 'subpartitions'
+        AND parent.relname = l_part_name_2nd
+    ) AS partitions
+    WHERE metric_timestamp >= lower_text::timestamptz 
+      AND metric_timestamp < upper_text::timestamptz
+    LIMIT 1;
+  END IF;
 
 END;
 $SQL$ LANGUAGE plpgsql;
 
--- GRANT EXECUTE ON FUNCTION admin.ensure_partition_metric_dbname_time(text,text,timestamp with time zone,integer) TO pgwatch;
+-- GRANT EXECUTE ON FUNCTION admin.ensure_partition_metric_dbname_time(text,text,timestamp with time zone,interval,integer) TO pgwatch;