Skip to content

Commit

Permalink
Fix to respect environment specific cloud storage
Browse files Browse the repository at this point in the history
  • Loading branch information
james03160927 committed Feb 2, 2025
1 parent 3aa3a99 commit 1189815
Show file tree
Hide file tree
Showing 3 changed files with 6 additions and 33 deletions.
29 changes: 0 additions & 29 deletions gateways/storage/files.go
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
package storage

import (
"bytes"
"context"
"fmt"
"io"
Expand All @@ -18,7 +17,6 @@ import (
// StorageService defines the interface for interacting with cloud storage.
type StorageService interface {
UploadFile(ctx context.Context, bucket, object, filePath string) (string, error)
StreamFileUpload(w io.Writer, objectName, blob string) (string, string, error)
GetFileUrl(ctx context.Context, bucketName, objectPath string) (string, error)
GenerateSignedURL(bucketName, objectName string) (string, error)
}
Expand Down Expand Up @@ -95,33 +93,6 @@ func (s *storageService) UploadFile(ctx context.Context, bucket, object, filePat
return publicURL, nil
}

// StreamFileUpload uploads an object via a stream to GCP storage.
func (s *storageService) StreamFileUpload(w io.Writer, objectName, blob string) (string, string, error) {
ctx := context.Background()

b := []byte(blob)
buf := bytes.NewBuffer(b)

ctx, cancel := context.WithTimeout(ctx, time.Second*50)
defer cancel()

// Upload the object as a stream
wc := s.client.Bucket(s.config.CloudStorageBucketName).Object(objectName).NewWriter(ctx)
wc.ChunkSize = 0 // Note: retries are not supported for chunk size 0

if _, err := io.Copy(wc, buf); err != nil {
return "", "", fmt.Errorf("io.Copy: %w", err)
}

if err := wc.Close(); err != nil {
return "", "", fmt.Errorf("Writer.Close: %w", err)
}

log.Ctx(ctx).Info().Msgf("%v uploaded to %v.\n", objectName, s.config.CloudStorageBucketName)

return s.config.CloudStorageBucketName, objectName, nil
}

// GetFileUrl gets the public URL of a file from GCP storage.
func (s *storageService) GetFileUrl(ctx context.Context, bucketName, objectPath string) (string, error) {
defer tracing.TraceDefaultSegment(ctx, "StorageService.GetFileUrl")()
Expand Down
8 changes: 5 additions & 3 deletions integration-tests/node_version_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,9 @@ func TestRegistryNodeVersion(t *testing.T) {
defer cleanup()

// Initialize server implementation and authorization middleware
impl := NewStrictServerImplementationWithMocks(client, &config.Config{})
impl := NewStrictServerImplementationWithMocks(client, &config.Config{
CloudStorageBucketName: "test-bucket",
})
authz := authorization.NewAuthorizationManager(client, impl.RegistryService, impl.NewRelicApp).AuthorizationMiddleware()

// Setup user context and publisher
Expand All @@ -47,7 +49,7 @@ func TestRegistryNodeVersion(t *testing.T) {
nodeVersion := randomNodeVersion(0)
signedUrl := "test-url"
downloadUrl := fmt.Sprintf(
"https://storage.googleapis.com/comfy-registry/%s/%s/%s/node.zip", publisherId, *node.Id, *nodeVersion.Version)
"https://storage.googleapis.com/test-bucket/%s/%s/%s/node.zip", publisherId, *node.Id, *nodeVersion.Version)

// Mock external service responses for storage and Discord
impl.mockStorageService.
Expand Down Expand Up @@ -388,7 +390,7 @@ func TestRegistryNodeVersion(t *testing.T) {
// Creating a new random node and version for scanning
node := randomNode()
nodeVersion := randomNodeVersion(0)
downloadUrl := fmt.Sprintf("https://storage.googleapis.com/comfy-registry/%s/%s/%s/node.zip", publisherId, *node.Id, *nodeVersion.Version)
downloadUrl := fmt.Sprintf("https://storage.googleapis.com/test-bucket/%s/%s/%s/node.zip", publisherId, *node.Id, *nodeVersion.Version)

// Mocking the behavior of services for URL generation and message sending
impl.mockStorageService.On("GenerateSignedURL", mock.Anything, mock.Anything).Return("test-url", nil)
Expand Down
2 changes: 1 addition & 1 deletion services/registry/registry_svc.go
Original file line number Diff line number Diff line change
Expand Up @@ -389,7 +389,7 @@ func (s *RegistryService) CreateNodeVersion(

log.Ctx(ctx).Info().Msgf(
"creating node version: %v for nodeId %v", nodeVersion, nodeID)
bucketName := "comfy-registry"
bucketName := s.config.CloudStorageBucketName
return db.WithTxResult(ctx, client, func(tx *ent.Tx) (*NodeVersionCreation, error) {
// If the node version is not provided, we will generate a new version
if nodeVersion.Version != nil {
Expand Down

0 comments on commit 1189815

Please sign in to comment.