From 997adcebe0d83e30cfa528e0a7785e72eeea3029 Mon Sep 17 00:00:00 2001
From: Konstantin Sykulev <konst@sykulev.com>
Date: Wed, 12 Mar 2025 14:48:52 -0500
Subject: [PATCH] Batched getPoliciesBySoftwareTitleIDs (#27062)

Mysql has a max of 65535 placeholders in a sql statement. When > 33k
title ids are passed to `getPoliciesBySoftwareTitleIDs` this causes a
`Prepared statement contains too many placeholders` error. Fixed this by
splitting up the query into multiple queries and aggregating the results
in memory.

https://github.com/fleetdm/fleet/issues/26753

- [x] Changes file added for user-visible changes in `changes/`,
`orbit/changes/` or `ee/fleetd-chrome/changes`.
See [Changes
files](https://github.com/fleetdm/fleet/blob/main/docs/Contributing/Committing-Changes.md#changes-files)
for more information.
- [x] Input data is properly validated, `SELECT *` is avoided, SQL
injection is prevented (using placeholders for values in statements)
- [x] Added/updated automated tests
- [x] A detailed QA plan exists on the associated ticket (if it isn't
there, work with the product group's QA engineer to add it)
- [x] Manual QA for all new/changed functionality
---
 changes/26753-batch-software-titles     |  1 +
 server/datastore/mysql/policies.go      | 26 ++++++++++++++++++-------
 server/datastore/mysql/policies_test.go | 11 +++++++++++
 3 files changed, 31 insertions(+), 7 deletions(-)
 create mode 100644 changes/26753-batch-software-titles

diff --git a/changes/26753-batch-software-titles b/changes/26753-batch-software-titles
new file mode 100644
index 0000000000..3ed8e2444a
--- /dev/null
+++ b/changes/26753-batch-software-titles
@@ -0,0 +1 @@
+* Fixed an error when requesting /fleet/software/titles endpoint unpaginated with > 33k software titles by batching the policies by software title id query
\ No newline at end of file
diff --git a/server/datastore/mysql/policies.go b/server/datastore/mysql/policies.go
index be1f642199..1c8d5432b0 100644
--- a/server/datastore/mysql/policies.go
+++ b/server/datastore/mysql/policies.go
@@ -1815,7 +1815,7 @@ func (ds *Datastore) getPoliciesBySoftwareTitleIDs(
 		return nil, nil
 	}
 
-	query := `
+	baseQuery := `
 	SELECT
 		p.id AS id,
 		p.name AS name,
@@ -1832,14 +1832,26 @@ func (ds *Datastore) getPoliciesBySoftwareTitleIDs(
 		tmID = *teamID
 	}
 
-	query, args, err := sqlx.In(query, softwareTitleIDs, softwareTitleIDs, tmID)
+	batchSize := 32000 // see https://github.com/fleetdm/fleet/issues/26753 on the math behind this number
+	var policies []fleet.AutomaticInstallPolicy
+	err := common_mysql.BatchProcessSimple(softwareTitleIDs, batchSize, func(softwareTitleIDsToProcess []uint) error {
+		query, args, err := sqlx.In(baseQuery, softwareTitleIDsToProcess, softwareTitleIDsToProcess, tmID)
+		if err != nil {
+			return ctxerr.Wrap(ctx, err, "build select get policies by software id query")
+		}
+
+		var policyBatch []fleet.AutomaticInstallPolicy
+		if err := sqlx.SelectContext(ctx, ds.reader(ctx), &policyBatch, query, args...); err != nil {
+			return ctxerr.Wrap(ctx, err, "get policies by software installer id")
+		}
+
+		policies = append(policies, policyBatch...)
+
+		return nil
+	})
 	if err != nil {
-		return nil, ctxerr.Wrap(ctx, err, "build select get policies by software id query")
+		return nil, err
 	}
 
-	var policies []fleet.AutomaticInstallPolicy
-	if err := sqlx.SelectContext(ctx, ds.reader(ctx), &policies, query, args...); err != nil {
-		return nil, ctxerr.Wrap(ctx, err, "get policies by software installer id")
-	}
 	return policies, nil
 }
diff --git a/server/datastore/mysql/policies_test.go b/server/datastore/mysql/policies_test.go
index 2089acd126..eac3a66a3c 100644
--- a/server/datastore/mysql/policies_test.go
+++ b/server/datastore/mysql/policies_test.go
@@ -5695,6 +5695,17 @@ func testPoliciesBySoftwareTitleID(t *testing.T, ds *Datastore) {
 		require.Equal(t, expected[got.ID], got)
 	}
 
+	// performance test for 50_000 title ids, ensure batching works
+	megaTitleIDs := make([]uint, 0, 50_000)
+	megaTitleIDs = append(megaTitleIDs, *installer3.TitleID)
+	for i := uint(0); i < (50_000 - 2); i++ {
+		megaTitleIDs = append(megaTitleIDs, *installer4.TitleID+i+1)
+	}
+	megaTitleIDs = append(megaTitleIDs, *installer4.TitleID)
+	policies, err = ds.getPoliciesBySoftwareTitleIDs(ctx, megaTitleIDs, nil)
+	require.NoError(t, err)
+	require.Len(t, policies, 2)
+
 	// "No team" titles should not have any policies when filtering by team 1
 	policies, err = ds.getPoliciesBySoftwareTitleIDs(ctx, []uint{*installer3.TitleID, *installer4.TitleID}, ptr.Uint(1))
 	require.NoError(t, err)