opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Enhance attachment upload process and content type handling

Browse source 
- Updated the upload process to cache metadata from the first chunk, ensuring consistency in content type, filename, and user path across chunked uploads.
- Modified content type validation in tests to allow for charset variations, improving robustness in file type checks.
- Refactored database save logic to differentiate between new records and updates for chunked uploads, enhancing data integrity and performance.
This commit is contained in:
Max 2025-11-06 11:31:24 +08:00
parent 80d0b80a6d
commit e435c03e0e
2 changed files with 83 additions and 37 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

104
attachment/manager.go
View file

@ -40,6 +40,10 @@ type UploadChunk struct {
Total int64
Chunksize int64
TotalChunks int64
// Cache metadata from first chunk to avoid inconsistencies
ContentType string
Filename string
UserPath string
}
// Parse parses an attachment wrapper string and returns uploader name and file ID
@ -395,6 +399,10 @@ func (manager Manager) Upload(ctx context.Context, fileheader *FileHeader, reade
Total: total,
Chunksize: chunksize,
TotalChunks: totalChunks,
// Cache metadata from first chunk
ContentType: file.ContentType,
Filename: file.Filename,
UserPath: file.UserPath,
})
}
@ -411,6 +419,11 @@ func (manager Manager) Upload(ctx context.Context, fileheader *FileHeader, reade
chunkIndex = chunkdata.Last + 1
chunkdata.Last = chunkIndex
uploadChunks.Store(file.ID, chunkdata)
// For non-first chunks, use cached metadata from first chunk
file.ContentType = chunkdata.ContentType
file.Filename = chunkdata.Filename
file.UserPath = chunkdata.UserPath
}
// Apply gzip compression if requested
@ -429,6 +442,15 @@ func (manager Manager) Upload(ctx context.Context, fileheader *FileHeader, reade
return nil, err
}
// Save to database on first chunk only
if start == 0 {
file.Status = "uploading"
err = manager.saveFileToDatabase(ctx, file, file.Path, option)
if err != nil {
return nil, fmt.Errorf("failed to create database record for chunked upload: %w", err)
}
}
// Fix the file size, the file size is the sum of all chunks
file.Bytes = chunkIndex * int(chunkdata.Chunksize)
file.Status = "uploading"
@ -451,14 +473,14 @@ func (manager Manager) Upload(ctx context.Context, fileheader *FileHeader, reade
// Remove the chunk data
uploadChunks.Delete(file.ID)
// Fix the file size
// Fix the file size and update status to uploaded
file.Bytes = int(chunkdata.Total)
file.Status = "uploaded"
// Save file information to database when chunked upload is complete
// Update only bytes and status for the last chunk
err = manager.saveFileToDatabase(ctx, file, file.Path, option)
if err != nil {
return nil, fmt.Errorf("failed to save chunked file to database: %w", err)
return nil, fmt.Errorf("failed to update chunked file status: %w", err)
}
}
@ -829,9 +851,6 @@ func (manager Manager) makeFile(file *FileHeader, option UploadOption) (*File, e
return nil, fmt.Errorf("file size %d exceeds the maximum size of %d", file.Size, manager.maxsize)
}
// Get the content type
contentType := file.Header.Get("Content-Type")
// Use original filename if provided, otherwise use the file header filename
filename := file.Filename
userPath := option.OriginalFilename
@ -842,6 +861,22 @@ func (manager Manager) makeFile(file *FileHeader, option UploadOption) (*File, e
extension := filepath.Ext(filename)
// Get the content type
// For chunked uploads, file.Header may have incorrect content-type (e.g., application/octet-stream for Blob)
// Try to detect from filename extension first, then fallback to header
contentType := file.Header.Get("Content-Type")
if extension != "" {
// Try to get content type from extension
detectedType := mime.TypeByExtension(extension)
if detectedType != "" {
// If detected type is not the generic octet-stream, use it
// This handles chunked uploads where the header has incorrect type
if detectedType != "application/octet-stream" || contentType == "application/octet-stream" {
contentType = detectedType
}
}
}
// Get the extension from the content type if not available from filename
if extension == "" {
// Special handling for common types
@ -1084,10 +1119,41 @@ func (manager Manager) Delete(ctx context.Context, fileID string) error {
}
// saveFileToDatabase saves file information to the database
// For chunked uploads, it only updates bytes/status/progress if record exists
func (manager Manager) saveFileToDatabase(ctx context.Context, file *File, storagePath string, option UploadOption) error {
m := model.Select("__yao.attachment")
// Check if record exists first
records, err := m.Get(model.QueryParam{
Select: []interface{}{"file_id"},
Wheres: []model.QueryWhere{
{Column: "file_id", Value: file.ID},
},
})
if err != nil {
return fmt.Errorf("failed to check existing record: %w", err)
}
if len(records) > 0 {
// Record exists - this is a chunked upload update
// Only update bytes, status, and progress (don't overwrite metadata)
updateData := map[string]interface{}{
"bytes": int64(file.Bytes),
"status": file.Status,
}
_, err = m.UpdateWhere(model.QueryParam{
Wheres: []model.QueryWhere{
{Column: "file_id", Value: file.ID},
},
}, updateData)
return err
}
// Record doesn't exist - create new record with full metadata
// Set default value for share if empty
share := option.Share
if share == "" {
@ -1124,30 +1190,8 @@ func (manager Manager) saveFileToDatabase(ctx context.Context, file *File, stora
data["__yao_tenant_id"] = option.YaoTenantID
}
// Check if record exists first
records, err := m.Get(model.QueryParam{
Select: []interface{}{"file_id"},
Wheres: []model.QueryWhere{
{Column: "file_id", Value: file.ID},
},
})
if err != nil {
return fmt.Errorf("failed to check existing record: %w", err)
}
if len(records) > 0 {
// Update existing record
_, err = m.UpdateWhere(model.QueryParam{
Wheres: []model.QueryWhere{
{Column: "file_id", Value: file.ID},
},
}, data)
} else {
// Create new record
_, err = m.Create(data)
}
// Create new record
_, err = m.Create(data)
return err
}

16
attachment/manager_test.go
View file

@ -67,7 +67,8 @@ func TestManagerUpload(t *testing.T) {
t.Errorf("Expected filename 'test.txt', got '%s'", file.Filename)
}
if file.ContentType != "text/plain" {
// Content type may include charset
if !strings.HasPrefix(file.ContentType, "text/plain") {
t.Errorf("Expected content type 'text/plain', got '%s'", file.ContentType)
}
@ -581,7 +582,8 @@ func TestInfo(t *testing.T) {
t.Errorf("Expected filename %s, got %s", uploadedFile.Filename, fileInfo.Filename)
}
if fileInfo.ContentType != "text/plain" {
// Content type may include charset
if !strings.HasPrefix(fileInfo.ContentType, "text/plain") {
t.Errorf("Expected content type 'text/plain', got %s", fileInfo.ContentType)
}
@ -731,9 +733,9 @@ func TestList(t *testing.T) {
// Test filtering by content type
t.Run("FilterByContentType", func(t *testing.T) {
result, err := manager.List(context.Background(), ListOption{
Filters: map[string]interface{}{
"uploader": managerName,
"content_type": "text/plain",
Wheres: []model.QueryWhere{
{Column: "uploader", Value: managerName},
{Column: "content_type", Value: "text/plain%", OP: "like"},
},
})
if err != nil {
@ -745,9 +747,9 @@ func TestList(t *testing.T) {
t.Errorf("Expected %d text files, got %d", expectedCount, len(result.Files))
}
// Verify all returned files are text/plain
// Verify all returned files are text/plain (may include charset)
for _, file := range result.Files {
if file.ContentType != "text/plain" {
if !strings.HasPrefix(file.ContentType, "text/plain") {
t.Errorf("Expected content type 'text/plain', got '%s'", file.ContentType)
}
}