Permalink
Browse files

Merge pull request #58 from manover/add_replicaset_support

Add replicaset support
  • Loading branch information...
2 parents 745c33c + 93282ad commit 810e11b8bca4393cd14815a3619467de02ca231d @jehiah jehiah committed Apr 3, 2013
Showing with 452 additions and 98 deletions.
  1. +6 −2 .travis.yml
  2. +2 −0 MANIFEST
  3. +191 −0 asyncmongo/asyncjobs.py
  4. +76 −93 asyncmongo/connection.py
  5. +3 −0 asyncmongo/errors.py
  6. +9 −0 asyncmongo/helpers.py
  7. +1 −0 test/test_authentication.py
  8. +161 −0 test/test_replica_set.py
  9. +3 −3 test/test_shunt.py
View
@@ -5,8 +5,12 @@ python:
- "2.7"
services: mongodb
install:
- - "pip install tornado --use-mirrors"
+ - "pip install tornado==2.4.1 --use-mirrors"
- "pip install pymongo --use-mirrors"
-script: py.test
+env:
+ - TEST_REPLICA_SETS="true"
+ - TEST_REPLICA_SETS="false"
+script:
+ - if [[ $TEST_REPLICA_SETS == 'true' ]]; then py.test test/test_replica_set.py; else py.test --ignore=test/test_replica_set.py; fi
notifications:
email: false
View
@@ -9,6 +9,7 @@ asyncmongo/errors.py
asyncmongo/helpers.py
asyncmongo/message.py
asyncmongo/pool.py
+asyncmongo/asyncjobs.py
asyncmongo/backends/__init__.py
asyncmongo/backends/glib2_backend.py
asyncmongo/backends/glib3_backend.py
@@ -28,5 +29,6 @@ test/test_query.py
test/test_safe_updates.py
test/test_shunt.py
test/test_slave_only.py
+test/test_replica_set.py
test/testgtk2/test.py
test/testgtk3/test.py
View
@@ -0,0 +1,191 @@
+#!/bin/env python
+#
+# Copyright 2013 bit.ly
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you may
+# not use this file except in compliance with the License. You may obtain
+# a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations
+# under the License.
+
+"""Tools for creating `messages
+<http://www.mongodb.org/display/DOCS/Mongo+Wire+Protocol>`_ to be sent to
+MongoDB.
+
+.. note:: This module is for internal use and is generally not needed by
+ application developers.
+"""
+
+import logging
+from bson import SON
+
+import message
+import helpers
+from errors import AuthenticationError, RSConnectionError, InterfaceError
+
+class AsyncMessage(object):
+ def __init__(self, connection, message, callback):
+ super(AsyncMessage, self).__init__()
+ self.connection = connection
+ self.message = message
+ self.callback = callback
+
+ def process(self, *args, **kwargs):
+ try:
+ self.connection._send_message(self.message, self.callback)
+ except Exception, e:
+ if self.callback is None:
+ logging.error("Error occurred in safe update mode: %s", e)
+ else:
+ self.callback(None, e)
+
+class AuthorizeJob(object):
+ def __init__(self, connection, dbuser, dbpass, pool):
+ super(AuthorizeJob, self).__init__()
+ self.connection = connection
+ self._state = "start"
+ self.dbuser = dbuser
+ self.dbpass = dbpass
+ self.pool = pool
+
+ def __repr__(self):
+ return "AuthorizeJob at 0x%X, state = %r" % (id(self), self._state)
+
+ def process(self, response=None, error=None):
+ if error:
+ logging.debug(error)
+ logging.debug(response)
+ raise AuthenticationError(error)
+
+ if self._state == "start":
+ self._state = "nonce"
+ logging.debug("Sending nonce")
+ msg = message.query(
+ 0,
+ "%s.$cmd" % self.pool._dbname,
+ 0,
+ 1,
+ SON({'getnonce': 1}),
+ SON({})
+ )
+ self.connection._send_message(msg, self.process)
+ elif self._state == "nonce":
+ # this is the nonce response
+ self._state = "finish"
+ nonce = response['data'][0]['nonce']
+ logging.debug("Nonce received: %r", nonce)
+ key = helpers._auth_key(nonce, self.dbuser, self.dbpass)
+
+ msg = message.query(
+ 0,
+ "%s.$cmd" % self.pool._dbname,
+ 0,
+ 1,
+ SON([('authenticate', 1),
+ ('user', self.dbuser),
+ ('nonce', nonce),
+ ('key', key)]),
+ SON({})
+ )
+ self.connection._send_message(msg, self.process)
+ elif self._state == "finish":
+ self._state = "done"
+ assert response['number_returned'] == 1
+ response = response['data'][0]
+ if response['ok'] != 1:
+ logging.debug('Failed authentication %s' % response['errmsg'])
+ raise AuthenticationError(response['errmsg'])
+ self.connection._next_job()
+ else:
+ raise ValueError("Unexpected state: %s" % self._state)
+
+class ConnectRSJob(object):
+ def __init__(self, connection, seed, rs):
+ self.connection = connection
+ self.known_hosts = set(seed)
+ self.rs = rs
+ self._tried_hosts = set()
+ self._state = "seed"
+ self._primary = None
+
+ def __repr__(self):
+ return "ConnectRSJob at 0x%X, state = %s" % (id(self), self._state)
+
+ def process(self, response=None, error=None):
+ if error:
+ logging.debug("Problem connecting: %s", error)
+
+ if self._state == "ismaster":
+ self._state = "seed"
+
+ if self._state == "seed":
+ fresh = self.known_hosts ^ self._tried_hosts
+ logging.debug("Working through the rest of the host list: %r", fresh)
+
+ while fresh:
+ if self._primary and self._primary not in self._tried_hosts:
+ # Try primary first
+ h = self._primary
+ else:
+ h = fresh.pop()
+
+ self._tried_hosts.add(h)
+
+ logging.debug("Connecting to %s:%s", *h)
+ self.connection._host, self.connection._port = h
+ try:
+ self.connection._socket_connect()
+ logging.debug("Connected to %s", h)
+ except InterfaceError, e:
+ logging.error("Failed to connect to the host: %s", e)
+ else:
+ break
+
+ else:
+ raise RSConnectionError("No more hosts to try, tried: %s" % self.known_hosts)
+
+ self._state = "ismaster"
+ msg = message.query(
+ options=0,
+ collection_name="admin.$cmd",
+ num_to_skip=0,
+ num_to_return=-1,
+ query=SON([("ismaster", 1)])
+ )
+ self.connection._send_message(msg, self.process)
+
+ elif self._state == "ismaster":
+ logging.debug("ismaster response: %r", response)
+
+ if len(response["data"]) == 1:
+ res = response["data"][0]
+ else:
+ raise RSConnectionError("Invalid response data: %r" % response["data"])
+
+ rs_name = res.get("setName")
+ if rs_name:
+ if rs_name != self.rs:
+ raise RSConnectionError("Wrong replica set: %s, expected: %s" %
+ (rs_name, self.rs))
+ hosts = res.get("hosts")
+ if hosts:
+ self.known_hosts.update(helpers._parse_host(h) for h in hosts)
+
+ ismaster = res.get("ismaster")
+ if ismaster:
+ logging.info("Connected to master")
+ self._state = "done"
+ self.connection._next_job()
+ else:
+ primary = res.get("primary")
+ if primary:
+ self._primary = helpers._parse_host(primary)
+
+ self._state = "seed"
+ self.process()
Oops, something went wrong.

0 comments on commit 810e11b

Please sign in to comment.