Merge pull request #805 from estolfo/RUBY-1111-custom-grid-file-id

RUBY-1111 Support custom GridFS file IDs

Author: estolfo

lib/mongo/grid/fs_bucket.rb

@@ -337,6 +337,7 @@ def download_to_stream_by_name(filename, io, opts = {})
       # @param [ String ] filename The filename of the file to upload.
       # @param [ Hash ] opts The options for the write stream.
       #
+      # @option opts [ Object ] :file_id An optional unique file id. An ObjectId is generated otherwise.
       # @option opts [ Integer ] :chunk_size Override the default chunk size.
       # @option opts [ Hash ] :write The write concern.
       # @option opts [ Hash ] :metadata User data for the 'metadata' field of the files
@@ -366,12 +367,13 @@ def open_upload_stream(filename, opts = {})
       # document for the filename in the files collection.
       #
       # @example Upload a file to the GridFS bucket.
-      #   fs.upload_from_stream('a-file.txt')
+      #   fs.upload_from_stream('a-file.txt', file)
       #
       # @param [ String ] filename The filename of the file to upload.
       # @param [ IO ] io The source io stream to upload from.
       # @param [ Hash ] opts The options for the write stream.
       #
+      # @option opts [ Object ] :file_id An optional unique file id. An ObjectId is generated otherwise.
       # @option opts [ Integer ] :chunk_size Override the default chunk size.
       # @option opts [ Hash ] :write The write concern.
       # @option opts [ Hash ] :metadata User data for the 'metadata' field of the files

lib/mongo/grid/stream/write.rb

@@ -27,7 +27,7 @@ class Write
           # @since 2.1.0
           attr_reader :fs
 
-          # @return [ BSON::ObjectId ] file_id The id of the file being uploaded.
+          # @return [ Object ] file_id The id of the file being uploaded.
           #
           # @since 2.1.0
           attr_reader :file_id
@@ -50,6 +50,7 @@ class Write
           # @param [ FSBucket ] fs The GridFS bucket object.
           # @param [ Hash ] options The write stream options.
           #
+          # @option opts [ Object ] :file_id The file id. An ObjectId is generated otherwise.
           # @option opts [ Integer ] :chunk_size Override the default chunk size.
           # @option opts [ Hash ] :write The write concern.
           # @option opts [ Hash ] :metadata User data for the 'metadata' field of the files collection document.
@@ -63,7 +64,7 @@ def initialize(fs, options)
             @fs = fs
             @length = 0
             @n = 0
-            @file_id = BSON::ObjectId.new
+            @file_id = options[:file_id] || BSON::ObjectId.new
             @options = options
             @filename = @options[:filename]
             @open = true

spec/mongo/grid/fs_bucket_spec.rb

@@ -462,6 +462,25 @@
     it 'removes the file from the db' do
       expect(from_db).to be_nil
     end
+
+    context 'when a custom file id is used' do
+
+      let(:custom_file_id) do
+        fs.upload_from_stream(filename, file, file_id: 'Custom ID')
+      end
+
+      before do
+        fs.delete(custom_file_id)
+      end
+
+      let(:from_db) do
+        fs.find_one(:filename => filename)
+      end
+
+      it 'removes the file from the db' do
+        expect(from_db).to be_nil
+      end
+    end
   end
 
   context 'when a read stream is opened' do
@@ -526,6 +545,51 @@
           expect(io.size).to eq(0)
         end
       end
+
+      context 'when a custom file id is provided' do
+
+        let(:file) do
+          File.open(__FILE__)
+        end
+
+        let!(:file_id) do
+          fs.open_upload_stream(filename, file_id: 'Custom ID') do |stream|
+            stream.write(file)
+          end.file_id
+        end
+
+        context 'when a block is provided' do
+
+          let!(:stream) do
+            fs.open_download_stream(file_id) do |stream|
+              io.write(stream.read)
+            end
+          end
+
+          it 'yields the stream to the block' do
+            expect(io.size).to eq(file.size)
+          end
+        end
+
+        context 'when a block is not provided' do
+
+          let!(:stream) do
+            fs.open_download_stream(file_id)
+          end
+
+          it 'returns a Stream::Read object' do
+            expect(stream).to be_a(Mongo::Grid::FSBucket::Stream::Read)
+          end
+
+          it 'does not close the stream' do
+            expect(stream.closed?).to be(false)
+          end
+
+          it 'does not yield the stream to the block' do
+            expect(io.size).to eq(0)
+          end
+        end
+      end
     end
 
     describe '#download_to_stream' do
