refactor: moving image listener to own function

backend/events.go

@@ -0,0 +1,79 @@
+package main
+
+import (
+	"context"
+	"database/sql"
+	"log"
+	"os"
+	"screenmark/screenmark/agents"
+	"screenmark/screenmark/models"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/lib/pq"
+)
+
+func ListenNewImageEvents(db *sql.DB) {
+	listener := pq.NewListener(os.Getenv("DB_CONNECTION"), time.Second, time.Second, func(event pq.ListenerEventType, err error) {
+		if err != nil {
+			panic(err)
+		}
+	})
+	defer listener.Close()
+
+	locationModel := models.NewLocationModel(db)
+	eventModel := models.NewEventModel(db)
+	noteModel := models.NewNoteModel(db)
+	imageModel := models.NewImageModel(db)
+	contactModel := models.NewContactModel(db)
+
+	err := listener.Listen("new_image")
+	if err != nil {
+		panic(err)
+	}
+
+	for {
+		select {
+		case parameters := <-listener.Notify:
+			imageId := uuid.MustParse(parameters.Extra)
+
+			ctx := context.Background()
+
+			go func() {
+				locationAgent, err := agents.NewLocationEventAgent(locationModel, eventModel, contactModel)
+				if err != nil {
+					panic(err)
+				}
+
+				noteAgent, err := agents.NewNoteAgent(noteModel)
+				if err != nil {
+					panic(err)
+				}
+
+				image, err := imageModel.GetToProcessWithData(ctx, imageId)
+				if err != nil {
+					log.Println("Failed to GetToProcessWithData")
+					log.Println(err)
+					return
+				}
+
+				_, err = imageModel.FinishProcessing(ctx, image.ID)
+				if err != nil {
+					log.Println("Failed to FinishProcessing")
+					log.Println(err)
+					return
+				}
+
+				orchestrator, err := agents.NewOrchestratorAgent(locationAgent, noteAgent, image.Image.ImageName, image.Image.Image)
+				if err != nil {
+					panic(err)
+				}
+
+				err = orchestrator.Orchestrate(image.UserID, image.ImageID, image.Image.ImageName, image.Image.Image)
+				if err != nil {
+					log.Println(err)
+				}
+			}()
+		}
+	}
+}

backend/main.go

@@ -2,26 +2,21 @@ package main
 
 import (
 	"bytes"
-	"context"
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
 	"io"
 	"log"
 	"net/http"
-	"os"
 	"path/filepath"
 	"screenmark/screenmark/.gen/haystack/haystack/model"
-	"screenmark/screenmark/agents"
 	"screenmark/screenmark/agents/client"
 	"screenmark/screenmark/models"
-	"time"
 
 	"github.com/go-chi/chi/v5"
 	"github.com/go-chi/chi/v5/middleware"
 	"github.com/google/uuid"
 	"github.com/joho/godotenv"
-	"github.com/lib/pq"
 )
 
 type TestAiClient struct {
@@ -38,84 +33,20 @@ func main() {
 		panic(err)
 	}
 
-	mode := os.Getenv("MODE")
-	log.Printf("Mode: %s\n", mode)
-
 	db, err := models.InitDatabase()
 	if err != nil {
 		panic(err)
 	}
 
 	imageModel := models.NewImageModel(db)
-	locationModel := models.NewLocationModel(db)
-	eventModel := models.NewEventModel(db)
 	userModel := models.NewUserModel(db)
-	contactModel := models.NewContactModel(db)
-	noteModel := models.NewNoteModel(db)
 
-	listener := pq.NewListener(os.Getenv("DB_CONNECTION"), time.Second, time.Second, func(event pq.ListenerEventType, err error) {
+	go ListenNewImageEvents(db)
-		if err != nil {
-			panic(err)
-		}
-	})
-	defer listener.Close()
-
-	go func() {
-		err := listener.Listen("new_image")
-		if err != nil {
-			panic(err)
-		}
-
-		for {
-
-			select {
-			case parameters := <-listener.Notify:
-				imageId := uuid.MustParse(parameters.Extra)
-
-				ctx := context.Background()
-
-				go func() {
-					locationAgent, err := agents.NewLocationEventAgent(locationModel, eventModel, contactModel)
-					if err != nil {
-						panic(err)
-					}
-
-					noteAgent, err := agents.NewNoteAgent(noteModel)
-					if err != nil {
-						panic(err)
-					}
-
-					image, err := imageModel.GetToProcessWithData(ctx, imageId)
-					if err != nil {
-						log.Println("Failed to GetToProcessWithData")
-						log.Println(err)
-						return
-					}
-
-					_, err = imageModel.FinishProcessing(ctx, image.ID)
-					if err != nil {
-						log.Println("Failed to FinishProcessing")
-						log.Println(err)
-						return
-					}
-
-					orchestrator, err := agents.NewOrchestratorAgent(locationAgent, noteAgent, image.Image.ImageName, image.Image.Image)
-					if err != nil {
-						panic(err)
-					}
-
-					err = orchestrator.Orchestrate(image.UserID, image.ImageID, image.Image.ImageName, image.Image.Image)
-					if err != nil {
-						fmt.Println(err)
-					}
-				}()
-			}
-		}
-	}()
 
 	r := chi.NewRouter()
 
 	r.Use(middleware.Logger)
+	r.Use(CorsMiddleware)
 	r.Use(func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			w.Header().Add("Content-Type", "application/json")
@@ -124,19 +55,9 @@ func main() {
 		})
 	})
 
-	r.Options("/*", func(w http.ResponseWriter, r *http.Request) {
-		w.Header().Add("Access-Control-Allow-Origin", "*")
-		w.Header().Add("Access-Control-Allow-Credentials", "*")
-		w.Header().Add("Access-Control-Allow-Headers", "*")
-	})
-
 	r.Get("/image", func(w http.ResponseWriter, r *http.Request) {
 		userId := r.Header.Get("userId")
 
-		w.Header().Add("Access-Control-Allow-Origin", "*")
-		w.Header().Add("Access-Control-Allow-Credentials", "*")
-		w.Header().Add("Access-Control-Allow-Headers", "*")
-
 		images, err := userModel.ListWithProperties(r.Context(), uuid.MustParse(userId))
 		if err != nil {
 			log.Println(err)
@@ -190,10 +111,6 @@ func main() {
 	r.Get("/image/{id}", func(w http.ResponseWriter, r *http.Request) {
 		imageId := r.PathValue("id")
 
-		w.Header().Add("Access-Control-Allow-Origin", "*")
-		w.Header().Add("Access-Control-Allow-Credentials", "*")
-		w.Header().Add("Access-Control-Allow-Headers", "*")
-
 		// TODO: really need authorization here!
 		image, err := imageModel.Get(r.Context(), uuid.MustParse(imageId))
 		if err != nil {

backend/middleware.go

@@ -0,0 +1,13 @@
+package main
+
+import "net/http"
+
+func CorsMiddleware(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Add("Access-Control-Allow-Origin", "*")
+		w.Header().Add("Access-Control-Allow-Credentials", "*")
+		w.Header().Add("Access-Control-Allow-Headers", "*")
+
+		next.ServeHTTP(w, r)
+	})
+}