From 29309dac9a926077962fc0778e5b2fbaf1d29cc2 Mon Sep 17 00:00:00 2001
From: Omer Katz <omer.drow@gamil.com>
Date: Wed, 25 Jun 2014 16:53:24 +0300
Subject: [PATCH] Mongo clients with the same settings should be shared since
 they manage a connection pool. Also, I removed old code that was supposed to
 support Pymongo<2.1 which we don't support anymore.

---
 mongoengine/connection.py | 33 ++++++++++++++++++---------------
 tests/test_connection.py  | 11 +++++++++++
 2 files changed, 29 insertions(+), 15 deletions(-)

diff --git a/mongoengine/connection.py b/mongoengine/connection.py
index 7cc626f4..d3efac62 100644
--- a/mongoengine/connection.py
+++ b/mongoengine/connection.py
@@ -93,20 +93,11 @@ def get_connection(alias=DEFAULT_CONNECTION_NAME, reconnect=False):
             raise ConnectionError(msg)
         conn_settings = _connection_settings[alias].copy()
 
-        if hasattr(pymongo, 'version_tuple'):  # Support for 2.1+
-            conn_settings.pop('name', None)
-            conn_settings.pop('slaves', None)
-            conn_settings.pop('is_slave', None)
-            conn_settings.pop('username', None)
-            conn_settings.pop('password', None)
-        else:
-            # Get all the slave connections
-            if 'slaves' in conn_settings:
-                slaves = []
-                for slave_alias in conn_settings['slaves']:
-                    slaves.append(get_connection(slave_alias))
-                conn_settings['slaves'] = slaves
-                conn_settings.pop('read_preference', None)
+        conn_settings.pop('name', None)
+        conn_settings.pop('slaves', None)
+        conn_settings.pop('is_slave', None)
+        conn_settings.pop('username', None)
+        conn_settings.pop('password', None)
 
         connection_class = MongoClient
         if 'replicaSet' in conn_settings:
@@ -119,7 +110,19 @@ def get_connection(alias=DEFAULT_CONNECTION_NAME, reconnect=False):
             connection_class = MongoReplicaSetClient
 
         try:
-            _connections[alias] = connection_class(**conn_settings)
+            connection = None
+            connection_settings_iterator = ((alias, settings.copy()) for alias, settings in _connection_settings.iteritems())
+            for alias, connection_settings in connection_settings_iterator:
+                connection_settings.pop('name', None)
+                connection_settings.pop('slaves', None)
+                connection_settings.pop('is_slave', None)
+                connection_settings.pop('username', None)
+                connection_settings.pop('password', None)
+                if conn_settings == connection_settings and _connections.get(alias, None):
+                    connection = _connections[alias]
+                    break
+
+            _connections[alias] = connection if connection else connection_class(**conn_settings)
         except Exception, e:
             raise ConnectionError("Cannot connect to database %s :\n%s" % (alias, e))
     return _connections[alias]
diff --git a/tests/test_connection.py b/tests/test_connection.py
index 96135bc5..bf615ceb 100644
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -34,6 +34,17 @@ class ConnectionTest(unittest.TestCase):
         conn = get_connection('testdb')
         self.assertTrue(isinstance(conn, pymongo.mongo_client.MongoClient))
 
+    def test_sharing_connections(self):
+        """Ensure that connections are shared when the connection settings are exactly the same
+        """
+        connect('mongoenginetest', alias='testdb1')
+
+        expected_connection = get_connection('testdb1')
+
+        connect('mongoenginetest', alias='testdb2')
+        actual_connection = get_connection('testdb2')
+        self.assertIs(expected_connection, actual_connection)
+
     def test_connect_uri(self):
         """Ensure that the connect() method works properly with uri's
         """