Haystack/backend/router.go

72 lines
1.7 KiB
Go

package main
import (
"database/sql"
"os"
"screenmark/screenmark/agents/client"
"screenmark/screenmark/auth"
"screenmark/screenmark/images"
"screenmark/screenmark/models"
"screenmark/screenmark/stacks"
ourmiddleware "screenmark/screenmark/middleware"
"github.com/go-chi/chi/v5"
"github.com/go-chi/chi/v5/middleware"
)
type TestAiClient struct {
ImageInfo client.ImageMessageContent
}
func (client TestAiClient) GetImageInfo(imageName string, imageData []byte) (client.ImageMessageContent, error) {
return client.ImageInfo, nil
}
func setupRouter(db *sql.DB) chi.Router {
imageModel := models.NewImageModel(db)
stackModel := models.NewListModel(db)
stackHandler := stacks.CreateStackHandler(db)
authHandler := auth.CreateAuthHandler(db)
imageHandler := images.CreateImageHandler(db)
notifier := NewNotifier[Notification](10)
// Only start event listeners if not in test environment
if os.Getenv("GO_TEST_ENVIRONMENT") != "true" {
// TODO: should extract these into a notification manager
// And actually make them the same code.
// The events are basically the same.
go ListenNewImageEvents(db)
go ListenProcessingImageStatus(db, imageModel, &notifier)
go ListenNewStackEvents(db)
go ListenProcessingStackStatus(db, stackModel, &notifier)
}
r := chi.NewRouter()
r.Use(middleware.Logger)
r.Use(ourmiddleware.CorsMiddleware)
r.Route("/stacks", stackHandler.CreateRoutes)
r.Route("/auth", authHandler.CreateRoutes)
r.Route("/images", imageHandler.CreateRoutes)
r.Route("/notifications", func(r chi.Router) {
r.Use(ourmiddleware.GetUserIdFromUrl)
r.Get("/", CreateEventsHandler(&notifier))
})
logWriter := DatabaseWriter{
dbPool: db,
}
r.Route("/logs", createLogHandler(&logWriter))
return r
}