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

@@ -237,7 +237,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),
@@ -416,16 +416,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) {
@@ -434,6 +433,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"),
@@ -1104,6 +1104,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")
+}