From c1ddc47c95e04c242263ee9c5074cac9a2422ca2 Mon Sep 17 00:00:00 2001
From: Julien Dessaux
Date: Tue, 17 Dec 2024 23:18:34 +0100
Subject: chore(tfstated): refactor database environment handling

---
 cmd/tfstated/main.go      | 19 +------------------
 cmd/tfstated/main_test.go | 31 +++++++++++++++----------------
 2 files changed, 16 insertions(+), 34 deletions(-)

(limited to 'cmd')

diff --git a/cmd/tfstated/main.go b/cmd/tfstated/main.go
index e6c6272..a1721e1 100644
--- a/cmd/tfstated/main.go
+++ b/cmd/tfstated/main.go
@@ -10,7 +10,6 @@ import (
 	"net/http"
 	"os"
 	"os/signal"
-	"strconv"
 	"sync"
 	"time"
 
@@ -36,22 +35,6 @@ func run(
 	ctx, cancel := signal.NotifyContext(ctx, os.Interrupt)
 	defer cancel()
 
-	dataEncryptionKey := getenv("DATA_ENCRYPTION_KEY")
-	if dataEncryptionKey == "" {
-		return fmt.Errorf("the DATA_ENCRYPTION_KEY environment variable is not set")
-	}
-	if err := db.SetDataEncryptionKey(dataEncryptionKey); err != nil {
-		return err
-	}
-	versionsHistoryLimit := getenv("VERSIONS_HISTORY_LIMIT")
-	if versionsHistoryLimit != "" {
-		n, err := strconv.Atoi(versionsHistoryLimit)
-		if err != nil {
-			return fmt.Errorf("failed to parse the VERSIONS_HISTORY_LIMIT environment variable: %w", err)
-		}
-		db.SetVersionsHistoryLimit(n)
-	}
-
 	if err := db.InitAdminAccount(); err != nil {
 		return err
 	}
@@ -107,7 +90,7 @@ func main() {
 		Port: "8080",
 	}
 
-	db, err := database.NewDB(ctx, "./tfstate.db?_txlock=immediate")
+	db, err := database.NewDB(ctx, "./tfstate.db?_txlock=immediate", os.Getenv)
 	if err != nil {
 		fmt.Fprintf(os.Stderr, "database init error: %+v\n", err)
 		os.Exit(1)
diff --git a/cmd/tfstated/main_test.go b/cmd/tfstated/main_test.go
index bd21918..dc19a3f 100644
--- a/cmd/tfstated/main_test.go
+++ b/cmd/tfstated/main_test.go
@@ -24,19 +24,10 @@ var adminPassword string
 var adminPasswordMutex sync.Mutex
 
 func TestMain(m *testing.M) {
-	ctx := context.Background()
-	ctx, cancel := context.WithCancel(ctx)
 	config := Config{
 		Host: "127.0.0.1",
 		Port: "8081",
 	}
-	_ = os.Remove("./test.db")
-	var err error
-	db, err = database.NewDB(ctx, "./test.db")
-	if err != nil {
-		fmt.Fprintf(os.Stderr, "%+v\n", err)
-		os.Exit(1)
-	}
 	getenv := func(key string) string {
 		switch key {
 		case "DATA_ENCRYPTION_KEY":
@@ -48,6 +39,15 @@ func TestMain(m *testing.M) {
 		}
 	}
 
+	ctx := context.Background()
+	ctx, cancel := context.WithCancel(ctx)
+	_ = os.Remove("./test.db")
+	var err error
+	db, err = database.NewDB(ctx, "./test.db", getenv)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "%+v\n", err)
+		os.Exit(1)
+	}
 	database.AdvertiseAdminPassword = func(password string) {
 		adminPasswordMutex.Lock()
 		defer adminPasswordMutex.Unlock()
@@ -125,14 +125,13 @@ func waitForReady(
 		resp, err := client.Do(req)
 		if err != nil {
 			fmt.Printf("Error making request: %s\n", err.Error())
-			continue
-		}
-		if resp.StatusCode == http.StatusOK {
-			fmt.Println("Endpoint is ready!")
-			resp.Body.Close()
-			return nil
+		} else {
+			_ = resp.Body.Close()
+			if resp.StatusCode == http.StatusOK {
+				fmt.Println("Endpoint is ready!")
+				return nil
+			}
 		}
-		resp.Body.Close()
 
 		select {
 		case <-ctx.Done():
-- 
cgit v1.2.3