feat(query): more typing

Makefile

@@ -74,6 +74,7 @@ fmt: ## Format and simplify the source code using `gofmt`
 
 .PHONY: generate
 generate: \
+	axiom/query/aggregation_string.go \
 	axiom/querylegacy/aggregation_string.go \
 	axiom/querylegacy/filter_string.go \
 	axiom/querylegacy/kind_string.go \

axiom/datasets_integration_test.go

@@ -230,7 +230,7 @@ func (s *DatasetsTestSuite) Test() {
 	startTime := now.Add(-time.Minute)
 	endTime := now.Add(time.Minute)
 
-	// Run a simple APL query.
+	// Run a simple APL query...
 	apl := fmt.Sprintf("['%s']", s.dataset.ID)
 	queryResult, err := s.client.Datasets.Query(s.ctx, apl,
 		query.SetStartTime(startTime),
@@ -401,16 +401,15 @@ func (s *DatasetsTestSuite) TestCursor() {
 	)
 	s.Require().NoError(err)
 
-	// FIXME(lukasmalkmus): Tabular results format is not yet returning the
-	// _rowID column.
-	s.T().Skip()
-
 	// HINT(lukasmalkmus): Expecting four columns: _time, _sysTime, _rowID, foo.
 	// This is only checked once for the first query result to verify the
 	// dataset scheme. The following queries will only check the results in the
 	// columns.
+	// FIXME(lukasmalkmus): Tabular results format is not yet returning the
+	// _rowID column.
 	s.Require().Len(queryResult.Tables, 1)
-	s.Require().Len(queryResult.Tables[0].Columns, 4)
+	s.Require().Len(queryResult.Tables[0].Columns, 3)
+	// s.Require().Len(queryResult.Tables[0].Columns, 4)
 	s.Require().Len(queryResult.Tables[0].Columns[0], 3)
 
 	if s.Len(queryResult.Tables, 1) {
@@ -419,6 +418,10 @@ func (s *DatasetsTestSuite) TestCursor() {
 		s.Equal("bar", queryResult.Tables[0].Columns[2][2])
 	}
 
+	// FIXME(lukasmalkmus): Tabular results format is not yet returning the
+	// _rowID column.
+	s.T().Skip()
+
 	// HINT(lukasmalkmus): In a real-world scenario, the cursor would be
 	// retrieved from the query status MinCursor or MaxCursor fields, depending
 	// on the queries sort order.

axiom/datasets_test.go

@@ -198,50 +198,50 @@ var (
 				Fields: []query.Field{
 					{
 						Name: "_time",
-						Type: "string",
+						Type: query.TypeString,
 					},
 					{
 						Name: "_sysTime",
-						Type: "string",
+						Type: query.TypeString,
 					},
 					{
 						Name: "_rowId",
-						Type: "string",
+						Type: query.TypeString,
 					},
 					{
 						Name: "agent",
-						Type: "string",
+						Type: query.TypeString,
 					},
 					{
 						Name: "bytes",
-						Type: "float64",
+						Type: query.TypeReal,
 					},
 					{
 						Name: "referrer",
-						Type: "string",
+						Type: query.TypeString,
 					},
 					{
 						Name: "remote_ip",
-						Type: "string",
+						Type: query.TypeString,
 					},
 					{
 						Name: "remote_user",
-						Type: "string",
+						Type: query.TypeString,
 					},
 					{
 						Name: "request",
-						Type: "string",
+						Type: query.TypeString,
 					},
 					{
 						Name: "response",
-						Type: "float64",
+						Type: query.TypeReal,
 					},
 					{
 						Name: "time",
-						Type: "string",
+						Type: query.TypeString,
 					},
 				},
-				Range: &query.RangeInfo{
+				Range: &query.Range{
 					Field: "_time",
 					Start: parseTimeOrPanic("2023-03-21T13:38:51.735448191Z"),
 					End:   parseTimeOrPanic("2023-03-28T13:38:51.735448191Z"),
@@ -1092,6 +1092,8 @@ func TestDatasetsService_Query(t *testing.T) {
 	assert.Equal(t, expQueryRes, res)
 }
 
+// TODO(lukasmalkmus): Add test for a query with an aggregation.
+
 func TestDatasetsService_QueryLegacy(t *testing.T) {
 	hf := func(w http.ResponseWriter, r *http.Request) {
 		assert.Equal(t, http.MethodPost, r.Method)

axiom/query/aggregation.go

@@ -0,0 +1,134 @@
+package query
+
+import (
+	"encoding/json"
+	"fmt"
+	"strings"
+)
+
+//go:generate go run golang.org/x/tools/cmd/stringer -type=AggregationOp -linecomment -output=aggregation_string.go
+
+// An AggregationOp describes the [Aggregation] operation applied on a [Field].
+type AggregationOp uint8
+
+// All available [Aggregation] operations.
+const (
+	OpUnknown AggregationOp = iota // unknown
+
+	OpCount               // count
+	OpCountIf             // countif
+	OpDistinct            // distinct
+	OpDistinctIf          // distinctif
+	OpSum                 // sum
+	OpSumIf               // sumif
+	OpAvg                 // avg
+	OpAvgIf               // avgif
+	OpMin                 // min
+	OpMinIf               // minif
+	OpMax                 // max
+	OpMaxIf               // maxif
+	OpTopk                // topk
+	OpPercentiles         // percentiles
+	OpHistogram           // histogram
+	OpStandardDeviation   // stdev
+	OpStandardDeviationIf // stdevif
+	OpVariance            // variance
+	OpVarianceIf          // varianceif
+	OpArgMin              // argmin
+	OpArgMax              // argmax
+	OpRate                // rate
+	OpPearson             // pearson_correlation
+	OpMakeSet             // makeset
+	OpMakeSetIf           // makesetif
+	OpMakeList            // makelist
+	OpMakeListIf          // makelistif
+)
+
+func aggregationOpFromString(s string) (op AggregationOp, err error) {
+	switch strings.ToLower(s) {
+	case OpCount.String():
+		op = OpCount
+	case OpCountIf.String():
+		op = OpCountIf
+	case OpDistinct.String():
+		op = OpDistinct
+	case OpDistinctIf.String():
+		op = OpDistinctIf
+	case OpSum.String():
+		op = OpSum
+	case OpSumIf.String():
+		op = OpSumIf
+	case OpAvg.String():
+		op = OpAvg
+	case OpAvgIf.String():
+		op = OpAvgIf
+	case OpMin.String():
+		op = OpMin
+	case OpMinIf.String():
+		op = OpMinIf
+	case OpMax.String():
+		op = OpMax
+	case OpMaxIf.String():
+		op = OpMaxIf
+	case OpTopk.String():
+		op = OpTopk
+	case OpPercentiles.String():
+		op = OpPercentiles
+	case OpHistogram.String():
+		op = OpHistogram
+	case OpStandardDeviation.String():
+		op = OpStandardDeviation
+	case OpStandardDeviationIf.String():
+		op = OpStandardDeviationIf
+	case OpVariance.String():
+		op = OpVariance
+	case OpVarianceIf.String():
+		op = OpVarianceIf
+	case OpArgMin.String():
+		op = OpArgMin
+	case OpArgMax.String():
+		op = OpArgMax
+	case OpRate.String():
+		op = OpRate
+	case OpPearson.String():
+		op = OpPearson
+	case OpMakeSet.String():
+		op = OpMakeSet
+	case OpMakeSetIf.String():
+		op = OpMakeSetIf
+	case OpMakeList.String():
+		op = OpMakeList
+	case OpMakeListIf.String():
+		op = OpMakeListIf
+	default:
+		return OpUnknown, fmt.Errorf("unknown aggregation operation: %s", s)
+	}
+
+	return op, nil
+}
+
+// UnmarshalJSON implements [json.Unmarshaler]. It is in place to unmarshal the
+// AggregationOp from the string representation the server returns.
+func (op *AggregationOp) UnmarshalJSON(b []byte) (err error) {
+	var s string
+	if err = json.Unmarshal(b, &s); err != nil {
+		return err
+	}
+
+	*op, err = aggregationOpFromString(s)
+
+	return err
+}
+
+// Aggregation that is applied to a [Field] in a [Table].
+type Aggregation struct {
+	// Op is the aggregation operation. If the aggregation is aliased, the alias
+	// is stored in the parent [Field.Name].
+	Op AggregationOp `json:"name"`
+	// Fields specifies the names of the fields this aggregation is computed on.
+	// E.g. ["players"] for "topk(players, 10)".
+	Fields []string `json:"fields"`
+	// Args are the non-field arguments of the aggregation, if any. E.g. "10"
+	// for "topk(players, 10)".
+	Args []any `json:"args"`
+}

axiom/query/aggregation_test.go

@@ -0,0 +1,47 @@
+package query
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestAggregationOp_Unmarshal(t *testing.T) {
+	var act struct {
+		Op AggregationOp `json:"name"`
+	}
+	err := json.Unmarshal([]byte(`{ "name": "count" }`), &act)
+	require.NoError(t, err)
+
+	assert.Equal(t, OpCount, act.Op)
+}
+
+func TestAggregationOp_String(t *testing.T) {
+	// Check outer bounds.
+	assert.Equal(t, OpUnknown, AggregationOp(0))
+	assert.Contains(t, (OpMakeListIf + 1).String(), "AggregationOp(")
+
+	for op := OpUnknown; op <= OpMakeListIf; op++ {
+		s := op.String()
+		assert.NotEmpty(t, s)
+		assert.NotContains(t, s, "AggregationOp(")
+	}
+}
+
+func TestAggregationOpFromString(t *testing.T) {
+	for op := OpCount; op <= OpMakeListIf; op++ {
+		s := op.String()
+
+		parsedOp, err := aggregationOpFromString(s)
+		if assert.NoError(t, err) {
+			assert.NotEmpty(t, s)
+			assert.Equal(t, op, parsedOp)
+		}
+	}
+
+	op, err := aggregationOpFromString("abc")
+	assert.Equal(t, OpUnknown, op)
+	assert.EqualError(t, err, "unknown aggregation operation: abc")
+}