Skip to content

Commit

Permalink
feat(lql): add limit option
Browse files Browse the repository at this point in the history 
ALLY-1102
  • Loading branch information
@hazedav
hazedav authored Aug 1, 2022
1 parent 810d005 commit d598def
Show file tree
Hide file tree
Showing 7 changed files with 172 additions and 41 deletions.
6 changes: 6 additions & 0 deletions api/lql_execute.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,18 +38,24 @@ const (
QueryEndTimeRange ExecuteQueryArgumentName = "EndTimeRange"
)

type ExecuteQueryOptions struct {
Limit *int `json:"limit,omitempty"`
}

type ExecuteQueryArgument struct {
Name ExecuteQueryArgumentName `json:"name"`
Value string `json:"value"`
}

type ExecuteQueryRequest struct {
Query ExecuteQuery `json:"query"`
Options ExecuteQueryOptions `json:"options"`
Arguments []ExecuteQueryArgument `json:"arguments"`
}

type ExecuteQueryByIDRequest struct {
QueryID string `json:"queryId,omitempty"`
Options ExecuteQueryOptions `json:"options"`
Arguments []ExecuteQueryArgument `json:"arguments"`
}

Expand Down
44 changes: 22 additions & 22 deletions api/lql_execute_internal_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,64 +34,64 @@ type validateQueryArgumentsTest struct {
}

