Skip to content

Commit a6eb919

Browse files
karakanbkarakanb
authored
Merge pull request #1105 from bruin-data/feat/pg-dry-run-with-tests
use test queries instead of select in postgres
2 parents 82c2173 + c5835ba commit a6eb919

File tree

2 files changed

+261
-17
lines changed

2 files changed

+261
-17
lines changed

‎pkg/postgres/db.go

Lines changed: 14 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -139,7 +139,12 @@ func (c *Client) Ping(ctx context.Context) error {
139139
}
140140

141141
func (c *Client) IsValid(ctx context.Context, query *query.Query) (bool, error) {
142-
rows, err := c.connection.Query(ctx, query.ToExplainQuery())
142+
// Wrap the query in a PostgreSQL test pattern that doesn't execute
143+
testQuery := fmt.Sprintf(`DO $TEST$ BEGIN RETURN;
144+
%s
145+
END; $TEST$;`, query.String())
146+
147+
rows, err := c.connection.Query(ctx, testQuery)
143148
if err == nil {
144149
err = rows.Err()
145150
}
@@ -249,7 +254,7 @@ func (c *Client) GetColumns(ctx context.Context, databaseName, tableName string)
249254
}
250255

251256
q := `
252-
SELECT
257+
SELECT
253258
column_name,
254259
data_type,
255260
is_nullable,
@@ -441,15 +446,15 @@ func (c *Client) GetTableSummary(ctx context.Context, tableName string, schemaOn
441446

442447
// Get table schema using information_schema
443448
schemaQuery := `
444-
SELECT
449+
SELECT
445450
column_name,
446451
data_type,
447452
is_nullable,
448453
column_default,
449454
character_maximum_length,
450455
numeric_precision,
451456
numeric_scale
452-
FROM information_schema.columns
457+
FROM information_schema.columns
453458
WHERE table_name = $1
454459
ORDER BY ordinal_position`
455460

@@ -558,7 +563,7 @@ func (c *Client) GetTableSummary(ctx context.Context, tableName string, schemaOn
558563
func (c *Client) fetchNumericalStats(ctx context.Context, tableName, columnName string) (*diff.NumericalStatistics, error) {
559564
stats := &diff.NumericalStatistics{}
560565
query := fmt.Sprintf(`
561-
SELECT
566+
SELECT
562567
COUNT(*) as count,
563568
COUNT(*) - COUNT(%s) as null_count,
564569
MIN(%s) as min_val,
@@ -595,7 +600,7 @@ func (c *Client) fetchNumericalStats(ctx context.Context, tableName, columnName
595600
func (c *Client) fetchStringStats(ctx context.Context, tableName, columnName string) (*diff.StringStatistics, error) {
596601
stats := &diff.StringStatistics{}
597602
query := fmt.Sprintf(`
598-
SELECT
603+
SELECT
599604
COUNT(*) as count,
600605
COUNT(*) - COUNT(%s) as null_count,
601606
COUNT(DISTINCT %s) as distinct_count,
@@ -629,7 +634,7 @@ func (c *Client) fetchStringStats(ctx context.Context, tableName, columnName str
629634
func (c *Client) fetchBooleanStats(ctx context.Context, tableName, columnName string) (*diff.BooleanStatistics, error) {
630635
stats := &diff.BooleanStatistics{}
631636
query := fmt.Sprintf(`
632-
SELECT
637+
SELECT
633638
COUNT(*) as count,
634639
COUNT(*) - COUNT(%s) as null_count,
635640
COUNT(CASE WHEN %s = true THEN 1 END) as true_count,
@@ -656,7 +661,7 @@ func (c *Client) fetchBooleanStats(ctx context.Context, tableName, columnName st
656661
func (c *Client) fetchDateTimeStats(ctx context.Context, tableName, columnName string) (*diff.DateTimeStatistics, error) {
657662
stats := &diff.DateTimeStatistics{}
658663
query := fmt.Sprintf(`
659-
SELECT
664+
SELECT
660665
COUNT(*) as count,
661666
COUNT(*) - COUNT(%s) as null_count,
662667
COUNT(DISTINCT %s) as unique_count,
@@ -698,7 +703,7 @@ func (c *Client) fetchDateTimeStats(ctx context.Context, tableName, columnName s
698703
func (c *Client) fetchJSONStats(ctx context.Context, tableName, columnName string) (*diff.JSONStatistics, error) {
699704
stats := &diff.JSONStatistics{}
700705
query := fmt.Sprintf(`
701-
SELECT
706+
SELECT
702707
COUNT(*) as count,
703708
COUNT(*) - COUNT(%s) as null_count
704709
FROM %s`,

‎pkg/postgres/db_test.go

Lines changed: 247 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -341,28 +341,267 @@ func TestDB_IsValid(t *testing.T) {
341341
name: "simple valid select query is handled",
342342
setupMock: func(mock pgxmock.PgxPoolIface) {
343343
rows := pgxmock.NewRowsWithColumnDefinition(
344-
pgconn.FieldDescription{Name: "id"},
345-
pgconn.FieldDescription{Name: "name"},
346-
).AddRow(1, "John Doe")
347-
mock.ExpectQuery("EXPLAIN SELECT 1;").WillReturnRows(rows)
344+
pgconn.FieldDescription{Name: "result"},
345+
).AddRow("DO")
346+
expectedQuery := `DO \$TEST\$ BEGIN RETURN;
347+
SELECT 1
348+
END; \$TEST\$;`
349+
mock.ExpectQuery(expectedQuery).WillReturnRows(rows)
348350
},
349351
query: query.Query{
350352
Query: "SELECT 1",
351353
},
352354
want: true,
353355
},
354356
{
355-
name: "invalid query is properly handled",
357+
name: "complex valid query with multiple statements",
358+
setupMock: func(mock pgxmock.PgxPoolIface) {
359+
rows := pgxmock.NewRowsWithColumnDefinition(
360+
pgconn.FieldDescription{Name: "result"},
361+
).AddRow("DO")
362+
expectedQuery := `DO \$TEST\$ BEGIN RETURN;
363+
SELECT id, name FROM users WHERE active = true
364+
END; \$TEST\$;`
365+
mock.ExpectQuery(expectedQuery).WillReturnRows(rows)
366+
},
367+
query: query.Query{
368+
Query: "SELECT id, name FROM users WHERE active = true",
369+
},
370+
want: true,
371+
},
372+
{
373+
name: "valid INSERT query",
374+
setupMock: func(mock pgxmock.PgxPoolIface) {
375+
rows := pgxmock.NewRowsWithColumnDefinition(
376+
pgconn.FieldDescription{Name: "result"},
377+
).AddRow("DO")
378+
expectedQuery := `DO \$TEST\$ BEGIN RETURN;
379+
INSERT INTO users \(name, email\) VALUES \('John', 'john@example\.com'\)
380+
END; \$TEST\$;`
381+
mock.ExpectQuery(expectedQuery).WillReturnRows(rows)
382+
},
383+
query: query.Query{
384+
Query: "INSERT INTO users (name, email) VALUES ('John', '[email protected]')",
385+
},
386+
want: true,
387+
},
388+
{
389+
name: "valid UPDATE query",
390+
setupMock: func(mock pgxmock.PgxPoolIface) {
391+
rows := pgxmock.NewRowsWithColumnDefinition(
392+
pgconn.FieldDescription{Name: "result"},
393+
).AddRow("DO")
394+
expectedQuery := `DO \$TEST\$ BEGIN RETURN;
395+
UPDATE users SET active = false WHERE last_login < NOW\(\) - INTERVAL '1 year'
396+
END; \$TEST\$;`
397+
mock.ExpectQuery(expectedQuery).WillReturnRows(rows)
398+
},
399+
query: query.Query{
400+
Query: "UPDATE users SET active = false WHERE last_login < NOW() - INTERVAL '1 year'",
401+
},
402+
want: true,
403+
},
404+
{
405+
name: "valid DELETE query",
406+
setupMock: func(mock pgxmock.PgxPoolIface) {
407+
rows := pgxmock.NewRowsWithColumnDefinition(
408+
pgconn.FieldDescription{Name: "result"},
409+
).AddRow("DO")
410+
expectedQuery := `DO \$TEST\$ BEGIN RETURN;
411+
DELETE FROM logs WHERE created_at < NOW\(\) - INTERVAL '30 days'
412+
END; \$TEST\$;`
413+
mock.ExpectQuery(expectedQuery).WillReturnRows(rows)
414+
},
415+
query: query.Query{
416+
Query: "DELETE FROM logs WHERE created_at < NOW() - INTERVAL '30 days'",
417+
},
418+
want: true,
419+
},
420+
{
421+
name: "invalid query with syntax error",
356422
setupMock: func(mock pgxmock.PgxPoolIface) {
357-
mock.ExpectQuery(`EXPLAIN some broken query;`).
358-
WillReturnError(errors.New("some actual error"))
423+
expectedQuery := `DO \$TEST\$ BEGIN RETURN;
424+
SELECT \* FORM users
425+
END; \$TEST\$;`
426+
mock.ExpectQuery(expectedQuery).
427+
WillReturnError(errors.New("syntax error at or near \"FORM\""))
428+
},
429+
query: query.Query{
430+
Query: "SELECT * FORM users",
431+
},
432+
want: false,
433+
wantErr: true,
434+
errorMessage: "syntax error at or near \"FORM\"",
435+
},
436+
{
437+
name: "invalid query with missing table",
438+
setupMock: func(mock pgxmock.PgxPoolIface) {
439+
expectedQuery := `DO \$TEST\$ BEGIN RETURN;
440+
SELECT \* FROM non_existent_table
441+
END; \$TEST\$;`
442+
mock.ExpectQuery(expectedQuery).
443+
WillReturnError(errors.New("relation \"non_existent_table\" does not exist"))
444+
},
445+
query: query.Query{
446+
Query: "SELECT * FROM non_existent_table",
447+
},
448+
want: false,
449+
wantErr: true,
450+
errorMessage: "relation \"non_existent_table\" does not exist",
451+
},
452+
{
453+
name: "invalid query with wrong column reference",
454+
setupMock: func(mock pgxmock.PgxPoolIface) {
455+
expectedQuery := `DO \$TEST\$ BEGIN RETURN;
456+
SELECT invalid_column FROM users
457+
END; \$TEST\$;`
458+
mock.ExpectQuery(expectedQuery).
459+
WillReturnError(errors.New("column \"invalid_column\" does not exist"))
460+
},
461+
query: query.Query{
462+
Query: "SELECT invalid_column FROM users",
463+
},
464+
want: false,
465+
wantErr: true,
466+
errorMessage: "column \"invalid_column\" does not exist",
467+
},
468+
{
469+
name: "query with CTE (WITH clause)",
470+
setupMock: func(mock pgxmock.PgxPoolIface) {
471+
rows := pgxmock.NewRowsWithColumnDefinition(
472+
pgconn.FieldDescription{Name: "result"},
473+
).AddRow("DO")
474+
expectedQuery := `DO \$TEST\$ BEGIN RETURN;
475+
WITH active_users AS \(SELECT \* FROM users WHERE active = true\) SELECT \* FROM active_users
476+
END; \$TEST\$;`
477+
mock.ExpectQuery(expectedQuery).WillReturnRows(rows)
478+
},
479+
query: query.Query{
480+
Query: "WITH active_users AS (SELECT * FROM users WHERE active = true) SELECT * FROM active_users",
481+
},
482+
want: true,
483+
},
484+
{
485+
name: "query with JOIN",
486+
setupMock: func(mock pgxmock.PgxPoolIface) {
487+
rows := pgxmock.NewRowsWithColumnDefinition(
488+
pgconn.FieldDescription{Name: "result"},
489+
).AddRow("DO")
490+
expectedQuery := `DO \$TEST\$ BEGIN RETURN;
491+
SELECT u\.name, p\.title FROM users u JOIN posts p ON u\.id = p\.user_id
492+
END; \$TEST\$;`
493+
mock.ExpectQuery(expectedQuery).WillReturnRows(rows)
494+
},
495+
query: query.Query{
496+
Query: "SELECT u.name, p.title FROM users u JOIN posts p ON u.id = p.user_id",
497+
},
498+
want: true,
499+
},
500+
{
501+
name: "invalid query with broken SQL",
502+
setupMock: func(mock pgxmock.PgxPoolIface) {
503+
expectedQuery := `DO \$TEST\$ BEGIN RETURN;
504+
some broken query
505+
END; \$TEST\$;`
506+
mock.ExpectQuery(expectedQuery).
507+
WillReturnError(errors.New("syntax error"))
359508
},
360509
query: query.Query{
361510
Query: "some broken query",
362511
},
363512
want: false,
364513
wantErr: true,
365-
errorMessage: "some actual error",
514+
errorMessage: "syntax error",
515+
},
516+
{
517+
name: "query with multiline formatting",
518+
setupMock: func(mock pgxmock.PgxPoolIface) {
519+
rows := pgxmock.NewRowsWithColumnDefinition(
520+
pgconn.FieldDescription{Name: "result"},
521+
).AddRow("DO")
522+
expectedQuery := `DO \$TEST\$ BEGIN RETURN;
523+
SELECT
524+
id,
525+
name,
526+
email
527+
FROM users
528+
WHERE active = true
529+
END; \$TEST\$;`
530+
mock.ExpectQuery(expectedQuery).WillReturnRows(rows)
531+
},
532+
query: query.Query{
533+
Query: `SELECT
534+
id,
535+
name,
536+
email
537+
FROM users
538+
WHERE active = true`,
539+
},
540+
want: true,
541+
},
542+
{
543+
name: "CREATE TABLE query validation",
544+
setupMock: func(mock pgxmock.PgxPoolIface) {
545+
rows := pgxmock.NewRowsWithColumnDefinition(
546+
pgconn.FieldDescription{Name: "result"},
547+
).AddRow("DO")
548+
expectedQuery := `DO \$TEST\$ BEGIN RETURN;
549+
CREATE TABLE test_table \(id SERIAL PRIMARY KEY, name VARCHAR\(100\)\)
550+
END; \$TEST\$;`
551+
mock.ExpectQuery(expectedQuery).WillReturnRows(rows)
552+
},
553+
query: query.Query{
554+
Query: "CREATE TABLE test_table (id SERIAL PRIMARY KEY, name VARCHAR(100))",
555+
},
556+
want: true,
557+
},
558+
{
559+
name: "ALTER TABLE query validation",
560+
setupMock: func(mock pgxmock.PgxPoolIface) {
561+
rows := pgxmock.NewRowsWithColumnDefinition(
562+
pgconn.FieldDescription{Name: "result"},
563+
).AddRow("DO")
564+
expectedQuery := `DO \$TEST\$ BEGIN RETURN;
565+
ALTER TABLE users ADD COLUMN age INTEGER
566+
END; \$TEST\$;`
567+
mock.ExpectQuery(expectedQuery).WillReturnRows(rows)
568+
},
569+
query: query.Query{
570+
Query: "ALTER TABLE users ADD COLUMN age INTEGER",
571+
},
572+
want: true,
573+
},
574+
{
575+
name: "DROP TABLE query validation",
576+
setupMock: func(mock pgxmock.PgxPoolIface) {
577+
rows := pgxmock.NewRowsWithColumnDefinition(
578+
pgconn.FieldDescription{Name: "result"},
579+
).AddRow("DO")
580+
expectedQuery := `DO \$TEST\$ BEGIN RETURN;
581+
DROP TABLE IF EXISTS temp_table
582+
END; \$TEST\$;`
583+
mock.ExpectQuery(expectedQuery).WillReturnRows(rows)
584+
},
585+
query: query.Query{
586+
Query: "DROP TABLE IF EXISTS temp_table",
587+
},
588+
want: true,
589+
},
590+
{
591+
name: "query with special characters in strings",
592+
setupMock: func(mock pgxmock.PgxPoolIface) {
593+
rows := pgxmock.NewRowsWithColumnDefinition(
594+
pgconn.FieldDescription{Name: "result"},
595+
).AddRow("DO")
596+
expectedQuery := `DO \$TEST\$ BEGIN RETURN;
597+
SELECT \* FROM users WHERE name = 'O''Brien' AND email LIKE '%@example\.com'
598+
END; \$TEST\$;`
599+
mock.ExpectQuery(expectedQuery).WillReturnRows(rows)
600+
},
601+
query: query.Query{
602+
Query: "SELECT * FROM users WHERE name = 'O''Brien' AND email LIKE '%@example.com'",
603+
},
604+
want: true,
366605
},
367606
}
368607
for _, tt := range tests {

0 commit comments

Comments
 (0)