alpha

cmd/root.go

@@ -1,13 +1,18 @@
 package cmd
 
 import (
+	"context"
+	"errors"
 	"fmt"
 	"log"
 	"net/http"
 	"os"
+	"os/signal"
 	"otpm/database"
 	"otpm/handlers"
 	"otpm/utils"
+	"syscall"
+	"time"
 
 	"github.com/jmoiron/sqlx"
 	"github.com/julienschmidt/httprouter"
@@ -57,20 +62,50 @@ func initConfig() {
 	}
 }
 
-func initApp(db *sqlx.DB) {
-	if err := database.MigrateDB(db); err != nil {
-		log.Fatalf("Error migrating the database: %v", err)
-	}
+type App struct {
+	db     *sqlx.DB
+	router http.Handler
+	port   int
 }
 
-func startApp() {
+func NewApp() (*App, error) {
+	db, err := connectDB()
+	if err != nil {
+		return nil, err
+	}
+
+	if err := runMigrations(db); err != nil {
+		return nil, err
+	}
+
+	router := setupRouter(db)
+
 	port := viper.GetInt("port")
+
+	return &App{
+		db:     db,
+		router: router,
+		port:   port,
+	}, nil
+}
+
+func connectDB() (*sqlx.DB, error) {
 	db, err := database.InitDB()
 	if err != nil {
-		log.Fatalf("Error connecting to the database: %v", err)
+		return nil, fmt.Errorf("failed to connect to the database: %v", err)
 	}
-	defer db.Close()
-	initApp(db)
+	return db, nil
+}
+
+func runMigrations(db *sqlx.DB) error {
+	if err := database.MigrateDB(db); err != nil {
+		log.Fatalf("Error migrating the database: %v", err)
+		return fmt.Errorf("error migrating the database: %w", err)
+	}
+	return nil
+}
+
+func setupRouter(db *sqlx.DB) http.Handler {
 	handler := &handlers.Handler{DB: db}
 
 	router := httprouter.New()
@@ -78,6 +113,39 @@ func startApp() {
 	router.POST("/set", utils.AdaptHandler(handler.UpdateOrCreateOtp))
 	router.GET("/get", utils.AdaptHandler(handler.GetOtp))
 
-	log.Println("Starting server on :8080")
-	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", port), router))
+	return router
+}
+
+func (a *App) Start() error {
+	server := &http.Server{Addr: fmt.Sprintf(":%d", a.port), Handler: a.router}
+	go func() {
+		if err := server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
+			log.Fatalf("Failed to start server: %v", err)
+		}
+	}()
+
+	quit := make(chan os.Signal, 1)
+	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
+	<-quit
+	log.Println("Shutting down server...")
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	return server.Shutdown(ctx)
+}
+
+func startApp() {
+	app, err := NewApp()
+	if err != nil {
+		log.Fatalf("Failed to initialize application: %v", err)
+	}
+	defer func() {
+		if err := app.db.Close(); err != nil {
+			log.Printf("Failed to close database connection: %v", err)
+		}
+	}()
+	if err := app.Start(); err != nil {
+		log.Fatalf("Failed to start application: %v", err)
+	}
 }