From ae80a66ff8d30bc9785d12a5f9d8fafd0bc37e2a Mon Sep 17 00:00:00 2001
From: Aaron <aaron@10gen.com>
Date: Mon, 22 Feb 2010 10:32:55 -0800
Subject: [PATCH] SERVER-305 record uri on connect

---
 scripting/engine.cpp |  3 +++
 scripting/engine.h   |  8 ++++++++
 scripting/sm_db.cpp  |  1 +
 scripting/v8_db.cpp  |  2 +-
 shell/dbshell.cpp    |  1 +
 shell/utils.cpp      | 12 +++++++-----
 shell/utils.h        |  1 +
 7 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/scripting/engine.cpp b/scripting/engine.cpp
index ab0d6dcc223..0dff413dbde 100644
--- a/scripting/engine.cpp
+++ b/scripting/engine.cpp
@@ -418,5 +418,8 @@ namespace mongo {
         }
     }
     
+    void ( *ScriptEngine::_connectCallback )( DBClientWithCommands & ) = 0;
+    
     ScriptEngine * globalScriptEngine;
 }
+    
\ No newline at end of file
diff --git a/scripting/engine.h b/scripting/engine.h
index 371c3ffb1da..9907d31244a 100644
--- a/scripting/engine.h
+++ b/scripting/engine.h
@@ -120,6 +120,8 @@ namespace mongo {
     
     void installGlobalUtils( Scope& scope );
 
+    class DBClientWithCommands;
+    
     class ScriptEngine : boost::noncopyable {
     public:
         ScriptEngine();
@@ -146,12 +148,18 @@ namespace mongo {
         virtual auto_ptr<Unlocker> newThreadUnlocker() { return auto_ptr< Unlocker >( new Unlocker ); }
         
         void setScopeInitCallback( void ( *func )( Scope & ) ) { _scopeInitCallback = func; }
+        static void setConnectCallback( void ( *func )( DBClientWithCommands& ) ) { _connectCallback = func; }
+        static void runConnectCallback( DBClientWithCommands &c ) {
+            if ( _connectCallback )
+                _connectCallback( c );
+        }
         
     protected:
         virtual Scope * createScope() = 0;
         
     private:
         void ( *_scopeInitCallback )( Scope & );
+        static void ( *_connectCallback )( DBClientWithCommands & );
     };
 
     extern ScriptEngine * globalScriptEngine;
diff --git a/scripting/sm_db.cpp b/scripting/sm_db.cpp
index ca4526a3c40..b15c6b8e416 100644
--- a/scripting/sm_db.cpp
+++ b/scripting/sm_db.cpp
@@ -155,6 +155,7 @@ namespace mongo {
                 JS_ReportError( cx , ((string)"couldn't connect: " + errmsg).c_str() );
                 return JS_FALSE;
             }
+            ScriptEngine::runConnectCallback( *c );
         }
         else { // paired
             int numCommas = 0;
diff --git a/scripting/v8_db.cpp b/scripting/v8_db.cpp
index c328a53def3..7ba5d3a535d 100644
--- a/scripting/v8_db.cpp
+++ b/scripting/v8_db.cpp
@@ -156,7 +156,7 @@ namespace mongo {
         if ( ! conn->connect( host , errmsg ) ){
             return v8::ThrowException( v8::String::New( "couldn't connect" ) );
         }
-
+        ScriptEngine::runConnectCallback( *conn );
         // NOTE I don't believe the conn object will ever be freed.
         args.This()->Set( CONN_STRING , External::New( conn ) );
         args.This()->Set( v8::String::New( "slaveOk" ) , Boolean::New( false ) );
diff --git a/shell/dbshell.cpp b/shell/dbshell.cpp
index cdd7c72a4dc..0013974a0bd 100644
--- a/shell/dbshell.cpp
+++ b/shell/dbshell.cpp
@@ -406,6 +406,7 @@ int _main(int argc, char* argv[]) {
 
     }
 
+    mongo::ScriptEngine::setConnectCallback( mongo::shellUtils::onConnect );
     mongo::ScriptEngine::setup();
     mongo::globalScriptEngine->setScopeInitCallback( mongo::shellUtils::initScope );
     auto_ptr< mongo::Scope > scope( mongo::globalScriptEngine->newScope() );    
diff --git a/shell/utils.cpp b/shell/utils.cpp
index f2db0e799b8..ef614e30a10 100644
--- a/shell/utils.cpp
+++ b/shell/utils.cpp
@@ -675,8 +675,6 @@ namespace mongo {
 #endif
         }
 
-        vector< string > _allMyUris;
-        
         void initScope( Scope &scope ) {
             scope.externalSetup();
             mongo::shellUtils::installShellUtils( scope );
@@ -684,14 +682,18 @@ namespace mongo {
             
             if ( !_dbConnect.empty() ) {
                 uassert( 12513, "connect failed", scope.exec( _dbConnect , "(connect)" , false , true , false ) );
-                uassert( 13010, "whatsmyuri failed", scope.exec( "__myuri = db.runCommand( {whatsmyuri:1} ).you", "(whatsmyuri)", false , true , false ) );
-                string uri = scope.getString( "__myuri" );
-                _allMyUris.push_back( uri );
                 if ( !_dbAuth.empty() ) {
                     installGlobalUtils( scope );
                     uassert( 12514, "login failed", scope.exec( _dbAuth , "(auth)" , true , true , false ) );
                 }
             }
         }
+        
+        vector< string > _allMyUris;        
+        void onConnect( DBClientWithCommands &c ) {
+            BSONObj info;
+            uassert( 13010, "whatsmyuri failed", c.runCommand( "admin", BSON( "whatsmyuri" << 1 ), info ) );
+            _allMyUris.push_back( info[ "you" ].str() );            
+        }
     }
 }
diff --git a/shell/utils.h b/shell/utils.h
index b3341589e7d..b4fea4226ba 100644
--- a/shell/utils.h
+++ b/shell/utils.h
@@ -40,5 +40,6 @@ namespace mongo {
         void KillMongoProgramInstances();
         
         void initScope( Scope &scope );
+        void onConnect( DBClientWithCommands &c );
     }
 }