Cherry pick: Make software batch endpoint asynchronous (#22258) (#22270)

Cherry pick for #22069 into `minor-fleet-v4.57.0`
fleetdm · Sep 20, 2024 · ce361d6 · ce361d6
1 parent 6d87096
commit ce361d6
Show file tree

Hide file tree

Showing 20 changed files with 757 additions and 155 deletions.
diff --git a/changes/22069-gitops-async-software-batch b/changes/22069-gitops-async-software-batch
@@ -0,0 +1 @@
+* Modified `POST /api/latest/fleet/software/batch` endpoint to be asynchronous and added a new endpoint `GET /api/latest/fleet/software/batch/{request_uuid}` to retrieve the result of the batch upload.
diff --git a/cmd/fleet/serve.go b/cmd/fleet/serve.go
@@ -49,6 +49,7 @@ import (
 	"github.com/fleetdm/fleet/v4/server/pubsub"
 	"github.com/fleetdm/fleet/v4/server/service"
 	"github.com/fleetdm/fleet/v4/server/service/async"
+	"github.com/fleetdm/fleet/v4/server/service/redis_key_value"
 	"github.com/fleetdm/fleet/v4/server/service/redis_lock"
 	"github.com/fleetdm/fleet/v4/server/service/redis_policy_set"
 	"github.com/fleetdm/fleet/v4/server/sso"
@@ -798,6 +799,7 @@ the way that the Fleet server works.
 					softwareInstallStore,
 					bootstrapPackageStore,
 					distributedLock,
+					redis_key_value.New(redisPool),
 				)
 				if err != nil {
 					initFatal(err, "initial Fleet Premium service")

diff --git a/cmd/fleetctl/get_test.go b/cmd/fleetctl/get_test.go
@@ -2320,8 +2320,8 @@ func TestGetTeamsYAMLAndApply(t *testing.T) {
 		declaration.DeclarationUUID = uuid.NewString()
 		return declaration, nil
 	}
-	ds.BatchSetSoftwareInstallersFunc = func(ctx context.Context, tmID *uint, installers []*fleet.UploadSoftwareInstallerPayload) ([]fleet.SoftwarePackageResponse, error) {
-		return nil, nil
+	ds.BatchSetSoftwareInstallersFunc = func(ctx context.Context, tmID *uint, installers []*fleet.UploadSoftwareInstallerPayload) error {
+		return nil
 	}
 
 	actualYaml := runAppForTest(t, []string{"get", "teams", "--yaml"})

diff --git a/cmd/fleetctl/gitops_test.go b/cmd/fleetctl/gitops_test.go
@@ -182,7 +182,8 @@ func TestGitOpsBasicGlobalPremium(t *testing.T) {
 	license := &fleet.LicenseInfo{Tier: fleet.TierPremium, Expiration: time.Now().Add(24 * time.Hour)}
 	_, ds := runServerWithMockedDS(
 		t, &service.TestServerOpts{
-			License: license,
+			License:       license,
+			KeyValueStore: newMemKeyValueStore(),
 		},
 	)
 
@@ -229,7 +230,10 @@ func TestGitOpsBasicGlobalPremium(t *testing.T) {
 	ds.NewJobFunc = func(ctx context.Context, job *fleet.Job) (*fleet.Job, error) {
 		return &fleet.Job{}, nil
 	}
-	ds.BatchSetSoftwareInstallersFunc = func(ctx context.Context, teamID *uint, installers []*fleet.UploadSoftwareInstallerPayload) ([]fleet.SoftwarePackageResponse, error) {
+	ds.BatchSetSoftwareInstallersFunc = func(ctx context.Context, teamID *uint, installers []*fleet.UploadSoftwareInstallerPayload) error {
+		return nil
+	}
+	ds.GetSoftwareInstallersFunc = func(ctx context.Context, tmID uint) ([]fleet.SoftwarePackageResponse, error) {
 		return nil, nil
 	}
 
@@ -285,7 +289,8 @@ func TestGitOpsBasicTeam(t *testing.T) {
 	license := &fleet.LicenseInfo{Tier: fleet.TierPremium, Expiration: time.Now().Add(24 * time.Hour)}
 	_, ds := runServerWithMockedDS(
 		t, &service.TestServerOpts{
-			License: license,
+			License:       license,
+			KeyValueStore: newMemKeyValueStore(),
 		},
 	)
 
@@ -373,7 +378,10 @@ func TestGitOpsBasicTeam(t *testing.T) {
 	ds.DeleteMDMAppleDeclarationByNameFunc = func(ctx context.Context, teamID *uint, name string) error {
 		return nil
 	}
-	ds.BatchSetSoftwareInstallersFunc = func(ctx context.Context, teamID *uint, installers []*fleet.UploadSoftwareInstallerPayload) ([]fleet.SoftwarePackageResponse, error) {
+	ds.BatchSetSoftwareInstallersFunc = func(ctx context.Context, teamID *uint, installers []*fleet.UploadSoftwareInstallerPayload) error {
+		return nil
+	}
+	ds.GetSoftwareInstallersFunc = func(ctx context.Context, tmID uint) ([]fleet.SoftwarePackageResponse, error) {
 		return nil, nil
 	}
 	ds.ApplyEnrollSecretsFunc = func(ctx context.Context, teamID *uint, secrets []*fleet.EnrollSecret) error {
@@ -644,6 +652,7 @@ func TestGitOpsFullTeam(t *testing.T) {
 			MDMPusher:        mockPusher{},
 			FleetConfig:      &fleetCfg,
 			NoCacheDatastore: true,
+			KeyValueStore:    newMemKeyValueStore(),
 		},
 	)
 
@@ -804,8 +813,11 @@ func TestGitOpsFullTeam(t *testing.T) {
 		return nil
 	}
 	var appliedSoftwareInstallers []*fleet.UploadSoftwareInstallerPayload
-	ds.BatchSetSoftwareInstallersFunc = func(ctx context.Context, teamID *uint, installers []*fleet.UploadSoftwareInstallerPayload) ([]fleet.SoftwarePackageResponse, error) {
+	ds.BatchSetSoftwareInstallersFunc = func(ctx context.Context, teamID *uint, installers []*fleet.UploadSoftwareInstallerPayload) error {
 		appliedSoftwareInstallers = installers
+		return nil
+	}
+	ds.GetSoftwareInstallersFunc = func(ctx context.Context, tmID uint) ([]fleet.SoftwarePackageResponse, error) {
 		return nil, nil
 	}
 	ds.SetTeamVPPAppsFunc = func(ctx context.Context, teamID *uint, adamIDs []fleet.VPPAppTeam) error {
@@ -937,7 +949,8 @@ func TestGitOpsBasicGlobalAndTeam(t *testing.T) {
 	license := &fleet.LicenseInfo{Tier: fleet.TierPremium, Expiration: time.Now().Add(24 * time.Hour)}
 	_, ds := runServerWithMockedDS(
 		t, &service.TestServerOpts{
-			License: license,
+			License:       license,
+			KeyValueStore: newMemKeyValueStore(),
 		},
 	)
 
@@ -1055,7 +1068,10 @@ func TestGitOpsBasicGlobalAndTeam(t *testing.T) {
 		savedTeam = team
 		return team, nil
 	}
-	ds.BatchSetSoftwareInstallersFunc = func(ctx context.Context, teamID *uint, installers []*fleet.UploadSoftwareInstallerPayload) ([]fleet.SoftwarePackageResponse, error) {
+	ds.BatchSetSoftwareInstallersFunc = func(ctx context.Context, teamID *uint, installers []*fleet.UploadSoftwareInstallerPayload) error {
+		return nil
+	}
+	ds.GetSoftwareInstallersFunc = func(ctx context.Context, tmID uint) ([]fleet.SoftwarePackageResponse, error) {
 		return nil, nil
 	}
 	ds.ListSoftwareTitlesFunc = func(ctx context.Context, opt fleet.SoftwareTitleListOptions, tmFilter fleet.TeamFilter) ([]fleet.SoftwareTitleListResult, int, *fleet.PaginationMetadata, error) {
@@ -1201,7 +1217,8 @@ func TestGitOpsBasicGlobalAndNoTeam(t *testing.T) {
 	license := &fleet.LicenseInfo{Tier: fleet.TierPremium, Expiration: time.Now().Add(24 * time.Hour)}
 	_, ds := runServerWithMockedDS(
 		t, &service.TestServerOpts{
-			License: license,
+			License:       license,
+			KeyValueStore: newMemKeyValueStore(),
 		},
 	)
 	// Mock appConfig
@@ -1317,7 +1334,10 @@ func TestGitOpsBasicGlobalAndNoTeam(t *testing.T) {
 		savedTeam = team
 		return team, nil
 	}
-	ds.BatchSetSoftwareInstallersFunc = func(ctx context.Context, teamID *uint, installers []*fleet.UploadSoftwareInstallerPayload) ([]fleet.SoftwarePackageResponse, error) {
+	ds.BatchSetSoftwareInstallersFunc = func(ctx context.Context, teamID *uint, installers []*fleet.UploadSoftwareInstallerPayload) error {
+		return nil
+	}
+	ds.GetSoftwareInstallersFunc = func(ctx context.Context, tmID uint) ([]fleet.SoftwarePackageResponse, error) {
 		return nil, nil
 	}
 	ds.ListSoftwareTitlesFunc = func(ctx context.Context, opt fleet.SoftwareTitleListOptions, tmFilter fleet.TeamFilter) ([]fleet.SoftwareTitleListResult, int, *fleet.PaginationMetadata, error) {
@@ -1634,9 +1654,9 @@ func TestGitOpsTeamSofwareInstallers(t *testing.T) {
 		file    string
 		wantErr string
 	}{
-		{"testdata/gitops/team_software_installer_not_found.yml", "Please make sure that URLs are publicy accessible to the internet."},
+		{"testdata/gitops/team_software_installer_not_found.yml", "Please make sure that URLs are reachable from your Fleet server."},
 		{"testdata/gitops/team_software_installer_unsupported.yml", "The file should be .pkg, .msi, .exe or .deb."},
-		{"testdata/gitops/team_software_installer_too_large.yml", "The maximum file size is 500 MB"},
+		{"testdata/gitops/team_software_installer_too_large.yml", "The maximum file size is 500 MiB"},
 		{"testdata/gitops/team_software_installer_valid.yml", ""},
 		{"testdata/gitops/team_software_installer_valid_apply.yml", ""},
 		{"testdata/gitops/team_software_installer_pre_condition_multiple_queries.yml", "should have only one query."},
@@ -1668,10 +1688,13 @@ func TestGitOpsTeamSoftwareInstallersQueryEnv(t *testing.T) {
 
 	t.Setenv("QUERY_VAR", "IT_WORKS")
 
-	ds.BatchSetSoftwareInstallersFunc = func(ctx context.Context, tmID *uint, installers []*fleet.UploadSoftwareInstallerPayload) ([]fleet.SoftwarePackageResponse, error) {
+	ds.BatchSetSoftwareInstallersFunc = func(ctx context.Context, tmID *uint, installers []*fleet.UploadSoftwareInstallerPayload) error {
 		if installers[0].PreInstallQuery != "select IT_WORKS" {
-			return nil, fmt.Errorf("Missing env var, got %s", installers[0].PreInstallQuery)
+			return fmt.Errorf("Missing env var, got %s", installers[0].PreInstallQuery)
 		}
+		return nil
+	}
+	ds.GetSoftwareInstallersFunc = func(ctx context.Context, tmID uint) ([]fleet.SoftwarePackageResponse, error) {
 		return nil, nil
 	}
 
@@ -1686,9 +1709,9 @@ func TestGitOpsNoTeamSoftwareInstallers(t *testing.T) {
 		noTeamFile string
 		wantErr    string
 	}{
-		{"testdata/gitops/no_team_software_installer_not_found.yml", "Please make sure that URLs are publicy accessible to the internet."},
+		{"testdata/gitops/no_team_software_installer_not_found.yml", "Please make sure that URLs are reachable from your Fleet server."},
 		{"testdata/gitops/no_team_software_installer_unsupported.yml", "The file should be .pkg, .msi, .exe or .deb."},
-		{"testdata/gitops/no_team_software_installer_too_large.yml", "The maximum file size is 500 MB"},
+		{"testdata/gitops/no_team_software_installer_too_large.yml", "The maximum file size is 500 MiB"},
 		{"testdata/gitops/no_team_software_installer_valid.yml", ""},
 		{"testdata/gitops/no_team_software_installer_pre_condition_multiple_queries.yml", "should have only one query."},
 		{"testdata/gitops/no_team_software_installer_pre_condition_not_found.yml", "no such file or directory"},
@@ -2050,6 +2073,7 @@ func setupFullGitOpsPremiumServer(t *testing.T) (*mock.Store, **fleet.AppConfig,
 			FleetConfig:      &fleetCfg,
 			License:          license,
 			NoCacheDatastore: true,
+			KeyValueStore:    newMemKeyValueStore(),
 		},
 	)
 
@@ -2181,7 +2205,10 @@ func setupFullGitOpsPremiumServer(t *testing.T) (*mock.Store, **fleet.AppConfig,
 		declaration.DeclarationUUID = uuid.NewString()
 		return declaration, nil
 	}
-	ds.BatchSetSoftwareInstallersFunc = func(ctx context.Context, teamID *uint, installers []*fleet.UploadSoftwareInstallerPayload) ([]fleet.SoftwarePackageResponse, error) {
+	ds.BatchSetSoftwareInstallersFunc = func(ctx context.Context, teamID *uint, installers []*fleet.UploadSoftwareInstallerPayload) error {
+		return nil
+	}
+	ds.GetSoftwareInstallersFunc = func(ctx context.Context, tmID uint) ([]fleet.SoftwarePackageResponse, error) {
 		return nil, nil
 	}
 
@@ -2890,3 +2917,26 @@ software:
 		})
 	}
 }
+
+type memKeyValueStore struct {
+	m map[string]string
+}
+
+func newMemKeyValueStore() *memKeyValueStore {
+	return &memKeyValueStore{
+		m: make(map[string]string),
+	}
+}
+
+func (m *memKeyValueStore) Set(ctx context.Context, key string, value string, expireTime time.Duration) error {
+	m.m[key] = value
+	return nil
+}
+
+func (m *memKeyValueStore) Get(ctx context.Context, key string) (*string, error) {
+	v, ok := m.m[key]
+	if !ok {
+		return nil, nil
+	}
+	return &v, nil
+}
diff --git a/ee/server/service/mdm_external_test.go b/ee/server/service/mdm_external_test.go
@@ -109,6 +109,7 @@ func setupMockDatastorePremiumService(t testing.TB) (*mock.Store, *eeservice.Ser
 		nil,
 		nil,
 		nil,
+		nil,
 	)
 	if err != nil {
 		panic(err)

diff --git a/ee/server/service/service.go b/ee/server/service/service.go
@@ -30,6 +30,7 @@ type Service struct {
 	softwareInstallStore  fleet.SoftwareInstallerStore
 	bootstrapPackageStore fleet.MDMBootstrapPackageStore
 	distributedLock       fleet.Lock
+	keyValueStore         fleet.KeyValueStore
 }
 
 func NewService(
@@ -46,6 +47,7 @@ func NewService(
 	softwareInstallStore fleet.SoftwareInstallerStore,
 	bootstrapPackageStore fleet.MDMBootstrapPackageStore,
 	distributedLock fleet.Lock,
+	keyValueStore fleet.KeyValueStore,
 ) (*Service, error) {
 	authorizer, err := authz.NewAuthorizer()
 	if err != nil {
@@ -67,6 +69,7 @@ func NewService(
 		softwareInstallStore:  softwareInstallStore,
 		bootstrapPackageStore: bootstrapPackageStore,
 		distributedLock:       distributedLock,
+		keyValueStore:         keyValueStore,
 	}
 
 	// Override methods that can't be easily overriden via
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		* Modified `POST /api/latest/fleet/software/batch` endpoint to be asynchronous and added a new endpoint `GET /api/latest/fleet/software/batch/{request_uuid}` to retrieve the result of the batch upload.
-Original file line number
+Diff line change
@@ Expand Up @@
     		nil,
     		nil,
     		nil,
+    		nil,
     	)
     	if err != nil {
     		panic(err)
@@ Expand Down @@