chore: maybe init workspace from server (#686)

Author: adityachoudhari26

apps/workspace-engine/pkg/db/workspaces.go

@@ -33,3 +33,22 @@ func GetWorkspaceIDs(ctx context.Context) ([]string, error) {
 	}
 	return workspaceIDs, nil
 }
+
+const WORKSPACE_EXISTS_QUERY = `
+	SELECT EXISTS(SELECT 1 FROM workspace WHERE id = $1)
+`
+
+func WorkspaceExists(ctx context.Context, workspaceId string) (bool, error) {
+	db, err := GetDB(ctx)
+	if err != nil {
+		return false, err
+	}
+	defer db.Release()
+
+	var exists bool
+	err = db.QueryRow(ctx, WORKSPACE_EXISTS_QUERY, workspaceId).Scan(&exists)
+	if err != nil {
+		return false, err
+	}
+	return exists, nil
+}

apps/workspace-engine/pkg/db/workspaces_test.go

@@ -2,8 +2,34 @@ package db
 
 import (
 	"testing"
+
+	"github.com/google/uuid"
 )
 
+func TestDBWorkspaces_WorkspaceExists(t *testing.T) {
+	workspaceID, conn := setupTestWithWorkspace(t)
+	_ = conn // Avoid "declared but not used" error
+
+	// Test that the workspace exists
+	exists, err := WorkspaceExists(t.Context(), workspaceID)
+	if err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+	if !exists {
+		t.Fatalf("expected workspace %s to exist", workspaceID)
+	}
+
+	// Test that a non-existent workspace returns false
+	nonExistentID := uuid.New().String()
+	exists, err = WorkspaceExists(t.Context(), nonExistentID)
+	if err != nil {
+		t.Fatalf("expected no error for non-existent workspace, got %v", err)
+	}
+	if exists {
+		t.Fatalf("expected workspace %s to not exist, but it does", nonExistentID)
+	}
+}
+
 func TestDBWorkspaces_GetWorkspaceIDs(t *testing.T) {
 	// Create multiple workspaces
 	workspaceID1, conn1 := setupTestWithWorkspace(t)

apps/workspace-engine/pkg/events/handler/handler.go

@@ -143,7 +143,7 @@ func (el *EventListener) ListenAndRoute(ctx context.Context, msg *kafka.Message)
 	changeSet := changeset.NewChangeSet[any]()
 	if !wsExists {
 		ws = workspace.New(rawEvent.WorkspaceID)
-		if err := loadWorkspaceWithInitialState(ctx, ws); err != nil {
+		if err := workspace.PopulateWorkspaceWithInitialState(ctx, ws); err != nil {
 			span.RecordError(err)
 			span.SetStatus(codes.Error, "failed to load workspace")
 			log.Error("Failed to load workspace", "error", err, "workspaceID", rawEvent.WorkspaceID)
@@ -154,7 +154,14 @@ func (el *EventListener) ListenAndRoute(ctx context.Context, msg *kafka.Message)
 	}
 	ctx = changeset.WithChangeSet(ctx, changeSet)
 
-	err := handler(ctx, ws, rawEvent)
+	if err := handler(ctx, ws, rawEvent); err != nil {
+		span.RecordError(err)
+		span.SetStatus(codes.Error, "handler failed")
+		log.Error("Handler failed to process event",
+			"eventType", rawEvent.EventType,
+			"error", err)
+		return fmt.Errorf("handler failed to process event %s: %w", rawEvent.EventType, err)
+	}
 
 	releaseTargetChanges, err := ws.ReleaseManager().ProcessChanges(ctx, changeSet)
 	if err != nil {
@@ -173,15 +180,6 @@ func (el *EventListener) ListenAndRoute(ctx context.Context, msg *kafka.Message)
 		return fmt.Errorf("failed to flush changeset: %w", err)
 	}
 
-	if err != nil {
-		span.RecordError(err)
-		span.SetStatus(codes.Error, "handler failed")
-		log.Error("Handler failed to process event",
-			"eventType", rawEvent.EventType,
-			"error", err)
-		return fmt.Errorf("handler failed to process event %s: %w", rawEvent.EventType, err)
-	}
-
 	span.SetStatus(codes.Ok, "event processed successfully")
 	log.Debug("Successfully processed event",
 		"eventType", rawEvent.EventType)

apps/workspace-engine/pkg/server/openapi/deployments/server.go

@@ -11,8 +11,11 @@ import (
 type Deployments struct{}
 
 func (s *Deployments) GetDeploymentResources(c *gin.Context, workspaceId string, deploymentId string) {
-	ws := utils.GetWorkspace(c, workspaceId)
-	if ws == nil {
+	ws, err := utils.GetWorkspace(c, workspaceId)
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{
+			"error": err.Error(),
+		})
 		return
 	}
 

apps/workspace-engine/pkg/server/openapi/environments/server.go

@@ -12,8 +12,11 @@ import (
 type Environments struct{}
 
 func (s *Environments) GetEnvironmentResources(c *gin.Context, workspaceId string, environmentId string) {
-	ws := utils.GetWorkspace(c, workspaceId)
-	if ws == nil {
+	ws, err := utils.GetWorkspace(c, workspaceId)
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{
+			"error": err.Error(),
+		})
 		return
 	}
 

apps/workspace-engine/pkg/server/openapi/policies/policies.go

@@ -12,10 +12,10 @@ import (
 type Policies struct{}
 
 func (p *Policies) GetReleaseTargetsForPolicy(c *gin.Context, workspaceId oapi.WorkspaceId, policyId oapi.PolicyId) {
-	ws := utils.GetWorkspace(c, workspaceId)
-	if ws == nil {
-		c.JSON(http.StatusNotFound, gin.H{
-			"error": "Workspace not found",
+	ws, err := utils.GetWorkspace(c, workspaceId)
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{
+			"error": "Failed to get workspace: " + err.Error(),
 		})
 		return
 	}

apps/workspace-engine/pkg/server/openapi/releasetargets/releasetargets.go

@@ -15,10 +15,10 @@ type ReleaseTargets struct {
 // EvaluateReleaseTarget implements oapi.ServerInterface.
 func (s *ReleaseTargets) EvaluateReleaseTarget(c *gin.Context, workspaceId oapi.WorkspaceId) {
 	// Get workspace
-	ws := utils.GetWorkspace(c, workspaceId)
-	if ws == nil {
-		c.JSON(http.StatusNotFound, gin.H{
-			"error": "Workspace not found",
+	ws, err := utils.GetWorkspace(c, workspaceId)
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{
+			"error": err.Error(),
 		})
 		return
 	}
@@ -69,10 +69,10 @@ func (s *ReleaseTargets) EvaluateReleaseTarget(c *gin.Context, workspaceId oapi.
 
 // GetPoliciesForReleaseTarget implements oapi.ServerInterface.
 func (s *ReleaseTargets) GetPoliciesForReleaseTarget(c *gin.Context, workspaceId oapi.WorkspaceId, releaseTargetId oapi.ReleaseTargetId) {
-	ws := utils.GetWorkspace(c, workspaceId)
-	if ws == nil {
-		c.JSON(http.StatusNotFound, gin.H{
-			"error": "Workspace not found",
+	ws, err := utils.GetWorkspace(c, workspaceId)
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{
+			"error": "Failed to get workspace: " + err.Error(),
 		})
 		return
 	}

apps/workspace-engine/pkg/server/openapi/utils/utils.go

@@ -1,20 +1,30 @@
 package utils
 
 import (
-	"net/http"
+	"fmt"
+	"workspace-engine/pkg/db"
 	"workspace-engine/pkg/workspace"
 
 	"github.com/gin-gonic/gin"
 )
 
-func GetWorkspace(c *gin.Context, workspaceId string) *workspace.Workspace {
-	ok := workspace.HasWorkspace(workspaceId)
-	if !ok {
-		c.JSON(http.StatusNotFound, gin.H{
-			"error": "Workspace not found",
-		})
-		return nil
+func GetWorkspace(c *gin.Context, workspaceId string) (*workspace.Workspace, error) {
+	existsInDB, err := db.WorkspaceExists(c.Request.Context(), workspaceId)
+	if err != nil {
+		return nil, err
+	}
+	if !existsInDB {
+		return nil, fmt.Errorf("workspace %s not found in database", workspaceId)
+	}
+
+	if workspace.Exists(workspaceId) {
+		return workspace.GetWorkspace(workspaceId), nil
 	}
 
-	return workspace.GetWorkspace(workspaceId)
+	ws := workspace.New(workspaceId)
+	if err := workspace.PopulateWorkspaceWithInitialState(c.Request.Context(), ws); err != nil {
+		return nil, fmt.Errorf("failed to populate workspace with initial state: %w", err)
+	}
+	workspace.Set(workspaceId, ws)
+	return ws, nil
 }

apps/workspace-engine/pkg/events/handler/workspace.go renamed to apps/workspace-engine/pkg/workspace/populate_workspace.go

@@ -1,12 +1,11 @@
-package handler
+package workspace
 
 import (
 	"context"
 	"workspace-engine/pkg/db"
-	"workspace-engine/pkg/workspace"
 )
 
-func loadWorkspaceWithInitialState(ctx context.Context, ws *workspace.Workspace) error {
+func PopulateWorkspaceWithInitialState(ctx context.Context, ws *Workspace) error {
 	initialWorkspaceState, err := db.LoadWorkspace(ctx, ws.ID)
 	if err != nil {
 		return err