Improve MongoDB support.

* Support configuration with MongoDB connection URI strings.
* Handle connecting to replica sets.
* Use standard notation database+collection rather than
  collection+sub_collection.

Note that these changes are NOT expected to be fully backwards
compatible for someone using mongodb with pysaml2 already.

1 files changed, 60 insertions, 15 deletions

diff --git a/src/saml2/mongo_store.py b/src/saml2/mongo_store.py
index 6fd38943..607678ee 100644
--- a/src/saml2/mongo_store.py
+++ b/src/saml2/mongo_store.py
@@ -2,6 +2,9 @@ from hashlib import sha1
 import logging
 
 from pymongo import MongoClient
+from pymongo.mongo_replica_set_client import MongoReplicaSetClient
+import pymongo.uri_parser
+import pymongo.errors
 from saml2.eptid import Eptid
 from saml2.mdstore import MetaData
 from saml2.s_utils import PolicyError
@@ -49,10 +52,9 @@ def context_match(cfilter, cntx):
 class SessionStorageMDB(object):
     """ Session information is stored in a MongoDB database"""
 
-    def __init__(self, collection=""):
-        connection = MongoClient()
-        db = connection[collection]
-        self.assertion = db.assertion
+    def __init__(self, database="", collection="assertion", **kwargs):
+        db = _mdb_get_database(database, **kwargs)
+        self.assertion = db[collection]
 
     def store_assertion(self, assertion, to_sign):
         name_id = assertion.subject.name_id
@@ -129,9 +131,9 @@ class SessionStorageMDB(object):
 
 
 class IdentMDB(IdentDB):
-    def __init__(self, collection="", domain="", name_qualifier=""):
+    def __init__(self, database="", collection="ident", domain="", name_qualifier=""):
         IdentDB.__init__(self, None, domain, name_qualifier)
-        self.mdb = MDB(collection, "ident")
+        self.mdb = MDB(database=database, collection=collection)
         self.mdb.primary_key = "user_id"
 
     def in_store(self, _id):
@@ -195,10 +197,9 @@ class IdentMDB(IdentDB):
 class MDB(object):
     primary_key = "mdb"
 
-    def __init__(self, collection="", sub_collection=""):
-        connection = MongoClient()
-        _db = connection[collection]
-        self.db = _db[sub_collection]
+    def __init__(self, database, collection, **kwargs):
+        _db = _mdb_get_database(database, **kwargs)
+        self.db = _db[collection]
 
     def store(self, value, **kwargs):
         if value:
@@ -250,11 +251,55 @@ class MDB(object):
         self.db.drop()
 
 
+
+def _mdb_get_database(uri, **kwargs):
+    """
+    Helper-function to connect to MongoDB and return a database object.
+
+    The `uri' argument should be either a full MongoDB connection URI string,
+    or just a database name in which case a connection to the default mongo
+    instance at mongodb://localhost:27017 will be made.
+
+    Performs explicit authentication if a username is provided in a connection
+    string URI, since PyMongo does not always seem to do that as promised.
+
+    :params database: name as string or (uri, name)
+    :returns: pymongo database object
+    """
+    connection_factory = MongoClient
+    _parsed_uri = {}
+    db_name = None
+    try:
+        _parsed_uri = pymongo.uri_parser.parse_uri(uri)
+    except pymongo.errors.InvalidURI:
+        # assume URI to be just the database name
+        db_name = uri
+        pass
+    else:
+        if "replicaset" in _parsed_uri["options"]:
+            connection_factory = MongoReplicaSetClient
+        db_name = _parsed_uri.get("database", "pysaml2")
+
+    if not "tz_aware" in kwargs:
+        # default, but not forced
+        kwargs["tz_aware"] = True
+
+    _conn = connection_factory(uri, **kwargs)
+    _db = _conn[db_name]
+
+    if "username" in _parsed_uri:
+        _db.authenticate(
+            _parsed_uri.get("username", None),
+            _parsed_uri.get("password", None)
+            )
+
+    return _db
+
 #------------------------------------------------------------------------------
 class EptidMDB(Eptid):
-    def __init__(self, secret, collection="", sub_collection="eptid"):
+    def __init__(self, secret, database="", collection="eptid"):
         Eptid.__init__(self, secret)
-        self.mdb = MDB(collection, sub_collection)
+        self.mdb = MDB(database, collection)
         self.mdb.primary_key = "eptid_key"
 
     def __getitem__(self, key):
@@ -330,9 +375,9 @@ def export_mdstore_to_mongo_db(mds, collection, sub_collection=""):
 
 
 class MetadataMDB(MetaData):
-    def __init__(self, onts, attrc, collection="", sub_collection=""):
+    def __init__(self, onts, attrc, database="", collection=""):
         MetaData.__init__(self, onts, attrc)
-        self.mdb = MDB(collection, sub_collection)
+        self.mdb = MDB(database, collection)
         self.mdb.primary_key = "entity_id"
 
     def _ext_service(self, entity_id, typ, service, binding):
@@ -381,4 +426,4 @@ class MetadataMDB(MetaData):
             raise CorruptDatabase("More then one document with key %s" % item)
 
     def bindings(self, entity_id, typ, service):
-        pass
\ No newline at end of file
+        pass