Update the way of creating distributed query in osctrl-api

api/handlers/queries.go

@@ -150,102 +150,91 @@ func (h *HandlersApi) QueriesRunHandler(w http.ResponseWriter, r *http.Request)
 		return
 	}
 
-	// Temporary list of UUIDs to calculate Expected
-	var expected []string
-	// Create targets
+	// List all the nodes that match the query
+	var expected []uint
+
+	targetNodesID := []uint{}
+	// Current logic is to select nodes meeting all criteria in the query
+	// TODO: I believe we should only allow to list nodes in one environment in URL paths
+	// We will refactor this part to be tag based queries and add more options to the query
 	if len(q.Environments) > 0 {
+		expected = []uint{}
 		for _, e := range q.Environments {
 			if (e != "") && h.Envs.Exists(e) {
-				if err := h.Queries.CreateTarget(newQuery.Name, queries.QueryTargetEnvironment, e); err != nil {
-					apiErrorResponse(w, "error creating query environment target", http.StatusInternalServerError, err)
-					h.Inc(metricAPIQueriesErr)
-					return
-				}
 				nodes, err := h.Nodes.GetByEnv(e, "active", h.Settings.InactiveHours(settings.NoEnvironmentID))
 				if err != nil {
 					apiErrorResponse(w, "error getting nodes by environment", http.StatusInternalServerError, err)
 					h.Inc(metricAPIQueriesErr)
 					return
 				}
 				for _, n := range nodes {
-					expected = append(expected, n.UUID)
+					expected = append(expected, n.ID)
 				}
 			}
 		}
+		targetNodesID = utils.Intersect(targetNodesID, expected)
 	}
 	// Create platform target
 	if len(q.Platforms) > 0 {
+		expected = []uint{}
 		platforms, _ := h.Nodes.GetAllPlatforms()
 		for _, p := range q.Platforms {
 			if (p != "") && checkValidPlatform(platforms, p) {
-				if err := h.Queries.CreateTarget(newQuery.Name, queries.QueryTargetPlatform, p); err != nil {
-					apiErrorResponse(w, "error creating query platform target", http.StatusInternalServerError, err)
-					h.Inc(metricAPIQueriesErr)
-					return
-				}
 				nodes, err := h.Nodes.GetByPlatform(p, "active", h.Settings.InactiveHours(settings.NoEnvironmentID))
 				if err != nil {
 					apiErrorResponse(w, "error getting nodes by platform", http.StatusInternalServerError, err)
 					h.Inc(metricAPIQueriesErr)
 					return
 				}
 				for _, n := range nodes {
-					expected = append(expected, n.UUID)
+					expected = append(expected, n.ID)
 				}
 			}
 		}
+		targetNodesID = utils.Intersect(targetNodesID, expected)
 	}
 	// Create UUIDs target
 	if len(q.UUIDs) > 0 {
+		expected = []uint{}
 		for _, u := range q.UUIDs {
-			if (u != "") && h.Nodes.CheckByUUID(u) {
-				if err := h.Queries.CreateTarget(newQuery.Name, queries.QueryTargetUUID, u); err != nil {
-					apiErrorResponse(w, "error creating query UUID target", http.StatusInternalServerError, err)
-					h.Inc(metricAPIQueriesErr)
-					return
+			if u != "" {
+				node, err := h.Nodes.GetByUUID(u)
+				if err != nil {
+					log.Err(err).Msgf("error getting node %s and failed to create node query for it", u)
+					continue
 				}
-				expected = append(expected, u)
+				expected = append(expected, node.ID)
 			}
 		}
+		targetNodesID = utils.Intersect(targetNodesID, expected)
 	}
 	// Create hostnames target
 	if len(q.Hosts) > 0 {
-		for _, _h := range q.Hosts {
-			if (_h != "") && h.Nodes.CheckByHost(_h) {
-				if err := h.Queries.CreateTarget(newQuery.Name, queries.QueryTargetLocalname, _h); err != nil {
-					apiErrorResponse(w, "error creating query hostname target", http.StatusInternalServerError, err)
-					h.Inc(metricAPIQueriesErr)
-					return
+		expected = []uint{}
+		for _, hostName := range q.Hosts {
+			if hostName != "" {
+				node, err := h.Nodes.GetByIdentifier(hostName)
+				if err != nil {
+					log.Err(err).Msgf("error getting node %s and failed to create node query for it", hostName)
+					continue
 				}
-				expected = append(expected, _h)
+				expected = append(expected, node.ID)
 			}
 		}
+		targetNodesID = utils.Intersect(targetNodesID, expected)
 	}
 
-	// Remove duplicates from expected
-	expectedClear := removeStringDuplicates(expected)
-
-	// Create new record for query list
-	nodesID := make([]uint, len(expectedClear))
-	for _, nodeUUID := range expectedClear {
-		node, err := h.Nodes.GetByUUID(nodeUUID)
-		if err != nil {
-			log.Err(err).Msgf("error getting node %s and failed to create node query for it", nodeUUID)
-			continue
-		}
-		nodesID = append(nodesID, node.ID)
-	}
 	// If the list is empty, we don't need to create node queries
-	if len(nodesID) != 0 {
-		if err := h.Queries.CreateNodeQueries(nodesID, newQuery.ID); err != nil {
+	if len(targetNodesID) != 0 {
+		if err := h.Queries.CreateNodeQueries(targetNodesID, newQuery.ID); err != nil {
 			log.Err(err).Msgf("error creating node queries for query %s", newQuery.Name)
 			apiErrorResponse(w, "error creating node queries", http.StatusInternalServerError, err)
 			return
 		}
 	}
 
 	// Update value for expected
-	if err := h.Queries.SetExpected(queryName, len(expectedClear), env.ID); err != nil {
+	if err := h.Queries.SetExpected(queryName, len(targetNodesID), env.ID); err != nil {
 		apiErrorResponse(w, "error setting expected", http.StatusInternalServerError, err)
 		h.Inc(metricAPICarvesErr)
 		return

nodes/nodes.go

@@ -347,7 +347,7 @@ func (n *NodeManager) UpdateMetadataByUUID(uuid string, metadata NodeMetadata) e
 		return fmt.Errorf("RecordUsername %v", err)
 	}
 	if metadata.Username != node.Username && metadata.Username != "" {
-		updates["username"] =metadata.Username
+		updates["username"] = metadata.Username
 	}
 	// Record hostname
 	if err := n.RecordHostname(metadata.Hostname, node); err != nil {

utils/utils.go

@@ -31,3 +31,28 @@ func RandomForNames() string {
 	_, _ = hasher.Write([]byte(fmt.Sprintf("%x", b)))
 	return hex.EncodeToString(hasher.Sum(nil))
 }
+
+func Intersect(slice1, slice2 []uint) []uint {
+	if len(slice1) == 0 {
+		return slice2
+	}
+	// If slice2 is empty, return slice1
+	if len(slice2) == 0 {
+		return slice1
+	}
+
+	set := make(map[uint]struct{})
+	for _, item := range slice1 {
+		set[item] = struct{}{} // Add items from slice1 to the set
+	}
+
+	intersection := []uint{}
+	for _, item := range slice2 {
+		if _, exists := set[item]; exists {
+			intersection = append(intersection, item)
+			delete(set, item) // Ensure uniqueness in the result
+		}
+	}
+
+	return intersection
+}

utils/utils_test.go

@@ -1,27 +1,44 @@
-package utils
+package utils_test
 
 import (
 	"testing"
 
+	"github.com/jmpsec/osctrl/utils"
 	"github.com/stretchr/testify/assert"
 )
 
 func TestBytesReceivedConversionBytes(t *testing.T) {
-	assert.NotEmpty(t, BytesReceivedConversion(123))
-	assert.Equal(t, "123 bytes", BytesReceivedConversion(123))
+	assert.NotEmpty(t, utils.BytesReceivedConversion(123))
+	assert.Equal(t, "123 bytes", utils.BytesReceivedConversion(123))
 }
 
 func TestBytesReceivedConversionKBytes(t *testing.T) {
-	assert.NotEmpty(t, BytesReceivedConversion(1024))
-	assert.Equal(t, "1.0 KB", BytesReceivedConversion(1024))
+	assert.NotEmpty(t, utils.BytesReceivedConversion(1024))
+	assert.Equal(t, "1.0 KB", utils.BytesReceivedConversion(1024))
 }
 
 func TestBytesReceivedConversionMBytes(t *testing.T) {
-	assert.NotEmpty(t, BytesReceivedConversion(1048576))
-	assert.Equal(t, "1.0 MB", BytesReceivedConversion(1048576))
+	assert.NotEmpty(t, utils.BytesReceivedConversion(1048576))
+	assert.Equal(t, "1.0 MB", utils.BytesReceivedConversion(1048576))
 }
 
 func TestRandomForNames(t *testing.T) {
-	assert.NotEmpty(t, RandomForNames())
-	assert.Equal(t, 32, len(RandomForNames()))
+	assert.NotEmpty(t, utils.RandomForNames())
+	assert.Equal(t, 32, len(utils.RandomForNames()))
+}
+
+func TestIntersect(t *testing.T) {
+	var slice1 = []uint{1, 2, 3, 4, 5}
+	var slice2 = []uint{3, 4, 5, 6, 7}
+	var expected = []uint{3, 4, 5}
+	assert.Equal(t, expected, utils.Intersect(slice1, slice2))
+	slice1 = utils.Intersect(slice1, slice2)
+	assert.Equal(t, expected, slice1)
+}
+
+func TestIntersectEmpty(t *testing.T) {
+	var slice1 = []uint{}
+	var slice2 = []uint{3, 4, 5, 6, 7}
+	var expected = []uint{3, 4, 5, 6, 7}
+	assert.Equal(t, expected, utils.Intersect(slice1, slice2))
 }