Skip to content

Commit

Permalink
fix: don't attempt to import tuples if none are present (#408)
Browse files Browse the repository at this point in the history
  • Loading branch information
ewanharris authored Oct 21, 2024
2 parents 4156d21 + 7e2a509 commit d05efa8
Showing 1 changed file with 61 additions and 44 deletions.
105 changes: 61 additions & 44 deletions cmd/store/import.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ import (

"github.com/schollz/progressbar/v3"

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

Expand Down Expand Up @@ -116,30 +117,74 @@ func importStore(
maxTuplesPerWrite, maxParallelRequests int32,
fileName string,
) (*CreateStoreAndModelResponse, error) {
var (
response *CreateStoreAndModelResponse
err error
)
response, err := createOrUpdateStore(clientConfig, fgaClient, storeData, format, storeID, fileName)
if err != nil {
return nil, err
}

if len(storeData.Tuples) == 0 {
return response, nil
}

err = importTuples(fgaClient, storeData.Tuples, maxTuplesPerWrite, maxParallelRequests)
if err != nil {
return nil, err
}

return response, nil
}

func createOrUpdateStore(
clientConfig *fga.ClientConfig,
fgaClient client.SdkClient,
storeData *storetest.StoreData,
format authorizationmodel.ModelFormat,
storeID string,
fileName string,
) (*CreateStoreAndModelResponse, error) {
if storeID == "" {
response, err = createStore(clientConfig, storeData, format, fileName)
if err != nil {
return nil, fmt.Errorf("failed to create store: %w", err)
return createStore(clientConfig, storeData, format, fileName)
}

return updateStore(clientConfig, fgaClient, storeData, format, storeID)
}

func importTuples(
fgaClient client.SdkClient,
tuples []openfga.TupleKey,
maxTuplesPerWrite, maxParallelRequests int32,
) error {
bar := createProgressBar(len(tuples))

for index := 0; index < len(tuples); index += int(maxTuplesPerWrite) {
end := index + int(maxTuplesPerWrite)
if end > len(tuples) {
end = len(tuples)
}
} else {
response, err = updateStore(clientConfig, fgaClient, storeData, format, storeID)
if err != nil {
return nil, fmt.Errorf("failed to update store: %w", err)

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

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

time.Sleep(progressBarUpdateDelay)
}

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

// Initialize progress bar
bar := progressbar.NewOptions(len(storeData.Tuples),
return nil
}

func createProgressBar(total int) *progressbar.ProgressBar {
return progressbar.NewOptions(total,
progressbar.OptionSetWriter(os.Stderr),
progressbar.OptionSetDescription("Importing tuples"),
progressbar.OptionShowCount(),
Expand All @@ -158,34 +203,6 @@ func importStore(
BarEnd: "]",
}),
)

for index := 0; index < len(storeData.Tuples); index += int(maxTuplesPerWrite) {
end := index + int(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)
}

// Ensure progress bar is completed and cleared
if err := bar.Finish(); err != nil {
return nil, fmt.Errorf("failed to finish progress bar: %w", err)
}

return response, nil
}

// importCmd represents the get command.
Expand Down

0 comments on commit d05efa8

Please sign in to comment.