From a4b3ac0ecd0df83ef44897416c6104173bd3a489 Mon Sep 17 00:00:00 2001 From: Michal Wasilewski Date: Wed, 24 Jan 2024 14:21:24 +0100 Subject: [PATCH] feat(usage-view): download usage data in csv format Signed-off-by: Michal Wasilewski --- client/client.go | 41 ++++++- client/interface.go | 4 + internal/cmd/profile/flags.go | 91 ++++++++++++++++ internal/cmd/profile/profile.go | 1 + .../cmd/profile/usage_view_csv_command.go | 103 ++++++++++++++++++ 5 files changed, 238 insertions(+), 2 deletions(-) create mode 100644 internal/cmd/profile/usage_view_csv_command.go diff --git a/client/client.go b/client/client.go index a872ba7..8b56af8 100644 --- a/client/client.go +++ b/client/client.go @@ -5,6 +5,7 @@ import ( "fmt" "net/http" "net/url" + "strings" "github.com/shurcooL/graphql" "golang.org/x/oauth2" @@ -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 } diff --git a/client/interface.go b/client/interface.go index 7065d1f..64e1d85 100644 --- a/client/interface.go +++ b/client/interface.go @@ -2,6 +2,7 @@ package client import ( "context" + "net/http" "github.com/shurcooL/graphql" ) @@ -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) } diff --git a/internal/cmd/profile/flags.go b/internal/cmd/profile/flags.go index e2409f9..aa3372b 100644 --- a/internal/cmd/profile/flags.go +++ b/internal/cmd/profile/flags.go @@ -2,6 +2,7 @@ package profile import ( "fmt" + "time" "github.com/urfave/cli/v2" @@ -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"}, +} diff --git a/internal/cmd/profile/profile.go b/internal/cmd/profile/profile.go index 2fdd508..f3ecd2e 100644 --- a/internal/cmd/profile/profile.go +++ b/internal/cmd/profile/profile.go @@ -27,6 +27,7 @@ func Command() *cli.Command { Subcommands: []*cli.Command{ currentCommand(), exportTokenCommand(), + usageViewCSVCommand(), listCommand(), loginCommand(), logoutCommand(), diff --git a/internal/cmd/profile/usage_view_csv_command.go b/internal/cmd/profile/usage_view_csv_command.go new file mode 100644 index 0000000..09297ba --- /dev/null +++ b/internal/cmd/profile/usage_view_csv_command.go @@ -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) + 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 +}