Skip to content

Commit e32f5a9

Browse files
sweatybridgesweatybridge
committed
chore: refactor setup database logic
1 parent 30ff8d7 commit e32f5a9

File tree

6 files changed

+119
-112
lines changed

6 files changed

+119
-112
lines changed

‎internal/db/diff/migra.go‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -141,11 +141,7 @@ func MigrateShadowDatabaseVersions(ctx context.Context, container string, migrat
141141
return err
142142
}
143143
defer conn.Close(context.Background())
144-
if err := start.InitDatabase(ctx, container[:12], os.Stderr, options...); err != nil {
145-
return err
146-
}
147-
// Apply user migrations
148-
if err := start.CreateCustomRoles(ctx, conn, os.Stderr, fsys); err != nil {
144+
if err := start.SetupDatabase(ctx, conn, container[:12], os.Stderr, fsys); err != nil {
149145
return err
150146
}
151147
return apply.MigrateUp(ctx, conn, migrations, fsys)

‎internal/db/diff/migra_test.go‎

Lines changed: 40 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,40 @@ func TestRunMigra(t *testing.T) {
137137
func TestMigrateShadow(t *testing.T) {
138138
utils.Config.Db.MajorVersion = 14
139139

140+
t.Run("migrates shadow database", func(t *testing.T) {
141+
utils.Config.Db.ShadowPort = 54320
142+
utils.GlobalsSql = "create schema public"
143+
utils.InitialSchemaSql = "create schema private"
144+
// Setup in-memory fs
145+
fsys := afero.NewMemMapFs()
146+
path := filepath.Join(utils.MigrationsDir, "0_test.sql")
147+
sql := "create schema test"
148+
require.NoError(t, afero.WriteFile(fsys, path, []byte(sql), 0644))
149+
// Setup mock postgres
150+
conn := pgtest.NewConn()
151+
defer conn.Close(t)
152+
conn.Query(utils.GlobalsSql).
153+
Reply("CREATE SCHEMA").
154+
Query(utils.InitialSchemaSql).
155+
Reply("CREATE SCHEMA").
156+
Query(repair.CREATE_VERSION_SCHEMA).
157+
Reply("CREATE SCHEMA").
158+
Query(repair.CREATE_VERSION_TABLE).
159+
Reply("CREATE TABLE").
160+
Query(repair.ADD_STATEMENTS_COLUMN).
161+
Reply("ALTER TABLE").
162+
Query(repair.ADD_NAME_COLUMN).
163+
Reply("ALTER TABLE").
164+
Query(sql).
165+
Reply("CREATE SCHEMA").
166+
Query(repair.INSERT_MIGRATION_VERSION, "0", "test", fmt.Sprintf("{%s}", sql)).
167+
Reply("INSERT 0 1")
168+
// Run test
169+
err := MigrateShadowDatabase(context.Background(), "test-shadow-db", fsys, conn.Intercept)
170+
// Check error
171+
assert.NoError(t, err)
172+
})
173+
140174
t.Run("throws error on timeout", func(t *testing.T) {
141175
utils.Config.Db.ShadowPort = 54320
142176
// Setup in-memory fs
@@ -150,37 +184,27 @@ func TestMigrateShadow(t *testing.T) {
150184
assert.ErrorContains(t, err, "operation was canceled")
151185
})
152186

153-
t.Run("throws error on globals schema", func(t *testing.T) {
154-
utils.Config.Db.ShadowPort = 54320
155-
utils.GlobalsSql = "create schema public"
187+
t.Run("throws error on permission denied", func(t *testing.T) {
156188
// Setup in-memory fs
157-
fsys := afero.NewMemMapFs()
158-
// Setup mock postgres
159-
conn := pgtest.NewConn()
160-
defer conn.Close(t)
161-
conn.Query(utils.GlobalsSql).
162-
ReplyError(pgerrcode.DuplicateSchema, `schema "public" already exists`)
189+
fsys := afero.NewReadOnlyFs(afero.NewMemMapFs())
163190
// Run test
164-
err := MigrateShadowDatabase(context.Background(), "", fsys, conn.Intercept)
191+
err := MigrateShadowDatabase(context.Background(), "", fsys)
165192
// Check error
166-
assert.ErrorContains(t, err, `ERROR: schema "public" already exists (SQLSTATE 42P06)`)
193+
assert.ErrorIs(t, err, os.ErrPermission)
167194
})
168195

169-
t.Run("throws error on initial schema", func(t *testing.T) {
196+
t.Run("throws error on globals schema", func(t *testing.T) {
170197
utils.Config.Db.ShadowPort = 54320
171198
utils.GlobalsSql = "create schema public"
172-
utils.InitialSchemaSql = "create schema private"
173199
// Setup in-memory fs
174200
fsys := afero.NewMemMapFs()
175201
// Setup mock postgres
176202
conn := pgtest.NewConn()
177203
defer conn.Close(t)
178204
conn.Query(utils.GlobalsSql).
179-
Reply("CREATE SCHEMA").
180-
Query(utils.InitialSchemaSql).
181205
ReplyError(pgerrcode.DuplicateSchema, `schema "public" already exists`)
182206
// Run test
183-
err := MigrateShadowDatabase(context.Background(), "", fsys, conn.Intercept)
207+
err := MigrateShadowDatabase(context.Background(), "test-shadow-db", fsys, conn.Intercept)
184208
// Check error
185209
assert.ErrorContains(t, err, `ERROR: schema "public" already exists (SQLSTATE 42P06)`)
186210
})

‎internal/db/start/start.go‎

Lines changed: 11 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -73,8 +73,6 @@ EOF
7373
return config
7474
}
7575

76-
var noBackupVolume bool = true
77-
7876
func StartDatabase(ctx context.Context, fsys afero.Fs, w io.Writer, options ...func(*pgx.ConnConfig)) error {
7977
config := NewContainerConfig()
8078
hostPort := strconv.FormatUint(uint64(utils.Config.Db.Port), 10)
@@ -93,7 +91,7 @@ func StartDatabase(ctx context.Context, fsys afero.Fs, w io.Writer, options ...f
9391
}
9492
// Creating volume will not override existing volume, so we must inspect explicitly
9593
_, err := utils.Docker.VolumeInspect(ctx, utils.DbId)
96-
noBackupVolume = client.IsErrNotFound(err)
94+
noBackupVolume := client.IsErrNotFound(err)
9795
if noBackupVolume {
9896
fmt.Fprintln(w, "Starting database...")
9997
} else {
@@ -107,7 +105,7 @@ func StartDatabase(ctx context.Context, fsys afero.Fs, w io.Writer, options ...f
107105
}
108106
// Initialise if we are on PG14 and there's no existing db volume
109107
if noBackupVolume {
110-
if err := InitDatabase(ctx, utils.DbId, w, options...); err != nil {
108+
if err := setupDatabase(ctx, fsys, w, options...); err != nil {
111109
return err
112110
}
113111
}
@@ -137,18 +135,12 @@ func initCurrentBranch(fsys afero.Fs) error {
137135
return afero.WriteFile(fsys, utils.CurrBranchPath, []byte("main"), 0644)
138136
}
139137

140-
func InitDatabase(ctx context.Context, host string, w io.Writer, options ...func(*pgx.ConnConfig)) error {
141-
// Initialise globals
142-
conn, err := utils.ConnectLocalPostgres(ctx, pgconn.Config{}, options...)
143-
if err != nil {
144-
return err
145-
}
146-
defer conn.Close(context.Background())
138+
func initSchema(ctx context.Context, conn *pgx.Conn, host string, w io.Writer) error {
147139
fmt.Fprintln(w, "Setting up initial schema...")
148140
if utils.Config.Db.MajorVersion <= 14 {
149141
return initSchema14(ctx, conn)
150142
}
151-
return initSchema15(ctx, conn, host)
143+
return initSchema15(ctx, host)
152144
}
153145

154146
func initSchema14(ctx context.Context, conn *pgx.Conn) error {
@@ -158,7 +150,7 @@ func initSchema14(ctx context.Context, conn *pgx.Conn) error {
158150
return apply.BatchExecDDL(ctx, conn, strings.NewReader(utils.InitialSchemaSql))
159151
}
160152

161-
func initSchema15(ctx context.Context, conn *pgx.Conn, host string) error {
153+
func initSchema15(ctx context.Context, host string) error {
162154
// Apply service migrations
163155
if err := utils.DockerRunOnceWithStream(ctx, utils.StorageImage, []string{
164156
"ANON_KEY=" + utils.Config.Auth.AnonKey,
@@ -183,26 +175,22 @@ func initSchema15(ctx context.Context, conn *pgx.Conn, host string) error {
183175
}, []string{"gotrue", "migrate"}, io.Discard, os.Stderr)
184176
}
185177

186-
func SetupDatabase(ctx context.Context, dbConfig pgconn.Config, fsys afero.Fs, w io.Writer, options ...func(*pgx.ConnConfig)) error {
187-
if !noBackupVolume {
188-
return nil
189-
}
190-
if dbConfig.Host != utils.DbId {
191-
return nil
192-
}
193-
// conn, err := utils.ConnectLocalPostgres(ctx, dbConfig, options...)
178+
func setupDatabase(ctx context.Context, fsys afero.Fs, w io.Writer, options ...func(*pgx.ConnConfig)) error {
194179
conn, err := utils.ConnectLocalPostgres(ctx, pgconn.Config{}, options...)
195180
if err != nil {
196181
return err
197182
}
198183
defer conn.Close(context.Background())
199-
if err := CreateCustomRoles(ctx, conn, w, fsys); err != nil {
184+
if err := SetupDatabase(ctx, conn, utils.DbId, w, fsys); err != nil {
200185
return err
201186
}
202187
return reset.InitialiseDatabase(ctx, conn, fsys)
203188
}
204189

205-
func CreateCustomRoles(ctx context.Context, conn *pgx.Conn, w io.Writer, fsys afero.Fs) error {
190+
func SetupDatabase(ctx context.Context, conn *pgx.Conn, host string, w io.Writer, fsys afero.Fs) error {
191+
if err := initSchema(ctx, conn, host, w); err != nil {
192+
return err
193+
}
206194
roles, err := fsys.Open(utils.CustomRolesPath)
207195
if errors.Is(err, os.ErrNotExist) {
208196
return nil

0 commit comments

Comments
 (0)