From f6d0dae6b66b621b4334880467823525e62df6b2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dawid=20Wysoki=C5=84ski?= <contact@dwysokinski.me>
Date: Sat, 16 Dec 2023 10:37:26 +0100
Subject: [PATCH] feat: db migrate/rollback - react to shutdown signals

---
 cmd/twhelp/cmd_db.go  | 32 +++++++++++++++++++++++++-------
 cmd/twhelp/cmd_job.go |  1 +
 cmd/twhelp/utils.go   | 15 +++++++++++++++
 3 files changed, 41 insertions(+), 7 deletions(-)

diff --git a/cmd/twhelp/cmd_db.go b/cmd/twhelp/cmd_db.go
index d34e52a..d0f9f1c 100644
--- a/cmd/twhelp/cmd_db.go
+++ b/cmd/twhelp/cmd_db.go
@@ -39,6 +39,9 @@ var cmdDB = &cli.Command{
 			Usage: "Migrate database",
 			Flags: dbFlags,
 			Action: func(c *cli.Context) error {
+				shutdownSignalCtx, stop := newShutdownSignalContext(c.Context)
+				defer stop()
+
 				logger := loggerFromCtx(c.Context)
 
 				bunDB, err := newBunDBFromFlags(c)
@@ -46,29 +49,35 @@ var cmdDB = &cli.Command{
 					return err
 				}
 				defer func() {
-					logger.Debug("closing db connections...")
+					logger.Debug("closing db connections...", slog.Int("db.openConnections", bunDB.Stats().OpenConnections))
 					if dbCloseErr := bunDB.Close(); dbCloseErr != nil {
 						logger.Warn("couldn't close db connections", slog.Any("error", dbCloseErr))
+					} else {
+						logger.Debug("db connections closed")
 					}
 				}()
 
 				migrator := migrations.NewMigrator(bunDB)
 
-				if err = migrator.Init(c.Context); err != nil {
+				if err = migrator.Init(shutdownSignalCtx); err != nil {
 					return fmt.Errorf("couldn't init migrator: %w", err)
 				}
 
-				if err = migrator.Lock(c.Context); err != nil {
+				if err = migrator.Lock(shutdownSignalCtx); err != nil {
 					return fmt.Errorf("couldn't lock db: %w", err)
 				}
 				defer func() {
 					logger.Debug("unlocking db...")
+					// we use c.Context instead of shutdownSignalCtx here because we still want to unlock the db
+					// even after one of the shutdown signals has been received
 					if unlockErr := migrator.Unlock(c.Context); unlockErr != nil {
 						logger.Warn("couldn't unlock db", slog.Any("error", unlockErr))
+					} else {
+						logger.Debug("db unlocked")
 					}
 				}()
 
-				group, err := migrator.Migrate(c.Context)
+				group, err := migrator.Migrate(shutdownSignalCtx)
 				if err != nil {
 					return fmt.Errorf("migration failed: %w", err)
 				}
@@ -92,6 +101,9 @@ var cmdDB = &cli.Command{
 			Usage: "Rollback the last migration group",
 			Flags: dbFlags,
 			Action: func(c *cli.Context) error {
+				shutdownSignalCtx, stop := newShutdownSignalContext(c.Context)
+				defer stop()
+
 				logger := loggerFromCtx(c.Context)
 
 				bunDB, err := newBunDBFromFlags(c)
@@ -99,25 +111,31 @@ var cmdDB = &cli.Command{
 					return err
 				}
 				defer func() {
-					logger.Debug("closing db connections...")
+					logger.Debug("closing db connections...", slog.Int("db.openConnections", bunDB.Stats().OpenConnections))
 					if dbCloseErr := bunDB.Close(); dbCloseErr != nil {
 						logger.Warn("couldn't close db connections", slog.Any("error", dbCloseErr))
+					} else {
+						logger.Debug("db connections closed")
 					}
 				}()
 
 				migrator := migrations.NewMigrator(bunDB)
 
-				if err = migrator.Lock(c.Context); err != nil {
+				if err = migrator.Lock(shutdownSignalCtx); err != nil {
 					return fmt.Errorf("couldn't lock db: %w", err)
 				}
 				defer func() {
 					logger.Debug("unlocking db...")
+					// we use c.Context instead of shutdownSignalCtx here because we still want to unlock the db
+					// even after one of the shutdown signals has been received
 					if unlockErr := migrator.Unlock(c.Context); unlockErr != nil {
 						logger.Warn("couldn't unlock db", slog.Any("error", unlockErr))
+					} else {
+						logger.Debug("db unlocked")
 					}
 				}()
 
-				group, err := migrator.Rollback(c.Context)
+				group, err := migrator.Rollback(shutdownSignalCtx)
 				if err != nil {
 					return fmt.Errorf("couldn't rollback last migration: %w", err)
 				}
diff --git a/cmd/twhelp/cmd_job.go b/cmd/twhelp/cmd_job.go
index 5c22789..8668240 100644
--- a/cmd/twhelp/cmd_job.go
+++ b/cmd/twhelp/cmd_job.go
@@ -28,6 +28,7 @@ var (
 								return err
 							}
 							defer func() {
+								// watermill has its own logs, so we don't need to add any debug logs here
 								if closeErr := amqpConn.Close(); closeErr != nil {
 									logger.Warn("couldn't close amqp connection", slog.Any("error", err))
 								}
diff --git a/cmd/twhelp/utils.go b/cmd/twhelp/utils.go
index 2c92173..c0094ed 100644
--- a/cmd/twhelp/utils.go
+++ b/cmd/twhelp/utils.go
@@ -1,5 +1,12 @@
 package main
 
+import (
+	"context"
+	"os"
+	"os/signal"
+	"syscall"
+)
+
 func concatSlices[T any](slices ...[]T) []T {
 	var totalLen int
 
@@ -17,3 +24,11 @@ func concatSlices[T any](slices ...[]T) []T {
 
 	return result
 }
+
+var shutdownSignals = []os.Signal{os.Interrupt, syscall.SIGTERM}
+
+// newShutdownSignalContext returns a copy of the parent context that is marked done
+// (its Done channel is closed) when one of the shutdownSignals arrives.
+func newShutdownSignalContext(parent context.Context) (context.Context, context.CancelFunc) {
+	return signal.NotifyContext(parent, shutdownSignals...)
+}