Skip to content

Commit

Permalink
feat(import): add progress bar for store import
Browse files Browse the repository at this point in the history
  • Loading branch information
Siddhant-K-code authored and ewanharris committed Jun 18, 2024
1 parent cea1edc commit 601021c
Show file tree
Hide file tree
Showing 4 changed files with 156 additions and 56 deletions.
1 change: 1 addition & 0 deletions .golangci.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ linters-settings:
- gopkg.in/yaml.v3
- github.com/hashicorp/go-multierror
- github.com/gocarina/gocsv
- github.com/schollz/progressbar/v3
test:
files:
- "$test"
Expand Down
193 changes: 137 additions & 56 deletions cmd/store/import.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,10 @@ import (
"os"
"path"
"strings"
"time"

"github.com/openfga/go-sdk/client"
"github.com/schollz/progressbar/v3"
"github.com/spf13/cobra"

"github.com/openfga/cli/cmd/model"
Expand All @@ -35,79 +37,154 @@ import (
"github.com/openfga/cli/internal/storetest"
)

func importStore(
clientConfig fga.ClientConfig,
fgaClient client.SdkClient,
const (
progressBarWidth = 40
progressBarSleepDelay = 10 // time.Millisecond
progressBarThrottleValue = 65
progressBarUpdateDelay = 5 * time.Millisecond
)

// createStore creates a new store with the given client configuration and store data.
func createStore(
clientConfig *fga.ClientConfig,
storeData *storetest.StoreData,
format authorizationmodel.ModelFormat,
storeID string,
maxTuplesPerWrite int,
maxParallelRequests int,
fileName string,
) (*CreateStoreAndModelResponse, error) {
response := &CreateStoreAndModelResponse{}
storeDataName := storeData.Name
if storeDataName == "" {
storeDataName = strings.TrimSuffix(path.Base(fileName), ".fga.yaml")
}

if storeID == "" { //nolint:nestif
storeDataName := storeData.Name
if storeDataName == "" {
storeDataName = strings.TrimSuffix(path.Base(fileName), ".fga.yaml")
}
createStoreAndModelResponse, err := CreateStoreWithModel(*clientConfig, storeDataName, storeData.Model, format)
if err != nil {
return nil, err
}

createStoreAndModelResponse, err := CreateStoreWithModel(
clientConfig,
storeDataName,
storeData.Model,
format,
)
if err != nil {
return nil, err
}
clientConfig.StoreID = createStoreAndModelResponse.Store.Id

response = createStoreAndModelResponse
clientConfig.StoreID = createStoreAndModelResponse.Store.Id
} else {
store, err := fgaClient.GetStore(context.Background()).Execute()
if err != nil {
return nil, err //nolint:wrapcheck
}
return createStoreAndModelResponse, nil
}

// updateStore updates an existing store with the given client configuration, store data, and store ID.
func updateStore(
clientConfig *fga.ClientConfig,
fgaClient client.SdkClient,
storeData *storetest.StoreData,
format authorizationmodel.ModelFormat,
storeID string,
) (*CreateStoreAndModelResponse, error) {
store, err := fgaClient.GetStore(context.Background()).Execute()
if err != nil {
return nil, fmt.Errorf("failed to get store: %w", err)
}

response.Store = client.ClientCreateStoreResponse{
response := &CreateStoreAndModelResponse{
Store: client.ClientCreateStoreResponse{
CreatedAt: store.GetCreatedAt(),
Id: store.GetId(),
Name: store.GetName(),
UpdatedAt: store.GetUpdatedAt(),
}
},
}

authModel := authorizationmodel.AuthzModel{}
clientConfig.StoreID = storeID

if err := authModel.ReadModelFromString(storeData.Model, format); err != nil {
return nil, fmt.Errorf("failed to read model: %w", err)
}

authModel := authorizationmodel.AuthzModel{}
clientConfig.StoreID = storeID
modelWriteRes, err := model.Write(fgaClient, authModel)
if err != nil {
return nil, fmt.Errorf("failed to write model: %w", err)
}

response.Model = modelWriteRes

return response, nil
}

// importStore imports store data, either creating a new store or updating an existing one.
func importStore(
clientConfig *fga.ClientConfig,
fgaClient client.SdkClient,
storeData *storetest.StoreData,
format authorizationmodel.ModelFormat,
storeID string,
maxTuplesPerWrite, maxParallelRequests int,
fileName string,
) (*CreateStoreAndModelResponse, error) {
var (
response *CreateStoreAndModelResponse
err error
)

err = authModel.ReadModelFromString(storeData.Model, format)
if storeID == "" {
response, err = createStore(clientConfig, storeData, format, fileName)
if err != nil {
return nil, err //nolint:wrapcheck
return nil, fmt.Errorf("failed to create store: %w", err)
}

modelWriteRes, err := model.Write(fgaClient, authModel)
} else {
response, err = updateStore(clientConfig, fgaClient, storeData, format, storeID)
if err != nil {
return nil, fmt.Errorf("failed to write model due to %w", err)
return nil, fmt.Errorf("failed to update store: %w", err)
}

response.Model = modelWriteRes
}

fgaClient, err := clientConfig.GetFgaClient()
fgaClient, err = clientConfig.GetFgaClient()
if err != nil {
return nil, fmt.Errorf("failed to initialize FGA Client due to %w", err)
return nil, fmt.Errorf("failed to initialize FGA Client: %w", err)
}

writeRequest := client.ClientWriteRequest{
Writes: storeData.Tuples,
// Initialize progress bar
bar := progressbar.NewOptions(len(storeData.Tuples),
progressbar.OptionSetDescription("Importing tuples"),
progressbar.OptionShowCount(),
progressbar.OptionSetWidth(progressBarWidth),
progressbar.OptionClearOnFinish(),
progressbar.OptionFullWidth(),
progressbar.OptionThrottle(progressBarThrottleValue*progressBarSleepDelay),
progressbar.OptionShowIts(),
progressbar.OptionSetItsString("tuples"),
progressbar.OptionSetPredictTime(true),
progressbar.OptionSetRenderBlankState(true),
progressbar.OptionSetTheme(progressbar.Theme{
Saucer: "#",
SaucerPadding: " ",
BarStart: "[",
BarEnd: "]",
}),
)

for index := 0; index < len(storeData.Tuples); index += maxTuplesPerWrite {
end := index + maxTuplesPerWrite
if end > len(storeData.Tuples) {
end = len(storeData.Tuples)
}

writeRequest := client.ClientWriteRequest{
Writes: storeData.Tuples[index:end],
}
if _, err := tuple.ImportTuples(fgaClient, writeRequest, maxTuplesPerWrite, maxParallelRequests); err != nil {
return nil, fmt.Errorf("failed to import tuples: %w", err)
}

if err := bar.Add(end - index); err != nil {
return nil, fmt.Errorf("failed to update progress bar: %w", err)
}

// Introduce a small delay to smooth out the progress bar rendering
time.Sleep(progressBarUpdateDelay)
}

_, err = tuple.ImportTuples(fgaClient, writeRequest, maxTuplesPerWrite, maxParallelRequests)
if err != nil {
return nil, err //nolint:wrapcheck
// Ensure progress bar is completed and cleared
if err := bar.Finish(); err != nil {
return nil, fmt.Errorf("failed to finish progress bar: %w", err)
}

fmt.Println("✅ Store imported")

return response, nil
}

Expand All @@ -118,46 +195,50 @@ var importCmd = &cobra.Command{
Long: `Import a store: updating the name, model and appending the global tuples`,
Example: "fga store import --file=model.fga.yaml",
RunE: func(cmd *cobra.Command, _ []string) error {
var createStoreAndModelResponse *CreateStoreAndModelResponse
clientConfig := cmdutils.GetClientConfig(cmd)

storeID, err := cmd.Flags().GetString("store-id")
if err != nil {
return fmt.Errorf("failed to get store-id %w", err)
return fmt.Errorf("failed to get store-id: %w", err)
}

maxTuplesPerWrite, err := cmd.Flags().GetInt("max-tuples-per-write")
if err != nil {
return fmt.Errorf("failed to parse max tuples per write due to %w", err)
return fmt.Errorf("failed to parse max tuples per write: %w", err)
}

maxParallelRequests, err := cmd.Flags().GetInt("max-parallel-requests")
if err != nil {
return fmt.Errorf("failed to parse parallel requests due to %w", err)
return fmt.Errorf("failed to parse parallel requests: %w", err)
}

fileName, err := cmd.Flags().GetString("file")
if err != nil {
return err //nolint:wrapcheck
return fmt.Errorf("failed to get file name: %w", err)
}

format, storeData, err := storetest.ReadFromFile(fileName, path.Dir(fileName))
if err != nil {
return err //nolint:wrapcheck
return fmt.Errorf("failed to read from file: %w", err)
}

fgaClient, err := clientConfig.GetFgaClient()
if err != nil {
return fmt.Errorf("failed to initialize FGA Client due to %w", err)
return fmt.Errorf("failed to initialize FGA Client: %w", err)
}

createStoreAndModelResponse, err = importStore(clientConfig, fgaClient, storeData, format,
createStoreAndModelResponse, err := importStore(&clientConfig, fgaClient, storeData, format,
storeID, maxTuplesPerWrite, maxParallelRequests, fileName)
if err != nil {
return err
return fmt.Errorf("failed to import store: %w", err)
}

// Print the response using output.Display without printing <nil>
if outputErr := output.Display(createStoreAndModelResponse); outputErr != nil {
return fmt.Errorf("failed to display output: %w", outputErr)
}

return output.Display(createStoreAndModelResponse)
return nil
},
}

Expand Down
7 changes: 7 additions & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,12 @@ require (
gopkg.in/yaml.v3 v3.0.1
)

require (
github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db // indirect
github.com/rivo/uniseg v0.4.7 // indirect
golang.org/x/term v0.21.0 // indirect
)

require (
github.com/antlr4-go/antlr/v4 v4.13.1 // indirect
github.com/beorn7/perks v1.0.1 // indirect
Expand Down Expand Up @@ -56,6 +62,7 @@ require (
github.com/prometheus/procfs v0.15.1 // indirect
github.com/sagikazarmark/locafero v0.6.0 // indirect
github.com/sagikazarmark/slog-shim v0.1.0 // indirect
github.com/schollz/progressbar/v3 v3.14.4
github.com/sourcegraph/conc v0.3.0 // indirect
github.com/spf13/afero v1.11.0 // indirect
github.com/spf13/cast v1.6.0 // indirect
Expand Down
11 changes: 11 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,7 @@ github.com/jackc/puddle/v2 v2.2.1 h1:RhxXJtFG022u4ibrCSMSiu5aOq1i77R3OHKNJj77OAk
github.com/jackc/puddle/v2 v2.2.1/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
github.com/jarcoal/httpmock v1.3.1 h1:iUx3whfZWVf3jT01hQTO/Eo5sAYtB2/rqaUuOtpInww=
github.com/jarcoal/httpmock v1.3.1/go.mod h1:3yb8rc4BI7TCBhFY8ng0gjuLKJNquuDNiPaZjnENuYg=
github.com/k0kubun/go-ansi v0.0.0-20180517002512-3bf9e2903213/go.mod h1:vNUNkEQ1e29fT/6vq2aBdFsgNPmy8qMdSay1npru+Sw=
github.com/karlseguin/ccache/v3 v3.0.5 h1:hFX25+fxzNjsRlREYsoGNa2LoVEw5mPF8wkWq/UnevQ=
github.com/karlseguin/ccache/v3 v3.0.5/go.mod h1:qxC372+Qn+IBj8Pe3KvGjHPj0sWwEF7AeZVhsNPZ6uY=
github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8=
Expand Down Expand Up @@ -151,6 +152,8 @@ github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWE
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
github.com/mfridman/interpolate v0.0.2 h1:pnuTK7MQIxxFz1Gr+rjSIx9u7qVjf5VOoM/u6BbAxPY=
github.com/mfridman/interpolate v0.0.2/go.mod h1:p+7uk6oE07mpE/Ik1b8EckO0O4ZXiGAfshKBWLUM9Xg=
github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db h1:62I3jR2EmQ4l5rM/4FEfDWcRD+abF5XlKShorW5LRoQ=
github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db/go.mod h1:l0dey0ia/Uv7NcFFVbCLtqEBQbrT4OCwCSKTEv6enCw=
github.com/mitchellh/mapstructure v1.5.0 h1:jeMsZIYE/09sWLaz43PL7Gy6RuMjD2eJVyuac5Z2hdY=
github.com/mitchellh/mapstructure v1.5.0/go.mod h1:bFUtVrKA4DC2yAKiSyO/QUcy7e+RRV2QTWOzhPopBRo=
github.com/moby/docker-image-spec v1.3.1 h1:jMKff3w6PgbfSa69GfNg+zN/XLhfXJGnEx3Nl2EsFP0=
Expand Down Expand Up @@ -214,13 +217,17 @@ github.com/prometheus/common v0.54.0 h1:ZlZy0BgJhTwVZUn7dLOkwCZHUkrAqd3WYtcFCWnM
github.com/prometheus/common v0.54.0/go.mod h1:/TQgMJP5CuVYveyT7n/0Ix8yLNNXy9yRSkhnLTHPDIQ=
github.com/prometheus/procfs v0.15.1 h1:YagwOFzUgYfKKHX6Dr+sHT7km/hxC76UB0learggepc=
github.com/prometheus/procfs v0.15.1/go.mod h1:fB45yRUv8NstnjriLhBQLuOUt+WW4BsoGhij/e3PBqk=
github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ=
github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=
github.com/rogpeppe/go-internal v1.12.0 h1:exVL4IDcn6na9z1rAb56Vxr+CgyK3nn3O+epU5NdKM8=
github.com/rogpeppe/go-internal v1.12.0/go.mod h1:E+RYuTGaKKdloAfM02xzb0FW3Paa99yedzYV+kq4uf4=
github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
github.com/sagikazarmark/locafero v0.6.0 h1:ON7AQg37yzcRPU69mt7gwhFEBwxI6P9T4Qu3N51bwOk=
github.com/sagikazarmark/locafero v0.6.0/go.mod h1:77OmuIc6VTraTXKXIs/uvUxKGUXjE1GbemJYHqdNjX0=
github.com/sagikazarmark/slog-shim v0.1.0 h1:diDBnUNK9N/354PgrxMywXnAwEr1QZcOr6gto+ugjYE=
github.com/sagikazarmark/slog-shim v0.1.0/go.mod h1:SrcSrq8aKtyuqEI1uvTDTK1arOWRIczQRv+GVI1AkeQ=
github.com/schollz/progressbar/v3 v3.14.4 h1:W9ZrDSJk7eqmQhd3uxFNNcTr0QL+xuGNI9dEMrw0r74=
github.com/schollz/progressbar/v3 v3.14.4/go.mod h1:aT3UQ7yGm+2ZjeXPqsjTenwL3ddUiuZ0kfQ/2tHlyNI=
github.com/sethvargo/go-retry v0.2.4 h1:T+jHEQy/zKJf5s95UkguisicE0zuF9y7+/vgz08Ocec=
github.com/sethvargo/go-retry v0.2.4/go.mod h1:1afjQuvh7s4gflMObvjLPaWgluLLyhA1wmVZ6KLpICw=
github.com/shirou/gopsutil/v3 v3.23.12 h1:z90NtUkp3bMtmICZKpC4+WaknU1eXtp5vtbQ11DgpE4=
Expand Down Expand Up @@ -347,8 +354,12 @@ golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7w
golang.org/x/sys v0.0.0-20211025201205-69cdffdb9359/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.20.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
golang.org/x/sys v0.21.0 h1:rF+pYz3DAGSQAxAu1CbC7catZg4ebC4UIeIhKxBZvws=
golang.org/x/sys v0.21.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
golang.org/x/term v0.20.0/go.mod h1:8UkIAJTvZgivsXaD6/pH6U9ecQzZ45awqEOzuCvwpFY=
golang.org/x/term v0.21.0 h1:WVXCp+/EBEHOj53Rvu+7KiT/iElMrO8ACK16SMZ3jaA=
golang.org/x/term v0.21.0/go.mod h1:ooXLefLobQVslOqselCNF4SxFAaoS6KujMbsGzSDmX0=
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
golang.org/x/text v0.16.0 h1:a94ExnEXNtEwYLGJSIUxnWoxoRz/ZcCsV63ROupILh4=
Expand Down

0 comments on commit 601021c

Please sign in to comment.