SERVER-4345 mongorestore will read from stdin when filename is "-"

            check eof only on fifo files. restored original logic of checking bytes read against file length.
            mongorestore requires both --db and --collection when using stdin/fifo

3 files changed, 37 insertions, 17 deletions

diff --git a/jstests/tool/dumprestore1.js b/jstests/tool/dumprestore1.js
index fd1e8789ea6..aabe441244f 100644
--- a/jstests/tool/dumprestore1.js
+++ b/jstests/tool/dumprestore1.js
@@ -21,3 +21,11 @@ assert.eq( 22 , c.findOne().a , "after restore 2" );
 var ret = t.runTool( "dump" , "--collection" , "col" );
 assert.neq( ret, 0, "mongodump should return failure code" );
 t.stop();
+
+// Ensure that --db and --collection are provided when filename is "-" (stdin).
+ret = t.runTool( "restore" , "--collection" , "coll", "--dir", "-" );
+assert.neq( ret, 0, "mongorestore should return failure code" );
+t.stop();
+ret = t.runTool( "restore" , "--db" , "db", "--dir", "-" );
+assert.neq( ret, 0, "mongorestore should return failure code" );
+t.stop();
diff --git a/src/mongo/tools/restore.cpp b/src/mongo/tools/restore.cpp
index 7ded379eb0b..d4e2f8339af 100644
--- a/src/mongo/tools/restore.cpp
+++ b/src/mongo/tools/restore.cpp
@@ -413,11 +413,19 @@ public:
         boost::filesystem::file_status fileStatus = boost::filesystem::status(root);
         if ( ! ( endsWith( root.string().c_str() , ".bson" ) ||
                  endsWith( root.string().c_str() , ".bin" ) ) &&
+             ! ( root.string() == "-" ) &&
              ! ( fileStatus.type() == boost::filesystem::fifo_file ) ) {
             toolError() << "don't know what to do with file [" << root.string() << "]" << std::endl;
             return;
         }
 
+        // Both --db and --collection have to be provided when using stdin or fifo.
+        if ((root.string() == "-" || fileStatus.type() == boost::filesystem::fifo_file) &&
+            !(use_db && use_coll)) {
+            toolError() << "Both --db and --collection have to be provided when using stding/fifo";
+            exit(EXIT_FAILURE);
+        }
+
         toolInfoLog() << root.string() << std::endl;
 
         if ( root.leaf() == "system.profile.bson" ) {
@@ -496,6 +504,7 @@ public:
         BSONObj metadataObject;
         boost::filesystem::file_status fileStatus = boost::filesystem::status(root);
         if (fileStatus.type() != boost::filesystem::fifo_file &&
+            root.string() != "-" &&
             (mongoRestoreGlobalParams.restoreOptions || mongoRestoreGlobalParams.restoreIndexes)) {
             string oldCollName = root.leaf().string(); // Name of collection that was dumped from
             oldCollName = oldCollName.substr( 0 , oldCollName.find_last_of( "." ) );
diff --git a/src/mongo/tools/tool.cpp b/src/mongo/tools/tool.cpp
index b5aabb09729..92daecb1860 100644
--- a/src/mongo/tools/tool.cpp
+++ b/src/mongo/tools/tool.cpp
@@ -33,6 +33,7 @@
 #include <boost/filesystem/operations.hpp>
 #include <fstream>
 #include <iostream>
+#include <limits>
 
 #include "mongo/base/initializer.h"
 #include "mongo/base/init.h"
@@ -295,28 +296,27 @@ namespace mongo {
 
     long long BSONTool::processFile( const boost::filesystem::path& root ) {
         bool isFifoFile = boost::filesystem::status(root).type() == boost::filesystem::fifo_file;
+        bool isStdin = root == "-";
 
         std::string fileName = root.string();
 
-        unsigned long long fileLength = 0;
-        if (!isFifoFile) {
-            fileLength = file_size( root );
+        unsigned long long fileLength = (isFifoFile || isStdin) ?
+            std::numeric_limits<unsigned long long>::max() : file_size(root);
 
-            if ( fileLength == 0 ) {
-                toolInfoOutput() << "file " << fileName << " empty, skipping" << std::endl;
-                return 0;
-            }
+        if ( fileLength == 0 ) {
+            toolInfoOutput() << "file " << fileName << " empty, skipping" << std::endl;
+            return 0;
         }
 
 
-        FILE* file = fopen( fileName.c_str() , "rb" );
+        FILE* file = isStdin ? stdin : fopen( fileName.c_str() , "rb" );
         if ( ! file ) {
             toolError() << "error opening file: " << fileName << " " << errnoWithDescription()
                       << std::endl;
             return 0;
         }
 
-        if (!isFifoFile) {
+        if (!isFifoFile && !isStdin) {
 #ifdef POSIX_FADV_SEQUENTIAL
             posix_fadvise(fileno(file), 0, fileLength, POSIX_FADV_SEQUENTIAL);
 #endif
@@ -326,6 +326,7 @@ namespace mongo {
             }
         }
 
+        unsigned long long read = 0;
         unsigned long long num = 0;
         unsigned long long processed = 0;
 
@@ -336,16 +337,15 @@ namespace mongo {
         // no progress is available for FIFO
         // only for regular files
         boost::scoped_ptr<ProgressMeter> m;
-        if (!toolGlobalParams.quiet && !isFifoFile) {
+        if (!toolGlobalParams.quiet && !isFifoFile && !isStdin) {
             m.reset(new ProgressMeter( fileLength ));
-            // boost::scoped_ptr<ProgressMeter> m( fileLength );
             m->setUnits( "bytes" );
         }
 
-        while ( true ) {
+        while ( read < fileLength ) {
             size_t amt = fread(buf, 1, 4, file);
-            // end of fifo/file
-            if ( feof(file) ) {
+            // end of fifo
+            if ((isFifoFile || isStdin) && ::feof(file)) {
                 break;
             }
             verify( amt == 4 );
@@ -382,6 +382,7 @@ namespace mongo {
                 processed++;
             }
 
+            read += o.objsize();
             num++;
 
             if (m.get()) {
@@ -389,10 +390,12 @@ namespace mongo {
             }
         }
 
-        fclose( file );
+        if (!isStdin) {
+            fclose(file);
+        }
 
-        if (m.get()) {
-            uassert(10265, "counts don't match", m->done() == fileLength);
+        if (!isFifoFile && !isStdin) {
+            uassert(10265, "counts don't match", read == fileLength);
         }
         toolInfoOutput() << num << ((num == 1) ? " document" : " documents")
                          << " found" << std::endl;