feat(RAIN-91334): allow Rego queries to be created and retrieved (#1465)

* feat(RAIN-91334): allow Rego queries to be created and retrieved

Author: incredibler

api/lql.go

@@ -29,8 +29,9 @@ import (
 )
 
 type NewQuery struct {
-	QueryID   string `json:"queryId" yaml:"queryId"`
-	QueryText string `json:"queryText" yaml:"queryText"`
+	QueryID       string  `json:"queryId" yaml:"queryId"`
+	QueryLanguage *string `json:"queryLanguage,omitempty" yaml:"queryLanguage,omitempty"`
+	QueryText     string  `json:"queryText" yaml:"queryText"`
 }
 
 func ParseNewQuery(s string) (NewQuery, error) {
@@ -49,6 +50,7 @@ func ParseNewQuery(s string) (NewQuery, error) {
 	if err == nil && !reflect.DeepEqual(query, NewQuery{}) { // empty string unmarshals w/o error
 		return query, nil
 	}
+
 	// invalid query
 	return query, errors.New("unable to parse query")
 }
@@ -59,6 +61,7 @@ type UpdateQuery struct {
 
 type Query struct {
 	QueryID        string                   `json:"queryId" yaml:"queryId"`
+	QueryLanguage  *string                  `json:"queryLanguage,omitempty" yaml:"queryLanguage,omitempty"`
 	QueryText      string                   `json:"queryText" yaml:"queryText"`
 	Owner          string                   `json:"owner"`
 	LastUpdateTime string                   `json:"lastUpdateTime"`

api/lql_test.go

@@ -22,13 +22,15 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"strings"
 	"testing"
 
 	"github.com/pkg/errors"
 	"github.com/stretchr/testify/assert"
 
 	"github.com/lacework/go-sdk/api"
 	"github.com/lacework/go-sdk/internal/lacework"
+	"github.com/lacework/go-sdk/internal/pointer"
 )
 
 var (
@@ -46,6 +48,23 @@ var (
 queryId: %s
 queryText: %s`, newQuery.QueryID, newQuery.QueryText)
 	lqlErrorReponse = `{ "message": "This is an error message" }`
+
+	regoQueryID      = "my_rego"
+	newRegoQueryText = "package clitest\n" +
+		"import future.keywords\n" +
+		"import data.lacework\n" +
+		"source := lacework.aws.cfg.list(\"s3\", \"list-buckets\")\n" +
+		"assess := assessment.violation(input, \"just because\")"
+	newRegoQuery = api.NewQuery{
+		QueryID:       regoQueryID,
+		QueryLanguage: pointer.Of("Rego"),
+		QueryText:     newRegoQueryText,
+	}
+	newRegoQueryJSON = fmt.Sprintf(`{
+	"queryId": "%s",
+	"queryLanguage": "Rego",
+	"queryText": %#v
+}`, queryID, newRegoQueryText)
 )
 
 func mockQueryDataResponse(data string) string {
@@ -129,8 +148,8 @@ func TestQueryCreateMethod(t *testing.T) {
 	assert.Nil(t, err)
 }
 
-func TestQueryCreateOK(t *testing.T) {
-	mockResponse := mockQueryDataResponse(newQueryJSON)
+func createQueryOKTestHelper(t *testing.T, expectedResponseData string, testQuery api.NewQuery) {
+	mockResponse := mockQueryDataResponse(expectedResponseData)
 
 	fakeServer := lacework.MockServer()
 	fakeServer.MockAPI(
@@ -151,10 +170,26 @@ func TestQueryCreateOK(t *testing.T) {
 	_ = json.Unmarshal([]byte(mockResponse), &createExpected)
 
 	var createActual api.QueryResponse
-	createActual, err = c.V2.Query.Create(newQuery)
+	createActual, err = c.V2.Query.Create(testQuery)
 	assert.Nil(t, err)
 
 	assert.Equal(t, createExpected, createActual)
+
+	if strings.Contains(expectedResponseData, "queryLanguage") {
+		assert.Equal(t, "Rego", *createActual.Data.QueryLanguage)
+	} else {
+		assert.Nil(t, createActual.Data.QueryLanguage)
+	}
+}
+
+func TestLQLQueryCreateOK(t *testing.T) {
+	// queryLanguage is not available
+	createQueryOKTestHelper(t, newQueryJSON, newQuery)
+}
+
+func TestRegoQueryCreateOK(t *testing.T) {
+	// queryLanguage is available
+	createQueryOKTestHelper(t, newRegoQueryJSON, newRegoQuery)
 }
 
 func TestQueryCreateError(t *testing.T) {
@@ -198,8 +233,8 @@ func TestQueryListMethod(t *testing.T) {
 	assert.Nil(t, err)
 }
 
-func TestQueryGetQueryByIDOK(t *testing.T) {
-	mockResponse := mockQueryDataResponse(newQueryJSON)
+func getQueryByIDTestHelper(t *testing.T, expectedResponseData string, queryId string) {
+	mockResponse := mockQueryDataResponse(expectedResponseData)
 
 	fakeServer := lacework.MockServer()
 	fakeServer.MockAPI(
@@ -225,6 +260,22 @@ func TestQueryGetQueryByIDOK(t *testing.T) {
 	assert.Nil(t, err)
 
 	assert.Equal(t, getExpected, getActual)
+
+	if strings.Contains(expectedResponseData, "queryLanguage") {
+		assert.Equal(t, "Rego", *getActual.Data.QueryLanguage)
+	} else {
+		assert.Nil(t, getActual.Data.QueryLanguage)
+	}
+}
+
+func TestLQLQueryGetQueryByIDOK(t *testing.T) {
+	// queryLanguage is not available
+	getQueryByIDTestHelper(t, newQueryJSON, queryID)
+}
+
+func TestRegoQueryGetQueryByIDOK(t *testing.T) {
+	// queryLanguage is available
+	getQueryByIDTestHelper(t, newRegoQueryJSON, regoQueryID)
 }
 
 func TestQueryGetNotFound(t *testing.T) {

cli/cmd/lql_show.go

@@ -71,8 +71,9 @@ func showQuery(_ *cobra.Command, args []string) error {
 
 	if cli.YAMLOutput() {
 		return cli.OutputYAML(&api.NewQuery{
-			QueryID:   queryResponse.Data.QueryID,
-			QueryText: queryResponse.Data.QueryText,
+			QueryID:       queryResponse.Data.QueryID,
+			QueryLanguage: queryResponse.Data.QueryLanguage,
+			QueryText:     queryResponse.Data.QueryText,
 		})
 	}
 

internal/pointer/bool.go

@@ -8,3 +8,8 @@ func CompareBoolPtr(ptr *bool, b bool) bool {
 	}
 	return *ptr == b
 }
+
+// Of is a helper to allow taking the address of a string literal
+func Of[E any](e E) *E {
+	return &e
+}