Refactor database shutdown procedures in server, ingester, and producer modules

- Removed direct dependencies on bun.DB and replaced with a unified GracefulDatabaseShutdown function to handle various database types.
- Enhanced shutdown logic to support both SQL and MongoDB, ensuring proper disconnection and resource cleanup.
- Updated logging for clearer shutdown status across modules.

Author: 0xfurai

apps/server/cmd/api/main.go

@@ -41,7 +41,6 @@ import (
 	"peekaping/internal/version"
 	"syscall"
 
-	"github.com/uptrace/bun"
 	"go.uber.org/dig"
 	"go.uber.org/zap"
 )
@@ -186,7 +185,6 @@ func main() {
 	// Start the server with graceful shutdown
 	err = container.Invoke(func(
 		server *internal.Server,
-		db *bun.DB,
 		eventBus events.EventBus,
 		logger *zap.SugaredLogger,
 	) error {
@@ -220,9 +218,9 @@ func main() {
 			logger.Errorw("Failed to close event bus", "error", err)
 		}
 
-		// Perform graceful SQLite shutdown if using SQLite
-		if err := infra.GracefulSQLiteShutdown(db, internalCfg.DBType, logger); err != nil {
-			logger.Errorw("Failed to gracefully shutdown database", "error", err)
+		// Perform graceful database shutdown
+		if err := infra.GracefulDatabaseShutdown(container, internalCfg, logger); err != nil {
+			logger.Errorw("Failed to shutdown database", "error", err)
 		}
 
 		logger.Info("Server stopped gracefully")

apps/server/cmd/ingester/main.go

@@ -22,7 +22,6 @@ import (
 	"syscall"
 
 	"github.com/hibiken/asynq"
-	"github.com/uptrace/bun"
 	"go.uber.org/dig"
 	"go.uber.org/zap"
 )
@@ -118,7 +117,6 @@ func main() {
 	// Start the ingester
 	err = container.Invoke(func(
 		ing *ingester.Ingester,
-		db *bun.DB,
 		eventBus events.EventBus,
 		logger *zap.SugaredLogger,
 	) error {
@@ -145,9 +143,9 @@ func main() {
 			logger.Errorw("Failed to close event bus", "error", err)
 		}
 
-		// Perform graceful SQLite shutdown if using SQLite
-		if err := infra.GracefulSQLiteShutdown(db, internalCfg.DBType, logger); err != nil {
-			logger.Errorw("Failed to gracefully shutdown database", "error", err)
+		// Perform graceful database shutdown
+		if err := infra.GracefulDatabaseShutdown(container, internalCfg, logger); err != nil {
+			logger.Errorw("Failed to shutdown database", "error", err)
 		}
 
 		logger.Info("Ingester stopped gracefully")

apps/server/cmd/producer/main.go

@@ -27,7 +27,6 @@ import (
 	"peekaping/internal/version"
 	"syscall"
 
-	"github.com/uptrace/bun"
 	"go.uber.org/dig"
 	"go.uber.org/zap"
 )
@@ -91,7 +90,6 @@ func main() {
 	err = container.Invoke(func(
 		prod *producer.Producer,
 		eventListener *producer.EventListener,
-		db *bun.DB,
 		eventBus events.EventBus,
 		logger *zap.SugaredLogger,
 	) error {
@@ -120,9 +118,9 @@ func main() {
 			logger.Errorw("Failed to close event bus", "error", err)
 		}
 
-		// Perform graceful SQLite shutdown if using SQLite
-		if err := infra.GracefulSQLiteShutdown(db, internalCfg.DBType, logger); err != nil {
-			logger.Errorw("Failed to gracefully shutdown database", "error", err)
+		// Perform graceful database shutdown
+		if err := infra.GracefulDatabaseShutdown(container, internalCfg, logger); err != nil {
+			logger.Errorw("Failed to shutdown database", "error", err)
 		}
 
 		logger.Info("Producer stopped gracefully")

apps/server/internal/infra/sql.go

@@ -1,6 +1,7 @@
 package infra
 
 import (
+	"context"
 	"database/sql"
 	"fmt"
 	"peekaping/internal/config"
@@ -12,6 +13,8 @@ import (
 	"github.com/uptrace/bun/driver/pgdriver"
 	"github.com/uptrace/bun/driver/sqliteshim"
 	"github.com/uptrace/bun/extra/bundebug"
+	"go.mongodb.org/mongo-driver/mongo"
+	"go.uber.org/dig"
 	"go.uber.org/zap"
 
 	_ "github.com/go-sql-driver/mysql"
@@ -203,3 +206,29 @@ func GracefulSQLiteShutdown(db *bun.DB, dbType string, logger *zap.SugaredLogger
 	logger.Info("SQLite database closed gracefully")
 	return nil
 }
+
+// GracefulDatabaseShutdown performs graceful shutdown for any database type
+// It uses the DI container to get the appropriate database connection
+func GracefulDatabaseShutdown(container *dig.Container, cfg *config.Config, logger *zap.SugaredLogger) error {
+	switch cfg.DBType {
+	case "postgres", "postgresql", "mysql", "sqlite":
+		// For SQL databases, perform graceful shutdown
+		return container.Invoke(func(db *bun.DB) {
+			if err := GracefulSQLiteShutdown(db, cfg.DBType, logger); err != nil {
+				logger.Errorw("Failed to gracefully shutdown database", "error", err)
+			}
+		})
+	case "mongo", "mongodb":
+		// For MongoDB, disconnect the client
+		return container.Invoke(func(client *mongo.Client) {
+			if err := client.Disconnect(context.Background()); err != nil {
+				logger.Errorw("Failed to disconnect MongoDB client", "error", err)
+			} else {
+				logger.Info("MongoDB client disconnected gracefully")
+			}
+		})
+	default:
+		logger.Warnf("Unknown database type for shutdown: %s", cfg.DBType)
+		return nil
+	}
+}