generated from TBD54566975/tbd-project-template
-
Notifications
You must be signed in to change notification settings - Fork 7
/
devel.go
103 lines (91 loc) · 2.39 KB
/
devel.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
package databasetesting
import (
"context"
"fmt"
"time"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
"github.com/TBD54566975/ftl/backend/controller/sql"
"github.com/TBD54566975/ftl/internal/log"
)
// CreateForDevel creates and migrates a new database for development or testing.
//
// If "recreate" is true, the database will be dropped and recreated.
func CreateForDevel(ctx context.Context, dsn string, recreate bool) (*pgxpool.Pool, error) {
logger := log.FromContext(ctx)
config, err := pgx.ParseConfig(dsn)
if err != nil {
return nil, err
}
noDBDSN := config.Copy()
noDBDSN.Database = ""
var conn *pgx.Conn
for range 10 {
conn, err = pgx.ConnectConfig(ctx, noDBDSN)
if err == nil {
defer conn.Close(ctx)
break
}
logger.Debugf("Waiting for database to be ready: %v", err)
select {
case <-ctx.Done():
return nil, ctx.Err()
case <-time.After(1 * time.Second):
}
}
if conn == nil {
return nil, fmt.Errorf("database not ready after 10 tries: %w", err)
}
if recreate {
// Terminate any dangling connections.
_, err = conn.Exec(ctx, `
SELECT pid, pg_terminate_backend(pid)
FROM pg_stat_activity
WHERE datname = $1 AND pid <> pg_backend_pid()`,
config.Database)
if err != nil {
return nil, err
}
_, err = conn.Exec(ctx, fmt.Sprintf("DROP DATABASE IF EXISTS %q", config.Database))
if err != nil {
return nil, err
}
}
_, _ = conn.Exec(ctx, fmt.Sprintf("CREATE DATABASE %q", config.Database)) //nolint:errcheck // PG doesn't support "IF NOT EXISTS" so instead we just ignore any error.
err = sql.Migrate(ctx, dsn)
if err != nil {
return nil, err
}
realConn, err := pgxpool.New(ctx, dsn)
if err != nil {
return nil, err
}
// Reset transient state in the database to a clean state for development purposes.
// This includes things like resetting the state of async calls, leases,
// controller/runner registration, etc. but not anything more.
if !recreate {
_, err = realConn.Exec(ctx, `
WITH deleted AS (
DELETE FROM async_calls
RETURNING 1
), deleted_fsm_instances AS (
DELETE FROM fsm_instances
RETURNING 1
), deleted_leases AS (
DELETE FROM leases
RETURNING 1
), deleted_controllers AS (
DELETE FROM controller
RETURNING 1
), deleted_runners AS (
DELETE FROM runners
RETURNING 1
)
SELECT COUNT(*)
`)
if err != nil {
return nil, err
}
}
return realConn, nil
}