1 package sources_test
2
3 import (
4 "fmt"
5 "os"
6 "path/filepath"
7 "sync"
8 "testing"
9 "time"
10
11 "github.com/cybertec-postgresql/pgwatch/v6/internal/sources"
12 "github.com/stretchr/testify/assert"
13 )
14
15
16 const sampleEntriesNumber = 5
17
18 const (
19 contribDir = "../../contrib/"
20 sampleFile = "../../contrib/sample.sources.yaml"
21 )
22
23 func TestNewYAMLSourcesReaderWriter(t *testing.T) {
24 a := assert.New(t)
25 yamlrw, err := sources.NewYAMLSourcesReaderWriter(ctx, sampleFile)
26 a.NoError(err)
27 a.NotNil(t, yamlrw)
28 }
29
30 func TestYAMLGetMonitoredDatabases(t *testing.T) {
31 a := assert.New(t)
32
33 t.Run("single file", func(*testing.T) {
34 yamlrw, err := sources.NewYAMLSourcesReaderWriter(ctx, sampleFile)
35 a.NoError(err)
36
37 dbs, err := yamlrw.GetSources()
38 a.NoError(err)
39 a.Len(dbs, sampleEntriesNumber)
40 })
41
42 t.Run("nonexistent file", func(*testing.T) {
43 yamlrw, err := sources.NewYAMLSourcesReaderWriter(ctx, "nonexistent.yaml")
44 a.NoError(err)
45 dbs, err := yamlrw.GetSources()
46 a.Error(err)
47 a.Nil(dbs)
48 })
49
50 t.Run("garbage file", func(*testing.T) {
51 yamlrw, err := sources.NewYAMLSourcesReaderWriter(ctx, filepath.Join(contribDir, "yaml.go"))
52 a.NoError(err)
53 dbs, err := yamlrw.GetSources()
54 a.Error(err)
55 a.Nil(dbs)
56 })
57
58 t.Run("duplicate in single file", func(t *testing.T) {
59 tmpFile := filepath.Join(t.TempDir(), "duplicate.yaml")
60 yamlContent := `
61 - name: test1
62 conn_str: postgresql://localhost/test1
63 - name: test2
64 conn_str: postgresql://localhost/test2
65 - name: test1
66 conn_str: postgresql://localhost/test1_duplicate
67 `
68 err := os.WriteFile(tmpFile, []byte(yamlContent), 0644)
69 a.NoError(err)
70 yamlrw, err := sources.NewYAMLSourcesReaderWriter(ctx, tmpFile)
71 a.NoError(err)
72
73 dbs, err := yamlrw.GetSources()
74 a.Error(err)
75 a.Nil(dbs)
76 })
77
78 t.Run("duplicates across files", func(t *testing.T) {
79 tmpDir := t.TempDir()
80 yamlContent1 := `
81 - name: test1
82 conn_str: postgresql://localhost/test1
83 - name: test2
84 conn_str: postgresql://localhost/test2
85 `
86 err := os.WriteFile(filepath.Join(tmpDir, "sources1.yaml"), []byte(yamlContent1), 0644)
87 a.NoError(err)
88
89 yamlContent2 := `
90 - name: test1
91 conn_str: postgresql://localhost/test1_duplicate
92 `
93 err = os.WriteFile(filepath.Join(tmpDir, "sources2.yaml"), []byte(yamlContent2), 0644)
94 a.NoError(err)
95 yamlrw, err := sources.NewYAMLSourcesReaderWriter(ctx, tmpDir)
96 a.NoError(err)
97
98 dbs, err := yamlrw.GetSources()
99 a.Error(err)
100 a.Nil(dbs)
101 })
102
103 t.Run("directory with yaml and yml files", func(t *testing.T) {
104 tmpDir := t.TempDir()
105 yamlContent1 := `
106 - name: dir_test1
107 conn_str: postgresql://localhost/test1
108 `
109 yamlContent2 := `
110 - name: dir_test2
111 conn_str: postgresql://localhost/test2
112 `
113 err := os.WriteFile(filepath.Join(tmpDir, "sources.yaml"), []byte(yamlContent1), 0644)
114 a.NoError(err)
115 err = os.WriteFile(filepath.Join(tmpDir, "sources.yml"), []byte(yamlContent2), 0644)
116 a.NoError(err)
117 yamlrw, err := sources.NewYAMLSourcesReaderWriter(ctx, tmpDir)
118 a.NoError(err)
119 dbs, err := yamlrw.GetSources()
120 a.NoError(err)
121 a.Len(dbs, 2)
122 })
123 }
124
125 func TestYAMLDeleteDatabase(t *testing.T) {
126 a := assert.New(t)
127
128 t.Run("happy path", func(*testing.T) {
129 data, err := os.ReadFile(sampleFile)
130 a.NoError(err)
131 tmpSampleFile := filepath.Join(t.TempDir(), "sample.sources.yaml")
132 err = os.WriteFile(tmpSampleFile, data, 0644)
133 a.NoError(err)
134 defer os.Remove(tmpSampleFile)
135
136 yamlrw, err := sources.NewYAMLSourcesReaderWriter(ctx, tmpSampleFile)
137 a.NoError(err)
138
139 err = yamlrw.DeleteSource("test1")
140 a.NoError(err)
141
142 dbs, err := yamlrw.GetSources()
143 a.NoError(err)
144 a.Len(dbs, sampleEntriesNumber-1)
145 })
146
147 t.Run("nonexistent file", func(*testing.T) {
148 yamlrw, err := sources.NewYAMLSourcesReaderWriter(ctx, "nonexistent.yaml")
149 a.NoError(err)
150 err = yamlrw.DeleteSource("test1")
151 a.Error(err)
152 })
153 }
154
155 func TestYAMLUpdateDatabase(t *testing.T) {
156 a := assert.New(t)
157
158 t.Run("happy path", func(*testing.T) {
159 data, err := os.ReadFile(sampleFile)
160 a.NoError(err)
161 tmpSampleFile := filepath.Join(t.TempDir(), "sample.sources.yaml")
162 err = os.WriteFile(tmpSampleFile, data, 0644)
163 a.NoError(err)
164 defer os.Remove(tmpSampleFile)
165
166 yamlrw, err := sources.NewYAMLSourcesReaderWriter(ctx, tmpSampleFile)
167 a.NoError(err)
168
169
170 md := sources.Source{}
171 md.Name = "test1"
172 md.ConnStr = "postgresql://localhost/test1"
173 err = yamlrw.UpdateSource(md)
174 a.NoError(err)
175
176
177 md = sources.Source{}
178 md.Name = "test5"
179 md.ConnStr = "postgresql://localhost/test5"
180 err = yamlrw.UpdateSource(md)
181 a.NoError(err)
182
183 dbs, err := yamlrw.GetSources()
184 a.NoError(err)
185 a.Len(dbs, sampleEntriesNumber+1)
186 dbs[0].ConnStr = "postgresql://localhost/test1"
187 dbs[sampleEntriesNumber].ConnStr = "postgresql://localhost/test5"
188 })
189
190 t.Run("nonexistent file", func(*testing.T) {
191 yamlrw, err := sources.NewYAMLSourcesReaderWriter(ctx, "")
192 a.NoError(err)
193 err = yamlrw.UpdateSource(sources.Source{})
194 a.Error(err)
195 })
196 }
197
198 func TestYAMLCreateSource(t *testing.T) {
199 a := assert.New(t)
200
201 t.Run("happy_path", func(*testing.T) {
202 data, err := os.ReadFile(sampleFile)
203 a.NoError(err)
204 tmpSampleFile := filepath.Join(t.TempDir(), "sample.sources.yaml")
205 err = os.WriteFile(tmpSampleFile, data, 0644)
206 a.NoError(err)
207 defer os.Remove(tmpSampleFile)
208
209 yamlrw, err := sources.NewYAMLSourcesReaderWriter(ctx, tmpSampleFile)
210 a.NoError(err)
211
212
213 md := sources.Source{
214 Name: "new_source",
215 ConnStr: "postgresql://localhost/new_db",
216 Kind: sources.SourcePostgres,
217 }
218 err = yamlrw.CreateSource(md)
219 a.NoError(err)
220
221
222 dbs, err := yamlrw.GetSources()
223 a.NoError(err)
224 a.Len(dbs, sampleEntriesNumber+1)
225
226
227 err = yamlrw.CreateSource(md)
228 a.Error(err)
229 a.ErrorIs(sources.ErrSourceExists, err)
230 })
231
232 t.Run("duplicate_source", func(*testing.T) {
233 data, err := os.ReadFile(sampleFile)
234 a.NoError(err)
235 tmpSampleFile := filepath.Join(t.TempDir(), "sample.sources.yaml")
236 err = os.WriteFile(tmpSampleFile, data, 0644)
237 a.NoError(err)
238 defer os.Remove(tmpSampleFile)
239
240 yamlrw, err := sources.NewYAMLSourcesReaderWriter(ctx, tmpSampleFile)
241 a.NoError(err)
242
243
244 md := sources.Source{
245 Name: "test1",
246 ConnStr: "postgresql://localhost/test1",
247 Kind: sources.SourcePostgres,
248 }
249 err = yamlrw.CreateSource(md)
250 a.Error(err)
251 a.ErrorIs(sources.ErrSourceExists, err)
252 })
253 }
254
255 func TestExpandEnvVars(t *testing.T) {
256 a := assert.New(t)
257
258
259 t.Setenv("PGW_TEST_NAME", "expanded_name")
260 t.Setenv("PGW_TEST_GROUP", "expanded_group")
261 t.Setenv("PGW_TEST_CONNSTR", "postgresql://localhost/expanded")
262 t.Setenv("PGW_TEST_KIND", "postgres")
263 t.Setenv("PGW_TEST_INCLUDE", "include_pattern")
264 t.Setenv("PGW_TEST_EXCLUDE", "exclude_pattern")
265 t.Setenv("PGW_TEST_PRESET", "exhaustive")
266 t.Setenv("PGW_TEST_PRESET_STANDBY", "standby_preset")
267
268 t.Run("all fields expanded", func(*testing.T) {
269 tmpFile := filepath.Join(t.TempDir(), "env_sources.yaml")
270 yamlContent := `
271 - name: $PGW_TEST_NAME
272 group: $PGW_TEST_GROUP
273 conn_str: $PGW_TEST_CONNSTR
274 kind: $PGW_TEST_KIND
275 include_pattern: $PGW_TEST_INCLUDE
276 exclude_pattern: $PGW_TEST_EXCLUDE
277 preset_metrics: $PGW_TEST_PRESET
278 preset_metrics_standby: $PGW_TEST_PRESET_STANDBY
279 `
280 err := os.WriteFile(tmpFile, []byte(yamlContent), 0644)
281 a.NoError(err)
282
283 yamlrw, err := sources.NewYAMLSourcesReaderWriter(ctx, tmpFile)
284 a.NoError(err)
285
286 dbs, err := yamlrw.GetSources()
287 a.NoError(err)
288 a.Len(dbs, 1)
289
290 src := dbs[0]
291 a.Equal("expanded_name", src.Name)
292 a.Equal("expanded_group", src.Group)
293 a.Equal("postgresql://localhost/expanded", src.ConnStr)
294 a.Equal(sources.SourcePostgres, src.Kind)
295 a.Equal("include_pattern", src.IncludePattern)
296 a.Equal("exclude_pattern", src.ExcludePattern)
297 a.Equal("exhaustive", src.PresetMetrics)
298 a.Equal("standby_preset", src.PresetMetricsStandby)
299 })
300
301 t.Run("no expansion without dollar prefix", func(*testing.T) {
302 tmpFile := filepath.Join(t.TempDir(), "no_env_sources.yaml")
303 yamlContent := `
304 - name: literal_name
305 group: literal_group
306 conn_str: postgresql://localhost/literal
307 kind: postgres
308 include_pattern: literal_include
309 exclude_pattern: literal_exclude
310 preset_metrics: basic
311 preset_metrics_standby: basic_standby
312 `
313 err := os.WriteFile(tmpFile, []byte(yamlContent), 0644)
314 a.NoError(err)
315
316 yamlrw, err := sources.NewYAMLSourcesReaderWriter(ctx, tmpFile)
317 a.NoError(err)
318
319 dbs, err := yamlrw.GetSources()
320 a.NoError(err)
321 a.Len(dbs, 1)
322
323 src := dbs[0]
324 a.Equal("literal_name", src.Name)
325 a.Equal("literal_group", src.Group)
326 a.Equal("postgresql://localhost/literal", src.ConnStr)
327 a.Equal(sources.SourcePostgres, src.Kind)
328 a.Equal("literal_include", src.IncludePattern)
329 a.Equal("literal_exclude", src.ExcludePattern)
330 a.Equal("basic", src.PresetMetrics)
331 a.Equal("basic_standby", src.PresetMetricsStandby)
332 })
333
334 t.Run("unset env var expands to empty", func(*testing.T) {
335 tmpFile := filepath.Join(t.TempDir(), "unset_env_sources.yaml")
336 yamlContent := `
337 - name: $PGW_UNSET_VAR
338 conn_str: postgresql://localhost/test
339 `
340 err := os.WriteFile(tmpFile, []byte(yamlContent), 0644)
341 a.NoError(err)
342
343 yamlrw, err := sources.NewYAMLSourcesReaderWriter(ctx, tmpFile)
344 a.NoError(err)
345
346 dbs, err := yamlrw.GetSources()
347 a.NoError(err)
348 a.Len(dbs, 1)
349 a.Equal("", dbs[0].Name)
350 })
351 }
352
353 func TestConcurrentSourceUpdates(t *testing.T) {
354 a := assert.New(t)
355 tempDir := t.TempDir()
356 tempFile := filepath.Join(tempDir, "sources.yaml")
357
358 yamlrw, err := sources.NewYAMLSourcesReaderWriter(ctx, tempFile)
359 a.NoError(err)
360
361 err = yamlrw.WriteSources(sources.Sources{})
362 a.NoError(err)
363
364 numGoroutines := 10
365 var wg sync.WaitGroup
366
367
368 for id := range numGoroutines {
369 wg.Go(func() {
370 testSource := sources.Source{
371 Name: fmt.Sprintf("source_%d", id),
372 ConnStr: fmt.Sprintf("postgresql://localhost/test_%d", id),
373 Kind: sources.SourcePostgres,
374 PresetMetrics: "basic",
375 }
376 time.Sleep(time.Millisecond * time.Duration(id%3))
377 err := yamlrw.UpdateSource(testSource)
378 a.NoError(err, "Error during concurrent update")
379 })
380 }
381
382 wg.Wait()
383
384 finalSources, err := yamlrw.GetSources()
385 a.NoError(err)
386 a.Equal(numGoroutines, len(finalSources), "Some updates were lost due to race condition!")
387 }
388
389 func TestGetSourcesDirWithInvalidYAML(t *testing.T) {
390 dir := t.TempDir()
391 err := os.WriteFile(filepath.Join(dir, "bad.yaml"), []byte("invalid: yaml: {unclosed"), 0644)
392 assert.NoError(t, err)
393 yamlrw, err := sources.NewYAMLSourcesReaderWriter(ctx, dir)
394 assert.NoError(t, err)
395 _, err = yamlrw.GetSources()
396 assert.Error(t, err)
397 }
398
399 func TestCreateSourceGetSourcesError(t *testing.T) {
400 nonExistent := filepath.Join(t.TempDir(), "does_not_exist", "sources.yaml")
401 yamlrw, err := sources.NewYAMLSourcesReaderWriter(ctx, nonExistent)
402 assert.NoError(t, err)
403 assert.Error(t, yamlrw.CreateSource(sources.Source{Name: "x"}))
404 }
405
406
407
408 func TestMutationsWriteError(t *testing.T) {
409
410
411 dir := t.TempDir()
412 yamlrw, err := sources.NewYAMLSourcesReaderWriter(ctx, dir)
413 assert.NoError(t, err)
414
415 assert.Error(t, yamlrw.UpdateSource(sources.Source{Name: "x"}))
416 assert.Error(t, yamlrw.DeleteSource("x"))
417 assert.Error(t, yamlrw.CreateSource(sources.Source{Name: "x"}))
418 }
419
420
421
422 func TestYAML_PrometheusREQ033Example(t *testing.T) {
423 a := assert.New(t)
424 tmpFile := filepath.Join(t.TempDir(), "req033.sources.yaml")
425 yamlContent := `
426 - name: postgres-exporter-prod
427 kind: prometheus
428 conn_str: "https://user:secret@localhost:9187/metrics?tlsrootcert=/etc/ssl/certs/my-ca.pem"
429 is_enabled: true
430 custom_metrics:
431 pg_stat_activity_count: 30
432 pg_stat_bgwriter_checkpoints_timed: 60
433 custom_tags:
434 env: production
435
436 - name: node-exporter-prod
437 kind: prometheus
438 conn_str: "http://localhost:9100/metrics"
439 is_enabled: true
440 `
441 err := os.WriteFile(tmpFile, []byte(yamlContent), 0644)
442 a.NoError(err)
443
444 yamlrw, err := sources.NewYAMLSourcesReaderWriter(ctx, tmpFile)
445 a.NoError(err)
446
447 srcs, err := yamlrw.GetSources()
448 a.NoError(err)
449
450 srcs, err = srcs.Validate()
451 a.NoError(err)
452 if !a.Len(srcs, 2) {
453 return
454 }
455
456 prod := srcs[0]
457 a.Equal(sources.SourcePrometheus, prod.Kind)
458 a.Equal("postgres-exporter-prod", prod.Name)
459 a.Equal("https://user:secret@localhost:9187/metrics?tlsrootcert=/etc/ssl/certs/my-ca.pem", prod.ConnStr)
460 a.Equal("production", prod.CustomTags["env"])
461 a.Equal(30, prod.Metrics["pg_stat_activity_count"])
462 a.Equal(60, prod.Metrics["pg_stat_bgwriter_checkpoints_timed"])
463
464 scrapeAll := srcs[1]
465 a.Equal(sources.SourcePrometheus, scrapeAll.Kind)
466 a.Equal("node-exporter-prod", scrapeAll.Name)
467 a.Empty(scrapeAll.Metrics)
468 }
469
470 func TestYAML_PrometheusSourceRoundTrip(t *testing.T) {
471 a := assert.New(t)
472 tmpFile := filepath.Join(t.TempDir(), "prometheus.sources.yaml")
473 yamlContent := `
474 - name: my-postgres-exporter
475 kind: prometheus
476 conn_str: "http://localhost:9187/metrics"
477 custom_metrics:
478 pg_stat_activity_count: 30
479 custom_tags:
480 env: production
481 is_enabled: true
482 `
483
484 err := os.WriteFile(tmpFile, []byte(yamlContent), 0644)
485 a.NoError(err)
486
487 yamlrw, err := sources.NewYAMLSourcesReaderWriter(ctx, tmpFile)
488 a.NoError(err)
489
490 srcs, err := yamlrw.GetSources()
491 a.NoError(err)
492
493 srcs, err = srcs.Validate()
494 a.NoError(err)
495 if !a.Len(srcs, 1) {
496 return
497 }
498
499 src := srcs[0]
500 a.Equal(sources.SourcePrometheus, src.Kind)
501 a.Equal("http://localhost:9187/metrics", src.ConnStr)
502 a.Contains(src.Metrics, "pg_stat_activity_count")
503 a.Equal(30, src.Metrics["pg_stat_activity_count"])
504 a.Equal("production", src.CustomTags["env"])
505 }
506