Node.js MongoDB Driver API

Source: lib/gridfs-stream/upload.js

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