Potential datastore optimizations for concurrent use of list mdm command API to poll results by host identifier (#30804)

2025-07-17 15:25:31 -05:00
parent 2a7547eb37
commit ce02856f85
3 changed files with 247 additions and 20 deletions
@@ -0,0 +1,2 @@
+- Modified backend for GET /api/v1/fleet/commands when filtering by `host_identifier` to address performance
+  concerns and exhausting database connections when API is called concurrently for many hosts.
@@ -92,6 +92,11 @@ func (ds *Datastore) ListMDMCommands(
 	tmFilter fleet.TeamFilter,
 	listOpts *fleet.MDMCommandListOptions,
 ) ([]*fleet.MDMCommand, error) {
+	if listOpts != nil && listOpts.Filters.HostIdentifier != "" {
+		// separate codepath for more performant query by host identifier
+		return ds.listMDMCommandsByHostIdentifier(ctx, tmFilter, listOpts)
+	}
+
 	jointStmt, params := getCombinedMDMCommandsQuery(ds, listOpts.Filters.HostIdentifier)
 	jointStmt += ds.whereFilterHostsByTeams(tmFilter, "h")
 	jointStmt, params = addRequestTypeFilter(jointStmt, &listOpts.Filters, params)
@@ -103,6 +108,193 @@ func (ds *Datastore) ListMDMCommands(
 	return results, nil
 }

+// listMDMCommandsByHostIdentifier retrieves MDM commands by host identifier. It is implemented as a
+// distinct code path to optimize the query for use cases where a client may be polling for the
+// status of commands for a specific host.
+//
+// TODO: Additional optimizations not implemented yet:
+//   - restrict ordering by date to a new sorted index (probably `nano_enrollment_queue (id,
+//     created_at DESC)` would be a good candidate)
+//   - only search by hostname as a fallback if no results are found for UUID or hardware serial
+func (ds *Datastore) listMDMCommandsByHostIdentifier(
+	ctx context.Context,
+	teamFilter fleet.TeamFilter,
+	listOpts *fleet.MDMCommandListOptions,
+) ([]*fleet.MDMCommand, error) {
+	if listOpts == nil || listOpts.Filters.HostIdentifier == "" {
+		return nil, ctxerr.Wrap(ctx, errors.New("listMDMCommandsByHostIdentifier requires non-empty listOpts.Filters.HostIdentifier"))
+	}
+
+	// First, search for host by identifier (hostname, uuid, or hardware_serial).
+	//
+	// NOTE: We're not using existing methods like ds.whereFilterHostsByIdentifier,
+	// ds.HostIDsByIdentifier, ds.HostLiteByIdentifier because those methods are poorly
+	// optimized for the indexes we currently have on the hosts table.
+	// They filter with disjunctive conditions like `hostname = ? OR uuid = ?` as well as
+	// `? IN(hostname, uuid)`. These existing queries aren't really suited for either composite
+	// indexes or indexes on individual columns, and the optimizer ends up with executions that
+	// resort full table scans or minimally filtered results (when the optimizer is using
+	// indexes on team id and the like. Full-text indexes might be an option, but we've had
+	// difficulties managing those for the hosts table in the past.
+	//
+	// So we're writing a custom query here that uses a UNION with three subqueries, each targeting
+	// a specific column index: hostname, uuid, and hardware_serial.
+
+	identifier := listOpts.Filters.HostIdentifier
+	whereTeam := ds.whereFilterHostsByTeams(teamFilter, "h")
+	columns := "id, uuid, hardware_serial, hostname, platform, team_id"
+
+	// TODO: Add index for `hostname` or remove query? If removing, we'd need to update API
+	// documentation? Breaking change? For now, adding a secondary team filter inside hostname part
+	// of the union subquery to narrow the scope somewhat
+	stmt := `
+SELECT ` + columns + ` FROM (
+	SELECT ` + columns + ` FROM hosts h WHERE hostname = ? AND ` + whereTeam + `
+	UNION SELECT ` + columns + ` FROM hosts WHERE uuid = ?
+	UNION SELECT ` + columns + ` FROM hosts WHERE hardware_serial = ? ) h
+WHERE ` + whereTeam
+
+	var dest []fleet.Host // NOTE: we're using the hosts struct for convenience, but it will not be fully populated
+	args := []any{identifier, identifier, identifier}
+	err := sqlx.SelectContext(ctx, ds.reader(ctx), &dest, stmt, args...)
+	switch {
+	case err != nil:
+		return nil, ctxerr.Wrap(ctx, err, "get host by identifier for mdm")
+	case len(dest) == 0:
+		// TODO: should we return an empty slice or an error?
+		return []*fleet.MDMCommand{}, nil
+	case len(dest) > 1:
+		// TODO: how should we handle this unexpected case?
+		level.Debug(ds.logger).Log("msg", "list mdm commands: multiple hosts found for identifier",
+			"identifier", identifier, "count", len(dest),
+		)
+	}
+
+	// Next, build the query to list MDM commands. If the found host(s) are on the same platform,
+	// we can optimize the query by skipping the UNION ALL and using a single query targeted to the
+	// platform.
+
+	var appleStmt, winStmt string
+	var appleParams, winParams []any
+	var appleUUIDs, winUUIDs []string
+	byUUID := make(map[string]fleet.Host, len(dest)) // map UUID to host so that we can loop over command results to add hostname and team info and avoid joining hosts to commands in DB
+	for _, h := range dest {
+		if prev, ok := byUUID[h.UUID]; ok {
+			// TODO: how should we handle this unexpected case?
+			level.Debug(ds.logger).Log("msg", "list mdm commands: multiple hosts found for identifier",
+				"keeping", fmt.Sprintf("id: %d uuid: %s serial: %s hostname: %s platform: %s team: %+v", h.ID, h.UUID, h.HardwareSerial, h.Hostname, h.Platform, h.TeamID),
+				"skipping", fmt.Sprintf("id: %d uuid: %s serial: %s hostname: %s platform: %s team: %+v", prev.ID, prev.UUID, prev.HardwareSerial, prev.Hostname, prev.Platform, prev.TeamID),
+			)
+		}
+		byUUID[h.UUID] = h
+		switch fleet.MDMPlatform(h.Platform) {
+		case "darwin":
+			appleUUIDs = append(appleUUIDs, h.UUID)
+		case "windows":
+			winUUIDs = append(winUUIDs, h.UUID)
+		}
+	}
+
+	if len(appleUUIDs) > 0 {
+		appleParams = []any{appleUUIDs}
+		appleStmt = `
+SELECT
+	nq.id AS host_uuid,
+	nc.command_uuid,
+	COALESCE(ncr.updated_at, nc.created_at) AS updated_at,
+	COALESCE(NULLIF(ncr.status, ''), 'Pending') AS status,
+	request_type
+FROM
+	nano_enrollment_queue nq
+	JOIN nano_commands nc ON nq.command_uuid = nc.command_uuid
+	LEFT JOIN nano_command_results ncr ON nq.id = ncr.id
+		AND nc.command_uuid = ncr.command_uuid
+WHERE
+	nq.id IN(?) AND nq.active = 1`
+
+		appleStmt, appleParams = addRequestTypeFilter(appleStmt, &listOpts.Filters, appleParams)
+		appleStmt, appleParams, err = sqlx.In(appleStmt, appleParams...)
+		if err != nil {
+			return nil, ctxerr.Wrap(ctx, err, "prepare query to list MDM commands for Apple devices")
+		}
+	}
+
+	if len(winUUIDs) > 0 {
+		winParams = []any{winUUIDs}
+		winStmt = `
+SELECT
+	mwe.host_uuid,
+	wq.command_uuid,
+	COALESCE(wcr.updated_at, wc.created_at) AS updated_at,
+	COALESCE(NULLIF(wcr.status_code, ''), 'Pending') AS status,
+	wc.target_loc_uri AS request_type
+FROM
+	windows_mdm_command_queue wq
+	JOIN mdm_windows_enrollments mwe ON mwe.id = wq.enrollment_id
+	JOIN windows_mdm_commands wc ON wc.command_uuid = wq.command_uuid
+	LEFT JOIN windows_mdm_command_results wcr ON wcr.command_uuid = wq.command_uuid
+		AND wcr.enrollment_id = wq.enrollment_id
+WHERE
+	mwe.host_uuid IN (?)`
+
+		winStmt, winParams = addRequestTypeFilter(winStmt, &listOpts.Filters, winParams)
+		winStmt, winParams, err = sqlx.In(winStmt, winParams...)
+		if err != nil {
+			return nil, ctxerr.Wrap(ctx, err, "prepare query to list MDM commands for Windows devices")
+		}
+	}
+
+	var listStmt string
+	var params []any
+	switch {
+	case len(appleUUIDs) > 0 && len(winUUIDs) > 0:
+		listStmt = fmt.Sprintf(`SELECT * FROM ((%s) UNION ALL (%s)) u`,
+			appleStmt, winStmt)
+		params = append(params, appleParams...)
+		params = append(params, winParams...)
+	case len(appleUUIDs) > 0:
+		listStmt = appleStmt
+		params = appleParams
+	case len(winUUIDs) > 0:
+		listStmt = winStmt
+		params = winParams
+	}
+
+	// TODO: Maybe move this to the service method? What about pagination metadata?
+	if listOpts.OrderKey == "" {
+		listOpts.OrderKey = "updated_at"
+	}
+	// // FIXME: We probably ought to modify how listOptionsFromRequest in transport.go applies the
+	// // default order direction. Defaulting to ascending doesn't make sense for date fields like
+	// // updated_at. List options are decoded by transport before the specific gets the request
+	// // struct so there's no way apply a different default because at that point we can't tell if
+	// // the direction was set by the user or not. One approach would be to have listOptionsFromRequest
+	// // check if the order key is a date field (i.e. it ends with "_at") and default to descending
+	// // in those cases.
+	// if listOpts.OrderDirection == "" {
+	// 	listOpts.OrderDirection = fleet.OrderDescending
+	// }
+	if listOpts.PerPage == 0 {
+		listOpts.PerPage = 10
+	}
+	listStmt, params = appendListOptionsWithCursorToSQL(listStmt, params, &listOpts.ListOptions)
+
+	var results []*fleet.MDMCommand
+	if err := sqlx.SelectContext(ctx, ds.reader(ctx), &results, listStmt, params...); err != nil {
+		return nil, ctxerr.Wrap(ctx, err, "list commands")
+	}
+
+	// Add hostname and team info to the results based on the host UUIDs.
+	for i := range results {
+		if host, ok := byUUID[results[i].HostUUID]; ok {
+			results[i].Hostname = host.Hostname
+			results[i].TeamID = host.TeamID
+		}
+	}
+
+	return results, nil
+}
+
 func addRequestTypeFilter(stmt string, filter *fleet.MDMCommandFilters, params []interface{}) (string, []interface{}) {
 	if filter.RequestType != "" {
 		stmt += " AND request_type = ?"
@@ -72,7 +72,7 @@ func testMDMCommands(t *testing.T, ds *Datastore) {

 	// enroll a windows device
 	windowsH, err := ds.NewHost(ctx, &fleet.Host{
-		Hostname:       "windows-test",
+		Hostname:       "test-host", // ambiguous hostname shared with macOS host
 		OsqueryHostID:  ptr.String("osquery-windows"),
 		NodeKey:        ptr.String("node-key-windows"),
 		UUID:           uuid.NewString(),
@@ -109,7 +109,7 @@ func testMDMCommands(t *testing.T, ds *Datastore) {

 	// enroll a macOS device
 	macH, err := ds.NewHost(ctx, &fleet.Host{
-		Hostname:       "macos-test",
+		Hostname:       "test-host", // ambiguous hostname shared with windows host
 		OsqueryHostID:  ptr.String("osquery-macos"),
 		NodeKey:        ptr.String("node-key-macos"),
 		UUID:           uuid.NewString(),
@@ -274,32 +274,65 @@ func testMDMCommands(t *testing.T, ds *Datastore) {
 	require.NoError(t, err)
 	require.Len(t, cmds, 0)

-	// filter by host Identifier
-	identifiers := map[string][]string{
-		windowsH.Hostname:       {winCmd.CommandUUID, winCmd2.CommandUUID, winCmd3.CommandUUID},
-		windowsH.UUID:           {winCmd.CommandUUID, winCmd2.CommandUUID, winCmd3.CommandUUID},
-		windowsH.HardwareSerial: {winCmd.CommandUUID, winCmd2.CommandUUID, winCmd3.CommandUUID},
-		macH.Hostname:           {appleCmdUUID, appleCmdUUID2, appleCmdUUID3},
-		macH.UUID:               {appleCmdUUID, appleCmdUUID2, appleCmdUUID3},
-		macH.HardwareSerial:     {appleCmdUUID, appleCmdUUID2, appleCmdUUID3},
-	}
-
-	for identifier, expected := range identifiers {
-		t.Run(identifier, func(t *testing.T) {
-			cmds, err = ds.ListMDMCommands(
+	for _, tc := range []struct {
+		name       string
+		identifier string
+		expected   []string
+	}{
+		{
+			name:       "windows host by hostname ambiguous with macOS host",
+			identifier: windowsH.Hostname,
+			expected: []string{
+				winCmd.CommandUUID, winCmd2.CommandUUID, winCmd3.CommandUUID,
+				appleCmdUUID, appleCmdUUID2, appleCmdUUID3,
+			},
+		},
+		{
+			name:       "windows host by UUID",
+			identifier: windowsH.UUID,
+			expected:   []string{winCmd.CommandUUID, winCmd2.CommandUUID, winCmd3.CommandUUID},
+		},
+		{
+			name:       "windows host by hardware serial",
+			identifier: windowsH.HardwareSerial,
+			expected:   []string{winCmd.CommandUUID, winCmd2.CommandUUID, winCmd3.CommandUUID},
+		},
+		{
+			name:       "macOS host by hostname ambiguous with windows host",
+			identifier: macH.Hostname,
+			expected: []string{
+				appleCmdUUID, appleCmdUUID2, appleCmdUUID3,
+				winCmd.CommandUUID, winCmd2.CommandUUID, winCmd3.CommandUUID,
+			},
+		},
+		{
+			name:       "macOS host by UUID",
+			identifier: macH.UUID,
+			expected:   []string{appleCmdUUID, appleCmdUUID2, appleCmdUUID3},
+		},
+		{
+			name:       "macOS host by hardware serial",
+			identifier: macH.HardwareSerial,
+			expected:   []string{appleCmdUUID, appleCmdUUID2, appleCmdUUID3},
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			cmds, err := ds.ListMDMCommands(
 				ctx,
 				fleet.TeamFilter{User: test.UserAdmin},
 				&fleet.MDMCommandListOptions{
 					Filters: fleet.MDMCommandFilters{
-						HostIdentifier: identifier,
+						HostIdentifier: tc.identifier,
 					},
 				},
 			)
 			require.NoError(t, err)
-			require.Len(t, cmds, 3)
+			require.Len(t, cmds, len(tc.expected))
+			var got []string
 			for _, cmd := range cmds {
-				require.Contains(t, expected, cmd.CommandUUID)
+				got = append(got, cmd.CommandUUID)
 			}
+			require.ElementsMatch(t, tc.expected, got)
 		})
 	}

@@ -334,8 +367,8 @@ func testMDMCommands(t *testing.T, ds *Datastore) {
 	)
 	require.NoError(t, err)
 	require.Len(t, cmds, 2)
-	require.Equal(t, appleCmdUUID2, cmds[0].CommandUUID)
-	require.Equal(t, appleCmdUUID4, cmds[1].CommandUUID)
+	require.Equal(t, appleCmdUUID4, cmds[0].CommandUUID)
+	require.Equal(t, appleCmdUUID2, cmds[1].CommandUUID)

 	// filter by request_type and host_identifier
 	cmds, err = ds.ListMDMCommands(