Skip to content
This repository was archived by the owner on Apr 11, 2023. It is now read-only.

Commit 1778b30

Browse files
author
jld3103
authored
Merge pull request #53 from stack11/feature/initial-reset-without-flag
Reset the database intially if it's empty without specifying the option
2 parents 35c6ebb + d01c43e commit 1778b30

6 files changed

Lines changed: 92 additions & 44 deletions

File tree

‎cmd/apply.go‎

Lines changed: 23 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,8 @@ func NewApplyCommand() *cobra.Command {
4242
return nil
4343
},
4444
Run: func(cmd *cobra.Command, args []string) {
45+
ctx := context.Background()
46+
4547
config, err := internal.ReadConfig()
4648
if err != nil {
4749
log.Fatalf("Failed to read config: %v\n", err)
@@ -52,7 +54,7 @@ func NewApplyCommand() *cobra.Command {
5254
log.Fatalf("Failed to get working directory: %v\n", err)
5355
}
5456

55-
conn, err := pgx.Connect(context.Background(), fmt.Sprintf(
57+
conn, err := pgx.Connect(ctx, fmt.Sprintf(
5658
"postgres://%s:%s@%s:%d/%s?sslmode=%s",
5759
postgresUser,
5860
postgresPassword,
@@ -74,40 +76,50 @@ func NewApplyCommand() *cobra.Command {
7476
}
7577

7678
_, err = conn.Exec(
77-
context.Background(),
79+
ctx,
7880
fmt.Sprintf("DROP DATABASE IF EXISTS %q WITH (FORCE)", config.DatabaseName),
7981
)
8082
if err != nil {
8183
log.Fatalf("Failed to drop database: %v", err)
8284
}
8385
}
8486

85-
databaseExists, err := internal.CheckDatabaseExists(conn, config.DatabaseName)
87+
databaseExists, err := internal.CheckDatabaseExists(ctx, conn, config.DatabaseName)
8688
if err != nil {
8789
log.Fatalf("Failed to check if database exists: %v", err)
8890
}
8991
if !databaseExists {
90-
_, err = conn.Exec(context.Background(), fmt.Sprintf("CREATE DATABASE %q", config.DatabaseName))
92+
_, err = conn.Exec(ctx, fmt.Sprintf("CREATE DATABASE %q", config.DatabaseName))
9193
if err != nil {
9294
log.Fatalf("Failed to create database: %v", err)
9395
}
9496
}
9597

98+
schemaMigrationsTableExists, err := internal.CheckTableExists(ctx, conn, "public", "schema_migrations")
99+
if err != nil {
100+
log.Fatalf("Failed to check if public.schema_migrations exists: %v", err)
101+
}
102+
103+
if schemaMigrationsTableExists != databaseExists {
104+
//nolint:lll
105+
log.Fatalf("Something is wrong, the database and the schema_migrations table should always exist or not exist together!")
106+
}
107+
96108
for _, u := range config.DatabaseUsers {
97109
var userExists bool
98-
userExists, err = internal.CheckUserExists(conn, u)
110+
userExists, err = internal.CheckUserExists(ctx, conn, u)
99111
if err != nil {
100112
log.Fatalf("Failed to check if user exists: %v", err)
101113
}
102114
if !userExists {
103-
_, err = conn.Exec(context.Background(), fmt.Sprintf("CREATE ROLE %q WITH LOGIN", u))
115+
_, err = conn.Exec(ctx, fmt.Sprintf("CREATE ROLE %q WITH LOGIN", u))
104116
if err != nil {
105117
log.Fatalf("Failed to create user: %v", err)
106118
}
107119
}
108120
}
109121

110-
err = conn.Close(context.Background())
122+
err = conn.Close(ctx)
111123
if err != nil {
112124
log.Fatalf("Failed to close connection: %v", err)
113125
}
@@ -127,7 +139,7 @@ func NewApplyCommand() *cobra.Command {
127139
log.Fatalln(err)
128140
}
129141

130-
if resetDatabase {
142+
if resetDatabase || (!databaseExists && !schemaMigrationsTableExists) {
131143
var files []os.DirEntry
132144
files, err = os.ReadDir(filepath.Join(wd, "migrations"))
133145
if err != nil {
@@ -178,19 +190,19 @@ func NewApplyCommand() *cobra.Command {
178190
}
179191
}
180192

181-
conn, err = pgx.Connect(context.Background(), dsn)
193+
conn, err = pgx.Connect(ctx, dsn)
182194
if err != nil {
183195
log.Fatalf("Unable to connect to database: %v", err)
184196
}
185197

186198
for _, u := range config.DatabaseUsers {
187-
_, err = conn.Exec(context.Background(), fmt.Sprintf("GRANT SELECT ON public.schema_migrations TO %q", u))
199+
_, err = conn.Exec(ctx, fmt.Sprintf("GRANT SELECT ON public.schema_migrations TO %q", u))
188200
if err != nil {
189201
log.Fatalf("Failed to grant select permission on schema_migrations to %q: %v", u, err)
190202
}
191203
}
192204

193-
err = conn.Close(context.Background())
205+
err = conn.Close(ctx)
194206
if err != nil {
195207
log.Fatalf("Failed to close connection: %v", err)
196208
}

‎cmd/generate.go‎

Lines changed: 32 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,8 @@ func NewGenerateCommand() *cobra.Command {
6969
return nil
7070
},
7171
Run: func(cmd *cobra.Command, args []string) {
72+
ctx := context.Background()
73+
7274
config, err := internal.ReadConfig()
7375
if err != nil {
7476
log.Fatalf("Failed to read config: %v\n", err)
@@ -91,7 +93,7 @@ func NewGenerateCommand() *cobra.Command {
9193
initial := migrationsCount == 0
9294

9395
runnerFunc = func() error {
94-
return runWithStdout(config, initial)
96+
return runWithStdout(ctx, config, initial)
9597
}
9698
} else {
9799
migrationName := args[0]
@@ -114,7 +116,7 @@ func NewGenerateCommand() *cobra.Command {
114116
}()
115117

116118
runnerFunc = func() error {
117-
return runWithFile(config, newMigrationFilePath, migrationNumber)
119+
return runWithFile(ctx, config, newMigrationFilePath, migrationNumber)
118120
}
119121
}
120122

@@ -143,13 +145,13 @@ func NewGenerateCommand() *cobra.Command {
143145
return generateCmd
144146
}
145147

146-
func setupDatabase(name string, port uint32) (*embeddedpostgres.EmbeddedPostgres, *pgx.Conn) {
148+
func setupDatabase(ctx context.Context, name string, port uint32) (*embeddedpostgres.EmbeddedPostgres, *pgx.Conn) {
147149
postgres, dsn := internal.NewPostgresDatabase(fmt.Sprintf("/tmp/trek/%s", name), port)
148150
err := postgres.Start()
149151
if err != nil {
150152
log.Fatalf("Failed to start %s database: %v", name, err)
151153
}
152-
conn, err := pgx.Connect(context.Background(), dsn)
154+
conn, err := pgx.Connect(ctx, dsn)
153155
if err != nil {
154156
log.Fatalf("Unable to connect to %s database: %v", name, err)
155157
}
@@ -160,6 +162,7 @@ func setupDatabase(name string, port uint32) (*embeddedpostgres.EmbeddedPostgres
160162
type RunnerFunc = func() error
161163

162164
func runWithStdout(
165+
ctx context.Context,
163166
config *internal.Config,
164167
initial bool,
165168
) error {
@@ -168,20 +171,21 @@ func runWithStdout(
168171
return fmt.Errorf("failed to check if model has been updated: %w", err)
169172
}
170173
if updated {
171-
targetPostgres, targetConn := setupDatabase("target", 5432)
174+
targetPostgres, targetConn := setupDatabase(ctx, "target", 5432)
172175
defer func() {
173-
_ = targetConn.Close(context.Background())
176+
_ = targetConn.Close(ctx)
174177
_ = targetPostgres.Stop()
175178
}()
176179

177-
migratePostgres, migrateConn := setupDatabase("migrate", 5433)
180+
migratePostgres, migrateConn := setupDatabase(ctx, "migrate", 5433)
178181
defer func() {
179-
_ = migrateConn.Close(context.Background())
182+
_ = migrateConn.Close(ctx)
180183
_ = migratePostgres.Stop()
181184
}()
182185

183186
var statements string
184187
statements, err = generateMigrationStatements(
188+
ctx,
185189
config,
186190
initial,
187191
targetConn,
@@ -236,6 +240,7 @@ func runWithStdout(
236240
}
237241

238242
func runWithFile(
243+
ctx context.Context,
239244
config *internal.Config,
240245
newMigrationFilePath string,
241246
migrationNumber uint,
@@ -252,20 +257,21 @@ func runWithFile(
252257
}
253258
}
254259

255-
targetPostgres, targetConn := setupDatabase("target", 5432)
260+
targetPostgres, targetConn := setupDatabase(ctx, "target", 5432)
256261
defer func() {
257-
_ = targetConn.Close(context.Background())
262+
_ = targetConn.Close(ctx)
258263
_ = targetPostgres.Stop()
259264
}()
260265

261-
migratePostgres, migrateConn := setupDatabase("migrate", 5433)
266+
migratePostgres, migrateConn := setupDatabase(ctx, "migrate", 5433)
262267
defer func() {
263-
_ = migrateConn.Close(context.Background())
268+
_ = migrateConn.Close(ctx)
264269
_ = migratePostgres.Stop()
265270
}()
266271

267272
var statements string
268273
statements, err = generateMigrationStatements(
274+
ctx,
269275
config,
270276
migrationNumber == 1,
271277
targetConn,
@@ -301,7 +307,7 @@ func runWithFile(
301307
return fmt.Errorf("failed to write template files: %w", err)
302308
}
303309

304-
updated, err = generateDiffLockFile(newMigrationFilePath, targetConn, migrateConn)
310+
updated, err = generateDiffLockFile(ctx, newMigrationFilePath, targetConn, migrateConn)
305311
if err != nil {
306312
return fmt.Errorf("failed to generate diff lock file: %w", err)
307313
}
@@ -335,12 +341,17 @@ func checkIfUpdated(config *internal.Config) (bool, error) {
335341
return true, nil
336342
}
337343

338-
func generateDiffLockFile(newMigrationFilePath string, targetConn, migrateConn *pgx.Conn) (bool, error) {
344+
func generateDiffLockFile(
345+
ctx context.Context,
346+
newMigrationFilePath string,
347+
targetConn,
348+
migrateConn *pgx.Conn,
349+
) (bool, error) {
339350
newMigrationFileContent, err := os.ReadFile(newMigrationFilePath)
340351
if err != nil {
341352
return false, fmt.Errorf("failed to read new migratio file: %w", err)
342353
}
343-
_, err = migrateConn.Exec(context.Background(), string(newMigrationFileContent))
354+
_, err = migrateConn.Exec(ctx, string(newMigrationFileContent))
344355
if err != nil {
345356
return false, fmt.Errorf("failed to apply generated migration: %w", err)
346357
}
@@ -549,6 +560,7 @@ var ErrInvalidModel = errors.New("invalid model")
549560

550561
//nolint:cyclop
551562
func generateMigrationStatements(
563+
ctx context.Context,
552564
config *internal.Config,
553565
initial bool,
554566
targetConn,
@@ -573,17 +585,17 @@ func generateMigrationStatements(
573585
return "", ErrInvalidModel
574586
}
575587

576-
err = internal.CreateUsers(migrateConn, config.DatabaseUsers)
588+
err = internal.CreateUsers(ctx, migrateConn, config.DatabaseUsers)
577589
if err != nil {
578590
return "", fmt.Errorf("failed to setup migrate database: %w", err)
579591
}
580592

581-
err = internal.CreateUsers(targetConn, config.DatabaseUsers)
593+
err = internal.CreateUsers(ctx, targetConn, config.DatabaseUsers)
582594
if err != nil {
583595
return "", fmt.Errorf("failed to setup target database: %w", err)
584596
}
585597

586-
err = executeTargetSQL(targetConn, config)
598+
err = executeTargetSQL(ctx, targetConn, config)
587599
if err != nil {
588600
log.Println(err)
589601

@@ -661,7 +673,7 @@ func executeMigrateSQL(migrateConn *pgx.Conn) error {
661673
return nil
662674
}
663675

664-
func executeTargetSQL(targetConn *pgx.Conn, config *internal.Config) error {
676+
func executeTargetSQL(ctx context.Context, targetConn *pgx.Conn, config *internal.Config) error {
665677
wd, err := os.Getwd()
666678
if err != nil {
667679
return fmt.Errorf("failed to get working directory: %w", err)
@@ -672,7 +684,7 @@ func executeTargetSQL(targetConn *pgx.Conn, config *internal.Config) error {
672684
return fmt.Errorf("failed to read target sql: %w", err)
673685
}
674686

675-
_, err = targetConn.Exec(context.Background(), string(targetSQL))
687+
_, err = targetConn.Exec(ctx, string(targetSQL))
676688
if err != nil {
677689
return fmt.Errorf("failed to execute target sql: %w", err)
678690
}

‎cmd/init.go‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package cmd
22

33
import (
4+
"context"
45
"errors"
56
"fmt"
67
"log"
@@ -36,6 +37,8 @@ func NewInitCommand() *cobra.Command {
3637
return nil
3738
},
3839
Run: func(cmd *cobra.Command, args []string) {
40+
ctx := context.Background()
41+
3942
var err error
4043

4144
if version == "" {
@@ -166,7 +169,7 @@ func NewInitCommand() *cobra.Command {
166169
log.Fatalf("Failed to get working directory: %v\n", wdErr)
167170
}
168171

169-
err = runWithFile(config, filepath.Join(wd, "migrations", "001_init.up.sql"), 1)
172+
err = runWithFile(ctx, config, filepath.Join(wd, "migrations", "001_init.up.sql"), 1)
170173
if err != nil {
171174
log.Fatalln(err)
172175
}

‎example/docker-compose.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ services:
2222
TREK_POSTGRES_PASSWORD: postgres
2323
TREK_POSTGRES_DATABASE: postgres
2424
TREK_POSTGRES_SSLMODE: disable
25-
TREK_RESET_DATABASE: "true"
25+
TREK_RESET_DATABASE: "false"
2626
TREK_INSERT_TEST_DATA: "true"
2727
volumes:
2828
- ./:/data

‎internal/embed/docker-compose.yaml.tmpl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ services:
2222
TREK_POSTGRES_PASSWORD: postgres
2323
TREK_POSTGRES_DATABASE: postgres
2424
TREK_POSTGRES_SSLMODE: disable
25-
TREK_RESET_DATABASE: "true"
25+
TREK_RESET_DATABASE: "false"
2626
TREK_INSERT_TEST_DATA: "true"
2727
volumes:
2828
- ./:/data

0 commit comments

Comments
 (0)