diff --git a/api/lql_execute.go b/api/lql_execute.go index b8aa17011..314db43b9 100644 --- a/api/lql_execute.go +++ b/api/lql_execute.go @@ -38,6 +38,10 @@ const ( QueryEndTimeRange ExecuteQueryArgumentName = "EndTimeRange" ) +type ExecuteQueryOptions struct { + Limit *int `json:"limit,omitempty"` +} + type ExecuteQueryArgument struct { Name ExecuteQueryArgumentName `json:"name"` Value string `json:"value"` @@ -45,11 +49,13 @@ type ExecuteQueryArgument struct { 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"` } diff --git a/api/lql_execute_internal_test.go b/api/lql_execute_internal_test.go index 1223f4e32..a4b2ec0e1 100644 --- a/api/lql_execute_internal_test.go +++ b/api/lql_execute_internal_test.go @@ -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, }, @@ -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), diff --git a/api/lql_execute_test.go b/api/lql_execute_test.go index 384c17ee2..199abacc8 100644 --- a/api/lql_execute_test.go +++ b/api/lql_execute_test.go @@ -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", }, @@ -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, @@ -55,6 +77,9 @@ var ( "INSERT_ID": "35308423" } ]` + limitZero = 0 + limitNeg = -1 + limitOne = 1 ) func TestQueryExecuteMethod(t *testing.T) { @@ -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() @@ -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() diff --git a/cli/cmd/lql.go b/cli/cmd/lql.go index fbf51449a..a1e21dc54 100644 --- a/cli/cmd/lql.go +++ b/cli/cmd/lql.go @@ -40,6 +40,7 @@ var ( queryCmdState = struct { End string File string + Limit int Range string Start string URL string @@ -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 { @@ -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')", ) } @@ -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), }, @@ -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) @@ -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() @@ -513,6 +539,7 @@ func runAdhocQuery(cmd *cobra.Command, args []api.ExecuteQueryArgument) ( Query: api.ExecuteQuery{ QueryText: newQuery.QueryText, }, + Options: opts, Arguments: args, } diff --git a/cli/cmd/lql_preview.go b/cli/cmd/lql_preview.go index 86814740b..270cf902d 100644 --- a/cli/cmd/lql_preview.go +++ b/cli/cmd/lql_preview.go @@ -64,12 +64,16 @@ func previewQuerySource(_ *cobra.Command, args []string) error { return errors.New("unable to parse datasource schema") } + // initialize limit + limit := 1 + // initialize query executeQuery := api.ExecuteQueryRequest{ Query: api.ExecuteQuery{ QueryText: fmt.Sprintf( queryPreviewSourceTemplate, args[0], strings.Join(returns, ",")), }, + Options: api.ExecuteQueryOptions{Limit: &limit}, } // initialize time attempts @@ -84,11 +88,11 @@ func previewQuerySource(_ *cobra.Command, args []string) error { end, _ := lwtime.ParseRelative(timeAttempt["end"]) executeQuery.Arguments = []api.ExecuteQueryArgument{ - api.ExecuteQueryArgument{ + { Name: api.QueryStartTimeRange, Value: start.UTC().Format(lwtime.RFC3339Milli), }, - api.ExecuteQueryArgument{ + { Name: api.QueryEndTimeRange, Value: end.UTC().Format(lwtime.RFC3339Milli), }, diff --git a/integration/lql_test.go b/integration/lql_test.go index 8cf54790c..6ccd8b387 100644 --- a/integration/lql_test.go +++ b/integration/lql_test.go @@ -21,6 +21,7 @@ package integration import ( "bytes" + "encoding/json" "fmt" "os" "testing" @@ -140,9 +141,24 @@ func TestQueryRunFile(t *testing.T) { } defer os.Remove(file.Name()) - // run (explicit times) + // run (bad limit) out, stderr, exitcode := LaceworkCLIWithTOMLConfig( - "query", "run", "-f", file.Name(), "--start", queryStart, "--end", queryEnd) + "query", "run", "-f", file.Name(), "--start", queryStart, "--end", queryEnd, "--limit", "-1") + assert.Contains(t, stderr.String(), "limit must be at least 1") + assert.Equal(t, 1, exitcode, "EXITCODE is not the expected one") + + // run (explicit times / options) + out, stderr, exitcode = LaceworkCLIWithTOMLConfig( + "query", "run", "-f", file.Name(), "--start", queryStart, "--end", queryEnd, "--limit", "1", "--json") + + // check limit + var results []interface{} + err = json.Unmarshal(out.Bytes(), &results) + if err != nil { + assert.FailNow(t, err.Error()) + } + assert.Equal(t, 1, len(results)) + assert.Contains(t, out.String(), `"INSERT_ID"`) assert.Empty(t, stderr.String(), "STDERR should be empty") assert.Equal(t, 0, exitcode, "EXITCODE is not the expected one") diff --git a/integration/test_resources/help/query_run b/integration/test_resources/help/query_run index 6da9eb8e0..0ce44e8ff 100644 --- a/integration/test_resources/help/query_run +++ b/integration/test_resources/help/query_run @@ -35,6 +35,7 @@ Flags: --fail_on_count string fail if the results from a query match the provided expression (e.g. '>0') -f, --file string path to a query to run -h, --help help for run + --limit int result limit for query (default 0) --range string natural time range for query --start string start time for query (default "-24h") -u, --url string url to a query to run