-
Notifications
You must be signed in to change notification settings - Fork 6
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[PSL-1239] cloud storage handler (#913)
* [PSL-1239] cloud storage handler
- Loading branch information
1 parent
167864a
commit d9ae328
Showing
9 changed files
with
371 additions
and
31 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,128 @@ | ||
| package cloud | ||
|
|
||
| import ( | ||
| "bytes" | ||
| "fmt" | ||
|
|
||
| "os" | ||
| "os/exec" | ||
| "path/filepath" | ||
| "sync" | ||
|
|
||
| "github.com/pastelnetwork/gonode/common/log" | ||
| ) | ||
|
|
||
| type Storage interface { | ||
| Store(key string, data []byte) (string, error) | ||
| Fetch(key string) ([]byte, error) | ||
| StoreBatch(data [][]byte) error | ||
| FetchBatch(keys []string) (map[string][]byte, error) | ||
| } | ||
|
|
||
| type RcloneStorage struct { | ||
| bucketName string | ||
| specName string | ||
| } | ||
|
|
||
| func NewRcloneStorage(bucketName, specName string) *RcloneStorage { | ||
| return &RcloneStorage{ | ||
| bucketName: bucketName, | ||
| specName: specName, | ||
| } | ||
| } | ||
|
|
||
| func (r *RcloneStorage) Store(key string, data []byte) (string, error) { | ||
| filePath := filepath.Join(os.TempDir(), key) | ||
|
|
||
| // Write data to a temporary file using os.WriteFile | ||
| if err := os.WriteFile(filePath, data, 0644); err != nil { | ||
| return "", fmt.Errorf("failed to write data to file: %w", err) | ||
| } | ||
|
|
||
| // Construct the remote path where the file will be stored | ||
| // This example places the file at the root of the remote, but you can modify the path as needed | ||
| remotePath := fmt.Sprintf("%s:%s/%s", r.specName, r.bucketName, key) | ||
|
|
||
| // Use rclone to copy the file to the remote | ||
| cmd := exec.Command("rclone", "copyto", filePath, remotePath) | ||
| if err := cmd.Run(); err != nil { | ||
| // Clean up the local file if the upload fails | ||
| os.Remove(filePath) | ||
| return "", fmt.Errorf("rclone command failed: %w", err) | ||
| } | ||
|
|
||
| // Delete the local file after successful upload | ||
| go func() { | ||
| if err := os.Remove(filePath); err != nil { | ||
| log.Error("failed to delete local file", "path", filePath, "error", err) | ||
| } | ||
| }() | ||
|
|
||
| // Return the remote path where the file was stored | ||
| return remotePath, nil | ||
| } | ||
|
|
||
| func (r *RcloneStorage) Fetch(key string) ([]byte, error) { | ||
| // Construct the rclone command to fetch the file | ||
| cmd := exec.Command("rclone", "cat", fmt.Sprintf("%s:%s/%s", r.specName, r.bucketName, key)) | ||
| var out bytes.Buffer | ||
| cmd.Stdout = &out | ||
| err := cmd.Run() | ||
| if err != nil { | ||
| return nil, fmt.Errorf("rclone command failed: %w - out %s", err, out.String()) | ||
| } | ||
|
|
||
| return out.Bytes(), nil | ||
| } | ||
|
|
||
| func (r *RcloneStorage) StoreBatch(data [][]byte) error { | ||
| // Placeholder for StoreBatch implementation | ||
| return nil | ||
| } | ||
|
|
||
| func (r *RcloneStorage) FetchBatch(keys []string) (map[string][]byte, error) { | ||
| results := make(map[string][]byte) | ||
| errs := make(map[string]error) | ||
| var mu sync.Mutex | ||
|
|
||
| semaphore := make(chan struct{}, 50) | ||
|
|
||
| var wg sync.WaitGroup | ||
| for _, key := range keys { | ||
| wg.Add(1) | ||
| semaphore <- struct{}{} // Acquire a token | ||
|
|
||
| go func(key string) { | ||
| defer wg.Done() | ||
| data, err := r.Fetch(key) | ||
|
|
||
| func() { | ||
| mu.Lock() | ||
| defer mu.Unlock() | ||
| if err != nil { | ||
| errs[key] = err | ||
| } else { | ||
| results[key] = data | ||
| } | ||
|
|
||
| }() | ||
| <-semaphore // Release the token | ||
| }(key) | ||
| } | ||
|
|
||
| wg.Wait() | ||
|
|
||
| if len(results) > 0 { | ||
| return results, nil | ||
| } | ||
|
|
||
| if len(errs) > 0 { | ||
| combinedError := fmt.Errorf("errors occurred in fetching keys") | ||
| for k, e := range errs { | ||
| combinedError = fmt.Errorf("%v; key %s error: %v", combinedError, k, e) | ||
| } | ||
| return nil, combinedError | ||
| } | ||
|
|
||
| return results, nil | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,102 @@ | ||
| package cloud | ||
|
|
||
| /* | ||
| import ( | ||
| "fmt" | ||
| "os" | ||
| "testing" | ||
| ) | ||
| func TestRcloneStorage_Fetch(t *testing.T) { | ||
| SetRcloneB2Env("0058a83f04e57580000000002", "K005lP8Wd0Gvr4JOdEI6e6BTAyt6iZA") | ||
| storage := NewRcloneStorage("pastel-test", "b2") | ||
| tests := map[string]struct { | ||
| key string | ||
| expected string | ||
| wantErr bool | ||
| }{ | ||
| "File Exists": { | ||
| key: "copy_45707.txt", | ||
| expected: "expected content of testfile.txt", | ||
| wantErr: false, | ||
| }, | ||
| "File Does Not Exist": { | ||
| key: "copy_45707_nonexistent.txt", | ||
| wantErr: true, | ||
| }, | ||
| } | ||
| for name, tc := range tests { | ||
| t.Run(name, func(t *testing.T) { | ||
| data, err := storage.Fetch(tc.key) | ||
| if (err != nil) != tc.wantErr { | ||
| t.Errorf("Fetch() error = %v, wantErr %v", err, tc.wantErr) | ||
| } | ||
| if !tc.wantErr && string(data) != tc.expected { | ||
| t.Errorf("Fetch() got = %v, want %v", string(data), tc.expected) | ||
| } | ||
| }) | ||
| } | ||
| } | ||
| // SetRcloneB2Env sets the environment variables for rclone configuration. | ||
| func SetRcloneB2Env(accountID, appKey string) error { | ||
| err := os.Setenv("RCLONE_CONFIG_B2_TYPE", "b2") | ||
| if err != nil { | ||
| return fmt.Errorf("failed to set RCLONE_CONFIG_B2_TYPE: %w", err) | ||
| } | ||
| err = os.Setenv("RCLONE_CONFIG_B2_ACCOUNT", accountID) | ||
| if err != nil { | ||
| return fmt.Errorf("failed to set RCLONE_CONFIG_B2_ACCOUNT: %w", err) | ||
| } | ||
| err = os.Setenv("RCLONE_CONFIG_B2_KEY", appKey) | ||
| if err != nil { | ||
| return fmt.Errorf("failed to set RCLONE_CONFIG_B2_KEY: %w", err) | ||
| } | ||
| return nil | ||
| } | ||
| func TestRcloneStorage_FetchBatch(t *testing.T) { | ||
| storage := NewRcloneStorage("mybucket") | ||
| tests := map[string]struct { | ||
| keys []string | ||
| expected map[string]string | ||
| wantErr bool | ||
| }{ | ||
| "Multiple Files Exist": { | ||
| keys: []string{"testfile1.txt", "testfile2.txt"}, // Ensure these files exist | ||
| expected: map[string]string{ | ||
| "testfile1.txt": "content of testfile1.txt", | ||
| "testfile2.txt": "content of testfile2.txt", | ||
| }, | ||
| wantErr: false, | ||
| }, | ||
| "Some Files Do Not Exist": { | ||
| keys: []string{"testfile1.txt", "doesNotExist.txt"}, | ||
| wantErr: true, | ||
| }, | ||
| } | ||
| for name, tc := range tests { | ||
| t.Run(name, func(t *testing.T) { | ||
| results, err := storage.FetchBatch(tc.keys) | ||
| if (err != nil) != tc.wantErr { | ||
| t.Errorf("FetchBatch() error = %v, wantErr %v", err, tc.wantErr) | ||
| } | ||
| if !tc.wantErr { | ||
| for k, v := range tc.expected { | ||
| if results[k] != v { | ||
| t.Errorf("FetchBatch() got = %v, want %v for key %v", string(results[k]), v, k) | ||
| } | ||
| } | ||
| } | ||
| }) | ||
| } | ||
| } | ||
| */ |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.