var validateQueryArgumentsTests = []validateQueryArgumentsTest{
validateQueryArgumentsTest{
{
name: "empty",
arguments: []ExecuteQueryArgument{},
retrn: nil,
//retrn: errors.New(`parsing time "" as "2006-01-02T15:04:05.000Z07:00": cannot parse "" as "2006"`),
},
validateQueryArgumentsTest{
{
name: "start-bad",
arguments: []ExecuteQueryArgument{
ExecuteQueryArgument{Name: QueryStartTimeRange, Value: ""},
{Name: QueryStartTimeRange, Value: ""},
},
retrn: errors.New(
`invalid StartTimeRange argument: parsing time "" as "2006-01-02T15:04:05.000Z": cannot parse "" as "2006"`),
},
validateQueryArgumentsTest{
{
name: "start-nonutc",
arguments: []ExecuteQueryArgument{
ExecuteQueryArgument{Name: QueryStartTimeRange, Value: "2021-07-11T00:00:00.123Z07:00"},
{Name: QueryStartTimeRange, Value: "2021-07-11T00:00:00.123Z07:00"},
},
retrn: errors.New(
`invalid StartTimeRange argument: parsing time "2021-07-11T00:00:00.123Z07:00": extra text: "07:00"`),
},
validateQueryArgumentsTest{
{
name: "start-good",
arguments: []ExecuteQueryArgument{
ExecuteQueryArgument{Name: QueryStartTimeRange, Value: "2021-07-12T00:00:00.000Z"},
{Name: QueryStartTimeRange, Value: "2021-07-12T00:00:00.000Z"},
},
retrn: nil,
},
validateQueryArgumentsTest{
{
name: "end-bad",
arguments: []ExecuteQueryArgument{
ExecuteQueryArgument{Name: "EndTimeRange", Value: ""},
{Name: "EndTimeRange", Value: ""},
},
retrn: errors.New(
`invalid EndTimeRange argument: parsing time "" as "2006-01-02T15:04:05.000Z": cannot parse "" as "2006"`),
},
validateQueryArgumentsTest{
{
name: "end-good",
arguments: []ExecuteQueryArgument{
ExecuteQueryArgument{Name: "EndTimeRange", Value: "2021-07-12T00:00:00.000Z"},
{Name: "EndTimeRange", Value: "2021-07-12T00:00:00.000Z"},
},
retrn: nil,
},
validateQueryArgumentsTest{
{
name: "range-bad",
arguments: []ExecuteQueryArgument{
ExecuteQueryArgument{Name: QueryStartTimeRange, Value: "2021-07-13T00:00:00.000Z"},
ExecuteQueryArgument{Name: "EndTimeRange", Value: "2021-07-12T00:00:00.000Z"},
{Name: QueryStartTimeRange, Value: "2021-07-13T00:00:00.000Z"},
{Name: "EndTimeRange", Value: "2021-07-12T00:00:00.000Z"},
},
retrn: errors.New(
"date range should have a start time before the end time"),
},
validateQueryArgumentsTest{
{
name: "range-good",
arguments: []ExecuteQueryArgument{
ExecuteQueryArgument{Name: QueryStartTimeRange, Value: "2021-07-12T00:00:00.000Z"},
ExecuteQueryArgument{Name: "EndTimeRange", Value: "2021-07-13T00:00:00.000Z"},
{Name: QueryStartTimeRange, Value: "2021-07-12T00:00:00.000Z"},
{Name: "EndTimeRange", Value: "2021-07-13T00:00:00.000Z"},
},
retrn: nil,
},
Expand All @@ -118,31 +118,31 @@ type validateQueryRangeTest struct {
}

var validateQueryRangeTests = []validateQueryRangeTest{
validateQueryRangeTest{
{
name: "ok",
startTimeRange: time.Unix(0, 0),
endTimeRange: time.Unix(1, 0),
retrn: nil,
},
validateQueryRangeTest{
{
name: "empty-start",
startTimeRange: time.Time{},
endTimeRange: time.Unix(1, 0),
retrn: nil,
},
validateQueryRangeTest{
{
name: "empty-end",
startTimeRange: time.Unix(1, 0),
endTimeRange: time.Time{},
retrn: errors.New("date range should have a start time before the end time"),
},
validateQueryRangeTest{
{
name: "start-after-end",
startTimeRange: time.Unix(1717333947, 0),
endTimeRange: time.Unix(1617333947, 0),
retrn: errors.New("date range should have a start time before the end time"),
},
validateQueryRangeTest{
{
name: "start-equal-end",
startTimeRange: time.Unix(1617333947, 0),
endTimeRange: time.Unix(1617333947, 0),
Expand Down
81 changes: 79 additions & 2 deletions api/lql_execute_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,11 +31,11 @@ import (

var (
executeQueryArguments = []api.ExecuteQueryArgument{
api.ExecuteQueryArgument{
{
Name: api.QueryStartTimeRange,
Value: "2021-07-11T00:00:00.000Z",
},
api.ExecuteQueryArgument{
{
Name: api.QueryEndTimeRange,
Value: "2021-07-12T00:00:00.000Z",
},
Expand All @@ -46,6 +46,28 @@ var (
},
Arguments: executeQueryArguments,
}
executeQueryBadOptions = api.ExecuteQueryRequest{
Query: api.ExecuteQuery{
QueryText: newQueryText,
},
Options: api.ExecuteQueryOptions{Limit: &limitZero},
Arguments: executeQueryArguments,
}
executeQueryBadArguments = api.ExecuteQueryRequest{
Query: api.ExecuteQuery{
QueryText: newQueryText,
},
Arguments: []api.ExecuteQueryArgument{
{
Name: api.QueryStartTimeRange,
Value: "2021-07-12T00:00:00.000Z",
},
{
Name: api.QueryEndTimeRange,
Value: "2021-07-11T00:00:00.000Z",
},
},
}
executeQueryByID = api.ExecuteQueryByIDRequest{
QueryID: queryID,
Arguments: executeQueryArguments,
Expand All @@ -55,6 +77,9 @@ var (
"INSERT_ID": "35308423"
}
]`
limitZero = 0
limitNeg = -1
limitOne = 1
)

func TestQueryExecuteMethod(t *testing.T) {
Expand Down Expand Up @@ -108,6 +133,32 @@ func TestQueryExecuteOK(t *testing.T) {
assert.Equal(t, runExpected, runActual)
}

func TestQueryExecuteBad(t *testing.T) {
mockResponse := mockQueryDataResponse(executeQueryData)

fakeServer := lacework.MockServer()
fakeServer.UseApiV2()
fakeServer.MockAPI(
"Queries/execute",
func(w http.ResponseWriter, r *http.Request) {
fmt.Fprint(w, mockResponse)
},
)
defer fakeServer.Close()

c, err := api.NewClient("test",
api.WithToken("TOKEN"),
api.WithURL(fakeServer.URL()),
)
assert.Nil(t, err)

var runExpected api.ExecuteQueryResponse
_ = json.Unmarshal([]byte(mockResponse), &runExpected)

_, err = c.V2.Query.Execute(executeQueryBadArguments)
assert.NotNil(t, err)
}

func TestQueryExecuteError(t *testing.T) {
fakeServer := lacework.MockServer()
fakeServer.UseApiV2()
Expand Down Expand Up @@ -180,6 +231,32 @@ func TestQueryExecuteByIDOK(t *testing.T) {
assert.Equal(t, runExpected, runActual)
}

func TestQueryExecuteByIDBad(t *testing.T) {
mockResponse := mockQueryDataResponse(executeQueryData)

fakeServer := lacework.MockServer()
fakeServer.UseApiV2()
fakeServer.MockAPI(
"Queries/execute",
func(w http.ResponseWriter, r *http.Request) {
fmt.Fprint(w, mockResponse)
},
)
defer fakeServer.Close()

c, err := api.NewClient("test",
api.WithToken("TOKEN"),
api.WithURL(fakeServer.URL()),
)
assert.Nil(t, err)

var runExpected api.ExecuteQueryResponse
_ = json.Unmarshal([]byte(mockResponse), &runExpected)

_, err = c.V2.Query.Execute(executeQueryBadArguments)
assert.NotNil(t, err)
}

func TestQueryExecuteByIDError(t *testing.T) {
fakeServer := lacework.MockServer()
fakeServer.UseApiV2()
Expand Down
53 changes: 40 additions & 13 deletions cli/cmd/lql.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ var (
queryCmdState = struct {
End string
File string
Limit int
Range string
Start string
URL string
Expand Down Expand Up @@ -114,6 +115,11 @@ Start and end times are required to run a query:
A. CLI flags take precedence over JSON specifications`,
Args: cobra.MaximumNArgs(1),
PreRunE: func(_ *cobra.Command, _ []string) error {
// default is 0 hence the '< 0' comparison
if queryCmdState.Limit < 0 {
return errors.New("limit must be at least 1")
}

if queryCmdState.FailOnCount != "" {
var co failon.CountOperation
if err := co.Parse(queryCmdState.FailOnCount); err != nil {
Expand Down Expand Up @@ -148,34 +154,41 @@ func init() {
// run specific flags
setQuerySourceFlags(queryRunCmd)

// since time flag
queryRunCmd.Flags().StringVarP(
// limit flag
queryRunCmd.Flags().IntVar(
&queryCmdState.Limit,
"limit", 0,
"result limit for query (default 0)",
)

// range time flag
queryRunCmd.Flags().StringVar(
&queryCmdState.Range,
"range", "", "",
"range", "",
"natural time range for query",
)

// start time flag
queryRunCmd.Flags().StringVarP(
queryRunCmd.Flags().StringVar(
&queryCmdState.Start,
"start", "", "-24h",
"start", "-24h",
"start time for query",
)
// end time flag
queryRunCmd.Flags().StringVarP(
queryRunCmd.Flags().StringVar(
&queryCmdState.End,
"end", "", "now",
"end", "now",
"end time for query",
)
queryRunCmd.Flags().BoolVarP(
queryRunCmd.Flags().BoolVar(
&queryCmdState.ValidateOnly,
"validate_only", "", false,
"validate_only", false,
"validate query only (do not run)",
)
// fail on count
queryRunCmd.Flags().StringVarP(
queryRunCmd.Flags().StringVar(
&queryCmdState.FailOnCount,
"fail_on_count", "", "",
"fail_on_count", "",
"fail if the results from a query match the provided expression (e.g. '>0')",
)
}
Expand Down Expand Up @@ -428,11 +441,11 @@ func runQuery(cmd *cobra.Command, args []string) error {
}

queryArgs := []api.ExecuteQueryArgument{
api.ExecuteQueryArgument{
{
Name: api.QueryStartTimeRange,
Value: start.UTC().Format(lwtime.RFC3339Milli),
},
api.ExecuteQueryArgument{
{
Name: api.QueryEndTimeRange,
Value: end.UTC().Format(lwtime.RFC3339Milli),
},
Expand Down Expand Up @@ -482,8 +495,15 @@ func runQueryByID(id string, args []api.ExecuteQueryArgument) (
cli.StartProgress(getRunStartProgressMessage(args))
defer cli.StopProgress()

opts := api.ExecuteQueryOptions{}
// only add limit if > 0
if queryCmdState.Limit > 0 {
opts.Limit = &queryCmdState.Limit
}

request := api.ExecuteQueryByIDRequest{
QueryID: id,
Options: opts,
Arguments: args,
}
return cli.LwApi.V2.Query.ExecuteByID(request)
Expand All @@ -505,6 +525,12 @@ func runAdhocQuery(cmd *cobra.Command, args []api.ExecuteQueryArgument) (
return
}

opts := api.ExecuteQueryOptions{}
// only add limit if > 0
if queryCmdState.Limit > 0 {
opts.Limit = &queryCmdState.Limit
}

cli.StartProgress(getRunStartProgressMessage(args))
defer cli.StopProgress()

Expand All @@ -513,6 +539,7 @@ func runAdhocQuery(cmd *cobra.Command, args []api.ExecuteQueryArgument) (
Query: api.ExecuteQuery{
QueryText: newQuery.QueryText,
},
Options: opts,
Arguments: args,
}

Expand Down
Loading

0 comments on commit d598def

Please sign in to comment.