Add support for tokens

axiom/client.go

@@ -97,6 +97,7 @@ type Client struct {
  Monitors *MonitorsService
  Notifiers *NotifiersService
  Annotations *AnnotationsService
+ Tokens *TokensService
 }
 
 // NewClient returns a new Axiom API client. It automatically takes its
@@ -133,6 +134,7 @@ func NewClient(options ...Option) (*Client, error) {
  client.Monitors = &MonitorsService{client: client, basePath: "/v2/monitors"}
  client.Notifiers = &NotifiersService{client: client, basePath: "/v2/notifiers"}
  client.Annotations = &AnnotationsService{client: client, basePath: "/v2/annotations"}
+ client.Tokens = &TokensService{client: client, basePath: "/v2/tokens/api"}
 
  // Apply supplied options.
  if err := client.Options(options...); err != nil {

axiom/tokens.go

@@ -0,0 +1,156 @@
+package axiom
+
+import (
+ "context"
+ "go.opentelemetry.io/otel/attribute"
+ "go.opentelemetry.io/otel/trace"
+ "net/http"
+ "net/url"
+ "time"
+)
+
+type Actions []string
+
+type APIToken struct {
+ ID string `json:"id"`
+ Name string `json:"name"`
+ Description string `json:"description"`
+ ExpiresAt time.Time `json:"expiresAt"`
+ DatasetCapabilities map[string]DatasetCapabilities `json:"datasetCapabilities"`
+ OrgCapabilities OrgCapabilities `json:"orgCapabilities"`
+}
+
+type DatasetCapabilities struct {
+ Ingest []string `json:"ingest"`
+ Query []string `json:"query"`
+ StarredQueries []string `json:"starredQueries"`
+ VirtualFields []string `json:"virtualFields"`
+}
+
+type OrgCapabilities struct {
+ Annotations []string `json:"annotations,omitempty"`
+ APITokens []string `json:"apiTokens,omitempty"`
+ Billing []string `json:"billing,omitempty"`
+ Dashboards []string `json:"dashboards,omitempty"`
+ Datasets []string `json:"datasets,omitempty"`
+ Endpoints []string `json:"endpoints,omitempty"`
+ Flows []string `json:"flows,omitempty"`
+ Integrations []string `json:"integrations,omitempty"`
+ Monitors []string `json:"monitors,omitempty"`
+ Notifiers []string `json:"notifiers,omitempty"`
+ Rbac []string `json:"rbac,omitempty"`
+ SharedAccessKeys []string `json:"sharedAccessKeys,omitempty"`
+ Users []string `json:"users,omitempty"`
+}
+
+type CreateTokenRequest struct {
+ Name string `json:"name"`
+ Description string `json:"description"`
+ ExpiresAt time.Time `json:"expiresAt"`
+ DatasetCapabilities map[string]DatasetCapabilities `json:"datasetCapabilities"`
+ OrgCapabilities OrgCapabilities `json:"orgCapabilities"`
+}
+
+type CreateTokenResponse struct {
+ APIToken
+ Token string `json:"token"`
+}
+
+type RegenerateTokenRequest struct {
+ ExistingTokenExpiresAt time.Time `json:"existingTokenExpiresAt"`
+ NewTokenExpiresAt time.Time `json:"newTokenExpiresAt"`
+}
+
+// Axiom API Reference: /v2/tokens
+type TokensService service
+
+// List all available tokens.
+func (s *TokensService) List(ctx context.Context) ([]*APIToken, error) {
+ ctx, span := s.client.trace(ctx, "Tokens.List")
+ defer span.End()
+
+ var res []*APIToken
+ if err := s.client.Call(ctx, http.MethodGet, s.basePath, nil, &res); err != nil {
+ return nil, spanError(span, err)
+ }
+
+ return res, nil
+}
+
+// Get a token by id.
+func (s *TokensService) Get(ctx context.Context, id string) (*APIToken, error) {
+ ctx, span := s.client.trace(ctx, "Tokens.Get", trace.WithAttributes(
+ attribute.String("axiom.token_id", id),
+ ))
+ defer span.End()
+
+ path, err := url.JoinPath(s.basePath, id)
+ if err != nil {
+ return nil, spanError(span, err)
+ }
+
+ var res APIToken
+ if err := s.client.Call(ctx, http.MethodGet, path, nil, &res); err != nil {
+ return nil, spanError(span, err)
+ }
+
+ return &res, nil
+}
+
+// Create a token with the given properties.
+func (s *TokensService) Create(ctx context.Context, req CreateTokenRequest) (*CreateTokenResponse, error) {
+ ctx, span := s.client.trace(ctx, "Tokens.Create", trace.WithAttributes(
+ attribute.String("axiom.param.name", req.Name),
+ ))
+ defer span.End()
+
+ var res CreateTokenResponse
+ if err := s.client.Call(ctx, http.MethodPost, s.basePath, req, &res); err != nil {
+ return nil, spanError(span, err)
+ }
+
+ return &res, nil
+}
+
+// Regenerate the token identified by the given id.
+func (s *TokensService) Regenerate(ctx context.Context, id string, req RegenerateTokenRequest) (*CreateTokenResponse, error) {
+ ctx, span := s.client.trace(ctx, "Tokens.Regenerate", trace.WithAttributes(
+ attribute.String("axiom.token_id", id),
+ ))
+ defer span.End()
+
+ path, err := url.JoinPath(s.basePath, id, "/regenerate")
+ if err != nil {
+ return nil, spanError(span, err)
+ }
+
+ var res CreateTokenResponse
+ if err := s.client.Call(ctx, http.MethodPost, path, req, &res); err != nil {
+ return nil, spanError(span, err)
+ }
+
+ return &res, nil
+}
+
+// Delete the token identified by the given id.
+func (s *TokensService) Delete(ctx context.Context, id string) error {
+ ctx, span := s.client.trace(ctx, "Tokens.Delete", trace.WithAttributes(
+ attribute.String("axiom.token_id", id),
+ ))
+ defer span.End()
+
+ path, err := url.JoinPath(s.basePath, "/", id)
+ if err != nil {
+ return spanError(span, err)
+ }
+
+ if err := s.client.Call(ctx, http.MethodDelete, path, nil, nil); err != nil {
+ return spanError(span, err)
+ }
+
+ return nil
+}
+
+func (t *CreateTokenResponse) AsAPIToken() *APIToken {
+ return &t.APIToken
+}

axiom/tokens_integration_test.go

@@ -0,0 +1,91 @@
+//go:build integration
+
+package axiom_test
+
+import (
+ "context"
+ "testing"
+ "time"
+
+ "github.com/stretchr/testify/suite"
+
+ "github.com/axiomhq/axiom-go/axiom"
+)
+
+// TokensTestSuite tests all methods of the Axiom Tokens API against a
+// live deployment.
+type TokensTestSuite struct {
+ IntegrationTestSuite
+
+ apiToken *axiom.APIToken
+}
+
+func TestTokensTestSuite(t *testing.T) {
+ suite.Run(t, new(TokensTestSuite))
+}
+
+func (s *TokensTestSuite) SetupSuite() {
+ s.IntegrationTestSuite.SetupSuite()
+}
+
+func (s *TokensTestSuite) TearDownSuite() {
+ // Teardown routines use their own context to avoid not being run at all
+ // when the suite gets cancelled or times out.
+ ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
+ defer cancel()
+
+ err := s.client.Tokens.Delete(ctx, s.apiToken.ID)
+ s.NoError(err)
+
+ s.IntegrationTestSuite.TearDownSuite()
+}
+
+func (s *TokensTestSuite) Test() {
+ createdToken, err := s.client.Tokens.Create(s.suiteCtx, axiom.CreateTokenRequest{
+ Name: "Test token",
+ ExpiresAt: time.Now().Add(24 * time.Hour),
+ DatasetCapabilities: map[string]axiom.DatasetCapabilities{
+ "*": {Ingest: []string{"create"}}},
+ OrgCapabilities: axiom.OrgCapabilities{
+ Users: []string{"create", "read", "update", "delete"},
+ }})
+ s.Require().NoError(err)
+ s.Require().NotNil(createdToken)
+
+ s.apiToken = createdToken.AsAPIToken()
+
+ // Get the token and make sure it matches what we have updated it to.
+ token, err := s.client.Tokens.Get(s.ctx, s.apiToken.ID)
+ s.Require().NoError(err)
+ s.Require().NotNil(token)
+
+ s.Equal(s.apiToken, token)
+
+ // List all tokens and make sure the created token is part of that
+ // list.
+ tokens, err := s.client.Tokens.List(s.ctx)
+ s.Require().NoError(err)
+ s.Require().NotEmpty(tokens)
+
+ s.Contains(tokens, s.apiToken)
+
+ // Regenerate the token and make sure the new token is part of the list.
+ regeneratedToken, err := s.client.Tokens.Regenerate(s.ctx, s.apiToken.ID, axiom.RegenerateTokenRequest{
+ ExistingTokenExpiresAt: time.Now().Add(-time.Second),
+ NewTokenExpiresAt: time.Now().Add(24 * time.Hour),
+ })
+ s.Require().NoError(err)
+ s.Require().NotEmpty(tokens)
+
+ oldToken := s.apiToken
+ s.apiToken = regeneratedToken.AsAPIToken()
+
+ // List all tokens and make sure the created token is part of that
+ // list.
+ tokens, err = s.client.Tokens.List(s.ctx)
+ s.Require().NoError(err)
+ s.Require().NotEmpty(tokens)
+
+ s.NotContains(tokens, oldToken)
+ s.Contains(tokens, regeneratedToken.AsAPIToken())
+}