▽

Source file src/github.com/cybertec-postgresql/pgwatch/v6/internal/sinks/prometheus_test.go

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

     1  package sinks
     2  
     3  import (
     4  	"strconv"
     5  	"strings"
     6  	"testing"
     7  	"time"
     8  
     9  	dto "github.com/prometheus/client_model/go"
    10  
    11  	"github.com/cybertec-postgresql/pgwatch/v6/internal/log"
    12  	"github.com/cybertec-postgresql/pgwatch/v6/internal/metrics"
    13  	"github.com/cybertec-postgresql/pgwatch/v6/internal/testutil"
    14  	"github.com/prometheus/client_golang/prometheus"
    15  	"github.com/stretchr/testify/assert"
    16  	"github.com/stretchr/testify/require"
    17  )
    18  
    19  func newTestPrometheusWriter(namespace string) *PrometheusWriter {
    20  	return &PrometheusWriter{
    21  		ctx:       testutil.TestContext,
    22  		logger:    log.GetLogger(testutil.TestContext),
    23  		Namespace: namespace,
    24  		Cache:     make(PromMetricCache),
    25  		lastScrapeErrors: prometheus.NewGauge(prometheus.GaugeOpts{
    26  			Namespace: namespace,
    27  			Name:      "test_last_scrape_errors",
    28  		}),
    29  		totalScrapes: prometheus.NewCounter(prometheus.CounterOpts{
    30  			Namespace: namespace,
    31  			Name:      "test_total_scrapes",
    32  		}),
    33  		totalScrapeFailures: prometheus.NewCounter(prometheus.CounterOpts{
    34  			Namespace: namespace,
    35  			Name:      "test_total_scrape_failures",
    36  		}),
    37  	}
    38  }
    39  
    40  // TestWriteAfterCollect verifies that Write() works after Collect().
    41  // Since Collect() now reads a snapshot without clearing the cache,
    42  // the cache should still contain the original data after collect.
    43  func TestWriteAfterCollect(t *testing.T) {
    44  	promw := newTestPrometheusWriter("test")
    45  
    46  	// Write initial data
    47  	msg := metrics.MeasurementEnvelope{
    48  		DBName:     "db1",
    49  		MetricName: "metric1",
    50  		Data: metrics.Measurements{
    51  			{metrics.EpochColumnName: time.Now().UnixNano(), "value": int64(100)},
    52  		},
    53  	}
    54  	require.NoError(t, promw.Write(msg))
    55  
    56  	// Collect reads a snapshot — cache is NOT cleared
    57  	ch := make(chan prometheus.Metric, 100)
    58  	promw.Collect(ch)
    59  	assert.NotEmpty(t, promw.Cache, "cache should still contain data after Collect (snapshot-based)")
    60  
    61  	// Write after Collect — must work
    62  	msg.Data[0]["value"] = int64(200)
    63  	require.NoError(t, promw.Write(msg))
    64  
    65  	assert.Contains(t, promw.Cache, "db1")
    66  	assert.Equal(t, int64(200), promw.Cache["db1"]["metric1"].Data[0]["value"])
    67  }
    68  
    69  // TestCollect_CachePreserved verifies Collect() does not consume the cache.
    70  // Parallel scrapes and back-to-back scrapes should see the same data.
    71  func TestCollect_CachePreserved(t *testing.T) {
    72  	promw := newTestPrometheusWriter("test")
    73  
    74  	// Populate cache with multiple databases
    75  	for _, db := range []string{"db1", "db2", "db3", "db4", "db5"} {
    76  		promw.Cache[db] = map[string]metrics.MeasurementEnvelope{
    77  			"metric": {
    78  				DBName:     db,
    79  				MetricName: "metric",
    80  				Data: metrics.Measurements{
    81  					{metrics.EpochColumnName: time.Now().UnixNano(), "value": int64(1)},
    82  				},
    83  			},
    84  		}
    85  	}
    86  	assert.Len(t, promw.Cache, 5)
    87  
    88  	// First Collect
    89  	ch := make(chan prometheus.Metric, 100)
    90  	promw.Collect(ch)
    91  
    92  	// Cache should still have all 5 databases
    93  	assert.Len(t, promw.Cache, 5, "cache should be preserved after Collect")
    94  
    95  	// Second Collect should also work (back-to-back scrapes)
    96  	ch2 := make(chan prometheus.Metric, 100)
    97  	promw.Collect(ch2)
    98  	assert.Len(t, promw.Cache, 5, "cache should be preserved after second Collect")
    99  }
   100  
   101  // TestCollect_DeterministicLabelOrdering verifies that metrics with the same
   102  // labels produce consistent output regardless of map insertion order.
   103  // This is the fix for the "collected metric was collected before" error.
   104  func TestCollect_DeterministicLabelOrdering(t *testing.T) {
   105  	promw := newTestPrometheusWriter("test")
   106  	promw.gauges = map[string][]string{promInstanceUpStateMetric: {"*"}}
   107  
   108  	// Create two data rows with the same labels but potentially different
   109  	// map iteration orders (Go randomizes map iteration).
   110  	promw.Cache["db1"] = map[string]metrics.MeasurementEnvelope{
   111  		"metric1": {
   112  			DBName:     "db1",
   113  			MetricName: promInstanceUpStateMetric,
   114  			Data: metrics.Measurements{
   115  				{
   116  					metrics.EpochColumnName: time.Now().UnixNano(),
   117  					"tag_host":              "server1",
   118  					"tag_port":              "5432",
   119  					"tag_region":            "us-east-1",
   120  					"value":                 int64(42),
   121  				},
   122  			},
   123  		},
   124  	}
   125  
   126  	// Collect multiple times — should never produce duplicate errors.
   127  	// Each Collect emits 3 self-instrumentation metrics + 1 data metric = 4 total.
   128  	const metaMetrics = 3 // totalScrapes + totalScrapeFailures + lastScrapeErrors
   129  	for i := range 100 {
   130  		ch := make(chan prometheus.Metric, 100)
   131  		promw.Collect(ch)
   132  		close(ch)
   133  
   134  		var collected []prometheus.Metric
   135  		for m := range ch {
   136  			collected = append(collected, m)
   137  		}
   138  		// 1 data metric + 3 meta-metrics
   139  		assert.Len(t, collected, 1+metaMetrics, "iteration %d: expected 1 data + 3 meta metrics", i)
   140  	}
   141  }
   142  
   143  // TestCollect_DeduplicateMetrics verifies that duplicate row data
   144  // (same metric name + same label values) is emitted only once per scrape.
   145  func TestCollect_DeduplicateMetrics(t *testing.T) {
   146  	promw := newTestPrometheusWriter("test")
   147  	promw.gauges = map[string][]string{"metric1": {"*"}}
   148  
   149  	// Two identical rows — same labels, same field name
   150  	promw.Cache["db1"] = map[string]metrics.MeasurementEnvelope{
   151  		"metric1": {
   152  			DBName:     "db1",
   153  			MetricName: "metric1",
   154  			CustomTags: map[string]string{"sys_id": "42"}, // custom tags should not affect identity
   155  			Data: metrics.Measurements{
   156  				{
   157  					metrics.EpochColumnName: time.Now().UnixNano(),
   158  					"tag_host":              "server1",
   159  					"value":                 int64(42),
   160  					"bool_val":              false,
   161  					"extra_field1":          "ignored", // extra fields should not affect identity
   162  				},
   163  				{
   164  					metrics.EpochColumnName: time.Now().UnixNano(),
   165  					"tag_host":              "server1",
   166  					"value":                 int64(99), // different values, same identity
   167  					"bool_val":              true,
   168  					"extra_field1":          "ignored", // extra fields should not affect identity
   169  				},
   170  			},
   171  		},
   172  	}
   173  
   174  	ch := make(chan prometheus.Metric, 100)
   175  	promw.Collect(ch)
   176  	close(ch)
   177  
   178  	var count int
   179  	for c := range ch {
   180  		t.Log(c.Desc())
   181  		count++
   182  	}
   183  	// Should deduplicate — only the first occurrence is emitted.
   184  	// 2 data metrics + 3 meta-metrics = 5 total.
   185  	assert.Equal(t, 2+3, count, "duplicate metric identity should be deduplicated (1 data + 3 meta)")
   186  }
   187  
   188  // TestCollect_InvalidMetricDoesNotPanic verifies that a malformed metric
   189  // (label count mismatch) does not cause a panic. This tests the
   190  // NewConstMetric error handling path.
   191  func TestCollect_InvalidMetricDoesNotPanic(t *testing.T) {
   192  	promw := newTestPrometheusWriter("test")
   193  
   194  	// This should not panic under any circumstances
   195  	assert.NotPanics(t, func() {
   196  		ch := make(chan prometheus.Metric, 100)
   197  		promw.Collect(ch)
   198  	})
   199  }
   200  
   201  // TestCollect_EmptyCache verifies that collecting from an empty cache
   202  // produces only the self-instrumentation metrics.
   203  func TestCollect_EmptyCache(t *testing.T) {
   204  	promw := newTestPrometheusWriter("test")
   205  
   206  	ch := make(chan prometheus.Metric, 100)
   207  	promw.Collect(ch)
   208  	close(ch)
   209  
   210  	var count int
   211  	for range ch {
   212  		count++
   213  	}
   214  	// 0 data metrics + 3 meta-metrics (totalScrapes, totalScrapeFailures, lastScrapeErrors)
   215  	assert.Equal(t, 3, count, "empty cache should produce only 3 meta-metrics")
   216  }
   217  
   218  // TestCollect_StaleMetricsDropped verifies that metrics older than
   219  // promScrapingStalenessHardDropLimit are dropped.
   220  func TestCollect_StaleMetricsDropped(t *testing.T) {
   221  	promw := newTestPrometheusWriter("test")
   222  	promw.gauges = map[string][]string{"metric1": {"*"}}
   223  
   224  	staleEpoch := time.Now().Add(-promCacheTTL - time.Minute).UnixNano()
   225  	promw.Cache["db1"] = map[string]metrics.MeasurementEnvelope{
   226  		"metric1": {
   227  			DBName:     "db1",
   228  			MetricName: "metric1",
   229  			Data: metrics.Measurements{
   230  				{
   231  					metrics.EpochColumnName: staleEpoch,
   232  					"value":                 int64(42),
   233  				},
   234  			},
   235  		},
   236  	}
   237  
   238  	ch := make(chan prometheus.Metric, 100)
   239  	promw.Collect(ch)
   240  	close(ch)
   241  
   242  	var count int
   243  	for range ch {
   244  		count++
   245  	}
   246  	// 0 data metrics + 3 meta-metrics
   247  	assert.Equal(t, 3, count, "stale metrics should be dropped, only meta-metrics remain")
   248  }
   249  
   250  func TestPrometheusWriteEmpty(t *testing.T) {
   251  	promw := newTestPrometheusWriter("test")
   252  	assert.NoError(t, promw.Write(metrics.MeasurementEnvelope{}))
   253  	ch := make(chan prometheus.Metric, 100)
   254  	written, errCount := promw.WritePromMetrics(metrics.MeasurementEnvelope{}, ch)
   255  	assert.Zero(t, errCount)
   256  	assert.Zero(t, written)
   257  	close(ch)
   258  }
   259  
   260  func TestPrometheusWRiteUnsupportedMetric(t *testing.T) {
   261  	promw := newTestPrometheusWriter("test")
   262  	assert.NoError(t, promw.Write(metrics.MeasurementEnvelope{
   263  		DBName:     "db1",
   264  		MetricName: "change_events", // unsupported metric
   265  		Data: metrics.Measurements{
   266  			{metrics.EpochColumnName: time.Now().UnixNano(), "value": int64(100)},
   267  		},
   268  	}))
   269  }
   270  
   271  // TestSourceKindPrometheus covers T041: SourceKind == "prometheus" is the sentinel
   272  // used to distinguish prom-sourced envelopes from DB-sourced ones.
   273  func TestSourceKindPrometheus(t *testing.T) {
   274  	tests := []struct {
   275  		name     string
   276  		envelope metrics.MeasurementEnvelope
   277  		want     bool
   278  	}{
   279  		{name: "prometheus", envelope: metrics.MeasurementEnvelope{SourceKind: "prometheus"}, want: true},
   280  		{name: "empty", envelope: metrics.MeasurementEnvelope{}, want: false},
   281  		{name: "postgresql", envelope: metrics.MeasurementEnvelope{SourceKind: "postgresql"}, want: false},
   282  		{name: "pgbouncer", envelope: metrics.MeasurementEnvelope{SourceKind: "pgbouncer"}, want: false},
   283  	}
   284  	for _, tt := range tests {
   285  		t.Run(tt.name, func(t *testing.T) {
   286  			assert.Equal(t, tt.want, tt.envelope.SourceKind == "prometheus")
   287  		})
   288  	}
   289  }
   290  
   291  // collectDataMetrics drains ch and returns only the metrics whose desc string
   292  // contains nameFilter.
   293  func collectDataMetrics(t *testing.T, ch <-chan prometheus.Metric, nameFilter string) []prometheus.Metric {
   294  	t.Helper()
   295  	var out []prometheus.Metric
   296  	for m := range ch {
   297  		if strings.Contains(m.Desc().String(), nameFilter) {
   298  			out = append(out, m)
   299  		}
   300  	}
   301  	return out
   302  }
   303  
   304  // TestPrometheusWriter_PromSourcedEnvelope covers T042: Write + Collect for a
   305  // prom-sourced envelope.
   306  func TestPrometheusWriter_PromSourcedEnvelope(t *testing.T) {
   307  	const namespace = "pgwatch"
   308  	const metricName = "pg_stat_activity_count"
   309  	epochNs := time.Now().UnixNano()
   310  
   311  	newEnv := func() metrics.MeasurementEnvelope {
   312  		return metrics.MeasurementEnvelope{
   313  			DBName:     "mydb",
   314  			MetricName: metricName,
   315  			SourceKind: "prometheus",
   316  			Data: metrics.Measurements{
   317  				{
   318  					"tag_datname":           "defaultdb",
   319  					metricName:              float64(42),
   320  					metrics.EpochColumnName: epochNs,
   321  				},
   322  			},
   323  		}
   324  	}
   325  
   326  	t.Run("metric name has no namespace prefix", func(t *testing.T) {
   327  		promw := newTestPrometheusWriter(namespace)
   328  		require.NoError(t, promw.Write(newEnv()))
   329  
   330  		ch := make(chan prometheus.Metric, 100)
   331  		promw.Collect(ch)
   332  		close(ch)
   333  
   334  		dataMetrics := collectDataMetrics(t, ch, metricName)
   335  		require.Len(t, dataMetrics, 1)
   336  
   337  		descStr := dataMetrics[0].Desc().String()
   338  		assert.Contains(t, descStr, `fqName: "`+metricName+`"`)
   339  		assert.NotContains(t, descStr, namespace+"_"+metricName)
   340  	})
   341  
   342  	t.Run("tag_* columns become labels", func(t *testing.T) {
   343  		promw := newTestPrometheusWriter(namespace)
   344  		require.NoError(t, promw.Write(newEnv()))
   345  
   346  		ch := make(chan prometheus.Metric, 100)
   347  		promw.Collect(ch)
   348  		close(ch)
   349  
   350  		dataMetrics := collectDataMetrics(t, ch, metricName)
   351  		require.Len(t, dataMetrics, 1)
   352  
   353  		var dtoMetric dto.Metric
   354  		require.NoError(t, dataMetrics[0].Write(&dtoMetric))
   355  
   356  		labels := make(map[string]string)
   357  		for _, lp := range dtoMetric.GetLabel() {
   358  			labels[lp.GetName()] = lp.GetValue()
   359  		}
   360  		assert.Equal(t, "defaultdb", labels["datname"], "tag_datname should become label datname")
   361  		assert.Equal(t, "mydb", labels["dbname"], "DBName should be added as dbname label")
   362  	})
   363  
   364  	t.Run("epoch_ns used as metric timestamp", func(t *testing.T) {
   365  		promw := newTestPrometheusWriter(namespace)
   366  		require.NoError(t, promw.Write(newEnv()))
   367  
   368  		ch := make(chan prometheus.Metric, 100)
   369  		promw.Collect(ch)
   370  		close(ch)
   371  
   372  		dataMetrics := collectDataMetrics(t, ch, metricName)
   373  		require.Len(t, dataMetrics, 1)
   374  
   375  		var dtoMetric dto.Metric
   376  		require.NoError(t, dataMetrics[0].Write(&dtoMetric))
   377  		require.NotNil(t, dtoMetric.TimestampMs)
   378  		assert.Equal(t, epochNs/1_000_000, dtoMetric.GetTimestampMs())
   379  	})
   380  
   381  	t.Run("duplicate label sets are deduplicated", func(t *testing.T) {
   382  		promw := newTestPrometheusWriter(namespace)
   383  
   384  		env := metrics.MeasurementEnvelope{
   385  			DBName:     "mydb",
   386  			MetricName: "pg_connections",
   387  			SourceKind: "prometheus",
   388  			Data: metrics.Measurements{
   389  				{
   390  					"tag_host":              "server1",
   391  					"pg_connections":        float64(10),
   392  					metrics.EpochColumnName: time.Now().UnixNano(),
   393  				},
   394  				{
   395  					"tag_host":              "server1", // same label set → duplicate
   396  					"pg_connections":        float64(20),
   397  					metrics.EpochColumnName: time.Now().UnixNano(),
   398  				},
   399  			},
   400  		}
   401  		require.NoError(t, promw.Write(env))
   402  
   403  		ch := make(chan prometheus.Metric, 100)
   404  		promw.Collect(ch)
   405  		close(ch)
   406  
   407  		dataMetrics := collectDataMetrics(t, ch, "pg_connections")
   408  		assert.Len(t, dataMetrics, 1, "duplicate (metric_name, label_set) pair should be emitted only once")
   409  	})
   410  }
   411  
   412  // TestPrometheusWriter_NonPromSourced_NamespacePrefix covers T043: the pgwatch
   413  // namespace IS still prepended for non-prometheus-sourced envelopes.
   414  func TestPrometheusWriter_NonPromSourced_NamespacePrefix(t *testing.T) {
   415  	const namespace = "pgwatch"
   416  	promw := newTestPrometheusWriter(namespace)
   417  	promw.gauges = map[string][]string{"pg_stat_activity": {"*"}}
   418  
   419  	env := metrics.MeasurementEnvelope{
   420  		DBName:     "mydb",
   421  		MetricName: "pg_stat_activity",
   422  		SourceKind: "", // not prometheus
   423  		Data: metrics.Measurements{
   424  			{
   425  				metrics.EpochColumnName: time.Now().UnixNano(),
   426  				"numbackends":           int64(5),
   427  			},
   428  		},
   429  	}
   430  	require.NoError(t, promw.Write(env))
   431  
   432  	ch := make(chan prometheus.Metric, 100)
   433  	promw.Collect(ch)
   434  	close(ch)
   435  
   436  	dataMetrics := collectDataMetrics(t, ch, "pg_stat_activity")
   437  	require.Len(t, dataMetrics, 1)
   438  
   439  	descStr := dataMetrics[0].Desc().String()
   440  	assert.Contains(t, descStr, `fqName: "`+namespace+`_pg_stat_activity_numbackends"`)
   441  }
   442  
   443  // BenchmarkWritePromMetrics measures the per-scrape conversion cost of one
   444  // cached envelope (500 rows, 3 tags, 8 numeric columns).
   445  func BenchmarkWritePromMetrics(b *testing.B) {
   446  	promw := newTestPrometheusWriter("pgwatch")
   447  	promw.gauges = map[string][]string{"bench": {"*"}}
   448  
   449  	const rows, fields = 500, 8
   450  	epoch := time.Now().UnixNano()
   451  	data := make(metrics.Measurements, rows)
   452  	for i := range data {
   453  		m := metrics.Measurement{
   454  			metrics.EpochColumnName: epoch,
   455  			"tag_schema":            "public",
   456  			"tag_table":             "table_" + strconv.Itoa(i),
   457  			"tag_state":             "active",
   458  		}
   459  		for f := range fields {
   460  			m["col_"+strconv.Itoa(f)] = int64(i * f)
   461  		}
   462  		data[i] = m
   463  	}
   464  	msg := metrics.MeasurementEnvelope{DBName: "db1", MetricName: "bench", Data: data}
   465  
   466  	ch := make(chan prometheus.Metric, rows*fields+1)
   467  	b.ReportAllocs()
   468  	for b.Loop() {
   469  		promw.WritePromMetrics(msg, ch)
   470  		for len(ch) > 0 {
   471  			<-ch
   472  		}
   473  	}
   474  }
   475