feat(RAIN-95114): include new request params in query endpoints (#1774)

* feat(RAIN-95114): include new request params in query endpoints

* style:refotmat

* feat(RAIN-95114): add test case

Author: qdfuxiao

api/lql.go

@@ -28,9 +28,42 @@ import (
 	"gopkg.in/yaml.v3"
 )
 
+type providerType int
+type providerTypes map[providerType]string
+
+const (
+	Agent providerType = iota
+	APA
+	AWS
+	Azure
+	CIEM
+	GCP
+	K8s
+	OCI
+)
+
+var ValidProviderTypes = providerTypes{
+	Agent: "Agent",
+	APA:   "APA",
+	AWS:   "AWS",
+	Azure: "Azure",
+	CIEM:  "CIEM",
+	GCP:   "GCP",
+	K8s:   "K8s",
+	OCI:   "OCI",
+}
+
+type UsageContext struct {
+	Providers []string `json:"providers,omitempty" yaml:"providers,omitempty"`
+	Policies  []string `json:"policies,omitempty" yaml:"policies,omitempty"`
+}
+
 type NewQuery struct {
-	QueryID   string `json:"queryId" yaml:"queryId"`
-	QueryText string `json:"queryText" yaml:"queryText"`
+	QueryID      string       `json:"queryId" yaml:"queryId"`
+	QueryText    string       `json:"queryText" yaml:"queryText"`
+	QueryName    string       `json:"queryName,omitempty" yaml:"queryName,omitempty"`
+	Description  string       `json:"description,omitempty" yaml:"description,omitempty"`
+	UsageContext UsageContext `json:"usageContext,omitempty" yaml:"usageContext,omitempty"`
 }
 
 func ParseNewQuery(s string) (NewQuery, error) {
@@ -54,7 +87,10 @@ func ParseNewQuery(s string) (NewQuery, error) {
 }
 
 type UpdateQuery struct {
-	QueryText string `json:"queryText"`
+	QueryText    string       `json:"queryText"`
+	QueryName    string       `json:"queryName,omitempty" yaml:"queryName,omitempty"`
+	Description  string       `json:"description,omitempty" yaml:"description,omitempty"`
+	UsageContext UsageContext `json:"usageContext,omitempty" yaml:"usageContext,omitempty"`
 }
 
 type Query struct {
@@ -64,6 +100,11 @@ type Query struct {
 	LastUpdateTime string                   `json:"lastUpdateTime"`
 	LastUpdateUser string                   `json:"lastUpdateUser"`
 	ResultSchema   []map[string]interface{} `json:"resultSchema"`
+	QueryName      string                   `json:"queryName,omitempty" yaml:"queryName,omitempty"`
+	Description    string                   `json:"description,omitempty" yaml:"description,omitempty"`
+	UsageContext   UsageContext             `json:"usageContext,omitempty" yaml:"usageContext,omitempty"`
+	Providers      []string                 `json:"providers,omitempty" yaml:"providers,omitempty"`
+	UsedByPolicies []string                 `json:"usedByPolicies,omitempty" yaml:"usedByPolicies,omitempty"`
 }
 
 type QueryResponse struct {

api/lql_test.go

@@ -34,6 +34,8 @@ import (
 var (
 	queryID      = "my_lql"
 	newQueryText = `my_lql { source { CloudTrailRawEvents } return { INSERT_ID } }`
+	queryName    = "Lql Query Name"
+	description  = "Lql Query Description"
 	newQuery     = api.NewQuery{
 		QueryID:   queryID,
 		QueryText: newQueryText,
@@ -45,6 +47,18 @@ var (
 	newQueryYAML = fmt.Sprintf(`---
 queryId: %s
 queryText: %s`, newQuery.QueryID, newQuery.QueryText)
+	newQueryWithNameDescriptionJSON = fmt.Sprintf(`{
+	"queryId": "%s",
+	"queryText": "%s",
+	"queryName": "%s",
+	"description": "%s"
+}`, queryID, newQueryText, queryName, description)
+	newQueryWithNameDescription = api.NewQuery{
+		QueryID:     queryID,
+		QueryText:   newQueryText,
+		QueryName:   queryName,
+		Description: description,
+	}
 	lqlErrorReponse = `{ "message": "This is an error message" }`
 )
 
@@ -92,6 +106,12 @@ var parseNewQueryTests = []parseNewQueryTest{
 		Expected: newQuery,
 		Error:    nil,
 	},
+	parseNewQueryTest{
+		Name:     "name-description-blobl",
+		Input:    newQueryWithNameDescriptionJSON,
+		Expected: newQueryWithNameDescription,
+		Error:    nil,
+	},
 }
 
 func TestParseNewQuery(t *testing.T) {

cli/cmd/lql_update.go

@@ -89,8 +89,10 @@ func updateQuery(cmd *cobra.Command, args []string) error {
 		}
 
 		queryYaml, err := yaml.Marshal(&api.NewQuery{
-			QueryID:   queryRes.Data.QueryID,
-			QueryText: queryRes.Data.QueryText,
+			QueryID:     queryRes.Data.QueryID,
+			QueryText:   queryRes.Data.QueryText,
+			QueryName:   queryRes.Data.QueryName,
+			Description: queryRes.Data.Description,
 		})
 		if err != nil {
 			return errors.Wrap(err, msg)
@@ -133,7 +135,9 @@ func updateQuery(cmd *cobra.Command, args []string) error {
 	cli.Log.Debugw("updating query", "query", queryString)
 	cli.StartProgress(" Updating query...")
 	update, err := cli.LwApi.V2.Query.Update(newQuery.QueryID, api.UpdateQuery{
-		QueryText: newQuery.QueryText,
+		QueryText:   newQuery.QueryText,
+		QueryName:   newQuery.QueryName,
+		Description: newQuery.Description,
 	})
 	cli.StopProgress()