Node.js MongoDB Driver API

Source: lib/gridfs-stream/upload.js

  1. 'use strict';
  3. var core = require('mongodb-core');
  4. var crypto = require('crypto');
  5. var stream = require('stream');
  6. var util = require('util');
  8. var ERROR_NAMESPACE_NOT_FOUND = 26;
  10. module.exports = GridFSBucketWriteStream;
  12. /**
  13.  * A writable stream that enables you to write buffers to GridFS.
  14.  *
  15.  * Do not instantiate this class directly. Use `openUploadStream()` instead.
  16.  *
  17.  * @class
  18.  * @param {GridFSBucket} bucket Handle for this stream's corresponding bucket
  19.  * @param {string} filename The value of the 'filename' key in the files doc
  20.  * @param {object} [options=null] Optional settings.
  21.  * @param {string|number|object} [options.id=null] Custom file id for the GridFS file.
  22.  * @param {number} [options.chunkSizeBytes=null] The chunk size to use, in bytes
  23.  * @param {number} [options.w=null] The write concern
  24.  * @param {number} [options.wtimeout=null] The write concern timeout
  25.  * @param {number} [options.j=null] The journal write concern
  26.  * @fires GridFSBucketWriteStream#error
  27.  * @fires GridFSBucketWriteStream#finish
  28.  * @return {GridFSBucketWriteStream} a GridFSBucketWriteStream instance.
  29.  */
  31. function GridFSBucketWriteStream(bucket, filename, options) {
  32.   options = options || {};
  33.   this.bucket = bucket;
  34.   this.chunks = bucket.s._chunksCollection;
  35.   this.filename = filename;
  36.   this.files = bucket.s._filesCollection;
  37.   this.options = options;
  38.   // Signals the write is all done
  39.   this.done = false;
  41.   this.id = options.id ? options.id : core.BSON.ObjectId();
  42.   this.chunkSizeBytes = this.options.chunkSizeBytes;
  43.   this.bufToStore = new Buffer(this.chunkSizeBytes);
  44.   this.length = 0;
  45.   this.md5 = crypto.createHash('md5');
  46.   this.n = 0;
  47.   this.pos = 0;
  48.   this.state = {
  49.     streamEnd: false,
  50.     outstandingRequests: 0,
  51.     errored: false,
  52.     aborted: false,
  53.     promiseLibrary: this.bucket.s.promiseLibrary
  54.   };
  56.   if (!this.bucket.s.calledOpenUploadStream) {
  57.     this.bucket.s.calledOpenUploadStream = true;
  59.     var _this = this;
  60.     checkIndexes(this, function() {
  61.       _this.bucket.s.checkedIndexes = true;
  62.       _this.bucket.emit('index');
  63.     });
  64.   }
  65. }
  67. util.inherits(GridFSBucketWriteStream, stream.Writable);
  69. /**
  70.  * An error occurred
  71.  *
  72.  * @event GridFSBucketWriteStream#error
  73.  * @type {Error}
  74.  */
  76. /**
  77.  * `end()` was called and the write stream successfully wrote the file
  78.  * metadata and all the chunks to MongoDB.
  79.  *
  80.  * @event GridFSBucketWriteStream#finish
  81.  * @type {object}
  82.  */
  84. /**
  85.  * Write a buffer to the stream.
  86.  *
  87.  * @method
  88.  * @param {Buffer} chunk Buffer to write
  89.  * @param {String} encoding Optional encoding for the buffer
  90.  * @param {Function} callback Function to call when the chunk was added to the buffer, or if the entire chunk was persisted to MongoDB if this chunk caused a flush.
  91.  * @return {Boolean} False if this write required flushing a chunk to MongoDB. True otherwise.
  92.  */
  94. GridFSBucketWriteStream.prototype.write = function(chunk, encoding, callback) {
  95.   var _this = this;
  96.   return waitForIndexes(this, function() {
  97.     return doWrite(_this, chunk, encoding, callback);
  98.   });
  99. };
  101. /**
  102.  * Places this write stream into an aborted state (all future writes fail)
  103.  * and deletes all chunks that have already been written.
  104.  *
  105.  * @method
  106.  * @param {GridFSBucket~errorCallback} callback called when chunks are successfully removed or error occurred
  107.  * @return {Promise} if no callback specified
  108.  */
  110. GridFSBucketWriteStream.prototype.abort = function(callback) {
  111.   if (this.state.streamEnd) {
  112.     var error = new Error('Cannot abort a stream that has already completed');
  113.     if (typeof callback === 'function') {
  114.       return callback(error);
  115.     }
  116.     return this.state.promiseLibrary.reject(error);
  117.   }
  118.   if (this.state.aborted) {
  119.     error = new Error('Cannot call abort() on a stream twice');
  120.     if (typeof callback === 'function') {
  121.       return callback(error);
  122.     }
  123.     return this.state.promiseLibrary.reject(error);
  124.   }
  125.   this.state.aborted = true;
  126.   this.chunks.deleteMany({ files_id: this.id }, function(error) {
  127.     if (typeof callback === 'function') callback(error);
  128.   });
  129. };
  131. /**
  132.  * Tells the stream that no more data will be coming in. The stream will
  133.  * persist the remaining data to MongoDB, write the files document, and
  134.  * then emit a 'finish' event.
  135.  *
  136.  * @method
  137.  * @param {Buffer} chunk Buffer to write
  138.  * @param {String} encoding Optional encoding for the buffer
  139.  * @param {Function} callback Function to call when all files and chunks have been persisted to MongoDB
  140.  */
  142. GridFSBucketWriteStream.prototype.end = function(chunk, encoding, callback) {
  143.   var _this = this;
  144.   if (typeof chunk === 'function') {
  145.     (callback = chunk), (chunk = null), (encoding = null);
  146.   } else if (typeof encoding === 'function') {
  147.     (callback = encoding), (encoding = null);
  148.   }
  150.   if (checkAborted(this, callback)) {
  151.     return;
  152.   }
  153.   this.state.streamEnd = true;
  155.   if (callback) {
  156.     this.once('finish', function(result) {
  157.       callback(null, result);
  158.     });
  159.   }
  161.   if (!chunk) {
  162.     waitForIndexes(this, function() {
  163.       writeRemnant(_this);
  164.     });
  165.     return;
  166.   }
  168.   this.write(chunk, encoding, function() {
  169.     writeRemnant(_this);
  170.   });
  171. };
  173. /**
  174.  * @ignore
  175.  */
  177. function __handleError(_this, error, callback) {
  178.   if (_this.state.errored) {
  179.     return;
  180.   }
  181.   _this.state.errored = true;
  182.   if (callback) {
  183.     return callback(error);
  184.   }
  185.   _this.emit('error', error);
  186. }
  188. /**
  189.  * @ignore
  190.  */
  192. function createChunkDoc(filesId, n, data) {
  193.   return {
  194.     _id: core.BSON.ObjectId(),
  195.     files_id: filesId,
  196.     n: n,
  197.     data: data
  198.   };
  199. }
  201. /**
  202.  * @ignore
  203.  */
  205. function checkChunksIndex(_this, callback) {
  206.   _this.chunks.listIndexes().toArray(function(error, indexes) {
  207.     if (error) {
  208.       // Collection doesn't exist so create index
  209.       if (error.code === ERROR_NAMESPACE_NOT_FOUND) {
  210.         var index = { files_id: 1, n: 1 };
  211.         _this.chunks.createIndex(index, { background: false, unique: true }, function(error) {
  212.           if (error) {
  213.             return callback(error);
  214.           }
  216.           callback();
  217.         });
  218.         return;
  219.       }
  220.       return callback(error);
  221.     }
  223.     var hasChunksIndex = false;
  224.     indexes.forEach(function(index) {
  225.       if (index.key) {
  226.         var keys = Object.keys(index.key);
  227.         if (keys.length === 2 && index.key.files_id === 1 && index.key.n === 1) {
  228.           hasChunksIndex = true;
  229.         }
  230.       }
  231.     });
  233.     if (hasChunksIndex) {
  234.       callback();
  235.     } else {
  236.       index = { files_id: 1, n: 1 };
  237.       var indexOptions = getWriteOptions(_this);
  239.       indexOptions.background = false;
  240.       indexOptions.unique = true;
  242.       _this.chunks.createIndex(index, indexOptions, function(error) {
  243.         if (error) {
  244.           return callback(error);
  245.         }
  247.         callback();
  248.       });
  249.     }
  250.   });
  251. }
  253. /**
  254.  * @ignore
  255.  */
  257. function checkDone(_this, callback) {
  258.   if (_this.done) return true;
  259.   if (_this.state.streamEnd && _this.state.outstandingRequests === 0 && !_this.state.errored) {
  260.     // Set done so we dont' trigger duplicate createFilesDoc
  261.     _this.done = true;
  262.     // Create a new files doc
  263.     var filesDoc = createFilesDoc(
  264.       _this.id,
  265.       _this.length,
  266.       _this.chunkSizeBytes,
  267.       _this.md5.digest('hex'),
  268.       _this.filename,
  269.       _this.options.contentType,
  270.       _this.options.aliases,
  271.       _this.options.metadata
  272.     );
  274.     if (checkAborted(_this, callback)) {
  275.       return false;
  276.     }
  278.     _this.files.insert(filesDoc, getWriteOptions(_this), function(error) {
  279.       if (error) {
  280.         return __handleError(_this, error, callback);
  281.       }
  282.       _this.emit('finish', filesDoc);
  283.     });
  285.     return true;
  286.   }
  288.   return false;
  289. }
  291. /**
  292.  * @ignore
  293.  */
  295. function checkIndexes(_this, callback) {
  296.   _this.files.findOne({}, { _id: 1 }, function(error, doc) {
  297.     if (error) {
  298.       return callback(error);
  299.     }
  300.     if (doc) {
  301.       return callback();
  302.     }
  304.     _this.files.listIndexes().toArray(function(error, indexes) {
  305.       if (error) {
  306.         // Collection doesn't exist so create index
  307.         if (error.code === ERROR_NAMESPACE_NOT_FOUND) {
  308.           var index = { filename: 1, uploadDate: 1 };
  309.           _this.files.createIndex(index, { background: false }, function(error) {
  310.             if (error) {
  311.               return callback(error);
  312.             }
  314.             checkChunksIndex(_this, callback);
  315.           });
  316.           return;
  317.         }
  318.         return callback(error);
  319.       }
  321.       var hasFileIndex = false;
  322.       indexes.forEach(function(index) {
  323.         var keys = Object.keys(index.key);
  324.         if (keys.length === 2 && index.key.filename === 1 && index.key.uploadDate === 1) {
  325.           hasFileIndex = true;
  326.         }
  327.       });
  329.       if (hasFileIndex) {
  330.         checkChunksIndex(_this, callback);
  331.       } else {
  332.         index = { filename: 1, uploadDate: 1 };
  334.         var indexOptions = getWriteOptions(_this);
  336.         indexOptions.background = false;
  338.         _this.files.createIndex(index, indexOptions, function(error) {
  339.           if (error) {
  340.             return callback(error);
  341.           }
  343.           checkChunksIndex(_this, callback);
  344.         });
  345.       }
  346.     });
  347.   });
  348. }
  350. /**
  351.  * @ignore
  352.  */
  354. function createFilesDoc(_id, length, chunkSize, md5, filename, contentType, aliases, metadata) {
  355.   var ret = {
  356.     _id: _id,
  357.     length: length,
  358.     chunkSize: chunkSize,
  359.     uploadDate: new Date(),
  360.     md5: md5,
  361.     filename: filename
  362.   };
  364.   if (contentType) {
  365.     ret.contentType = contentType;
  366.   }
  368.   if (aliases) {
  369.     ret.aliases = aliases;
  370.   }
  372.   if (metadata) {
  373.     ret.metadata = metadata;
  374.   }
  376.   return ret;
  377. }
  379. /**
  380.  * @ignore
  381.  */
  383. function doWrite(_this, chunk, encoding, callback) {
  384.   if (checkAborted(_this, callback)) {
  385.     return false;
  386.   }
  388.   var inputBuf = Buffer.isBuffer(chunk) ? chunk : new Buffer(chunk, encoding);
  390.   _this.length += inputBuf.length;
  392.   // Input is small enough to fit in our buffer
  393.   if (_this.pos + inputBuf.length < _this.chunkSizeBytes) {
  394.     inputBuf.copy(_this.bufToStore, _this.pos);
  395.     _this.pos += inputBuf.length;
  397.     callback && callback();
  399.     // Note that we reverse the typical semantics of write's return value
  400.     // to be compatible with node's `.pipe()` function.
  401.     // True means client can keep writing.
  402.     return true;
  403.   }
  405.   // Otherwise, buffer is too big for current chunk, so we need to flush
  406.   // to MongoDB.
  407.   var inputBufRemaining = inputBuf.length;
  408.   var spaceRemaining = _this.chunkSizeBytes - _this.pos;
  409.   var numToCopy = Math.min(spaceRemaining, inputBuf.length);
  410.   var outstandingRequests = 0;
  411.   while (inputBufRemaining > 0) {
  412.     var inputBufPos = inputBuf.length - inputBufRemaining;
  413.     inputBuf.copy(_this.bufToStore, _this.pos, inputBufPos, inputBufPos + numToCopy);
  414.     _this.pos += numToCopy;
  415.     spaceRemaining -= numToCopy;
  416.     if (spaceRemaining === 0) {
  417.       _this.md5.update(_this.bufToStore);
  418.       var doc = createChunkDoc(_this.id, _this.n, _this.bufToStore);
  419.       ++_this.state.outstandingRequests;
  420.       ++outstandingRequests;
  422.       if (checkAborted(_this, callback)) {
  423.         return false;
  424.       }
  426.       _this.chunks.insert(doc, getWriteOptions(_this), function(error) {
  427.         if (error) {
  428.           return __handleError(_this, error);
  429.         }
  430.         --_this.state.outstandingRequests;
  431.         --outstandingRequests;
  433.         if (!outstandingRequests) {
  434.           _this.emit('drain', doc);
  435.           callback && callback();
  436.           checkDone(_this);
  437.         }
  438.       });
  440.       spaceRemaining = _this.chunkSizeBytes;
  441.       _this.pos = 0;
  442.       ++_this.n;
  443.     }
  444.     inputBufRemaining -= numToCopy;
  445.     numToCopy = Math.min(spaceRemaining, inputBufRemaining);
  446.   }
  448.   // Note that we reverse the typical semantics of write's return value
  449.   // to be compatible with node's `.pipe()` function.
  450.   // False means the client should wait for the 'drain' event.
  451.   return false;
  452. }
  454. /**
  455.  * @ignore
  456.  */
  458. function getWriteOptions(_this) {
  459.   var obj = {};
  460.   if (_this.options.writeConcern) {
  461.     obj.w = _this.options.writeConcern.w;
  462.     obj.wtimeout = _this.options.writeConcern.wtimeout;
  463.     obj.j = _this.options.writeConcern.j;
  464.   }
  465.   return obj;
  466. }
  468. /**
  469.  * @ignore
  470.  */
  472. function waitForIndexes(_this, callback) {
  473.   if (_this.bucket.s.checkedIndexes) {
  474.     return callback(false);
  475.   }
  477.   _this.bucket.once('index', function() {
  478.     callback(true);
  479.   });
  481.   return true;
  482. }
  484. /**
  485.  * @ignore
  486.  */
  488. function writeRemnant(_this, callback) {
  489.   // Buffer is empty, so don't bother to insert
  490.   if (_this.pos === 0) {
  491.     return checkDone(_this, callback);
  492.   }
  494.   ++_this.state.outstandingRequests;
  496.   // Create a new buffer to make sure the buffer isn't bigger than it needs
  497.   // to be.
  498.   var remnant = new Buffer(_this.pos);
  499.   _this.bufToStore.copy(remnant, 0, 0, _this.pos);
  500.   _this.md5.update(remnant);
  501.   var doc = createChunkDoc(_this.id, _this.n, remnant);
  503.   // If the stream was aborted, do not write remnant
  504.   if (checkAborted(_this, callback)) {
  505.     return false;
  506.   }
  508.   _this.chunks.insert(doc, getWriteOptions(_this), function(error) {
  509.     if (error) {
  510.       return __handleError(_this, error);
  511.     }
  512.     --_this.state.outstandingRequests;
  513.     checkDone(_this);
  514.   });
  515. }
  517. /**
  518.  * @ignore
  519.  */
  521. function checkAborted(_this, callback) {
  522.   if (_this.state.aborted) {
  523.     if (typeof callback === 'function') {
  524.       callback(new Error('this stream has been aborted'));
  525.     }
  526.     return true;
  527.   }
  528.   return false;
  529. }