Skip to content

Commit

Permalink
feat(usage-view): download usage data in csv format
Browse files Browse the repository at this point in the history
Signed-off-by: Michal Wasilewski <[email protected]>
  • Loading branch information
mwasilew2 committed Jan 24, 2024
1 parent 16bff7c commit a4b3ac0
Show file tree
Hide file tree
Showing 5 changed files with 238 additions and 2 deletions.
41 changes: 39 additions & 2 deletions client/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (
"fmt"
"net/http"
"net/url"
"strings"

"github.com/shurcooL/graphql"
"golang.org/x/oauth2"
Expand Down Expand Up @@ -66,14 +67,50 @@ func (c *client) URL(format string, a ...interface{}) string {
}

func (c *client) apiClient(ctx context.Context) (*graphql.Client, error) {
httpC, err := c.httpClient(ctx)
if err != nil {
return nil, fmt.Errorf("graphql client creation failed at http client creation: %w", err)
}

return graphql.NewClient(c.session.Endpoint(), httpC), nil
}

func (c *client) Do(req *http.Request) (*http.Response, error) {
// get http client
httpC, err := c.httpClient(req.Context())
if err != nil {
return nil, fmt.Errorf("http client creation failed: %w", err)
}

// prepend request URL with spacelift endpoint
endpoint := strings.TrimRight(c.session.Endpoint(), "/graphql")
u, err := url.Parse(endpoint)
if err != nil {
return nil, fmt.Errorf("failed to parse endpoint: %w", err)
}
req.URL.Scheme = u.Scheme
req.URL.Host = u.Host

// execute request
resp, err := httpC.Do(req)
if err != nil {
return nil, fmt.Errorf("error executing request: %w", err)
}
if resp.StatusCode == http.StatusUnauthorized {
return nil, fmt.Errorf("unauthorized: you can re-login using `spacectl profile login`")
}
return resp, err
}

func (c *client) httpClient(ctx context.Context) (*http.Client, error) {
bearerToken, err := c.session.BearerToken(ctx)
if err != nil {
return nil, err
}

return graphql.NewClient(c.session.Endpoint(), oauth2.NewClient(
return oauth2.NewClient(
context.WithValue(ctx, oauth2.HTTPClient, c.wraps), oauth2.StaticTokenSource(
&oauth2.Token{AccessToken: bearerToken},
),
)), nil
), nil
}
4 changes: 4 additions & 0 deletions client/interface.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package client

import (
"context"
"net/http"

"github.com/shurcooL/graphql"
)
Expand All @@ -16,4 +17,7 @@ type Client interface {

// URL returns a full URL given a formatted path.
URL(string, ...interface{}) string

// Do executes an authenticated http request to the Spacelift API
Do(r *http.Request) (*http.Response, error)
}
91 changes: 91 additions & 0 deletions internal/cmd/profile/flags.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package profile

import (
"fmt"
"time"

"github.com/urfave/cli/v2"

Expand Down Expand Up @@ -65,3 +66,93 @@ var flagEndpoint = &cli.StringFlag{
Required: false,
EnvVars: []string{"SPACECTL_LOGIN_ENDPOINT"},
}

const (
usageViewCSVTimeFormat = "2006-01-02"
usageViewCSVDefaultRange = time.Duration(-1*30*24) * time.Hour
)

var flagUsageViewCSVSince = &cli.StringFlag{
Name: "since",
Usage: "[Optional] the start of the time range to query for usage data in format YYYY-MM-DD",
Required: false,
EnvVars: []string{"SPACECTL_USAGE_VIEW_CSV_SINCE"},
Value: time.Now().Add(usageViewCSVDefaultRange).Format(usageViewCSVTimeFormat),
Action: func(context *cli.Context, s string) error {
_, err := time.Parse(usageViewCSVTimeFormat, s)
if err != nil {
return err
}
return nil
},
}

var flagUsageViewCSVUntil = &cli.StringFlag{
Name: "until",
Usage: "[Optional] the end of the time range to query for usage data in format YYYY-MM-DD",
Required: false,
EnvVars: []string{"SPACECTL_USAGE_VIEW_CSV_UNTIL"},
Value: time.Now().Format(usageViewCSVTimeFormat),
Action: func(context *cli.Context, s string) error {
_, err := time.Parse(usageViewCSVTimeFormat, s)
if err != nil {
return err
}
return nil
},
}

const (
aspectRunMinutes = "run-minutes"
aspectWorkerCount = "worker-count"
)

var aspects = map[string]struct{}{
aspectRunMinutes: {},
aspectWorkerCount: {},
}

var flagUsageViewCSVAspect = &cli.StringFlag{
Name: "aspect",
Usage: "[Optional] the aspect to query for usage data",
Required: false,
EnvVars: []string{"SPACECTL_USAGE_VIEW_CSV_ASPECT"},
Value: aspectWorkerCount,
Action: func(context *cli.Context, s string) error {
if _, isValidAspect := aspects[s]; !isValidAspect {
return fmt.Errorf("invalid aspect: %s", s)
}
return nil
},
}

const (
groupByRunState = "run-state"
groupByRunType = "run-type"
)

var groupBys = map[string]struct{}{
groupByRunState: {},
groupByRunType: {},
}

var flagUsageViewCSVGroupBy = &cli.StringFlag{
Name: "group-by",
Usage: "[Optional] the aspect to group run minutes by",
Required: false,
EnvVars: []string{"SPACECTL_USAGE_VIEW_CSV_GROUP_BY"},
Value: groupByRunType,
Action: func(context *cli.Context, s string) error {
if _, isValidGroupBy := groupBys[s]; !isValidGroupBy {
return fmt.Errorf("invalid group-by: %s", s)
}
return nil
},
}

var flagUsageViewCSVFile = &cli.StringFlag{
Name: "file",
Usage: "[Optional] the file to save the CSV to",
Required: false,
EnvVars: []string{"SPACECTL_USAGE_VIEW_CSV_FILE"},
}
1 change: 1 addition & 0 deletions internal/cmd/profile/profile.go
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ func Command() *cli.Command {
Subcommands: []*cli.Command{
currentCommand(),
exportTokenCommand(),
usageViewCSVCommand(),
listCommand(),
loginCommand(),
logoutCommand(),
Expand Down
103 changes: 103 additions & 0 deletions internal/cmd/profile/usage_view_csv_command.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,103 @@
package profile

import (
"bufio"
"fmt"
"io"
"log"
"net/http"
"os"

"github.com/urfave/cli/v2"

"github.com/spacelift-io/spacectl/internal/cmd"
"github.com/spacelift-io/spacectl/internal/cmd/authenticated"
)

type queryArgs struct {
since string
until string
aspect string
groupBy string
}

func usageViewCSVCommand() *cli.Command {
return &cli.Command{
Name: "usage-view-csv",
Usage: "Prints CSV with usage data for the current account",
ArgsUsage: cmd.EmptyArgsUsage,
Flags: []cli.Flag{
flagUsageViewCSVSince,
flagUsageViewCSVUntil,
flagUsageViewCSVAspect,
flagUsageViewCSVGroupBy,
flagUsageViewCSVFile,
},
Before: authenticated.Ensure,
Action: func(ctx *cli.Context) error {
// prep http query
args := &queryArgs{
since: ctx.String(flagUsageViewCSVSince.Name),
until: ctx.String(flagUsageViewCSVUntil.Name),
aspect: ctx.String(flagUsageViewCSVAspect.Name),
groupBy: ctx.String(flagUsageViewCSVGroupBy.Name),
}
params := buildQueryParams(args)
req, err := http.NewRequestWithContext(ctx.Context, http.MethodGet, "/usageanalytics/csv", nil)
if err != nil {
return fmt.Errorf("failed to create an HTTP request: %w", err)
}
q := req.URL.Query()
for k, v := range params {
q.Add(k, v)
}
req.URL.RawQuery = q.Encode()

// execute http query
log.Println("Querying Spacelift for usage data...")
resp, err := authenticated.Client.Do(req)
if err != nil {
return err
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return fmt.Errorf("unexpected status code: %d", resp.StatusCode)
}

// save response to a file
var filename string
if !ctx.IsSet(flagUsageViewCSVFile.Name) {
filename = fmt.Sprintf("usage-%s-%s-%s.csv", args.aspect, args.since, args.until)
} else {
filename = ctx.String(flagUsageViewCSVFile.Name)
}
fd, err := os.OpenFile("./"+filename, os.O_CREATE|os.O_WRONLY|os.O_EXCL, 0644)

Check failure

Code scanning / gosec

Potential file inclusion via variable Error

Potential file inclusion via variable

Check failure

Code scanning / gosec

Expect file permissions to be 0600 or less Error

Expect file permissions to be 0600 or less
if err != nil {
return fmt.Errorf("failed to open a file descriptor: %w", err)
}
defer fd.Close()
bfd := bufio.NewWriter(fd)
defer bfd.Flush()
_, err = io.Copy(bfd, resp.Body)
if err != nil {
return fmt.Errorf("failed to write the response to a file: %w", err)
}
log.Println("Usage data saved to", filename)
return nil
},
}
}

func buildQueryParams(args *queryArgs) map[string]string {
params := make(map[string]string)

params["since"] = args.since
params["until"] = args.until
params["aspect"] = args.aspect

if args.aspect == "run-minutes" {
params["groupBy"] = args.groupBy
}

return params
}

0 comments on commit a4b3ac0

Please sign in to comment.