Refactor DB to remove restrictive query system

1 files changed, 18 insertions, 10 deletions

diff --git a/backend/cmd/main.go b/backend/cmd/main.go
index 6b49980..2a34780 100644
--- a/backend/cmd/main.go
+++ b/backend/cmd/main.go
@@ -54,8 +54,14 @@ func main() {
 		"logger.level", level,
 	)
 
+	// Create kms
+	kms, err := createKMS()
+	if err != nil {
+		log.Fatal("Unable to create KMS: ", err)
+	}
+
 	// Connect to the database
-	db, err := connectDB(logger, cfg)
+	db, err := connectDB(logger, cfg, kms)
 	defer func(db database.Database) {
 		_ = db.Close()
 	}(db)
@@ -78,7 +84,7 @@ func main() {
 	_ = auth
 
 	// Setup IBD client
-	client, err := setupIBDClient(cfg, db)
+	client, err := setupIBDClient(cfg, db, kms)
 	if err != nil {
 		log.Fatal("Unable to setup IBD client: ", err)
 	}
@@ -102,7 +108,7 @@ func main() {
 
 	// Start the gRPC server
 	go func() {
-		s, err := server.New(ctx, cfg.Server.Port, db, redisClient, client)
+		s, err := server.New(ctx, cfg.Server.Port, db, redisClient, client, kms, cfg.KMS.GCP.String())
 		if err != nil {
 			log.Fatal("Unable to create gRPC server: ", err)
 		}
@@ -120,6 +126,7 @@ func main() {
 			client,
 			redisClient,
 			db,
+			kms,
 			analyzer,
 		)
 		if err != nil {
@@ -146,7 +153,7 @@ func main() {
 	)
 }
 
-func setupIBDClient(cfg *config.Config, db database.Database) (*ibd.Client, error) {
+func setupIBDClient(cfg *config.Config, db database.Database, kms keys.KeyManagementService) (*ibd.Client, error) {
 	pUrl, err := url.Parse(cfg.IBD.ProxyURL)
 	if err != nil {
 		return nil, fmt.Errorf("unable to parse proxy URL: %w", err)
@@ -157,18 +164,19 @@ func setupIBDClient(cfg *config.Config, db database.Database) (*ibd.Client, erro
 		transport.NewStandardTransport(&http.Client{Transport: t}), // Default proxied transport
 		scrapfly.New(http.DefaultClient, cfg.IBD.APIKey),           // Scrapfly transport
 	}
-	client := ibd.NewClient(db, transports...)
+	client := ibd.NewClient(db, kms, transports...)
 	return client, nil
 }
 
-func connectDB(logger *slog.Logger, cfg *config.Config) (database.Database, error) {
+func createKMS() (keys.KeyManagementService, error) {
 	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
 	defer cancel()
+	return keys.NewGoogleKMS(ctx)
+}
 
-	kms, err := keys.NewGoogleKMS(ctx)
-	if err != nil {
-		return nil, fmt.Errorf("unable to create Google KMS Client: %w", err)
-	}
+func connectDB(logger *slog.Logger, cfg *config.Config, kms keys.KeyManagementService) (database.Database, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
 
 	db, err := database.New(ctx, logger, cfg.DB.URL, kms, cfg.KMS.GCP.String())
 	if err != nil {