Fix to respect environment specific cloud storage (#162)

* Fix to respect environment specific cloud storage

* fix integration test

---------

Co-authored-by: James Kwon <96548424+hongil0316@users.noreply.github.com>

Author: james03160927

gateways/storage/files.go

@@ -1,7 +1,6 @@
 package storage
 
 import (
-	"bytes"
 	"context"
 	"fmt"
 	"io"
@@ -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)
 }
@@ -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")()

integration-tests/comfy_node_test.go

@@ -21,7 +21,9 @@ func TestRegistryComfyNode(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 test user context and publisher

integration-tests/node_version_test.go

@@ -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
@@ -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.
@@ -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)

services/registry/registry_svc.go

@@ -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 {