APL support

Makefile

@@ -72,6 +72,7 @@ fmt: ## Format and simplify the source code using `gofmt`
 
 .PHONY: generate
 generate: $(STRINGER) \
+	axiom/apl/format_string.go \
 	axiom/query/kind_string.go \
 	axiom/query/result_string.go \
 	axiom/datasets_string.go \

axiom/apl/doc.go

@@ -0,0 +1,3 @@
+// Package apl provides the datatypes for construction APL. They usually extend
+// the functionality of existing types from the `query` package.
+package apl

axiom/apl/format.go

@@ -0,0 +1,22 @@
+package apl
+
+import (
+	"net/url"
+)
+
+//go:generate ../../bin/stringer -type=Format -linecomment -output=format_string.go
+
+// Format represents the format of an APL query.
+type Format uint8
+
+// All available query formats.
+const (
+	Legacy Format = iota // legacy
+)
+
+// EncodeValues implements `query.Encoder`. It is in place to encode the Format
+// into a string URL value because that's what the server expects.
+func (f Format) EncodeValues(key string, v *url.Values) error {
+	v.Set(key, f.String())
+	return nil
+}

axiom/apl/format_test.go

@@ -0,0 +1,41 @@
+package apl
+
+import (
+	"net/url"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestFormat_EncodeValues(t *testing.T) {
+	tests := []struct {
+		input Format
+		exp   string
+	}{
+		{Legacy, "legacy"},
+		// {0, "Format(0)"}, // HINT(lukasmalkmus): Maybe we want to sort this out by raising an error?
+	}
+	for _, tt := range tests {
+		t.Run(tt.input.String(), func(t *testing.T) {
+			v := &url.Values{}
+			err := tt.input.EncodeValues("test", v)
+			require.NoError(t, err)
+
+			assert.Equal(t, tt.exp, v.Get("test"))
+		})
+	}
+}
+
+func TestFormat_String(t *testing.T) {
+	// Check outer bounds.
+	// assert.Equal(t, Format(0).String(), "Format(0)")
+	// assert.Contains(t, (Legacy - 1).String(), "Format(")
+	assert.Contains(t, (Legacy + 1).String(), "Format(")
+
+	for c := Legacy; c <= Legacy; c++ {
+		s := c.String()
+		assert.NotEmpty(t, s)
+		assert.NotContains(t, s, "Format(")
+	}
+}

axiom/apl/options.go

@@ -0,0 +1,25 @@
+package apl
+
+import (
+	"time"
+)
+
+// Options specifies the optional parameters to APL query methods.
+type Options struct {
+	// StartTime of the query.
+	StartTime time.Time `url:"-"`
+	// EndTime of the query.
+	EndTime time.Time `url:"-"`
+
+	// NoCache omits the query cache.
+	NoCache bool `url:"nocache,omitempty"`
+	// Save the query on the server, if set to `true`. The ID of the saved query
+	// is returned with the query result as part of the response.
+	// HINT(lukasmalkmus): The server automatically sets the query kind to "apl"
+	// for queries going to the "/_apl" query endpoint. This allows us to set
+	// any value for the `saveAsKind` query param. For user experience, we use a
+	// bool here instead of forcing the user to set the value to `query.APL`.
+	Save bool `url:"saveAsKind,omitempty"`
+	// Format specifies the format of the APL query. Defaults to Legacy.
+	Format Format `url:"format"`
+}

axiom/apl/result.go

@@ -0,0 +1,12 @@
+package apl
+
+import "github.com/axiomhq/axiom-go/axiom/query"
+
+// Result is the result of an APL query. It adds the APL query request alongside
+// the query result it created, making it a superset of `query.Result`
+type Result struct {
+	// Request is the APL query request that created the result.
+	Request *query.Query `json:"request"`
+
+	*query.Result
+}

axiom/client.go

@@ -345,7 +345,7 @@ func (c *Client) newRequest(ctx context.Context, method, path string, body inter
 		req.Header.Set("Authorization", "Bearer "+c.accessToken)
 	}
 
-	// Set organization id header, if present.
+	// Set organization ID header, if present.
 	if c.orgID != "" {
 		req.Header.Set("X-Axiom-Org-Id", c.orgID)
 	}

axiom/client_integration_test.go

@@ -46,7 +46,7 @@ type IntegrationTestSuite struct {
 
 func (s *IntegrationTestSuite) SetupSuite() {
 	s.Require().NotEmpty(accessToken, "integration test needs a personal access token set")
-	s.Require().True(orgID != "" || deploymentURL != "", "integration test needs an organization id or deployment url set")
+	s.Require().True(orgID != "" || deploymentURL != "", "integration test needs an organization ID or deployment url set")
 
 	s.T().Logf("strict decoding is set to \"%t\"", strictDecoding)
 

axiom/datasets.go

@@ -12,6 +12,7 @@ import (
 	"time"
 	"unicode"
 
+	"github.com/axiomhq/axiom-go/axiom/apl"
 	"github.com/axiomhq/axiom-go/axiom/query"
 )
 
@@ -142,22 +143,6 @@ type TrimResult struct {
 	BlocksDeleted int `json:"numDeleted"`
 }
 
-// HistoryQuery represents a query stored inside the query history.
-type HistoryQuery struct {
-	// ID is the unique id of the starred query.
-	ID string `json:"id"`
-	// Kind of the starred query.
-	Kind query.Kind `json:"kind"`
-	// Dataset the starred query belongs to.
-	Dataset string `json:"dataset"`
-	// Owner is the ID of the starred queries owner. Can be a user or team ID.
-	Owner string `json:"who"`
-	// Query is the actual query.
-	Query query.Query `json:"query"`
-	// CreatedAt is the time the history query was created.
-	CreatedAt time.Time `json:"created"`
-}
-
 // IngestStatus is the status after an event ingestion operation.
 type IngestStatus struct {
 	// Ingested is the amount of events that have been ingested.
@@ -204,8 +189,17 @@ type datasetTrimRequest struct {
 	MaxDuration string `json:"maxDuration"`
 }
 
-// IngestOptions specifies the parameters for the Ingest and IngestEvents method
-// of the Datasets service.
+type aplQueryRequest struct {
+	// Raw is the raw APL query string.
+	Raw string `json:"apl"`
+	// StartTime of the query. Optional.
+	StartTime time.Time `json:"startTime"`
+	// EndTime of the query. Optional.
+	EndTime time.Time `json:"endTime"`
+}
+
+// IngestOptions specifies the optional parameters for the Ingest and
+// IngestEvents method of the Datasets service.
 type IngestOptions struct {
 	// TimestampField defines a custom field to extract the ingestion timestamp
 	// from. Defaults to `_time`.
@@ -321,10 +315,10 @@ func (s *DatasetsService) Trim(ctx context.Context, id string, maxDuration time.
 
 // History retrieves the query stored inside the query history dataset
 // identified by its id.
-func (s *DatasetsService) History(ctx context.Context, id string) (*HistoryQuery, error) {
+func (s *DatasetsService) History(ctx context.Context, id string) (*query.History, error) {
 	path := s.basePath + "/_history/" + id
 
-	var res HistoryQuery
+	var res query.History
 	if err := s.client.call(ctx, http.MethodGet, path, nil, &res); err != nil {
 		return nil, err
 	}
@@ -457,6 +451,35 @@ func (s *DatasetsService) Query(ctx context.Context, id string, q query.Query, o
 	return &res, nil
 }
 
+// APLQuery executes the given query specified using the Axiom Processing
+// Language (APL).
+func (s *DatasetsService) APLQuery(ctx context.Context, raw string, opts apl.Options) (*apl.Result, error) {
+	path, err := addOptions(s.basePath+"/_apl", opts)
+	if err != nil {
+		return nil, err
+	}
+
+	req, err := s.client.newRequest(ctx, http.MethodPost, path, aplQueryRequest{
+		Raw:       raw,
+		StartTime: opts.StartTime,
+		EndTime:   opts.EndTime,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	var (
+		res  apl.Result
+		resp *response
+	)
+	if resp, err = s.client.do(req, &res); err != nil {
+		return nil, err
+	}
+	res.SavedQueryID = resp.Header.Get("X-Axiom-History-Query-Id")
+
+	return &res, nil
+}
+
 // GZIPStreamer returns an io.Reader that gzip compresses the data it reads from
 // the provided reader using the specified compression level.
 func GZIPStreamer(r io.Reader, level int) (io.Reader, error) {

axiom/datasets_integration_test.go

@@ -5,6 +5,7 @@ package axiom_test
 import (
 	"bytes"
 	"context"
+	"fmt"
 	"io"
 	"strings"
 	"testing"
@@ -13,6 +14,7 @@ import (
 	"github.com/stretchr/testify/suite"
 
 	"github.com/axiomhq/axiom-go/axiom"
+	"github.com/axiomhq/axiom-go/axiom/apl"
 	"github.com/axiomhq/axiom-go/axiom/query"
 )
 
@@ -170,7 +172,7 @@ func (s *DatasetsTestSuite) Test() {
 			break
 		}
 	}
-	s.True(contains, "stats do not contain the dataset we created for this test")
+	s.True(contains, "stats do not contain the dataset created for this test")
 
 	// Run a query and make sure we see some results.
 	queryResult, err := s.client.Datasets.Query(s.ctx, s.dataset.ID, query.Query{
@@ -185,31 +187,57 @@ func (s *DatasetsTestSuite) Test() {
 	// This needs to pass in order for the history query test to have an input.
 	s.Require().NotEmpty(queryResult.SavedQueryID)
 
-	// FIXME(lukasmalkmus): For some reason we get "2" here?!
-	// s.EqualValues(1, queryResult.Status.BlocksExamined)
+	// s.EqualValues(1, queryResult.Status.BlocksExamined) // FIXME(lukasmalkmus): For some reason we get "2" here?!
 	s.EqualValues(4, queryResult.Status.RowsExamined)
 	s.EqualValues(4, queryResult.Status.RowsMatched)
 	s.Len(queryResult.Matches, 4)
 
+	// Run another query but using APL.
+	rawAPLQuery := fmt.Sprintf("['%s']", s.dataset.ID)
+	aplQueryResult, err := s.client.Datasets.APLQuery(s.ctx, rawAPLQuery, apl.Options{
+		Save: true,
+	})
+	s.Require().NoError(err)
+	s.Require().NotNil(aplQueryResult)
+
+	// This needs to pass in order for the history query test to have an input.
+	s.Require().NotEmpty(aplQueryResult.SavedQueryID)
+
+	// s.EqualValues(1, aplQueryResult.Status.BlocksExamined) // FIXME(lukasmalkmus): For some reason we get "2" here?!
+	s.EqualValues(4, aplQueryResult.Status.RowsExamined)
+	s.EqualValues(4, aplQueryResult.Status.RowsMatched)
+	s.Len(aplQueryResult.Matches, 4)
+
 	// HINT(lukasmalkmus): This test initializes a new client to make sure
 	// strict decoding is never set on this method. After this test, it gets
 	// set to its previous state.
 	// This is in place because the API returns a slightly different model with
 	// a lot of empty fields which are never set for a history query. Those are
 	// not part of the client side model for ease of use.
 	s.newClient()
-	defer func() {
-		if strictDecoding {
-			optsErr := s.client.Options(axiom.SetStrictDecoding())
-			s.Require().NoError(optsErr)
-		}
-	}()
 
-	query, err := s.client.Datasets.History(s.ctx, queryResult.SavedQueryID)
+	// Give the server some time to store the queries.
+	time.Sleep(time.Second)
+
+	historyQuery, err := s.client.Datasets.History(s.ctx, queryResult.SavedQueryID)
+	s.Require().NoError(err)
+	s.Require().NotNil(historyQuery)
+
+	s.Equal(queryResult.SavedQueryID, historyQuery.ID)
+	s.Equal(query.Analytics, historyQuery.Kind)
+
+	historyQuery, err = s.client.Datasets.History(s.ctx, aplQueryResult.SavedQueryID)
 	s.Require().NoError(err)
-	s.Require().NotNil(query)
+	s.Require().NotNil(historyQuery)
+
+	s.Equal(aplQueryResult.SavedQueryID, historyQuery.ID)
+	s.Equal(query.APL, historyQuery.Kind)
 
-	s.Equal(queryResult.SavedQueryID, query.ID)
+	// Revert to strict decoding.
+	if strictDecoding {
+		optsErr := s.client.Options(axiom.SetStrictDecoding())
+		s.Require().NoError(optsErr)
+	}
 
 	// Trim the dataset down to a minimum.
 	trimResult, err := s.client.Datasets.Trim(s.ctx, s.dataset.ID, time.Second)