diff --git a/backend/infra/contract/storage/option.go b/backend/infra/contract/storage/option.go index 831d2e17..127e7e84 100644 --- a/backend/infra/contract/storage/option.go +++ b/backend/infra/contract/storage/option.go @@ -38,6 +38,7 @@ type PutOption struct { ContentDisposition *string ContentLanguage *string Expires *time.Time + ObjectSize int64 } 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 { return func(o *PutOption) { o.ContentEncoding = &v diff --git a/backend/infra/contract/storage/storage.go b/backend/infra/contract/storage/storage.go index 252da4b3..70b09a87 100644 --- a/backend/infra/contract/storage/storage.go +++ b/backend/infra/contract/storage/storage.go @@ -16,11 +16,15 @@ 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 type Storage interface { 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) DeleteObject(ctx context.Context, objectKey string) error GetObjectUrl(ctx context.Context, objectKey string, opts ...GetOptFn) (string, error) diff --git a/backend/infra/impl/storage/minio/minio.go b/backend/infra/impl/storage/minio/minio.go index fd852020..12c7127f 100644 --- a/backend/infra/impl/storage/minio/minio.go +++ b/backend/infra/impl/storage/minio/minio.go @@ -138,6 +138,11 @@ func (m *minioClient) test() { } 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{} for _, opt := range opts { opt(&option) @@ -165,7 +170,7 @@ func (m *minioClient) PutObject(ctx context.Context, objectKey string, content [ } _, err := m.client.PutObject(ctx, m.bucketName, objectKey, - bytes.NewReader(content), int64(len(content)), minioOpts) + content, option.ObjectSize, minioOpts) if err != nil { return fmt.Errorf("PutObject failed: %v", err) } diff --git a/backend/infra/impl/storage/s3/s3.go b/backend/infra/impl/storage/s3/s3.go index 3e5025ee..1b848d20 100644 --- a/backend/infra/impl/storage/s3/s3.go +++ b/backend/infra/impl/storage/s3/s3.go @@ -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 { + 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 - body := bytes.NewReader(content) bucket := t.bucketName - // upload object - _, err := client.PutObject(ctx, &s3.PutObjectInput{ + option := storage.PutOption{} + for _, opt := range opts { + opt(&option) + } + + input := &s3.PutObjectInput{ Bucket: aws.String(bucket), 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 } diff --git a/backend/infra/impl/storage/tos/tos.go b/backend/infra/impl/storage/tos/tos.go index 1b9ecdfd..d7ea0168 100644 --- a/backend/infra/impl/storage/tos/tos.go +++ b/backend/infra/impl/storage/tos/tos.go @@ -140,22 +140,49 @@ func (t *tosClient) CheckAndCreateBucket(ctx context.Context) error { return err } - 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 - body := bytes.NewReader(content) bucketName := t.bucketName - _, err := client.PutObjectV2(ctx, &tos.PutObjectV2Input{ + option := storage.PutOption{} + for _, opt := range opts { + opt(&option) + } + + input := &tos.PutObjectV2Input{ PutObjectBasicInput: tos.PutObjectBasicInput{ Bucket: bucketName, 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 } diff --git a/backend/internal/mock/infra/contract/storage/storage_mock.go b/backend/internal/mock/infra/contract/storage/storage_mock.go index 35dda5e5..d1cc1862 100644 --- a/backend/internal/mock/infra/contract/storage/storage_mock.go +++ b/backend/internal/mock/infra/contract/storage/storage_mock.go @@ -11,6 +11,7 @@ package mock import ( context "context" + io "io" reflect "reflect" 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. -func (mr *MockStorageMockRecorder) DeleteObject(ctx, objectKey interface{}) *gomock.Call { +func (mr *MockStorageMockRecorder) DeleteObject(ctx, objectKey any) *gomock.Call { mr.mock.ctrl.T.Helper() 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. -func (mr *MockStorageMockRecorder) GetObject(ctx, objectKey interface{}) *gomock.Call { +func (mr *MockStorageMockRecorder) GetObject(ctx, objectKey any) *gomock.Call { mr.mock.ctrl.T.Helper() 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. func (m *MockStorage) GetObjectUrl(ctx context.Context, objectKey string, opts ...storage.GetOptFn) (string, error) { m.ctrl.T.Helper() - varargs := []interface{}{ctx, objectKey} + varargs := []any{ctx, objectKey} for _, a := range opts { 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. -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() - 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...) } -// 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. func (m *MockStorage) PutObject(ctx context.Context, objectKey string, content []byte, opts ...storage.PutOptFn) error { m.ctrl.T.Helper() - varargs := []interface{}{ctx, objectKey, content} + varargs := []any{ctx, objectKey, content} for _, a := range opts { 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. -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() - 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...) } + +// 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...) +}