From d984de41e70eabd589cc18117f2b81f679349ddc Mon Sep 17 00:00:00 2001
From: Tim Lee <timlee@fleetdm.com>
Date: Thu, 28 Mar 2024 11:39:27 -0600
Subject: [PATCH] Add batch gitops processing for policies and queries (#17714)

---
 changes/17264-batch-process-gitops |  1 +
 server/service/client.go           | 64 ++++++++++++++++++++++++------
 2 files changed, 52 insertions(+), 13 deletions(-)
 create mode 100644 changes/17264-batch-process-gitops

diff --git a/changes/17264-batch-process-gitops b/changes/17264-batch-process-gitops
new file mode 100644
index 0000000000..cfa7ce9776
--- /dev/null
+++ b/changes/17264-batch-process-gitops
@@ -0,0 +1 @@
+- `fleetctl gitops` now batch processes queries and policies
\ No newline at end of file
diff --git a/server/service/client.go b/server/service/client.go
index 2bb5042235..659c4bfffe 100644
--- a/server/service/client.go
+++ b/server/service/client.go
@@ -6,7 +6,6 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"golang.org/x/text/unicode/norm"
 	"io"
 	"net/http"
 	"os"
@@ -14,6 +13,8 @@ import (
 	"strings"
 	"time"
 
+	"golang.org/x/text/unicode/norm"
+
 	"github.com/fleetdm/fleet/v4/pkg/optjson"
 	"github.com/fleetdm/fleet/v4/pkg/spec"
 	"github.com/fleetdm/fleet/v4/server/contexts/ctxerr"
@@ -22,6 +23,8 @@ import (
 	kithttp "github.com/go-kit/kit/transport/http"
 )
 
+const batchSize = 100
+
 // Client is used to consume Fleet APIs from Go code
 type Client struct {
 	*baseClient
@@ -1099,11 +1102,19 @@ func (c *Client) doGitOpsPolicies(config *spec.GitOps, logFn func(format string,
 		numPolicies := len(config.Policies)
 		logFn("[+] syncing %d policies\n", numPolicies)
 		if !dryRun {
-			// Note: We are reusing the spec flow here for adding/updating policies, instead of creating a new flow for GitOps.
-			if err := c.ApplyPolicies(config.Policies); err != nil {
-				return fmt.Errorf("error applying policies: %w", err)
+			totalApplied := 0
+			for i := 0; i < len(config.Policies); i += batchSize {
+				end := i + batchSize
+				if end > len(config.Policies) {
+					end = len(config.Policies)
+				}
+				totalApplied += end - i
+				// Note: We are reusing the spec flow here for adding/updating policies, instead of creating a new flow for GitOps.
+				if err := c.ApplyPolicies(config.Policies[i:end]); err != nil {
+					return fmt.Errorf("error applying policies: %w", err)
+				}
+				logFn("[+] synced %d policies\n", totalApplied)
 			}
-			logFn("[+] synced %d policies\n", numPolicies)
 		}
 	}
 	var policiesToDelete []uint
@@ -1123,8 +1134,17 @@ func (c *Client) doGitOpsPolicies(config *spec.GitOps, logFn func(format string,
 	if len(policiesToDelete) > 0 {
 		logFn("[-] deleting %d policies\n", len(policiesToDelete))
 		if !dryRun {
-			if err := c.DeletePolicies(config.TeamID, policiesToDelete); err != nil {
-				return fmt.Errorf("error deleting policies: %w", err)
+			totalDeleted := 0
+			for i := 0; i < len(policiesToDelete); i += batchSize {
+				end := i + batchSize
+				if end > len(policiesToDelete) {
+					end = len(policiesToDelete)
+				}
+				totalDeleted += end - i
+				if err := c.DeletePolicies(config.TeamID, policiesToDelete[i:end]); err != nil {
+					return fmt.Errorf("error deleting policies: %w", err)
+				}
+				logFn("[-] deleted %d policies\n", totalDeleted)
 			}
 		}
 	}
@@ -1132,6 +1152,7 @@ func (c *Client) doGitOpsPolicies(config *spec.GitOps, logFn func(format string,
 }
 
 func (c *Client) doGitOpsQueries(config *spec.GitOps, logFn func(format string, args ...interface{}), dryRun bool) error {
+	batchSize := 100
 	// Get the ids and names of current queries to figure out which ones to delete
 	queries, err := c.GetQueries(config.TeamID, nil)
 	if err != nil {
@@ -1141,11 +1162,19 @@ func (c *Client) doGitOpsQueries(config *spec.GitOps, logFn func(format string,
 		numQueries := len(config.Queries)
 		logFn("[+] syncing %d queries\n", numQueries)
 		if !dryRun {
-			// Note: We are reusing the spec flow here for adding/updating queries, instead of creating a new flow for GitOps.
-			if err := c.ApplyQueries(config.Queries); err != nil {
-				return fmt.Errorf("error applying queries: %w", err)
+			appliedCount := 0
+			for i := 0; i < len(config.Queries); i += batchSize {
+				end := i + batchSize
+				if end > len(config.Queries) {
+					end = len(config.Queries)
+				}
+				appliedCount += end - i
+				// Note: We are reusing the spec flow here for adding/updating queries, instead of creating a new flow for GitOps.
+				if err := c.ApplyQueries(config.Queries[i:end]); err != nil {
+					return fmt.Errorf("error applying queries: %w", err)
+				}
+				logFn("[+] synced %d queries\n", appliedCount)
 			}
-			logFn("[+] synced %d queries\n", numQueries)
 		}
 	}
 	var queriesToDelete []uint
@@ -1165,8 +1194,17 @@ func (c *Client) doGitOpsQueries(config *spec.GitOps, logFn func(format string,
 	if len(queriesToDelete) > 0 {
 		logFn("[-] deleting %d queries\n", len(queriesToDelete))
 		if !dryRun {
-			if err := c.DeleteQueries(queriesToDelete); err != nil {
-				return fmt.Errorf("error deleting queries: %w", err)
+			deleteCount := 0
+			for i := 0; i < len(queriesToDelete); i += batchSize {
+				end := i + batchSize
+				if end > len(queriesToDelete) {
+					end = len(queriesToDelete)
+				}
+				deleteCount += end - i
+				if err := c.DeleteQueries(queriesToDelete[i:end]); err != nil {
+					return fmt.Errorf("error deleting queries: %w", err)
+				}
+				logFn("[-] deleted %d queries\n", deleteCount)
 			}
 		}
 	}