GODRIVER-1770 Use chunkSize from files doc for download streams (#523)

Author: Divjot Arora

mongo/gridfs/bucket.go

@@ -33,6 +33,9 @@ const DefaultChunkSize int32 = 255 * 1024 // 255 KiB
 // ErrFileNotFound occurs if a user asks to download a file with a file ID that isn't found in the files collection.
 var ErrFileNotFound = errors.New("file with given parameters not found")
 
+// ErrMissingChunkSize occurs when downloading a file if the files collection document is missing the "chunkSize" field.
+var ErrMissingChunkSize = errors.New("files collection document does not contain a 'chunkSize' field")
+
 // Bucket represents a GridFS bucket.
 type Bucket struct {
 	db         *mongo.Database
@@ -381,14 +384,21 @@ func (b *Bucket) openDownloadStream(filter interface{}, opts ...*options.FindOpt
 	}
 
 	if foundFile.Length == 0 {
-		return newDownloadStream(nil, b.chunkSize, &foundFile), nil
+		return newDownloadStream(nil, foundFile.ChunkSize, &foundFile), nil
+	}
+
+	// For a file with non-zero length, chunkSize must exist so we know what size to expect when downloading chunks.
+	if _, err := cursor.Current.LookupErr("chunkSize"); err != nil {
+		return nil, ErrMissingChunkSize
 	}
 
 	chunksCursor, err := b.findChunks(ctx, foundFile.ID)
 	if err != nil {
 		return nil, err
 	}
-	return newDownloadStream(chunksCursor, b.chunkSize, &foundFile), nil
+	// The chunk size can be overridden for individual files, so the expected chunk size should be the "chunkSize"
+	// field from the files collection document, not the bucket's chunk size.
+	return newDownloadStream(chunksCursor, foundFile.ChunkSize, &foundFile), nil
 }
 
 func deadlineContext(deadline time.Time) (context.Context, context.CancelFunc) {

mongo/integration/gridfs_test.go

@@ -15,6 +15,7 @@ import (
 	"time"
 
 	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/bson/primitive"
 	"go.mongodb.org/mongo-driver/event"
 	"go.mongodb.org/mongo-driver/internal/testutil/assert"
 	"go.mongodb.org/mongo-driver/internal/testutil/israce"
@@ -248,6 +249,48 @@ func TestGridFS(x *testing.T) {
 				})
 			}
 		})
+		mt.Run("chunk size determined by files collection document", func(mt *mtest.T) {
+			// Test that the chunk size for a file download is determined by the chunkSize field in the files
+			// collection document, not the bucket's chunk size.
+
+			bucket, err := gridfs.NewBucket(mt.DB)
+			assert.Nil(mt, err, "NewBucket error: %v", err)
+			defer func() { _ = bucket.Drop() }()
+
+			fileData := []byte("hello world")
+			uploadOpts := options.GridFSUpload().SetChunkSizeBytes(4)
+			fileID, err := bucket.UploadFromStream("file", bytes.NewReader(fileData), uploadOpts)
+			assert.Nil(mt, err, "UploadFromStream error: %v", err)
+
+			// If the bucket's chunk size was used, this would error because the actual chunk size is 4 and the bucket
+			// chunk size is 255 KB.
+			var downloadBuffer bytes.Buffer
+			_, err = bucket.DownloadToStream(fileID, &downloadBuffer)
+			assert.Nil(mt, err, "DownloadToStream error: %v", err)
+
+			downloadedBytes := downloadBuffer.Bytes()
+			assert.Equal(mt, fileData, downloadedBytes, "expected bytes %s, got %s", fileData, downloadedBytes)
+		})
+		mt.Run("error if files collection document does not have a chunkSize field", func(mt *mtest.T) {
+			// Test that opening a download returns ErrMissingChunkSize if the files collection document has no
+			// chunk size field.
+
+			oid := primitive.NewObjectID()
+			filesDoc := bson.D{
+				{"_id", oid},
+				{"length", 10},
+				{"filename", "filename"},
+			}
+			_, err := mt.DB.Collection("fs.files").InsertOne(mtest.Background, filesDoc)
+			assert.Nil(mt, err, "InsertOne error for files collection: %v", err)
+
+			bucket, err := gridfs.NewBucket(mt.DB)
+			assert.Nil(mt, err, "NewBucket error: %v", err)
+			defer func() { _ = bucket.Drop() }()
+
+			_, err = bucket.OpenDownloadStream(oid)
+			assert.Equal(mt, gridfs.ErrMissingChunkSize, err, "expected error %v, got %v", gridfs.ErrMissingChunkSize, err)
+		})
 	})
 
 	mt.RunOpts("bucket collection accessors", noClientOpts, func(mt *mtest.T) {