feat(infra): support uploading files via io.Reader (#1793)

This commit is contained in:
Ryo 2025-08-25 18:25:12 +08:00 committed by GitHub
parent 6fa2acf05a
commit 14ce6bc112
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 115 additions and 36 deletions

View File

@ -38,6 +38,7 @@ type PutOption struct {
ContentDisposition *string ContentDisposition *string
ContentLanguage *string ContentLanguage *string
Expires *time.Time Expires *time.Time
ObjectSize int64
} }
type PutOptFn func(option *PutOption) type PutOptFn func(option *PutOption)
@ -48,6 +49,12 @@ func WithContentType(v string) PutOptFn {
} }
} }
func WithObjectSize(v int64) PutOptFn {
return func(o *PutOption) {
o.ObjectSize = v
}
}
func WithContentEncoding(v string) PutOptFn { func WithContentEncoding(v string) PutOptFn {
return func(o *PutOption) { return func(o *PutOption) {
o.ContentEncoding = &v o.ContentEncoding = &v

View File

@ -16,11 +16,15 @@
package storage package storage
import "context" import (
"context"
"io"
)
//go:generate mockgen -destination ../../../internal/mock/infra/contract/storage/storage_mock.go -package mock -source storage.go Factory //go:generate mockgen -destination ../../../internal/mock/infra/contract/storage/storage_mock.go -package mock -source storage.go Factory
type Storage interface { type Storage interface {
PutObject(ctx context.Context, objectKey string, content []byte, opts ...PutOptFn) error PutObject(ctx context.Context, objectKey string, content []byte, opts ...PutOptFn) error
PutObjectWithReader(ctx context.Context, objectKey string, content io.Reader, opts ...PutOptFn) error
GetObject(ctx context.Context, objectKey string) ([]byte, error) GetObject(ctx context.Context, objectKey string) ([]byte, error)
DeleteObject(ctx context.Context, objectKey string) error DeleteObject(ctx context.Context, objectKey string) error
GetObjectUrl(ctx context.Context, objectKey string, opts ...GetOptFn) (string, error) GetObjectUrl(ctx context.Context, objectKey string, opts ...GetOptFn) (string, error)

View File

@ -138,6 +138,11 @@ func (m *minioClient) test() {
} }
func (m *minioClient) PutObject(ctx context.Context, objectKey string, content []byte, opts ...storage.PutOptFn) error { func (m *minioClient) PutObject(ctx context.Context, objectKey string, content []byte, opts ...storage.PutOptFn) error {
opts = append(opts, storage.WithObjectSize(int64(len(content))))
return m.PutObjectWithReader(ctx, objectKey, bytes.NewReader(content), opts...)
}
func (m *minioClient) PutObjectWithReader(ctx context.Context, objectKey string, content io.Reader, opts ...storage.PutOptFn) error {
option := storage.PutOption{} option := storage.PutOption{}
for _, opt := range opts { for _, opt := range opts {
opt(&option) opt(&option)
@ -165,7 +170,7 @@ func (m *minioClient) PutObject(ctx context.Context, objectKey string, content [
} }
_, err := m.client.PutObject(ctx, m.bucketName, objectKey, _, err := m.client.PutObject(ctx, m.bucketName, objectKey,
bytes.NewReader(content), int64(len(content)), minioOpts) content, option.ObjectSize, minioOpts)
if err != nil { if err != nil {
return fmt.Errorf("PutObject failed: %v", err) return fmt.Errorf("PutObject failed: %v", err)
} }

View File

@ -152,16 +152,47 @@ func (t *s3Client) CheckAndCreateBucket(ctx context.Context) error {
} }
func (t *s3Client) PutObject(ctx context.Context, objectKey string, content []byte, opts ...storage.PutOptFn) error { func (t *s3Client) PutObject(ctx context.Context, objectKey string, content []byte, opts ...storage.PutOptFn) error {
opts = append(opts, storage.WithObjectSize(int64(len(content))))
return t.PutObjectWithReader(ctx, objectKey, bytes.NewReader(content), opts...)
}
func (t *s3Client) PutObjectWithReader(ctx context.Context, objectKey string, content io.Reader, opts ...storage.PutOptFn) error {
client := t.client client := t.client
body := bytes.NewReader(content)
bucket := t.bucketName bucket := t.bucketName
// upload object option := storage.PutOption{}
_, err := client.PutObject(ctx, &s3.PutObjectInput{ for _, opt := range opts {
opt(&option)
}
input := &s3.PutObjectInput{
Bucket: aws.String(bucket), Bucket: aws.String(bucket),
Key: aws.String(objectKey), Key: aws.String(objectKey),
Body: body, Body: content,
}) }
if option.ContentType != nil {
input.ContentType = option.ContentType
}
if option.ContentEncoding != nil {
input.ContentEncoding = option.ContentEncoding
}
if option.ContentDisposition != nil {
input.ContentDisposition = option.ContentDisposition
}
if option.ContentLanguage != nil {
input.ContentLanguage = option.ContentLanguage
}
if option.Expires != nil {
input.Expires = option.Expires
}
if option.ObjectSize > 0 {
input.ContentLength = aws.Int64(option.ObjectSize)
}
// upload object
_, err := client.PutObject(ctx, input)
return err return err
} }

View File

@ -140,22 +140,49 @@ func (t *tosClient) CheckAndCreateBucket(ctx context.Context) error {
return err return err
} }
func (t *tosClient) PutObject(ctx context.Context, objectKey string, content []byte, opts ...storage.PutOptFn) error { func (t *tosClient) PutObject(ctx context.Context, objectKey string, content []byte, opts ...storage.PutOptFn) error {
opts = append(opts, storage.WithObjectSize(int64(len(content))))
return t.PutObjectWithReader(ctx, objectKey, bytes.NewReader(content), opts...)
}
func (t *tosClient) PutObjectWithReader(ctx context.Context, objectKey string, content io.Reader, opts ...storage.PutOptFn) error {
client := t.client client := t.client
body := bytes.NewReader(content)
bucketName := t.bucketName bucketName := t.bucketName
_, err := client.PutObjectV2(ctx, &tos.PutObjectV2Input{ option := storage.PutOption{}
for _, opt := range opts {
opt(&option)
}
input := &tos.PutObjectV2Input{
PutObjectBasicInput: tos.PutObjectBasicInput{ PutObjectBasicInput: tos.PutObjectBasicInput{
Bucket: bucketName, Bucket: bucketName,
Key: objectKey, Key: objectKey,
}, },
Content: body, Content: content,
}) }
// logs.CtxDebugf(ctx, "PutObject resp: %v, err: %v", conv.DebugJsonToStr(output), err) if option.ContentType != nil {
input.ContentType = *option.ContentType
}
if option.ContentEncoding != nil {
input.ContentEncoding = *option.ContentEncoding
}
if option.ContentDisposition != nil {
input.ContentDisposition = *option.ContentDisposition
}
if option.ContentLanguage != nil {
input.ContentLanguage = *option.ContentLanguage
}
if option.Expires != nil {
input.Expires = *option.Expires
}
if option.ObjectSize > 0 {
input.ContentLength = option.ObjectSize
}
_, err := client.PutObjectV2(ctx, input)
return err return err
} }

View File

@ -11,6 +11,7 @@ package mock
import ( import (
context "context" context "context"
io "io"
reflect "reflect" reflect "reflect"
storage "github.com/coze-dev/coze-studio/backend/infra/contract/storage" storage "github.com/coze-dev/coze-studio/backend/infra/contract/storage"
@ -50,7 +51,7 @@ func (m *MockStorage) DeleteObject(ctx context.Context, objectKey string) error
} }
// DeleteObject indicates an expected call of DeleteObject. // DeleteObject indicates an expected call of DeleteObject.
func (mr *MockStorageMockRecorder) DeleteObject(ctx, objectKey interface{}) *gomock.Call { func (mr *MockStorageMockRecorder) DeleteObject(ctx, objectKey any) *gomock.Call {
mr.mock.ctrl.T.Helper() mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteObject", reflect.TypeOf((*MockStorage)(nil).DeleteObject), ctx, objectKey) return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteObject", reflect.TypeOf((*MockStorage)(nil).DeleteObject), ctx, objectKey)
} }
@ -65,7 +66,7 @@ func (m *MockStorage) GetObject(ctx context.Context, objectKey string) ([]byte,
} }
// GetObject indicates an expected call of GetObject. // GetObject indicates an expected call of GetObject.
func (mr *MockStorageMockRecorder) GetObject(ctx, objectKey interface{}) *gomock.Call { func (mr *MockStorageMockRecorder) GetObject(ctx, objectKey any) *gomock.Call {
mr.mock.ctrl.T.Helper() mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetObject", reflect.TypeOf((*MockStorage)(nil).GetObject), ctx, objectKey) return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetObject", reflect.TypeOf((*MockStorage)(nil).GetObject), ctx, objectKey)
} }
@ -73,7 +74,7 @@ func (mr *MockStorageMockRecorder) GetObject(ctx, objectKey interface{}) *gomock
// GetObjectUrl mocks base method. // GetObjectUrl mocks base method.
func (m *MockStorage) GetObjectUrl(ctx context.Context, objectKey string, opts ...storage.GetOptFn) (string, error) { func (m *MockStorage) GetObjectUrl(ctx context.Context, objectKey string, opts ...storage.GetOptFn) (string, error) {
m.ctrl.T.Helper() m.ctrl.T.Helper()
varargs := []interface{}{ctx, objectKey} varargs := []any{ctx, objectKey}
for _, a := range opts { for _, a := range opts {
varargs = append(varargs, a) varargs = append(varargs, a)
} }
@ -84,31 +85,16 @@ func (m *MockStorage) GetObjectUrl(ctx context.Context, objectKey string, opts .
} }
// GetObjectUrl indicates an expected call of GetObjectUrl. // GetObjectUrl indicates an expected call of GetObjectUrl.
func (mr *MockStorageMockRecorder) GetObjectUrl(ctx, objectKey interface{}, opts ...interface{}) *gomock.Call { func (mr *MockStorageMockRecorder) GetObjectUrl(ctx, objectKey any, opts ...any) *gomock.Call {
mr.mock.ctrl.T.Helper() mr.mock.ctrl.T.Helper()
varargs := append([]interface{}{ctx, objectKey}, opts...) varargs := append([]any{ctx, objectKey}, opts...)
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetObjectUrl", reflect.TypeOf((*MockStorage)(nil).GetObjectUrl), varargs...) return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetObjectUrl", reflect.TypeOf((*MockStorage)(nil).GetObjectUrl), varargs...)
} }
// GetUploadAuth mocks base method.
func (m *MockStorage) GetUploadAuth(ctx context.Context) (*storage.SecurityToken, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetUploadAuth", ctx)
ret0, _ := ret[0].(*storage.SecurityToken)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetUploadAuth indicates an expected call of GetUploadAuth.
func (mr *MockStorageMockRecorder) GetUploadAuth(ctx interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUploadAuth", reflect.TypeOf((*MockStorage)(nil).GetUploadAuth), ctx)
}
// PutObject mocks base method. // PutObject mocks base method.
func (m *MockStorage) PutObject(ctx context.Context, objectKey string, content []byte, opts ...storage.PutOptFn) error { func (m *MockStorage) PutObject(ctx context.Context, objectKey string, content []byte, opts ...storage.PutOptFn) error {
m.ctrl.T.Helper() m.ctrl.T.Helper()
varargs := []interface{}{ctx, objectKey, content} varargs := []any{ctx, objectKey, content}
for _, a := range opts { for _, a := range opts {
varargs = append(varargs, a) varargs = append(varargs, a)
} }
@ -118,8 +104,27 @@ func (m *MockStorage) PutObject(ctx context.Context, objectKey string, content [
} }
// PutObject indicates an expected call of PutObject. // PutObject indicates an expected call of PutObject.
func (mr *MockStorageMockRecorder) PutObject(ctx, objectKey, content interface{}, opts ...interface{}) *gomock.Call { func (mr *MockStorageMockRecorder) PutObject(ctx, objectKey, content any, opts ...any) *gomock.Call {
mr.mock.ctrl.T.Helper() mr.mock.ctrl.T.Helper()
varargs := append([]interface{}{ctx, objectKey, content}, opts...) varargs := append([]any{ctx, objectKey, content}, opts...)
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PutObject", reflect.TypeOf((*MockStorage)(nil).PutObject), varargs...) return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PutObject", reflect.TypeOf((*MockStorage)(nil).PutObject), varargs...)
} }
// PutObjectWithReader mocks base method.
func (m *MockStorage) PutObjectWithReader(ctx context.Context, objectKey string, content io.Reader, opts ...storage.PutOptFn) error {
m.ctrl.T.Helper()
varargs := []any{ctx, objectKey, content}
for _, a := range opts {
varargs = append(varargs, a)
}
ret := m.ctrl.Call(m, "PutObjectWithReader", varargs...)
ret0, _ := ret[0].(error)
return ret0
}
// PutObjectWithReader indicates an expected call of PutObjectWithReader.
func (mr *MockStorageMockRecorder) PutObjectWithReader(ctx, objectKey, content any, opts ...any) *gomock.Call {
mr.mock.ctrl.T.Helper()
varargs := append([]any{ctx, objectKey, content}, opts...)
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PutObjectWithReader", reflect.TypeOf((*MockStorage)(nil).PutObjectWithReader), varargs...)
}