diff --git a/.github/workflows/test_examples.yaml b/.github/workflows/test_examples.yaml index 51d4d14d..6e8567dd 100644 --- a/.github/workflows/test_examples.yaml +++ b/.github/workflows/test_examples.yaml @@ -40,6 +40,7 @@ jobs: - oteltraces - query - querylegacy + - sas # HINT(lukasmalkmus): This test would require Go 1.21 (but uses Go # 1.20 as specified in go.mod) but for the sake of simplicity we do # not test it and trust the slogx test which is the same @@ -100,7 +101,7 @@ jobs: go-version-file: go.mod - name: Setup test dataset run: | - curl -sL $(curl -s https://api.github.com/repos/axiomhq/cli/releases/tags/v0.10.0 | grep "http.*linux_amd64.tar.gz" | awk '{print $2}' | sed 's|[\"\,]*||g') | tar xzvf - --strip-components=1 --wildcards -C /usr/local/bin "axiom_*_linux_amd64/axiom" + curl -sL $(curl -s https://api.github.com/repos/axiomhq/cli/releases/tags/v0.11.1 | grep "http.*linux_amd64.tar.gz" | awk '{print $2}' | sed 's|[\"\,]*||g') | tar xzvf - --strip-components=1 --wildcards -C /usr/local/bin "axiom_*_linux_amd64/axiom" axiom dataset create -n=$AXIOM_DATASET -d="Axiom Go ${{ matrix.example }} example test" - name: Setup example if: matrix.setup diff --git a/axiom/client.go b/axiom/client.go index 7313da2b..f74b4882 100644 --- a/axiom/client.go +++ b/axiom/client.go @@ -45,7 +45,7 @@ const ( otelTracerName = "github.com/axiomhq/axiom-go/axiom" ) -var validOnlyAPITokenPaths = regexp.MustCompile(`^/v1/datasets/([^/]+/(ingest|query)|_apl)(\?.+)?$`) +var validAPITokenPaths = regexp.MustCompile(`^/v1/datasets/([^/]+/(ingest|query)|_apl)(\?.+)?$`) // service is the base service used by all Axiom API services. type service struct { @@ -191,7 +191,7 @@ func (c *Client) NewRequest(ctx context.Context, method, path string, body any) } endpoint := c.config.BaseURL().ResolveReference(rel) - if config.IsAPIToken(c.config.Token()) && !validOnlyAPITokenPaths.MatchString(endpoint.Path) { + if config.IsAPIToken(c.config.Token()) && !validAPITokenPaths.MatchString(endpoint.Path) { return nil, ErrUnprivilegedToken } diff --git a/axiom/client_test.go b/axiom/client_test.go index 44c519e8..65a4bb91 100644 --- a/axiom/client_test.go +++ b/axiom/client_test.go @@ -680,7 +680,7 @@ func TestAPITokenPathRegex(t *testing.T) { } for _, tt := range tests { t.Run(tt.input, func(t *testing.T) { - assert.Equal(t, tt.match, validOnlyAPITokenPaths.MatchString(tt.input)) + assert.Equal(t, tt.match, validAPITokenPaths.MatchString(tt.input)) }) } } diff --git a/axiom/datasets.go b/axiom/datasets.go index 25f1c8ec..37a8397f 100644 --- a/axiom/datasets.go +++ b/axiom/datasets.go @@ -20,6 +20,8 @@ import ( "github.com/axiomhq/axiom-go/axiom/ingest" "github.com/axiomhq/axiom-go/axiom/query" "github.com/axiomhq/axiom-go/axiom/querylegacy" + "github.com/axiomhq/axiom-go/axiom/sas" + "github.com/axiomhq/axiom-go/internal/config" ) //go:generate go run golang.org/x/tools/cmd/stringer -type=ContentType,ContentEncoding -linecomment -output=datasets_string.go @@ -109,7 +111,7 @@ type DatasetUpdateRequest struct { } type wrappedDataset struct { - Dataset + *Dataset // HINT(lukasmalkmus) This is some future stuff we don't yet support in this // package so we just ignore it for now. @@ -190,7 +192,7 @@ func (s *DatasetsService) List(ctx context.Context) ([]*Dataset, error) { datasets := make([]*Dataset, len(res)) for i, r := range res { - datasets[i] = &r.Dataset + datasets[i] = r.Dataset } return datasets, nil @@ -213,7 +215,7 @@ func (s *DatasetsService) Get(ctx context.Context, id string) (*Dataset, error) return nil, spanError(span, err) } - return &res.Dataset, nil + return res.Dataset, nil } // Create a dataset with the given properties. @@ -229,7 +231,7 @@ func (s *DatasetsService) Create(ctx context.Context, req DatasetCreateRequest) return nil, spanError(span, err) } - return &res.Dataset, nil + return res.Dataset, nil } // Update the dataset identified by the given id with the given properties. @@ -250,7 +252,7 @@ func (s *DatasetsService) Update(ctx context.Context, id string, req DatasetUpda return nil, spanError(span, err) } - return &res.Dataset, nil + return res.Dataset, nil } // Delete the dataset identified by the given id. @@ -260,7 +262,7 @@ func (s *DatasetsService) Delete(ctx context.Context, id string) error { )) defer span.End() - path, err := url.JoinPath(s.basePath, "/", id) + path, err := url.JoinPath(s.basePath, id) if err != nil { return spanError(span, err) } @@ -584,20 +586,31 @@ func (s *DatasetsService) Query(ctx context.Context, apl string, options ...quer ctx, span := s.client.trace(ctx, "Datasets.Query", trace.WithAttributes( attribute.String("axiom.param.apl", apl), - attribute.String("axiom.param.start_time", opts.StartTime.String()), - attribute.String("axiom.param.end_time", opts.EndTime.String()), + attribute.String("axiom.param.start_time", opts.StartTime), + attribute.String("axiom.param.end_time", opts.EndTime), attribute.String("axiom.param.cursor", opts.Cursor), + attribute.Bool("axiom.param.include_cursor", opts.IncludeCursor), )) defer span.End() // The only query parameters supported can be hardcoded as they are not // configurable as of now. queryParams := struct { + *sas.Options + Format string `url:"format"` }{ Format: "legacy", // Hardcode legacy APL format for now. } + if t := s.client.config.Token(); config.IsSharedAccessSignature(t) { + options, err := sas.Decode(t) + if err != nil { + return nil, spanError(span, err) + } + queryParams.Options = &options + } + path, err := url.JoinPath(s.basePath, "_apl") if err != nil { return nil, spanError(span, err) @@ -614,6 +627,11 @@ func (s *DatasetsService) Query(ctx context.Context, apl string, options ...quer return nil, spanError(span, err) } + if config.IsSharedAccessSignature(s.client.config.Token()) { + req.Header.Del(headerAuthorization) + req.Header.Del(headerOrganizationID) + } + var res aplQueryResponse if _, err = s.client.Do(req, &res); err != nil { return nil, spanError(span, err) @@ -632,7 +650,10 @@ func (s *DatasetsService) Query(ctx context.Context, apl string, options ...quer // the future. Use [DatasetsService.Query] instead. func (s *DatasetsService) QueryLegacy(ctx context.Context, id string, q querylegacy.Query, opts querylegacy.Options) (*querylegacy.Result, error) { ctx, span := s.client.trace(ctx, "Datasets.QueryLegacy", trace.WithAttributes( - attribute.String("axiom.dataset_id", id), + attribute.String("axiom.param.dataset_id", id), + attribute.String("axiom.param.streaming_duration", opts.StreamingDuration.String()), + attribute.Bool("axiom.param.no_cache", opts.NoCache), + attribute.String("axiom.param.save_kind", opts.SaveKind.String()), )) defer span.End() diff --git a/axiom/doc.go b/axiom/doc.go index f822446a..bae2a356 100644 --- a/axiom/doc.go +++ b/axiom/doc.go @@ -7,6 +7,7 @@ // import "github.com/axiomhq/axiom-go/axiom/otel" // When using OpenTelemetry // import "github.com/axiomhq/axiom-go/axiom/query" // When constructing APL queries // import "github.com/axiomhq/axiom-go/axiom/querylegacy" // When constructing legacy queries +// import "github.com/axiomhq/axiom-go/axiom/sas" // When using shared access // // Construct a new Axiom client, then use the various services on the client to // access different parts of the Axiom API. The package automatically takes its diff --git a/axiom/limit.go b/axiom/limit.go index 18a9e513..1a2e173a 100644 --- a/axiom/limit.go +++ b/axiom/limit.go @@ -69,12 +69,12 @@ func limitScopeFromString(s string) (ls LimitScope, err error) { type Limit struct { // Scope a limit is enforced for. Only present on rate limited requests. Scope LimitScope - // The maximum limit a client is limited to for a specified time window + // The maximum limit a client is limited to for a specified time range // which resets at the time indicated by [Limit.Reset]. Limit uint64 // The remaining count towards the maximum limit. Remaining uint64 - // The time at which the current limit time window will reset. + // The time at which the current limit time range will reset. Reset time.Time limitType limitType diff --git a/axiom/orgs.go b/axiom/orgs.go index 8b34d014..8321936a 100644 --- a/axiom/orgs.go +++ b/axiom/orgs.go @@ -188,6 +188,17 @@ func (l *License) UnmarshalJSON(b []byte) error { return nil } +// SigningKeys are the signing keys used to sign shared access tokens that +// can be used by a third party to run queries on behalf of the organization. +// They can be rotated. +type SigningKeys struct { + // Primary signing key. Gets rotated to the secondary signing key after + // rotation. + Primary string `json:"primary"` + // Secondary signing key. Gets rotated out. + Secondary string `json:"secondary"` +} + // Organization represents an organization. type Organization struct { // ID is the unique ID of the organization. @@ -220,7 +231,7 @@ type Organization struct { } type wrappedOrganization struct { - Organization + *Organization // HINT(lukasmalkmus): Ignore these fields because they do not provide any // value to the user. @@ -247,7 +258,7 @@ func (s *OrganizationsService) List(ctx context.Context) ([]*Organization, error organizations := make([]*Organization, len(res)) for i, r := range res { - organizations[i] = &r.Organization + organizations[i] = r.Organization } return organizations, nil @@ -256,7 +267,7 @@ func (s *OrganizationsService) List(ctx context.Context) ([]*Organization, error // Get an organization by id. func (s *OrganizationsService) Get(ctx context.Context, id string) (*Organization, error) { ctx, span := s.client.trace(ctx, "Organizations.Get", trace.WithAttributes( - attribute.String("axiom.dataset_id", id), + attribute.String("axiom.organization_id", id), )) defer span.End() @@ -270,5 +281,47 @@ func (s *OrganizationsService) Get(ctx context.Context, id string) (*Organizatio return nil, spanError(span, err) } - return &res.Organization, nil + return res.Organization, nil +} + +// ViewSigningKeys views the shared access token signing keys for the +// organization identified by the given id. +func (s *OrganizationsService) ViewSigningKeys(ctx context.Context, id string) (*SigningKeys, error) { + ctx, span := s.client.trace(ctx, "Organizations.ViewSigningKeys", trace.WithAttributes( + attribute.String("axiom.organization_id", id), + )) + defer span.End() + + path, err := url.JoinPath(s.basePath, id, "keys") + if err != nil { + return nil, spanError(span, err) + } + + var res SigningKeys + if err := s.client.Call(ctx, http.MethodGet, path, nil, &res); err != nil { + return nil, spanError(span, err) + } + + return &res, nil +} + +// RotateSigningKeys rotates the shared access token signing keys for the +// organization identified by the given id. +func (s *OrganizationsService) RotateSigningKeys(ctx context.Context, id string) (*SigningKeys, error) { + ctx, span := s.client.trace(ctx, "Organizations.RotateSigningKeys", trace.WithAttributes( + attribute.String("axiom.organization_id", id), + )) + defer span.End() + + path, err := url.JoinPath(s.basePath, id, "rotate-keys") + if err != nil { + return nil, spanError(span, err) + } + + var res SigningKeys + if err := s.client.Call(ctx, http.MethodPut, path, nil, &res); err != nil { + return nil, spanError(span, err) + } + + return &res, nil } diff --git a/axiom/orgs_integration_test.go b/axiom/orgs_integration_test.go index e73aee3c..7fa71f43 100644 --- a/axiom/orgs_integration_test.go +++ b/axiom/orgs_integration_test.go @@ -31,4 +31,24 @@ func (s *OrganizationsTestSuite) Test() { s.Require().NotNil(organization) s.Contains(organizations, organization) + + keys, err := s.client.Organizations.ViewSigningKeys(s.ctx, organization.ID) + s.Require().NoError(err) + s.Require().NotNil(keys) + + s.NotEmpty(keys.Primary) + s.NotEmpty(keys.Secondary) + s.NotEqual(keys.Primary, keys.Secondary) + + // Rotate the signing keys on the organization and make sure the new keys + // are returned. + oldPrimaryKey, oldSecondaryKey := keys.Primary, keys.Secondary + keys, err = s.client.Organizations.RotateSigningKeys(s.ctx, organization.ID) + s.Require().NoError(err) + s.Require().NotNil(keys) + + s.NotEqual(oldPrimaryKey, keys.Primary) + s.NotEqual(oldSecondaryKey, keys.Secondary) + s.NotEqual(oldSecondaryKey, keys.Primary) + s.Equal(oldPrimaryKey, keys.Secondary) } diff --git a/axiom/orgs_test.go b/axiom/orgs_test.go index 1abec3be..c42cdded 100644 --- a/axiom/orgs_test.go +++ b/axiom/orgs_test.go @@ -194,6 +194,56 @@ func TestOrganizationsService_Get(t *testing.T) { assert.Equal(t, exp, res) } +func TestOrganizationsService_ViewSigningKeys(t *testing.T) { + exp := &SigningKeys{ + Primary: "75bb5815-8459-4b6e-a08f-1eb8058db44e", + Secondary: "6205e228-f8ed-4265-bee8-058a9b1091db", + } + + hf := func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, http.MethodGet, r.Method) + + w.Header().Set("Content-Type", mediaTypeJSON) + _, err := fmt.Fprint(w, `{ + "primary": "75bb5815-8459-4b6e-a08f-1eb8058db44e", + "secondary": "6205e228-f8ed-4265-bee8-058a9b1091db" + }`) + assert.NoError(t, err) + } + + client := setup(t, "/v1/orgs/axiom/keys", hf) + + res, err := client.Organizations.ViewSigningKeys(context.Background(), "axiom") + require.NoError(t, err) + + assert.Equal(t, exp, res) +} + +func TestOrganizationsService_RotateSigningKeys(t *testing.T) { + exp := &SigningKeys{ + Primary: "75bb5815-8459-4b6e-a08f-1eb8058db44e", + Secondary: "6205e228-f8ed-4265-bee8-058a9b1091db", + } + + hf := func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, http.MethodPut, r.Method) + + w.Header().Set("Content-Type", mediaTypeJSON) + _, err := fmt.Fprint(w, `{ + "primary": "75bb5815-8459-4b6e-a08f-1eb8058db44e", + "secondary": "6205e228-f8ed-4265-bee8-058a9b1091db" + }`) + assert.NoError(t, err) + } + + client := setup(t, "/v1/orgs/axiom/rotate-keys", hf) + + res, err := client.Organizations.RotateSigningKeys(context.Background(), "axiom") + require.NoError(t, err) + + assert.Equal(t, exp, res) +} + func TestPlan_Marshal(t *testing.T) { exp := `{ "plan": "personal" diff --git a/axiom/query/options.go b/axiom/query/options.go index 242bbc13..7cbd5f02 100644 --- a/axiom/query/options.go +++ b/axiom/query/options.go @@ -1,13 +1,15 @@ package query -import "time" +import ( + "time" +) // Options specifies the optional parameters for a query. type Options struct { // StartTime for the interval to query. - StartTime time.Time `json:"startTime,omitempty"` + StartTime string `json:"startTime,omitempty"` // EndTime of the interval to query. - EndTime time.Time `json:"endTime,omitempty"` + EndTime string `json:"endTime,omitempty"` // Cursor to use for pagination. When used, don't specify new start and end // times but rather use the start and end times of the query that returned // the cursor that will be used. @@ -27,15 +29,15 @@ type Option func(*Options) // SetStartTime specifies the start time of the query interval. When also using // [SetCursor], please make sure to use the start time of the query that // returned the cursor that will be used. -func SetStartTime(startTime time.Time) Option { - return func(o *Options) { o.StartTime = startTime } +func SetStartTime[T time.Time | string](startTime T) Option { + return func(o *Options) { o.StartTime = timeOrStringToString(startTime) } } // SetEndTime specifies the end time of the query interval. When also using // [SetCursor], please make sure to use the end time of the query that returned // the cursor that will be used. -func SetEndTime(endTime time.Time) Option { - return func(o *Options) { o.EndTime = endTime } +func SetEndTime[T time.Time | string](endTime T) Option { + return func(o *Options) { o.EndTime = timeOrStringToString(endTime) } } // SetCursor specifies the cursor of the query. If include is set to true the @@ -65,3 +67,13 @@ func SetVariable(name string, value any) Option { func SetVariables(variables map[string]any) Option { return func(o *Options) { o.Variables = variables } } + +func timeOrStringToString[T time.Time | string](t T) string { + switch t := any(t).(type) { + case time.Time: + return t.Format(time.RFC3339Nano) + case string: + return t + } + panic("time is neither time.Time nor string") +} diff --git a/axiom/sas/doc.go b/axiom/sas/doc.go new file mode 100644 index 00000000..f63222d1 --- /dev/null +++ b/axiom/sas/doc.go @@ -0,0 +1,20 @@ +// Package sas implements functionality for creating and verifying shared access +// signatures (SAS) and shared access tokens (SAT) as well as using them to +// query Axiom datasets. A SAS grants querying capabilities to a dataset for a +// given time range and with a global filter applied on behalf of an +// organization. A SAS is an URL query string composed of a set of query +// parameters that make up the payload for a signature and the cryptographic +// signature itself. That cryptographic signature is called SAT. +// +// Usage: +// +// import "github.com/axiomhq/axiom-go/axiom/sas" +// +// To create a SAS string, that can be attached to a query request, use the +// high-level [Create] function. The returned string is an already url encoded +// query string. +// +// To create a SAT string for a set of values that make up a signature, use +// the low-level [CreateToken] function. The returned string is an already +// base64 encoded string. +package sas diff --git a/axiom/sas/options.go b/axiom/sas/options.go new file mode 100644 index 00000000..6836d70d --- /dev/null +++ b/axiom/sas/options.go @@ -0,0 +1,74 @@ +package sas + +import ( + "errors" + "net/url" + + "github.com/google/go-querystring/query" +) + +// The parameter names for the shared access signature query string. +const ( + queryOrgID = "oi" + queryDataset = "dt" + queryFilter = "fl" + queryMinStartTime = "mst" + queryMaxEndTime = "met" + queryToken = "tk" +) + +// Options are the url query parameters used to authenticate a query request. +type Options struct { + Params + + // Token is the signature created from the other fields in the options. + Token string `url:"tk"` +} + +// Decode decodes the given signature into a set of options. +func Decode(signature string) (Options, error) { + q, err := url.ParseQuery(signature) + if err != nil { + return Options{}, err + } + + options := Options{ + Params: Params{ + OrganizationID: q.Get(queryOrgID), + Dataset: q.Get(queryDataset), + Filter: q.Get(queryFilter), + MinStartTime: q.Get(queryMinStartTime), + MaxEndTime: q.Get(queryMaxEndTime), + }, + Token: q.Get(queryToken), + } + + // Validate that the params are valid and the token is present. + if err := options.Params.Validate(); err != nil { + return options, err + } else if options.Token == "" { + return options, errors.New("missing token") + } + + return options, nil +} + +// Encode encodes the options into a url query string. +func (o Options) Encode() (string, error) { + q, err := query.Values(o) + if err != nil { + return "", err + } + + // Although officially there is no limit specified by RFC 2616, many + // security protocols and recommendations state that maxQueryStrings on a + // server should be set to a maximum character limit of 1024. While the + // entire URL, including the querystring, should be set to a max of 2048 + // characters. + s := q.Encode() + if len(s) > 1023 { // 1024 - 1 for '?' + return "", errors.New("signature too long") + } + + return s, nil +} diff --git a/axiom/sas/options_test.go b/axiom/sas/options_test.go new file mode 100644 index 00000000..16146bcb --- /dev/null +++ b/axiom/sas/options_test.go @@ -0,0 +1,44 @@ +package sas + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestOptions_Decode(t *testing.T) { + options, err := Decode("dt=logs&fl=customer+%3D%3D+%22vercel%22&met=now&mst=ago%281h%29&oi=axiom&tk=0M41vwyiTVtAqW_aw8ZaIgayOlxnSwtFoFbywuQ-VBc%3D") + require.NoError(t, err) + require.NotEmpty(t, options) + + assert.Equal(t, Options{ + Params: Params{ + OrganizationID: "axiom", + Dataset: "logs", + Filter: `customer == "vercel"`, + MinStartTime: "ago(1h)", + MaxEndTime: "now", + }, + Token: "0M41vwyiTVtAqW_aw8ZaIgayOlxnSwtFoFbywuQ-VBc=", + }, options) +} + +func TestOptions_Encode(t *testing.T) { + options := Options{ + Params: Params{ + OrganizationID: "axiom", + Dataset: "logs", + Filter: `customer == "vercel"`, + MinStartTime: "ago(1h)", + MaxEndTime: "now", + }, + Token: "0M41vwyiTVtAqW_aw8ZaIgayOlxnSwtFoFbywuQ-VBc=", + } + + s, err := options.Encode() + require.NoError(t, err) + require.NotEmpty(t, s) + + assert.Equal(t, "dt=logs&fl=customer+%3D%3D+%22vercel%22&met=now&mst=ago%281h%29&oi=axiom&tk=0M41vwyiTVtAqW_aw8ZaIgayOlxnSwtFoFbywuQ-VBc%3D", s) +} diff --git a/axiom/sas/params.go b/axiom/sas/params.go new file mode 100644 index 00000000..409b74a3 --- /dev/null +++ b/axiom/sas/params.go @@ -0,0 +1,79 @@ +package sas + +import ( + "crypto/hmac" + "crypto/sha256" + "encoding/base64" + "errors" + "fmt" + "strings" + + "github.com/google/uuid" +) + +// Params represents the parameters for creating a shared access token or a +// shared access signatur for a query request. +type Params struct { + // OrganizationID is the ID of the organization the token and signature is + // valid for. + OrganizationID string `url:"oi"` + // Dataset name the token and signature is valid for. + Dataset string `url:"dt"` + // Filter is the top-level query filter to apply to all query requests the + // token and signature is valid for. Must be a valid APL filter expression. + Filter string `url:"fl"` + // MinStartTime is the earliest query start time the token and signature is + // valid for. Can be a timestamp or an APL expression. + MinStartTime string `url:"mst"` + // MaxEndTime is the latest query end time the token and signature is valid + // for. Can be a timestamp or an APL expression. + MaxEndTime string `url:"met"` +} + +// Validate makes sure that all query parameters are provided. +func (p Params) Validate() error { + if p.OrganizationID == "" { + return errors.New("organization ID is required") + } else if p.Dataset == "" { + return errors.New("dataset is required") + } else if p.Filter == "" { + return errors.New("filter is required") + } else if p.MinStartTime == "" { + return errors.New("minimum start time is required") + } else if p.MaxEndTime == "" { + return errors.New("maximum end time is required") + } + return nil +} + +// sign the parameters with the given key and returns a base64 encoded token. +func (p Params) sign(key string) (string, error) { + k, err := uuid.Parse(key) + if err != nil { + return "", fmt.Errorf("invalid key: %s", err) + } + + var ( + pl = buildSignaturePayload(p) + h = hmac.New(sha256.New, k[:]) + ) + if _, err = h.Write([]byte(pl)); err != nil { + return "", fmt.Errorf("computing hmac: %s", err) + } + + return base64.URLEncoding.EncodeToString(h.Sum(nil)), nil +} + +// buildSignaturePayload builds the payload for a shared access token. The +// format is a simple, newline decoded string composed of the following values +// in that order: organization ID, dataset name, filter, minimum start time, +// maximum end time. +func buildSignaturePayload(params Params) string { + return strings.Join([]string{ + params.OrganizationID, + params.Dataset, + params.Filter, + params.MinStartTime, + params.MaxEndTime, + }, "\n") +} diff --git a/axiom/sas/params_test.go b/axiom/sas/params_test.go new file mode 100644 index 00000000..56a28d25 --- /dev/null +++ b/axiom/sas/params_test.go @@ -0,0 +1,56 @@ +package sas + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestParams_Validate(t *testing.T) { + var params Params + + err := params.Validate() + assert.EqualError(t, err, "organization ID is required") + + params.OrganizationID = "axiom" + + err = params.Validate() + assert.EqualError(t, err, "dataset is required") + + params.Dataset = "logs" + + err = params.Validate() + assert.EqualError(t, err, "filter is required") + + params.Filter = `customer == "vercel"` + + err = params.Validate() + assert.EqualError(t, err, "minimum start time is required") + + params.MinStartTime = "ago(1h)" + + err = params.Validate() + assert.EqualError(t, err, "maximum end time is required") + + params.MaxEndTime = "now()" + + err = params.Validate() + assert.NoError(t, err) +} + +func TestParams_sign(t *testing.T) { + params := Params{ + OrganizationID: "axiom", + Dataset: "logs", + Filter: `customer == "vercel"`, + MinStartTime: "ago(1h)", + MaxEndTime: "now", + } + + token, err := params.sign(testKeyStr) + require.NoError(t, err) + require.NotEmpty(t, token) + + assert.Equal(t, "0M41vwyiTVtAqW_aw8ZaIgayOlxnSwtFoFbywuQ-VBc=", token) +} diff --git a/axiom/sas/sas.go b/axiom/sas/sas.go new file mode 100644 index 00000000..c51cc5da --- /dev/null +++ b/axiom/sas/sas.go @@ -0,0 +1,65 @@ +package sas + +import ( + "fmt" + "net/http" + + "github.com/google/go-querystring/query" +) + +// Create creates a shared access signature using the given signing key and +// valid for the given parameters. The returned string is a query string that +// can be attached to a URL. +func Create(key string, params Params) (string, error) { + options, err := CreateOptions(key, params) + if err != nil { + return "", err + } + return options.Encode() +} + +// CreateOptions creates the options that compose a shared access signature +// signed with the given key and valid for the given parameters. It can be +// encoded into a query string by calling [Options.Encode] and attached to a +// URL. This operation would be equivalent to calling [Create]. +func CreateOptions(key string, params Params) (Options, error) { + token, err := CreateToken(key, params) + if err != nil { + return Options{}, err + } + + return Options{ + Params: params, + Token: token, + }, nil +} + +// CreateToken creates a shared access token signed with the given key and valid +// for the given parameters. +// +// This function is only useful if the intention is to create the shared access +// signature manually and without the help of [Create]. +func CreateToken(key string, params Params) (string, error) { + if err := params.Validate(); err != nil { + return "", fmt.Errorf("invalid parameters: %s", err) + } + return params.sign(key) +} + +// Attach attaches the given options to the given request as a query string. It +// retains existing query parameters unless they are overwritten by the key of +// one of the options. +func Attach(req *http.Request, options Options) error { + q, err := query.Values(options) + if err != nil { + return err + } + + qc := req.URL.Query() + for k := range q { + qc.Set(k, q.Get(k)) + } + req.URL.RawQuery = qc.Encode() + + return nil +} diff --git a/axiom/sas/sas_integration_test.go b/axiom/sas/sas_integration_test.go new file mode 100644 index 00000000..2c4560ad --- /dev/null +++ b/axiom/sas/sas_integration_test.go @@ -0,0 +1,161 @@ +package sas_test + +import ( + "context" + "fmt" + "os" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/axiomhq/axiom-go/axiom" + "github.com/axiomhq/axiom-go/axiom/ingest" + "github.com/axiomhq/axiom-go/axiom/query" + "github.com/axiomhq/axiom-go/axiom/sas" + "github.com/axiomhq/axiom-go/internal/config" + "github.com/axiomhq/axiom-go/internal/test/testhelper" +) + +const ingestData = `[ + { + "time": "17/May/2015:08:05:30 +0000", + "remote_ip": "93.180.71.1", + "remote_user": "-", + "request": "GET /downloads/product_1 HTTP/1.1", + "response": 304, + "bytes": 0, + "referrer": "-", + "agent": "Debian APT-HTTP/1.3 (0.8.16~exp12ubuntu10.21)" + }, + { + "time": "17/May/2015:08:05:31 +0000", + "remote_ip": "93.180.71.2", + "remote_user": "-", + "request": "GET /downloads/product_1 HTTP/1.1", + "response": 304, + "bytes": 0, + "referrer": "-", + "agent": "Debian APT-HTTP/1.3 (0.8.16~exp12ubuntu10.21)" + } +]` + +func TestSAS(t *testing.T) { + cfg := config.Default() + if err := cfg.IncorporateEnvironment(); err != nil { + t.Fatal(err) + } else if err = cfg.Validate(); err != nil { + t.Fatal(err) + } + + datasetSuffix := os.Getenv("AXIOM_DATASET_SUFFIX") + if datasetSuffix == "" { + datasetSuffix = "local" + } + + // Clear the environment to avoid unexpected behavior. + testhelper.SafeClearEnv(t) + + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + t.Cleanup(cancel) + + userAgent := fmt.Sprintf("axiom-go-sas-integration-test/%s", datasetSuffix) + client, err := axiom.NewClient( + axiom.SetNoEnv(), + axiom.SetURL(cfg.BaseURL().String()), + axiom.SetToken(cfg.Token()), + axiom.SetOrganizationID(cfg.OrganizationID()), + axiom.SetUserAgent(userAgent), + ) + require.NoError(t, err) + + // Get some info on the user that runs the test. + testUser, err := client.Users.Current(ctx) + require.NoError(t, err) + + t.Logf("using account %q", testUser.Name) + + // Create the dataset to use. + dataset, err := client.Datasets.Create(ctx, axiom.DatasetCreateRequest{ + Name: fmt.Sprintf("test-axiom-go-sas-%s", datasetSuffix), + Description: "This is a test dataset for adapter integration tests.", + }) + require.NoError(t, err) + t.Cleanup(func() { + // Restore token. + optsErr := client.Options(axiom.SetToken(cfg.Token())) + require.NoError(t, optsErr) + + teardownCtx := teardownContext(t, time.Second*15) + deleteErr := client.Datasets.Delete(teardownCtx, dataset.ID) + assert.NoError(t, deleteErr) + }) + + // Ingest some test data. + ingestRes, err := client.Ingest(ctx, dataset.ID, strings.NewReader(ingestData), axiom.JSON, axiom.Identity) + require.NoError(t, err) + require.EqualValues(t, 2, ingestRes.Ingested) + + // Ingest one event each with an earlier timestamp that will break the query + // test if the signatures time range is not respected. + now := time.Now() + then := now.Add(-time.Hour) + ingestRes, err = client.IngestEvents(ctx, dataset.ID, []axiom.Event{ + { + ingest.TimestampField: then.Format(time.RFC3339Nano), + "remote_ip": "93.180.71.1", + }, + { + ingest.TimestampField: then.Format(time.RFC3339Nano), + "remote_ip": "93.180.71.2", + }, + }) + require.NoError(t, err) + require.EqualValues(t, 2, ingestRes.Ingested) + + // List the keys we're going to use for creating the SAS. + keys, err := client.Organizations.ViewSigningKeys(ctx, cfg.OrganizationID()) + require.NoError(t, err) + + signature, err := sas.Create(keys.Primary, sas.Params{ + OrganizationID: cfg.OrganizationID(), + Dataset: dataset.ID, + Filter: `remote_ip == "93.180.71.1"`, + MinStartTime: "ago(5m)", + MaxEndTime: "now", + }) + require.NoError(t, err) + require.NotEmpty(t, signature) + + // Now use the SAS for authentication. + err = client.Options(axiom.SetToken(signature)) + require.NoError(t, err) + + queryRes, err := client.Query(ctx, fmt.Sprintf("['%s'] | count", dataset.ID), + query.SetStartTime("ago(1m)"), + query.SetEndTime("now"), + ) + require.NoError(t, err) + require.NotEmpty(t, queryRes) + + assert.EqualValues(t, 1, queryRes.Buckets.Totals[0].Aggregations[0].Value) + + // Now try to query and bypass the timerange via an APL 'where' statement. + queryRes, err = client.Query(ctx, fmt.Sprintf("['%s'] | where _time > ago(1d) | count", dataset.ID), + query.SetStartTime("ago(1m)"), + query.SetEndTime("now"), + ) + require.NoError(t, err) + + assert.EqualValues(t, 1, queryRes.Buckets.Totals[0].Aggregations[0].Value) +} + +func teardownContext(t *testing.T, timeout time.Duration) context.Context { + t.Helper() + + ctx, cancel := context.WithTimeout(context.Background(), timeout) + t.Cleanup(cancel) + return ctx +} diff --git a/axiom/sas/sas_test.go b/axiom/sas/sas_test.go new file mode 100644 index 00000000..46f012fd --- /dev/null +++ b/axiom/sas/sas_test.go @@ -0,0 +1,84 @@ +package sas + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +var testKeyStr = "aba84eee-3935-4b51-8aae-2c41b8693016" + +func TestCreate(t *testing.T) { + signature, err := Create(testKeyStr, Params{ + OrganizationID: "axiom", + Dataset: "logs", + Filter: `customer == "vercel"`, + MinStartTime: "ago(1h)", + MaxEndTime: "now", + }) + require.NoError(t, err) + require.NotEmpty(t, signature) + + assert.Equal(t, "dt=logs&fl=customer+%3D%3D+%22vercel%22&met=now&mst=ago%281h%29&oi=axiom&tk=0M41vwyiTVtAqW_aw8ZaIgayOlxnSwtFoFbywuQ-VBc%3D", signature) +} + +func TestCreateOptions(t *testing.T) { + options, err := CreateOptions(testKeyStr, Params{ + OrganizationID: "axiom", + Dataset: "logs", + Filter: `customer == "vercel"`, + MinStartTime: "ago(1h)", + MaxEndTime: "now", + }) + require.NoError(t, err) + require.NotEmpty(t, options) + + assert.Equal(t, Options{ + Params: Params{ + OrganizationID: "axiom", + Dataset: "logs", + Filter: `customer == "vercel"`, + MinStartTime: "ago(1h)", + MaxEndTime: "now", + }, + Token: "0M41vwyiTVtAqW_aw8ZaIgayOlxnSwtFoFbywuQ-VBc=", + }, options) +} + +func TestCreateToken(t *testing.T) { + token, err := CreateToken(testKeyStr, Params{ + OrganizationID: "axiom", + Dataset: "logs", + Filter: `customer == "vercel"`, + MinStartTime: "ago(1h)", + MaxEndTime: "now", + }) + require.NoError(t, err) + require.NotEmpty(t, token) + + assert.Equal(t, "0M41vwyiTVtAqW_aw8ZaIgayOlxnSwtFoFbywuQ-VBc=", token) +} + +func TestAttach(t *testing.T) { + options, err := CreateOptions(testKeyStr, Params{ + OrganizationID: "axiom", + Dataset: "logs", + Filter: `customer == "vercel"`, + MinStartTime: "ago(1h)", + MaxEndTime: "now", + }) + require.NoError(t, err) + + req := httptest.NewRequest(http.MethodPost, "/v1/datasets/_apl", nil) + + err = Attach(req, options) + require.NoError(t, err) + + parsedOptions, err := Decode(req.URL.RawQuery) + require.NoError(t, err) + + assert.Equal(t, options, parsedOptions) +} diff --git a/doc.go b/doc.go index 7f772880..40736254 100644 --- a/doc.go +++ b/doc.go @@ -7,6 +7,7 @@ // import "github.com/axiomhq/axiom-go/axiom/otel" // When using OpenTelemetry // import "github.com/axiomhq/axiom-go/axiom/query" // When constructing APL queries // import "github.com/axiomhq/axiom-go/axiom/querylegacy" // When constructing legacy queries +// import "github.com/axiomhq/axiom-go/axiom/sas" // When using shared access // // Construct a new Axiom client, then use the various services on the client to // access different parts of the Axiom API. The package automatically takes its diff --git a/examples/sas/main.go b/examples/sas/main.go new file mode 100644 index 00000000..776e997c --- /dev/null +++ b/examples/sas/main.go @@ -0,0 +1,121 @@ +// The purpose of this example is to show how to create a shared access +// signature (SAS) for a dataset and use it to query that dataset via an +// ordinary HTTP request. +package main + +import ( + "context" + "encoding/json" + "fmt" + "log" + "net/http" + "net/url" + "os" + "strings" + + "github.com/axiomhq/axiom-go/axiom" + "github.com/axiomhq/axiom-go/axiom/query" + "github.com/axiomhq/axiom-go/axiom/sas" +) + +func main() { + // Export "AXIOM_DATASET" in addition to the required environment variables. + + dataset := os.Getenv("AXIOM_DATASET") + if dataset == "" { + log.Fatal("AXIOM_DATASET is required") + } + + // 1. Initialize the Axiom API client. + client, err := axiom.NewClient() + if err != nil { + log.Fatal(err) + } + + // 2. Ingest some events with different values for the "team_id" field. + events := []axiom.Event{ + {"team_id": "a", "value": 1}, + {"team_id": "a", "value": 2}, + {"team_id": "b", "value": 4}, + {"team_id": "b", "value": 5}, + } + ingestRes, err := client.IngestEvents(context.Background(), dataset, events) + if err != nil { + log.Fatal(err) + } else if fails := len(ingestRes.Failures); fails > 0 { + log.Fatalf("Ingestion of %d events failed", fails) + } + + // 3. List the organizations signing keys. + keys, err := client.Organizations.ViewSigningKeys(context.Background(), os.Getenv("AXIOM_ORG_ID")) + if err != nil { + log.Fatal(err) + } + + // 4. Create a shared access signature that limits query access to events + // by the "team_id" field to only those with the value "a". The queries time + // range is limited to the last 5 minutes. + options, err := sas.CreateOptions(keys.Primary, sas.Params{ + OrganizationID: os.Getenv("AXIOM_ORG_ID"), + Dataset: dataset, + Filter: `team_id == "a"`, + MinStartTime: "ago(5m)", + MaxEndTime: "now", + }) + if err != nil { + log.Fatal(err) + } + + // ❗From here on, assume the code is executed by a non-Axiom user that is + // delegated query access via the SAS handed to him on behalf of the + // organization. + + // 5. Construct the Axiom API URL for the APL query endpoint. + u := os.Getenv("AXIOM_URL") + if u == "" { + u = "https://api.axiom.co" + } + queryURL, err := url.JoinPath(u, "/v1/datasets/_apl") + if err != nil { + log.Fatal(err) + } + queryURL += "?format=legacy" // Currently, must be set to "legacy". + + // 6. Construct the APL query request. + r := fmt.Sprintf(`{ + "apl": "['%s'] | count", + "startTime": "ago(1m)", + "endTime": "now" + }`, dataset) + req, err := http.NewRequest(http.MethodPost, queryURL, strings.NewReader(r)) + if err != nil { + log.Fatal(err) + } + req.Header.Set("Content-Type", "application/json") + + // 7. Attach the SAS to the request. + if err = sas.Attach(req, options); err != nil { + log.Fatal(err) + } + + // 8. Execute the request. + resp, err := http.DefaultClient.Do(req) + if err != nil { + log.Fatal(err) + } + defer resp.Body.Close() + + // 9. Check the response status code. + if code := resp.StatusCode; code != http.StatusOK { + log.Fatalf("unexpected status code: %d (%s)", code, http.StatusText(code)) + } + + // 10. Decode the response. + var res query.Result + if err = json.NewDecoder(resp.Body).Decode(&res); err != nil { + log.Fatal(err) + } + + // 11. Print the count, which should be "3". + fmt.Println(res.Buckets.Totals[0].Aggregations[0].Value) +} diff --git a/go.mod b/go.mod index 2fda22dd..297b39c9 100644 --- a/go.mod +++ b/go.mod @@ -7,6 +7,7 @@ require ( github.com/cenkalti/backoff/v4 v4.2.1 github.com/golangci/golangci-lint v1.54.2 github.com/google/go-querystring v1.1.0 + github.com/google/uuid v1.3.0 github.com/klauspost/compress v1.16.7 github.com/schollz/progressbar/v3 v3.13.1 github.com/sirupsen/logrus v1.9.3 diff --git a/go.sum b/go.sum index 53b4d359..1b853b0b 100644 --- a/go.sum +++ b/go.sum @@ -303,6 +303,8 @@ github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 h1:El6M4kTTCOh6aBiKaU github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510/go.mod h1:pupxD2MaaD3pAXIBCelhxNneeOaAeabZDe5s4K6zSpQ= github.com/google/uuid v1.1.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/google/uuid v1.1.2/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I= +github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/googleapis/gax-go/v2 v2.0.4/go.mod h1:0Wqv26UfaUD9n4G6kQubkQ+KchISgw+vpHVxEJEs9eg= github.com/googleapis/gax-go/v2 v2.0.5/go.mod h1:DWXyrwAJ9X0FpwwEdw+IPEYBICEFu5mhpdKc/us6bOk= github.com/googleapis/google-cloud-go-testing v0.0.0-20200911160855-bcd43fbb19e8/go.mod h1:dvDLG8qkwmyD9a/MJJN3XJcT3xFxOKAvTZGvuZmac9g= diff --git a/internal/config/config.go b/internal/config/config.go index 7f5f96f0..aa17fc42 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -8,7 +8,7 @@ import ( // Config is the configuration for Axiom related functionality. It should never // be created manually but always via the [Default] function. type Config struct { - // baseURL of the Axiom instance. Defaults to [CloudURL]. + // baseURL of the Axiom instance. Defaults to [APIURL]. baseURL *url.URL // token is the authentication token that will be set as 'Bearer' on the // 'Authorization' header. It must be an api or a personal token. @@ -103,7 +103,7 @@ func (c Config) Validate() error { if c.token == "" { return ErrMissingToken - } else if !IsValidToken(c.token) { + } else if !IsValidCredential(c.token) { return ErrInvalidToken } diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 7343aae6..60caae1d 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -11,11 +11,12 @@ import ( ) const ( - endpoint = "http://api.axiom.local" - apiToken = "xaat-XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX" - personalToken = "xapt-XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX" //nolint:gosec // Chill, it's just testing. - unspecifiedToken = "XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX" - organizationID = "awkward-identifier-c3po" + endpoint = "http://api.axiom.local" + apiToken = "xaat-XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX" + personalToken = "xapt-XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX" //nolint:gosec // Chill, it's just testing. + sharedAccessSignature = "dt=logs&fl=customer+%3D%3D+%22peter%22&met=now&mst=ago%281h%29&oi=axiom&tk=0M41vwyiTVtAqW_aw8ZaIgeyOlxnSwtFoFbywuQ-VCc%3D" + unspecifiedToken = "XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX" + organizationID = "awkward-identifier-c3po" ) func TestConfig_IncorporateEnvironment(t *testing.T) { diff --git a/internal/config/option.go b/internal/config/option.go index 242845d5..13d79706 100644 --- a/internal/config/option.go +++ b/internal/config/option.go @@ -22,7 +22,7 @@ func SetURL(baseURL string) Option { // SetToken specifies the token to use. func SetToken(token string) Option { return func(config *Config) error { - if !IsValidToken(token) { + if !IsValidCredential(token) { return ErrInvalidToken } diff --git a/internal/config/token.go b/internal/config/token.go index 00c37ceb..dcf7ee0d 100644 --- a/internal/config/token.go +++ b/internal/config/token.go @@ -1,18 +1,35 @@ package config -import "strings" +import ( + "strings" -// IsAPIToken returns true if the given token is an api token. + "github.com/axiomhq/axiom-go/axiom/sas" +) + +// IsAPIToken returns true if the given token is an api token. It does not +// validate the token itself. func IsAPIToken(token string) bool { return strings.HasPrefix(token, "xaat-") } -// IsPersonalToken returns true if the given token is a personal token. +// IsPersonalToken returns true if the given token is a personal token. It does +// not validate the token itself. func IsPersonalToken(token string) bool { return strings.HasPrefix(token, "xapt-") } -// IsValidToken returns true if the given token is a valid Axiom token. -func IsValidToken(token string) bool { - return IsAPIToken(token) || IsPersonalToken(token) +// IsSharedAccessSignature returns true if the given signature is a shared +// access signature. It does not validate the signature itself. +func IsSharedAccessSignature(signature string) bool { + _, err := sas.Decode(signature) + return err == nil +} + +// IsValidCredential returns true if the given credential is a valid Axiom +// token or shared access signature. It does not validate the token or signature +// itself. +func IsValidCredential(credential string) bool { + return IsAPIToken(credential) || + IsPersonalToken(credential) || + IsSharedAccessSignature(credential) } diff --git a/internal/config/token_test.go b/internal/config/token_test.go index 00e718e0..44d08668 100644 --- a/internal/config/token_test.go +++ b/internal/config/token_test.go @@ -9,17 +9,27 @@ import ( func TestIsAPIToken(t *testing.T) { assert.True(t, IsAPIToken(apiToken)) assert.False(t, IsAPIToken(personalToken)) + assert.False(t, IsAPIToken(sharedAccessSignature)) assert.False(t, IsAPIToken(unspecifiedToken)) } func TestIsPersonalToken(t *testing.T) { assert.False(t, IsPersonalToken(apiToken)) assert.True(t, IsPersonalToken(personalToken)) + assert.False(t, IsPersonalToken(sharedAccessSignature)) assert.False(t, IsPersonalToken(unspecifiedToken)) } +func TestIsSharedAccessSignature(t *testing.T) { + assert.False(t, IsSharedAccessSignature(apiToken)) + assert.False(t, IsSharedAccessSignature(personalToken)) + assert.True(t, IsSharedAccessSignature(sharedAccessSignature)) + assert.False(t, IsSharedAccessSignature(unspecifiedToken)) +} + func TestIsValidToken(t *testing.T) { - assert.True(t, IsValidToken(apiToken)) - assert.True(t, IsValidToken(personalToken)) - assert.False(t, IsValidToken(unspecifiedToken)) + assert.True(t, IsValidCredential(apiToken)) + assert.True(t, IsValidCredential(personalToken)) + assert.True(t, IsValidCredential(sharedAccessSignature)) + assert.False(t, IsValidCredential(unspecifiedToken)) } diff --git a/internal/test/adapters/integration.go b/internal/test/adapters/integration.go index df443c04..96178f53 100644 --- a/internal/test/adapters/integration.go +++ b/internal/test/adapters/integration.go @@ -26,6 +26,10 @@ type IntegrationTestFunc func(ctx context.Context, dataset string, client *axiom // IntegrationTest tests the given adapter with the given test function. It // takes care of setting up all surroundings for the integration test. func IntegrationTest(t *testing.T, adapterName string, testFunc IntegrationTestFunc) { + if adapterName == "" { + t.Fatal("adapter integration test needs the name of the adapter") + } + cfg := config.Default() if err := cfg.IncorporateEnvironment(); err != nil { t.Fatal(err) @@ -33,18 +37,14 @@ func IntegrationTest(t *testing.T, adapterName string, testFunc IntegrationTestF t.Fatal(err) } - // Clear the environment to avoid unexpected behavior. - testhelper.SafeClearEnv(t) - - if adapterName == "" { - t.Fatal("adapter integration test needs the name of the adapter") - } - datasetSuffix := os.Getenv("AXIOM_DATASET_SUFFIX") if datasetSuffix == "" { datasetSuffix = "local" } + // Clear the environment to avoid unexpected behavior. + testhelper.SafeClearEnv(t) + deadline := time.Minute ctx, cancel := context.WithTimeout(context.Background(), deadline) t.Cleanup(cancel) @@ -86,8 +86,6 @@ func IntegrationTest(t *testing.T, adapterName string, testFunc IntegrationTestF // Run the test function with the test client. testFunc(ctx, dataset.ID, client) - // time.Sleep(time.Second * 30) - // Make sure the dataset is not empty. res, err := client.Datasets.QueryLegacy(ctx, dataset.ID, querylegacy.Query{ StartTime: startTime,