@@ -858,6 +922,21 @@
         it 'creates an ObjectId for the file' do
           expect(stream.file_id).to be_a(BSON::ObjectId)
         end
+
+        context 'when a custom file ID is provided' do
+
+          let(:stream) do
+            fs.open_upload_stream(filename, file_id: 'Custom ID')
+          end
+
+          it 'returns a Stream::Write object' do
+            expect(stream).to be_a(Mongo::Grid::FSBucket::Stream::Write)
+          end
+
+          it 'creates an ObjectId for the file' do
+            expect(stream.file_id).to eq('Custom ID')
+          end
+        end
       end
 
       context 'when a block is provided' do
@@ -958,6 +1037,17 @@
         fs.open_upload_stream(filename, stream_options)
       end
 
+      context 'when a custom file id is provided' do
+
+        let(:stream_options) do
+          { file_id: 'Custom ID' }
+        end
+
+        it 'sets the file id on the stream' do
+          expect(stream.file_id).to eq('Custom ID')
+        end
+      end
+
       context 'when a write option is specified' do
 
         let(:stream_options) do

spec/mongo/grid/stream/write_spec.rb

@@ -83,7 +83,7 @@
 
         let(:extra_options) do
           {
-              write: { w: (WRITE_CONCERN[:w] + 1) }
+            write: { w: (WRITE_CONCERN[:w] + 1) }
           }
         end
 
@@ -174,6 +174,19 @@
           expect(stream.send(:file_info).document[:aliases]).to eq(options[:aliases])
         end
       end
+
+      context 'when provided a file_id option' do
+
+        let(:options) do
+          {
+            file_id: 'Custom ID'
+          }
+        end
+
+        it 'assigns the stream the file id' do
+          expect(stream.file_id).to eq(options[:file_id])
+        end
+      end
     end
   end
 
@@ -285,25 +298,63 @@
 
     context 'when provided an io stream' do
 
-      before do
-        stream.write(file)
-        stream.close
-      end
+      context 'when no file id is specified' do
 
-      it 'writes the contents of the stream' do
-        expect(file_from_db.data.size).to eq(file.size)
-      end
+        before do
+          stream.write(file)
+          stream.close
+        end
 
-      it 'updates the length written' do
-        expect(stream.send(:file_info).document['length']).to eq(file.size)
+        it 'writes the contents of the stream' do
+          expect(file_from_db.data.size).to eq(file.size)
+        end
+
+        it 'updates the length written' do
+          expect(stream.send(:file_info).document['length']).to eq(file.size)
+        end
+
+        it 'updates the position (n)' do
+          expect(stream.instance_variable_get(:@n)).to eq(1)
+        end
       end
 
-      it 'updates the position (n)' do
-        expect(stream.instance_variable_get(:@n)).to eq(1)
+      context 'when a custom file id is provided' do
+
+        let(:extra_options) do
+          {
+            file_id: 'Custom ID'
+          }
+        end
+
+        let!(:id) do
+          stream.write(file)
+          stream.close
+        end
+
+        it 'writes the contents of the stream' do
+          expect(file_from_db.data.size).to eq(file.size)
+        end
+
+        it 'updates the length written' do
+          expect(stream.send(:file_info).document['length']).to eq(file.size)
+        end
+
+        it 'updates the position (n)' do
+          expect(stream.instance_variable_get(:@n)).to eq(1)
+        end
+
+        it 'uses the custom file id' do
+          expect(id).to eq(options[:file_id])
+        end
       end
 
       context 'when the user file contains no data' do
 
+        before do
+          stream.write(file)
+          stream.close
+        end
+
         let(:file) do
           StringIO.new('')
         end