Skip to content

Commit

Permalink
Further improvements
Browse files Browse the repository at this point in the history 
- Address review comments
- Use zerolog everywhere
- Make the lambda mode be basically identical to normal CLI mode
  • Loading branch information
@vprus
vprus committed Jan 24, 2024
1 parent 9011f77 commit 2892be3
Show file tree
Hide file tree
Showing 6 changed files with 179 additions and 166 deletions.
26 changes: 16 additions & 10 deletions tools/storage-advisor/src/api.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,13 +5,13 @@ import (
resty "github.com/go-resty/resty/v2"
)

type Api struct {
type JoomCloudAPI struct {
url string
token string
client *resty.Client
}

func NewApi(url string, token *string) *Api {
func NewApi(url string, token *string) *JoomCloudAPI {
client := resty.New()
normalized := url
if normalized[len(normalized)-1] != '/' {
Expand All @@ -21,25 +21,31 @@ func NewApi(url string, token *string) *Api {
if token != nil {
theToken = *token
}
return &Api{url: normalized, token: theToken, client: client}
return &JoomCloudAPI{url: normalized, token: theToken, client: client}
}

func (api *Api) IsConfigured() bool {
func (api *JoomCloudAPI) IsConfigured() bool {
return api.token != ""
}

func (api *Api) Post(endpoint string, payload []byte) error {
resp, err := api.client.R().
SetAuthToken(api.token).SetBody(payload).Post(api.url + endpoint)
func (api *JoomCloudAPI) Post(endpoint string, payload []byte) error {
fullURL := api.url + endpoint
resp, err := api.client.R().SetAuthToken(api.token).SetBody(payload).Post(fullURL)

if err != nil {
fmt.Printf("Error: could not post to %s: %v\n", endpoint, err.Error())
return err
return fmt.Errorf("could not post to %s: %v", fullURL, err)
}

if resp.StatusCode() < 200 || resp.StatusCode() > 299 {
return fmt.Errorf("unexpected status code: %d", resp.StatusCode())
return fmt.Errorf("unexpected status code calling %s: %d", fullURL, resp.StatusCode())
}

return nil
}

func (api *JoomCloudAPI) PostS3Buckets(payload []byte) error {
return api.Post("storage-advisor/buckets", payload)
}
func (api *JoomCloudAPI) PostS3Inventory(payload []byte) error {
return api.Post("storage-advisor/s3-inventory", payload)
}
1 change: 1 addition & 0 deletions tools/storage-advisor/src/go.mod
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ require (
github.com/manifoldco/promptui v0.9.0 // indirect
github.com/mattn/go-colorable v0.1.13 // indirect
github.com/mattn/go-isatty v0.0.20 // indirect
github.com/rs/zerolog v1.31.0 // indirect
golang.org/x/net v0.17.0 // indirect
golang.org/x/sys v0.14.0 // indirect
)
8 changes: 8 additions & 0 deletions tools/storage-advisor/src/go.sum
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,18 +40,25 @@ github.com/chzyer/logex v1.1.10/go.mod h1:+Ywpsq7O8HXn0nuIou7OrIPyXbp3wmkHB+jjWR
github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e h1:fY5BOSpyZCqRo5OhCuC+XN+r/bBCmeuuJtjz+bCNIf8=
github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e/go.mod h1:nSuG5e5PlCu98SY8svDHJxuZscDgtXS6KTTbou5AhLI=
github.com/chzyer/test v0.0.0-20180213035817-a1ea475d72b1/go.mod h1:Q3SI9o4m/ZMnBNeIyt5eFwwo7qiLfzFZmjNmxjkiQlU=
github.com/coreos/go-systemd/v22 v22.5.0/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
github.com/fatih/color v1.16.0 h1:zmkK9Ngbjj+K0yRhTVONQh1p/HknKYSlNT+vZCzyokM=
github.com/fatih/color v1.16.0/go.mod h1:fL2Sau1YI5c0pdGEVCbKQbLXB6edEj1ZgiY4NijnWvE=
github.com/go-resty/resty/v2 v2.11.0 h1:i7jMfNOJYMp69lq7qozJP+bjgzfAzeOhuGlyDrqxT/8=
github.com/go-resty/resty/v2 v2.11.0/go.mod h1:iiP/OpA0CkcL3IGt1O0+/SIItFUbkkyw5BGXiVdTu+A=
github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
github.com/google/go-cmp v0.5.8 h1:e6P7q2lk1O+qJJb4BtCQXlK8vWEO8V1ZeuEdJNOqZyg=
github.com/manifoldco/promptui v0.9.0 h1:3V4HzJk1TtXW1MTZMP7mdlwbBpIinw3HztaIlYthEiA=
github.com/manifoldco/promptui v0.9.0/go.mod h1:ka04sppxSGFAtxX0qhlYQjISsg9mR4GWtQEhdbn6Pgg=
github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg=
github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
github.com/rs/xid v1.5.0/go.mod h1:trrq9SKmegXys3aeAKXMUTdJsYXVwGY3RLcfgqegfbg=
github.com/rs/zerolog v1.31.0 h1:FcTR3NnLWW+NnTwwhFWiJSZr4ECLpqCm6QsEnyvbV4A=
github.com/rs/zerolog v1.31.0/go.mod h1:/7mN4D5sKwJLZQ2b/znpjC3/GQWY/xaDXUM0kKWRHss=
github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
Expand All @@ -78,6 +85,7 @@ golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBc
golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.13.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.14.0 h1:Vz7Qs629MkJkGyHxUlRHizWJRG2j8fbQKjELVSNhy7Q=
golang.org/x/sys v0.14.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
Expand Down
27 changes: 12 additions & 15 deletions tools/storage-advisor/src/inventory.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@ import (
"encoding/json"
"github.com/aws/aws-sdk-go-v2/aws"
"github.com/aws/aws-sdk-go-v2/service/s3"
"github.com/rs/zerolog/log"
"io"
"log"
"strings"
"time"
)
Expand Down Expand Up @@ -37,13 +37,13 @@ func presignFile(presignClient *s3.PresignClient, bucket string, fileKey string)
return presignGetObject.URL
}

func uploadInventoryForBucket(client *s3.Client, api *Api, bucket string, prefix string) {
func uploadInventoryForBucket(client *s3.Client, api *JoomCloudAPI, bucket string, prefix string) {

datePrefixies, err := listCommonPrefixes(client, bucket, prefix)
datePrefixes, err := listCommonPrefixes(client, bucket, prefix)
dates := make([]time.Time, 0, 20)

dateFormat := "2006-01-02T15-04Z"
for _, datePrefix := range datePrefixies {
for _, datePrefix := range datePrefixes {
dateStr := strings.TrimPrefix(datePrefix, prefix)
dateStr = strings.TrimPrefix(dateStr, "/")
dateStr = strings.TrimSuffix(dateStr, "/")
Expand All @@ -53,7 +53,7 @@ func uploadInventoryForBucket(client *s3.Client, api *Api, bucket string, prefix
}
}
if len(dates) == 0 {
log.Println("Error: No dates found")
log.Warn().Msgf("no inventory data found at s3://%s/%s", bucket, prefix)
return
}

Expand All @@ -65,19 +65,19 @@ func uploadInventoryForBucket(client *s3.Client, api *Api, bucket string, prefix
}
}

log.Printf("Max date: %s\n", maxDate.Format(time.DateTime))
partitionDate := maxDate.Format(dateFormat)
content, err := readObjectContent(client, bucket, prefix+"/"+partitionDate+"/manifest.json")
manifestPath := prefix + "/" + partitionDate + "/manifest.json"
content, err := readObjectContent(client, bucket, manifestPath)
if err != nil {
log.Printf(err.Error())
}

var inputManifest InventoryManifest
err = json.Unmarshal(content, &inputManifest)
if err != nil {
log.Fatalf("Could not unmarsal inputManifest: %s\n", err.Error())
log.Err(err).Msgf("could not unmarshal manifest %s", manifestPath)
}
log.Printf("Manifest has %d files for bucket %s\n", len(inputManifest.Files), inputManifest.Bucket)
log.Debug().Msgf("manifest has %d files for bucket %s\n", len(inputManifest.Files), inputManifest.Bucket)

presignClient := s3.NewPresignClient(client)
outputManifest := InventoryManifest{PartitionDate: partitionDate, Version: manifestVersion, Bucket: inputManifest.Bucket}
Expand All @@ -96,12 +96,12 @@ func uploadInventoryForBucket(client *s3.Client, api *Api, bucket string, prefix
panic(err)
}

err = api.Post("storage-advisor/s3-inventory", marshal)
err = api.PostS3Inventory(marshal)
if err != nil {
log.Printf("Error: could not upload inventory: %s\n", err.Error())
log.Err(err).Msgf("could not upload inventory to API")
}

log.Printf("Uploaded data for s3://%s/%s", bucket, prefix)
log.Debug().Msgf("uploaded data for s3://%s/%s", bucket, prefix)
}

func listCommonPrefixes(client *s3.Client, bucket string, prefix string) ([]string, error) {
Expand All @@ -116,14 +116,11 @@ func listCommonPrefixes(client *s3.Client, bucket string, prefix string) ([]stri
Delimiter: aws.String("/"),
}

log.Println("Listing prefix", bucket, prefix)

paginator := s3.NewListObjectsV2Paginator(client, params)
var commonPrefixes []string
for paginator.HasMorePages() {
output, err := paginator.NextPage(context.Background())
if err != nil {
log.Println("Error:", err)
return nil, err
}
for _, commonPrefix := range output.CommonPrefixes {
Expand Down
140 changes: 73 additions & 67 deletions tools/storage-advisor/src/main.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,14 +2,18 @@ package main

import (
"context"
"errors"
"flag"
"fmt"
"github.com/aws/aws-lambda-go/events"
"github.com/aws/aws-lambda-go/lambda"
"github.com/aws/aws-sdk-go-v2/config"
"github.com/aws/aws-sdk-go-v2/service/s3"
"github.com/aws/aws-sdk-go-v2/service/sts"
"log"
"github.com/aws/smithy-go"
"github.com/fatih/color"
"github.com/rs/zerolog"
"github.com/rs/zerolog/log"
"os"
"strings"
)
Expand All @@ -21,96 +25,102 @@ type Config struct {
}

var mode = flag.String("mode", "cli", "Use as a CLI tool")
var region = flag.String("region", "", "AWS region to use")

var inventoryPrefix = flag.String("prefix", "", "Root inventory prefix")
var inventoryBucket = flag.String("inventory_bucket", "", "Bucket the inventory is stored in")
var apiToken = flag.String("api-token", "", "API token for Joom Cloud")

var apiEndpoint = flag.String("api-endpoint", "https://api.cloud.joom.ai/v1", "API endpoint URL")

func main() {
flag.Parse()
rootCtx := context.TODO()
rootCtx := context.Background()
cfg, err := config.LoadDefaultConfig(rootCtx)
if err != nil {
fmt.Printf("Error: could not initialize AWS config: %v\n", err.Error())
fmt.Printf("Error: could not initialize AWS config: %v\n", err)
return
}

client := s3.NewFromConfig(cfg)
api := NewApi(*apiEndpoint, apiToken)

switch *mode {
case "cli":
if *region != "" {
cfg.Region = *region
// For CLI, use a simplified logger
output := zerolog.ConsoleWriter{Out: os.Stdout, TimeFormat: ""}
output.FormatTimestamp = func(i interface{}) string {
return ""
}
output.FormatLevel = func(i interface{}) string {
switch i.(string) {
case "err":
return "Error: "
case "warn":
return "Warning: "
default:
return ""
}
}
log.Logger = zerolog.New(output).Level(zerolog.InfoLevel)

stsOutput, err := sts.NewFromConfig(cfg).GetCallerIdentity(rootCtx, &sts.GetCallerIdentityInput{})
if err != nil {
fmt.Printf("Error: unable to get your AWS indentity\n\n")

fmt.Printf(
"\nPossibly, you did not login to AWS or your token has expired.\n" +
"Please try to login again. Then, verify your identity using\n\n" +
" aws sts get-caller-identity\n\n" +
"If you are logged in, but still see this error, you might not\n" +
"have permissions to list S3 bucket. Double-check that IAM role\n" +
"associated with the role printed by the above command.\n\n")
fmt.Printf("Error: unable to get your AWS indentity\n\n" +
"Possibly, you did not login to AWS or your token has expired.\n" +
"Please try to login again. Then, verify your identity using\n\n" +
" aws sts get-caller-identity\n\n" +
"If you are logged in, but still see this error, you might not\n" +
"have permissions to list S3 bucket. Double-check that IAM role\n" +
"associated with the role printed by the above command.\n\n")

return
}
identity := *stsOutput.Arn
fmt.Printf("Info: your AWS identity is %s\n", identity)

buckets := runS3Checks(cfg, api, identity)
buckets, issues, err := runS3Checks(rootCtx, client)
var ae smithy.APIError
if errors.As(err, &ae) {
color.Red("Could not list buckets: %s: %s", ae.ErrorCode(), ae.ErrorMessage())
return
} else {
color.Red("Could not list buckets: %s", err.Error())
return
}

client := s3.NewFromConfig(cfg)
writeBasicIssues(issues, identity)

for _, bucket := range buckets {
if bucket.UsableInventory != "" {
fmt.Printf("Trying to process inventory s3://%s\n", bucket.UsableInventory)
i := strings.Index(bucket.UsableInventory, "/")
inventoryBucket := bucket.UsableInventory[:i]
inventoryPrefix := bucket.UsableInventory[i+1:]
uploadInventoryForBucket(client, api, inventoryBucket, inventoryPrefix)
}
if api.IsConfigured() {
uploadBasicIssues(api, buckets)
count := uploadInventories(api, client, buckets)
fmt.Printf(green("Uploaded inventory for %d buckets\n", count))
}
return

case "aws":
log.Println("Running in AWS lambda mode")

*inventoryPrefix = os.Getenv("prefix")
*region = os.Getenv("region")
*inventoryBucket = os.Getenv("inventoryBucket")
*apiToken = os.Getenv("jwt")

validateInputs()

cfg.Region = *region

client := s3.NewFromConfig(cfg)
if apiToken == nil {
log.Fatal().Msg("api token is not provided; either provide it or run in CLI mode")
return
}

handler := func(ctx context.Context) (events.APIGatewayProxyResponse, error) {
prefixList, err := listCommonPrefixes(client, *inventoryBucket, *inventoryPrefix)

stsOutput, err := sts.NewFromConfig(cfg).GetCallerIdentity(rootCtx, &sts.GetCallerIdentityInput{})
if err != nil {
response := events.APIGatewayProxyResponse{
StatusCode: 500,
Body: fmt.Sprintf("Failed to list prefix %s", err),
}

return response, err
log.Err(err).Msg("could not determine AWS identity, this usually means AWS auth is missing")
} else {
log.Info().Msgf("running with AWS identity %s", *stsOutput.Arn)
}

for _, prefix := range prefixList {
uploadInventoryForBucket(client, api, *inventoryBucket, prefix+"default")
buckets, _, err := runS3Checks(rootCtx, client)
if err != nil {
log.Fatal().Err(err).Msg("could not list buckets")
}

uploadBasicIssues(api, buckets)

count := uploadInventories(api, client, buckets)

response := events.APIGatewayProxyResponse{
StatusCode: 200,
Body: fmt.Sprintf("Processed prefixes %s for bucket %s", prefixList, *inventoryBucket),
Body: fmt.Sprintf("uploaded invetories for %d buckets", count),
}

return response, nil
Expand All @@ -121,22 +131,18 @@ func main() {

}

func validateInputs() {
if *region == "" {
log.Fatal("region is empty")
}

if *inventoryBucket == "" {
log.Fatal("inventoryBucket is empty")
}

if *inventoryPrefix == "" {
log.Fatal("prefix is empty")
}

if *apiToken == "" {
log.Fatal("jwt token is empty")
func uploadInventories(api *JoomCloudAPI, client *s3.Client, buckets []Bucket) int {
count := 0
for _, bucket := range buckets {
if bucket.UsableInventory == "" {
continue
}
log.Info().Msgf("Processing inventory s3://%s", bucket.UsableInventory)
i := strings.Index(bucket.UsableInventory, "/")
inventoryBucket := bucket.UsableInventory[:i]
inventoryPrefix := bucket.UsableInventory[i+1:]
uploadInventoryForBucket(client, api, inventoryBucket, inventoryPrefix)
count = count + 1
}

log.Println("Processing", *inventoryPrefix, *inventoryBucket, *region)
return count
}
Loading

0 comments on commit 2892be3

Please sign in to comment.