Ensure we return sockets if authentication fails.

Author: behackett

pymongo/mongo_client.py

@@ -351,11 +351,13 @@ def _cache_credentials(self, source, credentials):
  'to this database. You must logout first.')
 
  sock_info = self.__socket()
- # Logout any previous user for this source database
- self.__check_auth(sock_info)
- auth.authenticate(credentials, sock_info, self.__simple_command)
- sock_info.authset.add(credentials)
- self.__pool.maybe_return_socket(sock_info)
+ try:
+ # Since __check_auth was called in __socket
+ # there is no need to call it here.
+ auth.authenticate(credentials, sock_info, self.__simple_command)
+ sock_info.authset.add(credentials)
+ finally:
+ self.__pool.maybe_return_socket(sock_info)
 
  self.__auth_credentials[source] = credentials
 
@@ -662,7 +664,11 @@ def __socket(self):
  host_details = "%s:%d:" % (host, port)
  raise AutoReconnect("could not connect to "
  "%s %s" % (host_details, str(why)))
- self.__check_auth(sock_info)
+ try:
+ self.__check_auth(sock_info)
+ except OperationFailure:
+ self.__pool.maybe_return_socket(sock_info)
+ raise
  return sock_info
 
  def disconnect(self):

pymongo/mongo_replica_set_client.py

@@ -558,11 +558,13 @@ def _cache_credentials(self, source, credentials):
  members = self.__members.copy().values()
  member = select_member(members, ReadPreference.PRIMARY_PREFERRED)
  sock_info = self.__socket(member)
- # Logout any previous user for this source database
- self.__check_auth(sock_info)
- auth.authenticate(credentials, sock_info, self.__simple_command)
- sock_info.authset.add(credentials)
- member.pool.maybe_return_socket(sock_info)
+ try:
+ # Since __check_auth was called in __socket
+ # there is no need to call it here.
+ auth.authenticate(credentials, sock_info, self.__simple_command)
+ sock_info.authset.add(credentials)
+ finally:
+ member.pool.maybe_return_socket(sock_info)
 
  self.__auth_credentials[source] = credentials
 
@@ -917,7 +919,11 @@ def __socket(self, member):
 
  sock_info = member.pool.get_socket()
 
- self.__check_auth(sock_info)
+ try:
+ self.__check_auth(sock_info)
+ except OperationFailure:
+ member.pool.maybe_return_socket(sock_info)
+ raise
  return sock_info
 
  def disconnect(self):
@@ -1123,8 +1129,6 @@ def __send_and_receive(self, member, msg, **kwargs):
  member.pool.maybe_return_socket(sock_info)
 
  return response
- except OperationFailure:
- raise
  except (ConnectionFailure, socket.error), why:
  host, port = member.pool.pair
  member.pool.discard_socket(sock_info)