1 package sources_test
2
3 import (
4 "errors"
5 "testing"
6
7 "github.com/pashagolub/pgxmock/v5"
8 "github.com/stretchr/testify/assert"
9
10 "github.com/cybertec-postgresql/pgwatch/v6/internal/metrics"
11 "github.com/cybertec-postgresql/pgwatch/v6/internal/sources"
12 )
13
14 func TestNewPostgresSourcesReaderWriter(t *testing.T) {
15 a := assert.New(t)
16 t.Run("ConnectionError", func(*testing.T) {
17 pgrw, err := sources.NewPostgresSourcesReaderWriter(ctx, "postgres://user:pass@foohost:5432/db1")
18 a.Error(err)
19 a.NotNil(t, pgrw)
20 })
21 t.Run("InvalidConnStr", func(*testing.T) {
22 pgrw, err := sources.NewPostgresSourcesReaderWriter(ctx, "invalid_connstr")
23 a.Error(err)
24 a.Nil(pgrw)
25 })
26 }
27
28 func TestNewPostgresSourcesReaderWriterConn(t *testing.T) {
29 a := assert.New(t)
30 conn, err := pgxmock.NewPool()
31 a.NoError(err)
32 conn.ExpectQuery(`SELECT EXISTS`).WithArgs("pgwatch").
33 WillReturnRows(pgxmock.NewRows([]string{"exists"}).AddRow(true))
34 conn.ExpectPing()
35
36 pgrw, err := sources.NewPostgresSourcesReaderWriterConn(ctx, conn)
37 a.NoError(err)
38 a.NotNil(t, pgrw)
39 a.NoError(conn.ExpectationsWereMet())
40 }
41
42 func TestGetMonitoredDatabases(t *testing.T) {
43 a := assert.New(t)
44 conn, err := pgxmock.NewPool()
45 a.NoError(err)
46 conn.ExpectQuery(`SELECT EXISTS`).WithArgs("pgwatch").
47 WillReturnRows(pgxmock.NewRows([]string{"exists"}).AddRow(true))
48 conn.ExpectPing()
49 conn.ExpectQuery(`select \/\* pgwatch_generated \*\/`).WillReturnRows(pgxmock.NewRows([]string{
50 "name", "group", "dbtype", "connstr", "config", "config_standby", "preset_config",
51 "preset_config_standby", "include_pattern", "exclude_pattern",
52 "custom_tags", "only_if_master", "is_enabled",
53 }).AddRow(
54 "db1", "group1", sources.Kind("postgres"), "postgres://user:pass@localhost:5432/db1",
55 metrics.MetricIntervals{"metric": 60}, metrics.MetricIntervals{"standby_metric": 60}, "exhaustive", "exhaustive",
56 ".*", `\_.+`, map[string]string{"tag": "value"}, true, true,
57 ))
58 pgrw, err := sources.NewPostgresSourcesReaderWriterConn(ctx, conn)
59 a.NoError(err)
60
61 dbs, err := pgrw.GetSources()
62 a.NoError(err)
63 a.Len(dbs, 1)
64 a.NoError(conn.ExpectationsWereMet())
65
66
67 conn.ExpectQuery(`select \/\* pgwatch_generated \*\/`).WillReturnError(errors.New("failed query"))
68 dbs, err = pgrw.GetSources()
69 a.Error(err)
70 a.Nil(dbs)
71 a.NoError(conn.ExpectationsWereMet())
72 }
73
74 func TestDeleteDatabase(t *testing.T) {
75 a := assert.New(t)
76 conn, err := pgxmock.NewPool()
77 a.NoError(err)
78 conn.ExpectQuery(`SELECT EXISTS`).WithArgs("pgwatch").
79 WillReturnRows(pgxmock.NewRows([]string{"exists"}).AddRow(true))
80 conn.ExpectPing()
81 conn.ExpectExec(`delete from pgwatch\.source where name = \$1`).WithArgs("db1").WillReturnResult(pgxmock.NewResult("DELETE", 1))
82 pgrw, err := sources.NewPostgresSourcesReaderWriterConn(ctx, conn)
83 a.NoError(err)
84
85 err = pgrw.DeleteSource("db1")
86 a.NoError(err)
87 a.NoError(conn.ExpectationsWereMet())
88 }
89
90 func TestUpdateDatabase(t *testing.T) {
91 a := assert.New(t)
92 conn, err := pgxmock.NewPool()
93 a.NoError(err)
94
95 md := sources.Source{
96 Name: "db1",
97 Group: "group1",
98 Kind: sources.Kind("postgres"),
99 ConnStr: "postgres://user:pass@localhost:5432/db1",
100 Metrics: metrics.MetricIntervals{"metric": 60},
101 MetricsStandby: metrics.MetricIntervals{"standby_metric": 60},
102 IncludePattern: ".*",
103 ExcludePattern: `\_.+`,
104 CustomTags: map[string]string{"tag": "value"},
105 }
106 conn.ExpectQuery(`SELECT EXISTS`).WithArgs("pgwatch").
107 WillReturnRows(pgxmock.NewRows([]string{"exists"}).AddRow(true))
108 conn.ExpectPing()
109 conn.ExpectExec(`insert into pgwatch\.source`).WithArgs(
110 md.Name, md.Group, md.Kind,
111 md.ConnStr, `{"metric":60}`, `{"standby_metric":60}`,
112 md.PresetMetrics, md.PresetMetricsStandby,
113 md.IncludePattern, md.ExcludePattern, `{"tag":"value"}`,
114 md.OnlyIfMaster, md.IsEnabled,
115 ).WillReturnResult(pgxmock.NewResult("UPDATE", 1))
116
117 pgrw, err := sources.NewPostgresSourcesReaderWriterConn(ctx, conn)
118 a.NoError(err)
119 err = pgrw.UpdateSource(md)
120 a.NoError(err)
121 a.NoError(conn.ExpectationsWereMet())
122 }
123
124 func TestWriteMonitoredDatabases(t *testing.T) {
125 var (
126 pgrw sources.ReaderWriter
127 err error
128 )
129 a := assert.New(t)
130 conn, err := pgxmock.NewPool()
131 a.NoError(err)
132 md := sources.Source{
133 Name: "db1",
134 Group: "group1",
135 Kind: sources.Kind("postgres"),
136 ConnStr: "postgres://user:pass@localhost:5432/db1",
137 Metrics: metrics.MetricIntervals{"metric": 60},
138 MetricsStandby: metrics.MetricIntervals{"standby_metric": 60},
139 IncludePattern: ".*",
140 ExcludePattern: `\_.+`,
141 CustomTags: map[string]string{"tag": "value"},
142 }
143 mds := sources.Sources{md}
144
145 t.Run("happy path", func(*testing.T) {
146 conn.ExpectQuery(`SELECT EXISTS`).WithArgs("pgwatch").
147 WillReturnRows(pgxmock.NewRows([]string{"exists"}).AddRow(true))
148 conn.ExpectPing()
149 conn.ExpectBegin()
150 conn.ExpectExec(`truncate pgwatch\.source`).WillReturnResult(pgxmock.NewResult("TRUNCATE", 1))
151 conn.ExpectExec(`insert into pgwatch\.source`).WithArgs(
152 md.Name, md.Group, md.Kind,
153 md.ConnStr, `{"metric":60}`, `{"standby_metric":60}`, md.PresetMetrics, md.PresetMetricsStandby,
154 md.IncludePattern, md.ExcludePattern, `{"tag":"value"}`,
155 md.OnlyIfMaster, md.IsEnabled,
156 ).WillReturnResult(pgxmock.NewResult("INSERT", 1))
157 conn.ExpectCommit()
158 conn.ExpectRollback()
159
160 pgrw, err = sources.NewPostgresSourcesReaderWriterConn(ctx, conn)
161 a.NoError(err)
162 err = pgrw.WriteSources(mds)
163 a.NoError(err)
164 a.NoError(conn.ExpectationsWereMet())
165 })
166
167 t.Run("failed transaction begin", func(*testing.T) {
168 conn.ExpectBegin().WillReturnError(errors.New("failed transaction begin"))
169
170 err = pgrw.WriteSources(mds)
171 a.Error(err)
172 a.NoError(conn.ExpectationsWereMet())
173 })
174
175 t.Run("failed truncate", func(*testing.T) {
176 conn.ExpectBegin()
177 conn.ExpectExec(`truncate pgwatch\.source`).WillReturnError(errors.New("failed truncate"))
178
179 err = pgrw.WriteSources(mds)
180 a.Error(err)
181 a.NoError(conn.ExpectationsWereMet())
182 })
183
184 t.Run("failed insert", func(*testing.T) {
185 conn.ExpectBegin()
186 conn.ExpectExec(`truncate pgwatch\.source`).WillReturnResult(pgxmock.NewResult("TRUNCATE", 1))
187 conn.ExpectExec(`insert into pgwatch\.source`).WithArgs(
188 md.Name, md.Group, md.Kind,
189 md.ConnStr, `{"metric":60}`, `{"standby_metric":60}`, md.PresetMetrics, md.PresetMetricsStandby,
190 md.IncludePattern, md.ExcludePattern, `{"tag":"value"}`,
191 md.OnlyIfMaster, md.IsEnabled,
192 ).WillReturnError(errors.New("failed insert"))
193 conn.ExpectRollback()
194
195 err = pgrw.WriteSources(mds)
196 a.Error(err)
197 a.NoError(conn.ExpectationsWereMet())
198 })
199 }
200
201
202 func anyArgs(n int) []any {
203 args := make([]any, n)
204 for i := range args {
205 args[i] = pgxmock.AnyArg()
206 }
207 return args
208 }
209
210 func TestNewPostgresSourcesReaderWriterConn_Bootstrap(t *testing.T) {
211 a := assert.New(t)
212
213 t.Run("FullBootstrap", func(*testing.T) {
214 df := metrics.GetDefaultMetrics()
215 metricsCount := len(df.MetricDefs)
216 presetsCount := len(df.PresetDefs)
217
218 conn, err := pgxmock.NewPool()
219 a.NoError(err)
220 conn.ExpectQuery(`SELECT EXISTS`).WithArgs("pgwatch").
221 WillReturnRows(pgxmock.NewRows([]string{"exists"}).AddRow(false))
222 conn.ExpectBegin()
223 conn.ExpectExec("CREATE SCHEMA IF NOT EXISTS pgwatch").
224 WillReturnResult(pgxmock.NewResult("CREATE", 1))
225 conn.ExpectBegin()
226 conn.ExpectExec(`INSERT.+metric`).WithArgs(anyArgs(8)...).
227 WillReturnResult(pgxmock.NewResult("INSERT", 1)).Times(uint(metricsCount))
228 conn.ExpectExec(`INSERT.+preset`).WithArgs(anyArgs(3)...).
229 WillReturnResult(pgxmock.NewResult("INSERT", 1)).Times(uint(presetsCount))
230 conn.ExpectCommit()
231 conn.ExpectCommit()
232 conn.ExpectPing()
233
234 rw, err := sources.NewPostgresSourcesReaderWriterConn(ctx, conn)
235 a.NoError(err)
236 a.NotNil(rw)
237 a.NoError(conn.ExpectationsWereMet())
238 })
239
240 t.Run("SchemaQueryFail", func(*testing.T) {
241 conn, err := pgxmock.NewPool()
242 a.NoError(err)
243 conn.ExpectQuery(`SELECT EXISTS`).WithArgs("pgwatch").
244 WillReturnError(assert.AnError)
245 rw, err := sources.NewPostgresSourcesReaderWriterConn(ctx, conn)
246 a.Error(err)
247 a.Nil(rw)
248 a.NoError(conn.ExpectationsWereMet())
249 })
250 }
251
252 func TestSourcesNeedsMigration(t *testing.T) {
253 a := assert.New(t)
254 conn, err := pgxmock.NewPool()
255 a.NoError(err)
256 conn.ExpectQuery(`SELECT EXISTS`).WithArgs("pgwatch").
257 WillReturnRows(pgxmock.NewRows([]string{"exists"}).AddRow(true))
258 conn.ExpectPing()
259 conn.ExpectQuery(`SELECT to_regclass`).
260 WithArgs("pgwatch.migration").
261 WillReturnRows(pgxmock.NewRows([]string{"to_regclass"}).AddRow(true))
262 conn.ExpectQuery(`SELECT count`).WillReturnRows(pgxmock.NewRows([]string{"count"}).AddRow(0))
263
264 rw, err := sources.NewPostgresSourcesReaderWriterConn(ctx, conn)
265 a.NoError(err)
266
267 needs, err := rw.(interface {
268 NeedsMigration() (bool, error)
269 }).NeedsMigration()
270 a.NoError(err)
271 a.True(needs)
272 a.NoError(conn.ExpectationsWereMet())
273 }
274