source: Add on_prem flag to source store model

Add `on_premises` flag on store model to know if the source is created from an inventory file or not. Also, add filter for this flat at store level. Signed-off-by: Cosmin Tupangiu <[email protected]>
kubev2v · Jan 16, 2025 · f631700 · f631700
1 parent 5381dfc
commit f631700
Show file tree

Hide file tree

Showing 12 changed files with 182 additions and 77 deletions.
diff --git a/api/v1alpha1/agent/spec.gen.go b/api/v1alpha1/agent/spec.gen.go
diff --git a/api/v1alpha1/openapi.yaml b/api/v1alpha1/openapi.yaml
@@ -361,6 +361,8 @@ components:
           format: date-time
         name:
           type: string
+        onPremises:
+          type: boolean
         agents:
           type: array
           items:
@@ -369,6 +371,7 @@ components:
         - id
         - createdAt
         - updatedAt
+        - onPremises
 
     SourceCreate:
       type: object

diff --git a/api/v1alpha1/spec.gen.go b/api/v1alpha1/spec.gen.go
diff --git a/api/v1alpha1/types.gen.go b/api/v1alpha1/types.gen.go
diff --git a/internal/service/agent/handler.go b/internal/service/agent/handler.go
@@ -64,7 +64,7 @@ func (h *AgentServiceHandler) ReplaceSourceStatus(ctx context.Context, request a
 
 	associated := false
 	if source == nil {
-		source, err = h.store.Source().Create(ctx, mappers.SourceFromApi(request.Id, username, orgID, nil))
+		source, err = h.store.Source().Create(ctx, mappers.SourceFromApi(request.Id, username, orgID, nil, false))
 		if err != nil {
 			return agentServer.ReplaceSourceStatus400JSONResponse{}, nil
 		}
@@ -90,7 +90,7 @@ func (h *AgentServiceHandler) ReplaceSourceStatus(ctx context.Context, request a
 		return agentServer.ReplaceSourceStatus400JSONResponse{}, nil
 	}
 
-	newSource := mappers.SourceFromApi(request.Id, username, "", &request.Body.Inventory)
+	newSource := mappers.SourceFromApi(request.Id, username, "", &request.Body.Inventory, false)
 	result, err := h.store.Source().Update(ctx, newSource)
 	if err != nil {
 		_, _ = store.Rollback(ctx)

diff --git a/internal/service/mappers/inbound.go b/internal/service/mappers/inbound.go
@@ -19,11 +19,12 @@ func AgentFromApi(username string, orgID string, resource *apiAgent.AgentStatusU
 	}
 }
 
-func SourceFromApi(id uuid.UUID, username string, orgID string, inventory *api.Inventory) model.Source {
+func SourceFromApi(id uuid.UUID, username string, orgID string, inventory *api.Inventory, onPremises bool) model.Source {
 	source := model.Source{
-		ID:       id,
-		Username: username,
-		OrgID:    orgID,
+		ID:         id,
+		Username:   username,
+		OrgID:      orgID,
+		OnPremises: onPremises,
 	}
 
 	if inventory != nil {

diff --git a/internal/service/mappers/outbound.go b/internal/service/mappers/outbound.go
@@ -8,10 +8,11 @@ import (
 
 func SourceToApi(s model.Source) api.Source {
 	source := api.Source{
-		Id:        s.ID,
-		Inventory: nil,
-		CreatedAt: s.CreatedAt,
-		UpdatedAt: s.UpdatedAt,
+		Id:         s.ID,
+		Inventory:  nil,
+		CreatedAt:  s.CreatedAt,
+		UpdatedAt:  s.UpdatedAt,
+		OnPremises: s.OnPremises,
 	}
 
 	if len(s.Agents) > 0 {

diff --git a/internal/service/source.go b/internal/service/source.go
@@ -107,7 +107,15 @@ func (h *ServiceHandler) CreateSource(ctx context.Context, request server.Create
 		return server.CreateSource500JSONResponse{}, nil
 	}
 
-	if _, err = h.store.Source().Create(ctx, mappers.SourceFromApi(id, username, orgID, &inventory)); err != nil {
+	source, err := h.store.Source().Get(ctx, id)
+	if err == nil && source != nil {
+		if _, err = h.store.Source().Update(ctx, mappers.SourceFromApi(id, username, orgID, &inventory, true)); err != nil {
+			return server.CreateSource500JSONResponse{}, nil
+		}
+		return server.CreateSource201JSONResponse{}, nil
+	}
+
+	if _, err = h.store.Source().Create(ctx, mappers.SourceFromApi(id, username, orgID, &inventory, true)); err != nil {
 		return server.CreateSource500JSONResponse{}, nil
 	}
 

diff --git a/internal/service/source_test.go b/internal/service/source_test.go
@@ -7,6 +7,7 @@ import (
 
 	cloudevents "github.com/cloudevents/sdk-go/v2"
 	"github.com/google/uuid"
+	"github.com/kubev2v/migration-planner/api/v1alpha1"
 	"github.com/kubev2v/migration-planner/internal/api/server"
 	"github.com/kubev2v/migration-planner/internal/auth"
 	"github.com/kubev2v/migration-planner/internal/config"
@@ -22,6 +23,7 @@ const (
 	insertAgentWithSourceStm    = "INSERT INTO agents (id, source_id,associated) VALUES ('%s', '%s',  TRUE);"
 	insertSourceStm             = "INSERT INTO sources (id) VALUES ('%s');"
 	insertSourceWithUsernameStm = "INSERT INTO sources (id,username, org_id) VALUES ('%s', '%s', '%s');"
+	insertSourceOnPremisesStm   = "INSERT INTO sources (id,username, org_id, on_premises) VALUES ('%s', '%s', '%s', TRUE);"
 )
 
 var _ = Describe("source handler", Ordered, func() {
@@ -79,6 +81,32 @@ var _ = Describe("source handler", Ordered, func() {
 			Expect(resp).To(HaveLen(1))
 		})
 
+		It("successfully list all the sources -- on premises", func() {
+			tx := gormdb.Exec(fmt.Sprintf(insertSourceWithUsernameStm, uuid.NewString(), "admin", "admin"))
+			Expect(tx.Error).To(BeNil())
+			tx = gormdb.Exec(fmt.Sprintf(insertSourceOnPremisesStm, uuid.NewString(), "admin", "admin"))
+			Expect(tx.Error).To(BeNil())
+
+			eventWriter := newTestWriter()
+
+			user := auth.User{
+				Username:     "admin",
+				Organization: "admin",
+			}
+			ctx := auth.NewUserContext(context.TODO(), user)
+
+			srv := service.NewServiceHandler(s, events.NewEventProducer(eventWriter))
+			resp, err := srv.ListSources(ctx, server.ListSourcesRequestObject{})
+			Expect(err).To(BeNil())
+			Expect(reflect.TypeOf(resp)).To(Equal(reflect.TypeOf(server.ListSources200JSONResponse{})))
+			Expect(resp).To(HaveLen(2))
+
+			count := 0
+			tx = gormdb.Raw("SELECT count(*) from sources where on_premises IS TRUE;").Scan(&count)
+			Expect(tx.Error).To(BeNil())
+			Expect(count).To(Equal(1))
+		})
+
 		AfterEach(func() {
 			gormdb.Exec("DELETE FROM agents;")
 			gormdb.Exec("DELETE FROM sources;")
@@ -232,6 +260,40 @@ var _ = Describe("source handler", Ordered, func() {
 			gormdb.Exec("DELETE FROM agents;")
 			gormdb.Exec("DELETE FROM sources;")
 		})
+
+		Context("create", func() {
+			It("successfully creates a source on prem", func() {
+				eventWriter := newTestWriter()
+				srv := service.NewServiceHandler(s, events.NewEventProducer(eventWriter))
+
+				resp, err := srv.CreateSource(context.TODO(), server.CreateSourceRequestObject{
+					Body: &v1alpha1.SourceCreate{
+						Inventory: v1alpha1.Inventory{
+							Vcenter: v1alpha1.VCenter{
+								Id: uuid.NewString(),
+							},
+						},
+					},
+				})
+				Expect(err).To(BeNil())
+				Expect(reflect.TypeOf(resp)).To(Equal(reflect.TypeOf(server.CreateSource201JSONResponse{})))
+
+				count := 0
+				tx := gormdb.Raw("SELECT count(*) from sources;").Scan(&count)
+				Expect(tx.Error).To(BeNil())
+				Expect(count).To(Equal(1))
+
+				onPrem := false
+				tx = gormdb.Raw("SELECT on_premises from sources LIMIT 1;").Scan(&onPrem)
+				Expect(tx.Error).To(BeNil())
+				Expect(onPrem).To(BeTrue())
+			})
+		})
+
+		AfterEach(func() {
+			gormdb.Exec("DELETE FROM agents;")
+			gormdb.Exec("DELETE FROM sources;")
+		})
 	})
 })
 

diff --git a/internal/store/model/source.go b/internal/store/model/source.go
@@ -10,14 +10,15 @@ import (
 )
 
 type Source struct {
-	ID        openapi_types.UUID `json:"id" gorm:"primaryKey"`
-	Username  string
-	OrgID     string
-	CreatedAt time.Time
-	UpdatedAt time.Time
-	DeletedAt gorm.DeletedAt            `gorm:"index"`
-	Inventory *JSONField[api.Inventory] `gorm:"type:jsonb"`
-	Agents    []Agent                   `gorm:"constraint:OnDelete:SET NULL;"`
+	ID         openapi_types.UUID `json:"id" gorm:"primaryKey"`
+	Username   string
+	OrgID      string
+	CreatedAt  time.Time
+	UpdatedAt  time.Time
+	DeletedAt  gorm.DeletedAt            `gorm:"index"`
+	Inventory  *JSONField[api.Inventory] `gorm:"type:jsonb"`
+	OnPremises bool
+	Agents     []Agent `gorm:"constraint:OnDelete:SET NULL;"`
 }
 
 type SourceList []Source

diff --git a/internal/store/options.go b/internal/store/options.go
@@ -118,3 +118,13 @@ func (sf *SourceQueryFilter) ByDefaultInventory() *SourceQueryFilter {
 	})
 	return sf
 }
+
+func (qf *SourceQueryFilter) ByOnPremises(isOnPremises bool) *SourceQueryFilter {
+	qf.QueryFn = append(qf.QueryFn, func(tx *gorm.DB) *gorm.DB {
+		if isOnPremises {
+			return tx.Where("on_premises IS TRUE")
+		}
+		return tx.Where("on_premises IS NOT TRUE")
+	})
+	return qf
+}
diff --git a/internal/store/source_test.go b/internal/store/source_test.go
@@ -16,6 +16,7 @@ import (
 const (
 	insertSourceStm             = "INSERT INTO sources (id) VALUES ('%s');"
 	insertSourceWithUsernameStm = "INSERT INTO sources (id, username, org_id) VALUES ('%s', '%s', '%s');"
+	insertSourceOnPremisesStm   = "INSERT INTO sources (id,username, org_id, on_premises) VALUES ('%s', '%s', '%s', TRUE);"
 )
 
 var _ = Describe("source store", Ordered, func() {
@@ -84,6 +85,23 @@ var _ = Describe("source store", Ordered, func() {
 			Expect(sources[0].Username).To(Equal("admin"))
 		})
 
+		It("successfully list the source on prem", func() {
+			tx := gormdb.Exec(fmt.Sprintf(insertSourceWithUsernameStm, uuid.NewString(), "admin", "admin"))
+			Expect(tx.Error).To(BeNil())
+			tx = gormdb.Exec(fmt.Sprintf(insertSourceWithUsernameStm, uuid.NewString(), "admin", "admin"))
+			Expect(tx.Error).To(BeNil())
+			tx = gormdb.Exec(fmt.Sprintf(insertSourceOnPremisesStm, uuid.NewString(), "admin", "admin"))
+			Expect(tx.Error).To(BeNil())
+
+			sources, err := s.Source().List(context.TODO(), store.NewSourceQueryFilter().ByUsername("admin").ByOnPremises(false))
+			Expect(err).To(BeNil())
+			Expect(sources).To(HaveLen(2))
+
+			sources, err = s.Source().List(context.TODO(), store.NewSourceQueryFilter().ByUsername("admin").ByOnPremises(true))
+			Expect(err).To(BeNil())
+			Expect(sources).To(HaveLen(1))
+		})
+
 		AfterEach(func() {
 			gormdb.Exec("DELETE from agents;")
 			gormdb.Exec("DELETE from sources;")