chore: add gh entities on workspace load (#687)

Author: adityachoudhari26

apps/workspace-engine/pkg/workspace/jobdispatch/github.go

@@ -9,7 +9,7 @@ import (
 	"strconv"
 
 	"workspace-engine/pkg/oapi"
-	"workspace-engine/pkg/workspace/store/repository"
+	"workspace-engine/pkg/workspace/store"
 
 	"github.com/bradleyfalzon/ghinstallation/v2"
 	"github.com/google/go-github/v66/github"
@@ -48,21 +48,21 @@ func (r *realGithubClient) DispatchWorkflow(ctx context.Context, owner, repo str
 }
 
 type GithubDispatcher struct {
-	repo          *repository.Repository
+	store         *store.Store
 	clientFactory func(installationID int) (GithubClient, error)
 }
 
-func NewGithubDispatcher(repo *repository.Repository) *GithubDispatcher {
+func NewGithubDispatcher(store *store.Store) *GithubDispatcher {
 	return &GithubDispatcher{
-		repo:          repo,
+		store:         store,
 		clientFactory: nil, // will use default
 	}
 }
 
 // NewGithubDispatcherWithClientFactory creates a dispatcher with a custom client factory (useful for testing)
-func NewGithubDispatcherWithClientFactory(repo *repository.Repository, clientFactory func(installationID int) (GithubClient, error)) *GithubDispatcher {
+func NewGithubDispatcherWithClientFactory(store *store.Store, clientFactory func(installationID int) (GithubClient, error)) *GithubDispatcher {
 	return &GithubDispatcher{
-		repo:          repo,
+		store:         store,
 		clientFactory: clientFactory,
 	}
 }
@@ -91,16 +91,6 @@ func (d *GithubDispatcher) parseConfig(job *oapi.Job) (githubJobConfig, error) {
 	return parsed, nil
 }
 
-func (d *GithubDispatcher) getGithubEntity(cfg githubJobConfig) *oapi.GithubEntity {
-	ghEntities := d.repo.GithubEntities.IterBuffered()
-	for ghEntity := range ghEntities {
-		if ghEntity.Val.InstallationId == cfg.InstallationId && ghEntity.Val.Slug == cfg.Owner {
-			return ghEntity.Val
-		}
-	}
-	return nil
-}
-
 func (d *GithubDispatcher) getEnv(key string) string {
 	return os.Getenv(key)
 }
@@ -177,8 +167,8 @@ func (d *GithubDispatcher) DispatchJob(ctx context.Context, job *oapi.Job) error
 		return err
 	}
 
-	ghEntity := d.getGithubEntity(cfg)
-	if ghEntity == nil {
+	ghEntity, exists := d.store.GithubEntities.Get(cfg.Owner, cfg.InstallationId)
+	if !exists {
 		return fmt.Errorf("github entity not found for job %s", job.Id)
 	}
 

apps/workspace-engine/pkg/workspace/jobdispatch/github_test.go

@@ -5,9 +5,10 @@ import (
 	"os"
 	"testing"
 
-	"workspace-engine/pkg/cmap"
 	"workspace-engine/pkg/oapi"
-	"workspace-engine/pkg/workspace/store/repository"
+	"workspace-engine/pkg/workspace/store"
+
+	"github.com/stretchr/testify/require"
 )
 
 // mockGithubClient stores dispatched workflows for validation
@@ -36,19 +37,19 @@ func (m *mockGithubClient) DispatchWorkflow(ctx context.Context, owner, repo str
 
 func TestGithubDispatcher_DispatchJob_Success(t *testing.T) {
 	// Setup mock repository with GitHub entity
-	mockRepo := &repository.Repository{
-		GithubEntities: cmap.New[*oapi.GithubEntity](),
-	}
-	mockRepo.GithubEntities.Set("test-entity", &oapi.GithubEntity{
+	mockStore := store.New()
+	if err := mockStore.GithubEntities.Upsert(context.Background(), &oapi.GithubEntity{
 		InstallationId: 12345,
 		Slug:           "test-owner",
-	})
+	}); err != nil {
+		t.Fatalf("Failed to upsert GitHub entity: %v", err)
+	}
 
 	// Setup mock client
 	mockClient := &mockGithubClient{}
 
 	// Create dispatcher with mock client factory
-	dispatcher := NewGithubDispatcherWithClientFactory(mockRepo, func(installationID int) (GithubClient, error) {
+	dispatcher := NewGithubDispatcherWithClientFactory(mockStore, func(installationID int) (GithubClient, error) {
 		return mockClient, nil
 	})
 
@@ -93,17 +94,17 @@ func TestGithubDispatcher_DispatchJob_Success(t *testing.T) {
 }
 
 func TestGithubDispatcher_DispatchJob_WithCustomRef(t *testing.T) {
-	mockRepo := &repository.Repository{
-		GithubEntities: cmap.New[*oapi.GithubEntity](),
-	}
-	mockRepo.GithubEntities.Set("test-entity", &oapi.GithubEntity{
+	mockStore := store.New()
+	if err := mockStore.GithubEntities.Upsert(context.Background(), &oapi.GithubEntity{
 		InstallationId: 12345,
 		Slug:           "test-owner",
-	})
+	}); err != nil {
+		t.Fatalf("Failed to upsert GitHub entity: %v", err)
+	}
 
 	mockClient := &mockGithubClient{}
 
-	dispatcher := NewGithubDispatcherWithClientFactory(mockRepo, func(installationID int) (GithubClient, error) {
+	dispatcher := NewGithubDispatcherWithClientFactory(mockStore, func(installationID int) (GithubClient, error) {
 		return mockClient, nil
 	})
 
@@ -135,13 +136,11 @@ func TestGithubDispatcher_DispatchJob_WithCustomRef(t *testing.T) {
 }
 
 func TestGithubDispatcher_DispatchJob_EntityNotFound(t *testing.T) {
-	mockRepo := &repository.Repository{
-		GithubEntities: cmap.New[*oapi.GithubEntity](),
-	}
+	mockStore := store.New()
 
 	mockClient := &mockGithubClient{}
 
-	dispatcher := NewGithubDispatcherWithClientFactory(mockRepo, func(installationID int) (GithubClient, error) {
+	dispatcher := NewGithubDispatcherWithClientFactory(mockStore, func(installationID int) (GithubClient, error) {
 		return mockClient, nil
 	})
 
@@ -405,30 +404,31 @@ func TestGithubDispatcher_GetGithubEntity(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			mockRepo := &repository.Repository{
-				GithubEntities: cmap.New[*oapi.GithubEntity](),
+			mockStore := store.New()
+			if err := mockStore.GithubEntities.Upsert(context.Background(), &oapi.GithubEntity{
+				InstallationId: tt.entities[0].installationID,
+				Slug:           tt.entities[0].slug,
+			}); err != nil {
+				t.Fatalf("Failed to upsert GitHub entity: %v", err)
 			}
 
 			// Add test entities
-			for i, e := range tt.entities {
-				mockRepo.GithubEntities.Set(
-					string(rune(i)),
-					&oapi.GithubEntity{
-						InstallationId: e.installationID,
-						Slug:           e.slug,
-					},
-				)
+			for _, e := range tt.entities {
+				if err := mockStore.GithubEntities.Upsert(context.Background(), &oapi.GithubEntity{
+					InstallationId: e.installationID,
+					Slug:           e.slug,
+				}); err != nil {
+					t.Fatalf("Failed to upsert GitHub entity: %v", err)
+				}
 			}
 
-			dispatcher := NewGithubDispatcher(mockRepo)
-			result := dispatcher.getGithubEntity(tt.searchCfg)
-
-			if tt.expectFound && result == nil {
-				t.Error("Expected to find entity, got nil")
-			}
-			if !tt.expectFound && result != nil {
-				t.Errorf("Expected nil, got entity: %+v", result)
+			dispatcher := NewGithubDispatcher(mockStore)
+			result, exists := dispatcher.store.GithubEntities.Get(tt.searchCfg.Owner, tt.searchCfg.InstallationId)
+			if !exists {
+				t.Errorf("Expected to find entity, got nil")
 			}
+			require.Equal(t, tt.searchCfg.InstallationId, result.InstallationId)
+			require.Equal(t, tt.searchCfg.Owner, result.Slug)
 		})
 	}
 }

apps/workspace-engine/pkg/workspace/populate_workspace.go

@@ -55,6 +55,11 @@ func PopulateWorkspaceWithInitialState(ctx context.Context, ws *Workspace) error
 			return err
 		}
 	}
+	for _, githubEntity := range initialWorkspaceState.GithubEntities() {
+		if err := ws.Store().GithubEntities.Upsert(ctx, githubEntity); err != nil {
+			return err
+		}
+	}
 
 	return nil
 }

apps/workspace-engine/pkg/workspace/releasemanager/deployment/jobs/dispatcher.go

@@ -32,7 +32,7 @@ func (d *Dispatcher) DispatchJob(ctx context.Context, job *oapi.Job) error {
 	}
 
 	if jobAgent.Type == string(jobdispatch.JobAgentTypeGithub) {
-		return jobdispatch.NewGithubDispatcher(d.store.Repo()).DispatchJob(ctx, job)
+		return jobdispatch.NewGithubDispatcher(d.store).DispatchJob(ctx, job)
 	}
 
 	return ErrUnsupportedJobAgent

apps/workspace-engine/pkg/workspace/store/github_entities.go

@@ -0,0 +1,54 @@
+package store
+
+import (
+	"context"
+	"fmt"
+	"workspace-engine/pkg/changeset"
+	"workspace-engine/pkg/oapi"
+	"workspace-engine/pkg/workspace/store/repository"
+)
+
+type GithubEntities struct {
+	repo *repository.Repository
+}
+
+func NewGithubEntities(store *Store) *GithubEntities {
+	return &GithubEntities{
+		repo: store.repo,
+	}
+}
+
+func (g *GithubEntities) key(slug string, installationId int) string {
+	return fmt.Sprintf("%s-%d", slug, installationId)
+}
+
+func (g *GithubEntities) Upsert(ctx context.Context, githubEntity *oapi.GithubEntity) error {
+	key := g.key(githubEntity.Slug, githubEntity.InstallationId)
+	g.repo.GithubEntities.Set(key, githubEntity)
+	if cs, ok := changeset.FromContext[any](ctx); ok {
+		cs.Record(changeset.ChangeTypeUpsert, githubEntity)
+	}
+	return nil
+}
+
+func (g *GithubEntities) Get(slug string, installationId int) (*oapi.GithubEntity, bool) {
+	key := g.key(slug, installationId)
+	return g.repo.GithubEntities.Get(key)
+}
+
+func (g *GithubEntities) Remove(ctx context.Context, slug string, installationId int) {
+	key := g.key(slug, installationId)
+	githubEntity, ok := g.repo.GithubEntities.Get(key)
+	if !ok {
+		return
+	}
+
+	g.repo.GithubEntities.Remove(key)
+	if cs, ok := changeset.FromContext[any](ctx); ok {
+		cs.Record(changeset.ChangeTypeDelete, githubEntity)
+	}
+}
+
+func (g *GithubEntities) Items() map[string]*oapi.GithubEntity {
+	return g.repo.GithubEntities.Items()
+}

apps/workspace-engine/pkg/workspace/store/store.go

@@ -31,6 +31,7 @@ func New() *Store {
 	store.Variables = NewVariables(store)
 	store.ResourceVariables = NewResourceVariables(store)
 	store.ResourceProviders = NewResourceProviders(store)
+	store.GithubEntities = NewGithubEntities(store)
 
 	return store
 }
@@ -54,6 +55,7 @@ type Store struct {
 	UserApprovalRecords *UserApprovalRecords
 	Relationships       *RelationshipRules
 	Variables           *Variables
+	GithubEntities      *GithubEntities
 }
 
 func (s *Store) Repo() *repository.Repository {
@@ -101,6 +103,7 @@ func (s *Store) GobDecode(data []byte) error {
 	s.UserApprovalRecords = NewUserApprovalRecords(s)
 	s.Relationships = NewRelationshipRules(s)
 	s.Variables = NewVariables(s)
+	s.GithubEntities = NewGithubEntities(s)
 
 	// Reinitialize materialized views for environments and deployments
 	s.Environments.ReinitializeMaterializedViews()