feat: shared access support

.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

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
 	}
 

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))
 		})
 	}
 }

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()
 

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

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

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
 }

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)
 }

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"