导出系统

pkg/storage/s3.go

@@ -73,7 +73,7 @@ func (p *S3Provider) Upload(ctx context.Context, key string, reader io.Reader, c
 	return nil
 }
 
-func (p *S3Provider) Download(ctx context.Context, key string, writer io.Writer) error {
+func (p *S3Provider) Download(ctx context.Context, key string) (io.ReadCloser, error) {
 	input := &s3.GetObjectInput{
 		Bucket: aws.String(p.bucket),
 		Key:    aws.String(key),
@@ -82,17 +82,11 @@ func (p *S3Provider) Download(ctx context.Context, key string, writer io.Writer)
 	result, err := p.client.GetObjectWithContext(ctx, input)
 	if err != nil {
 		if strings.Contains(err.Error(), "NoSuchKey") {
-			return fmt.Errorf("文件不存在: %s", key)
+			return nil, fmt.Errorf("文件不存在: %s", key)
 		}
-		return fmt.Errorf("下载文件失败: %w", err)
+		return nil, fmt.Errorf("下载文件失败: %w", err)
 	}
-	defer result.Body.Close()
-
-	_, err = io.Copy(writer, result.Body)
-	if err != nil {
-		return fmt.Errorf("写入文件内容失败: %w", err)
-	}
-	return nil
+	return result.Body, nil
 }
 
 func (p *S3Provider) DownloadToTemp(ctx context.Context, key string) (string, func(), error) {
@@ -111,11 +105,19 @@ func (p *S3Provider) DownloadToTemp(ctx context.Context, key string) (string, fu
 		_ = os.Remove(tempPath)
 	}
 
-	if err := p.Download(ctx, key, tempFile); err != nil {
+	reader, err := p.Download(ctx, key)
+	if err != nil {
 		tempFile.Close()
 		cleanup()
 		return "", nil, err
 	}
+	defer reader.Close()
+
+	if _, err := io.Copy(tempFile, reader); err != nil {
+		tempFile.Close()
+		cleanup()
+		return "", nil, fmt.Errorf("写入文件内容失败: %w", err)
+	}
 
 	if err := tempFile.Close(); err != nil {
 		cleanup()