feat(service): add file extension in output filename

instill-ai · Nov 25, 2024 · 20bab6f · 20bab6f
1 parent 5086140
commit 20bab6f
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 6 deletions.
diff --git a/pkg/data/file.go b/pkg/data/file.go
@@ -8,6 +8,7 @@ import (
 	"unicode/utf8"
 
 	"github.com/gabriel-vasile/mimetype"
+	"github.com/gofrs/uuid"
 	"google.golang.org/protobuf/types/known/structpb"
 
 	"github.com/instill-ai/pipeline-backend/pkg/data/format"
@@ -43,6 +44,11 @@ func NewFileFromBytes(b []byte, contentType, filename string) (bin *fileData, er
 		contentType = strings.Split(mimetype.Detect(b).String(), ";")[0]
 	}
 
+	if filename == "" {
+		fileUID, _ := uuid.NewV4()
+		filename = fmt.Sprintf("%s.%s", fileUID, strings.ToLower(mimetype.Detect(b).Extension()))
+	}
+
 	f := &fileData{
 		raw:         b,
 		contentType: contentType,

diff --git a/pkg/worker/workflow.go b/pkg/worker/workflow.go
@@ -669,7 +669,7 @@ func (w *worker) uploadFileAndReplaceWithURL(ctx context.Context, param *Compone
 	}
 	switch v := (*value).(type) {
 	case format.File:
-		downloadURL, err := w.uploadBlobDataAndGetDownloadURL(ctx, param, &v)
+		downloadURL, err := w.uploadBlobDataAndGetDownloadURL(ctx, param, v)
 		if err != nil || downloadURL == "" {
 			logger.Warn("uploading blob data", zap.Error(err))
 			return v
@@ -692,15 +692,15 @@ func (w *worker) uploadFileAndReplaceWithURL(ctx context.Context, param *Compone
 	}
 }
 
-func (w *worker) uploadBlobDataAndGetDownloadURL(ctx context.Context, param *ComponentActivityParam, value *format.File) (string, error) {
+func (w *worker) uploadBlobDataAndGetDownloadURL(ctx context.Context, param *ComponentActivityParam, value format.File) (string, error) {
 	artifactClient := w.artifactPublicServiceClient
 	requesterID := param.SystemVariables.PipelineRequesterID
 
 	sysVarJSON := utils.StructToMap(param.SystemVariables, "json")
 
 	ctx = metadata.NewOutgoingContext(ctx, getRequestMetadata(sysVarJSON))
 
-	objectName := fmt.Sprintf("%s/%s", requesterID, uuid.Must(uuid.NewV4()).String())
+	objectName := fmt.Sprintf("%s/%s", requesterID, value.Filename())
 	resp, err := artifactClient.GetObjectUploadURL(ctx, &artifactpb.GetObjectUploadURLRequest{
 		NamespaceId:      requesterID,
 		ObjectName:       objectName,
@@ -729,7 +729,7 @@ func (w *worker) uploadBlobDataAndGetDownloadURL(ctx context.Context, param *Com
 	return respDownloadURL.GetDownloadUrl(), nil
 }
 
-func uploadBlobData(ctx context.Context, uploadURL string, value *format.File, logger *zap.Logger) error {
+func uploadBlobData(ctx context.Context, uploadURL string, value format.File, logger *zap.Logger) error {
 	if uploadURL == "" {
 		return fmt.Errorf("empty upload URL provided")
 	}
@@ -745,8 +745,8 @@ func uploadBlobData(ctx context.Context, uploadURL string, value *format.File, l
 	}
 	parsedURL.Host = fmt.Sprintf("%s:%d", config.Config.APIGateway.Host, config.Config.APIGateway.PublicPort)
 	fullURL := parsedURL.String()
-	contentType := (*value).ContentType().String()
-	fileBytes, err := (*value).Binary()
+	contentType := value.ContentType().String()
+	fileBytes, err := value.Binary()
 
 	if err != nil {
 		return fmt.Errorf("getting file bytes: %w", err)