Update mongodb config (#15)

ecmwf · Nov 20, 2023 · bf9d082 · bf9d082
1 parent a833364
commit bf9d082
Show file tree

Hide file tree

Showing 11 changed files with 57 additions and 126 deletions.
diff --git a/polytope_server/common/authentication/mongoapikey_authentication.py b/polytope_server/common/authentication/mongoapikey_authentication.py
@@ -40,22 +40,17 @@ class ApiKeyMongoAuthentication(authentication.Authentication):
 
     def __init__(self, name, realm, config):
         self.config = config
-        host = config.get("host", "localhost")
-        port = config.get("port", "27017")
+        uri = config.get("uri", "mongodb://localhost:27017")
         collection = config.get("collection", "keys")
         username = config.get("username")
         password = config.get("password")
-        srv = bool(config.get("srv", False))
-        tls = bool(config.get("tls", False))
-        tlsCAFile = config.get("tlsCAFile", None)
 
-        endpoint = "{}:{}".format(host, port)
-        self.mongo_client = mongo_client_factory.create_client(host, port, username, password, srv, tls, tlsCAFile)
+        self.mongo_client = mongo_client_factory.create_client(uri, username, password)
         self.database = self.mongo_client.keys
         self.keys = self.database[collection]
         assert realm == "polytope"
 
-        self.storage_metric_collector = MongoStorageMetricCollector(endpoint, self.mongo_client, "keys", collection)
+        self.storage_metric_collector = MongoStorageMetricCollector(uri, self.mongo_client, "keys", collection)
 
         super().__init__(name, realm, config)
 

diff --git a/polytope_server/common/authentication/mongodb_authentication.py b/polytope_server/common/authentication/mongodb_authentication.py
@@ -32,22 +32,17 @@
 class MongoAuthentication(authentication.Authentication):
     def __init__(self, name, realm, config):
         self.config = config
-        host = config.get("host", "localhost")
-        port = config.get("port", "27017")
+        uri = config.get("uri", "mongodb://localhost:27017")
         collection = config.get("collection", "users")
         username = config.get("username")
         password = config.get("password")
-        srv = bool(config.get("srv", False))
-        tls = bool(config.get("tls", False))
-        tlsCAFile = config.get("tlsCAFile", None)
 
-        endpoint = "{}:{}".format(host, port)
-        self.mongo_client = mongo_client_factory.create_client(host, port, username, password, srv, tls, tlsCAFile)
+        self.mongo_client = mongo_client_factory.create_client(uri, username, password)
         self.database = self.mongo_client.authentication
         self.users = self.database[collection]
 
         self.storage_metric_collector = MongoStorageMetricCollector(
-            endpoint, self.mongo_client, "authentication", collection
+            uri, self.mongo_client, "authentication", collection
         )
 
         super().__init__(name, realm, config)

diff --git a/polytope_server/common/authorization/mongodb_authorization.py b/polytope_server/common/authorization/mongodb_authorization.py
@@ -28,24 +28,17 @@ class MongoDBAuthorization(authorization.Authorization):
     def __init__(self, name, realm, config):
         self.config = config
         assert self.config["type"] == "mongodb"
-        self.host = config.get("host", "localhost")
-        self.port = config.get("port", "27017")
+        self.uri = config.get("uri", "mongodb://localhost:27017")
         self.collection = config.get("collection", "users")
         username = config.get("username")
         password = config.get("password")
-        srv = bool(config.get("srv", False))
-        tls = bool(config.get("tls", False))
-        tlsCAFile = config.get("tlsCAFile", None)
 
-        endpoint = "{}:{}".format(self.host, self.port)
-        self.mongo_client = mongo_client_factory.create_client(
-            self.host, self.port, username, password, srv, tls, tlsCAFile
-        )
+        self.mongo_client = mongo_client_factory.create_client(self.uri, username, password)
         self.database = self.mongo_client.authentication
         self.users = self.database[self.collection]
 
         self.storage_metric_collector = MongoStorageMetricCollector(
-            endpoint, self.mongo_client, "authentication", self.collection
+            self.uri, self.mongo_client, "authentication", self.collection
         )
 
         super().__init__(name, realm, config)

diff --git a/polytope_server/common/caching/caching.py b/polytope_server/common/caching/caching.py
@@ -195,23 +195,20 @@ def collect_metric_info(self):
 class MongoDBCaching(Caching):
     def __init__(self, cache_config):
         super().__init__(cache_config)
-        host = cache_config.get("host", "localhost")
-        port = cache_config.get("port", 27017)
+        uri = cache_config.get("uri", "mongodb://localhost:27017")
+
         username = cache_config.get("username")
         password = cache_config.get("password")
-        srv = bool(cache_config.get("srv", False))
-        tls = bool(cache_config.get("tls", False))
-        tlsCAFile = cache_config.get("tlsCAFile", None)
-        endpoint = "{}:{}".format(host, port)
+
         collection = cache_config.get("collection", "cache")
-        self.client = mongo_client_factory.create_client(host, port, username, password, srv, tls, tlsCAFile)
+        self.client = mongo_client_factory.create_client(uri, username, password,)
 
         self.database = self.client.cache
         self.collection = self.database[collection]
         self.collection.create_index("expire_at", expireAfterSeconds=0)
         self.collection.update_one({"_id": "hits"}, {"$setOnInsert": {"n": 0}}, upsert=True)
         self.collection.update_one({"_id": "misses"}, {"$setOnInsert": {"n": 0}}, upsert=True)
-        self.storage_metric_collector = MongoStorageMetricCollector(endpoint, self.client, "cache", collection)
+        self.storage_metric_collector = MongoStorageMetricCollector(uri, self.client, "cache", collection)
         self.cache_metric_collector = MongoCacheMetricCollector(self.client, "cache", collection)
 
     def get_type(self):

diff --git a/polytope_server/common/config/schema.yaml b/polytope_server/common/config/schema.yaml
@@ -68,7 +68,7 @@ mapping:
             desc: point to a hosted mongodb
             type: map
             mapping:
-              endpoint:
+              uri:
                 desc: host and port
                 example: localhost:27017
                 type: str
@@ -116,7 +116,7 @@ mapping:
             desc: point to a hosted mongodb
             type: map
             mapping:
-              endpoint:
+              uri:
                 desc: host and port
                 example: localhost:27017
                 type: str

diff --git a/polytope_server/common/identity/mongodb_identity.py b/polytope_server/common/identity/mongodb_identity.py
@@ -28,18 +28,16 @@
 class MongoDBIdentity(identity.Identity):
     def __init__(self, config):
         self.config = config
-        self.host = config.get("host", "localhost")
-        self.port = config.get("port", "27017")
+        self.uri = config.get("uri", "mongodb://localhost:27017")
+
         self.collection = config.get("collection", "users")
         username = config.get("username")
         password = config.get("password")
-        srv = bool(config.get("srv", False))
-        tls = bool(config.get("tls", False))
-        tlsCAFile = config.get("tlsCAFile", None)
 
-        endpoint = "{}:{}".format(self.host, self.port)
         self.mongo_client = mongo_client_factory.create_client(
-            self.host, self.port, username, password, srv, tls, tlsCAFile
+            self.uri,
+            username,
+            password,
         )
         self.database = self.mongo_client.authentication
         self.users = self.database[self.collection]
@@ -53,7 +51,7 @@ def __init__(self, config):
                 pass
 
         self.storage_metric_collector = MongoStorageMetricCollector(
-            endpoint, self.mongo_client, "authentication", self.collection
+            self.uri, self.mongo_client, "authentication", self.collection
         )
         self.identity_metric_collector = MetricCollector()
 

diff --git a/polytope_server/common/keygenerator/mongodb_keygenerator.py b/polytope_server/common/keygenerator/mongodb_keygenerator.py
@@ -33,23 +33,17 @@ class MongoKeyGenerator(keygenerator.KeyGenerator):
     def __init__(self, config):
         self.config = config
         assert self.config["type"] == "mongodb"
-        host = config.get("host", "localhost")
-        port = config.get("port", "27017")
+        uri = config.get("uri", "mongodb://localhost:27017")
         collection = config.get("collection", "keys")
         username = config.get("username")
         password = config.get("password")
-        srv = bool(config.get("srv", False))
-        tls = bool(config.get("tls", False))
-        tlsCAFile = config.get("tlsCAFile", None)
 
-        endpoint = "{}:{}".format(host, port)
-
-        self.mongo_client = mongo_client_factory.create_client(host, port, username, password, srv, tls, tlsCAFile)
+        self.mongo_client = mongo_client_factory.create_client(uri, username, password)
         self.database = self.mongo_client.keys
         self.keys = self.database[collection]
         self.realms = config.get("allowed_realms")
 
-        self.storage_metric_collector = MongoStorageMetricCollector(endpoint, self.mongo_client, "keys", collection)
+        self.storage_metric_collector = MongoStorageMetricCollector(uri, self.mongo_client, "keys", collection)
 
     def create_key(self, user: User) -> ApiKey:
         if user.realm not in self.realms:

diff --git a/polytope_server/common/metric_store/mongodb_metric_store.py b/polytope_server/common/metric_store/mongodb_metric_store.py
@@ -39,19 +39,13 @@
 
 class MongoMetricStore(MetricStore):
     def __init__(self, config=None):
-        host = config.get("host", "localhost")
-        port = config.get("port", "27017")
+        uri = config.get("uri", "mongodb://localhost:27017")
         metric_collection = config.get("collection", "metrics")
 
         username = config.get("username")
         password = config.get("password")
-        srv = bool(config.get("srv", False))
-        tls = bool(config.get("tls", False))
-        tlsCAFile = config.get("tlsCAFile", None)
 
-        endpoint = "{}:{}".format(host, port)
-
-        self.mongo_client = mongo_client_factory.create_client(host, port, username, password, srv, tls, tlsCAFile)
+        self.mongo_client = mongo_client_factory.create_client(uri, username, password)
         self.database = self.mongo_client.metric_storeg
         self.store = self.database[metric_collection]
 
@@ -65,10 +59,10 @@ def __init__(self, config=None):
         }
 
         self.storage_metric_collector = MongoStorageMetricCollector(
-            endpoint, self.mongo_client, "metric_store", metric_collection
+            uri, self.mongo_client, "metric_store", metric_collection
         )
 
-        logging.info("MongoClient configured to open at {}".format(endpoint))
+        logging.info("MongoClient configured to open at {}".format(uri))
 
     def get_type(self):
         return "mongodb"

diff --git a/polytope_server/common/mongo_client_factory.py b/polytope_server/common/mongo_client_factory.py
@@ -1,27 +1,14 @@
 import typing
-import urllib.parse
 
 import pymongo
 
 
 def create_client(
-    host: str,
-    port: str,
+    uri: str,
     username: typing.Optional[str] = None,
     password: typing.Optional[str] = None,
-    srv: bool = False,
-    tls: bool = False,
-    tlsCAFile: typing.Optional[str] = None,
 ) -> pymongo.MongoClient:
-    protocol = "mongodb"
-    if srv:
-        protocol = "mongodb+srv"
-
-    endpoint = f"{protocol}://{host}:{port}"
-
     if username and password:
-        encoded_username = urllib.parse.quote_plus(username)
-        encoded_password = urllib.parse.quote_plus(password)
-        endpoint = f"{protocol}://{encoded_username}:{encoded_password}@{host}:{port}"
-
-    return pymongo.MongoClient(endpoint, journal=True, connect=False, tls=tls, tlsCAFile=tlsCAFile)
+        return pymongo.MongoClient(host=uri, journal=True, connect=False, username=username, password=password)
+    else:
+        return pymongo.MongoClient(host=uri, journal=True, connect=False)
diff --git a/polytope_server/common/request_store/mongodb_request_store.py b/polytope_server/common/request_store/mongodb_request_store.py
@@ -35,19 +35,12 @@
 
 class MongoRequestStore(request_store.RequestStore):
     def __init__(self, config=None, metric_store_config=None):
-        host = config.get("host", "localhost")
-        port = config.get("port", "27017")
+        uri = config.get("uri", "mongodb://localhost:27017")
         request_collection = config.get("collection", "requests")
-
         username = config.get("username")
         password = config.get("password")
-        srv = bool(config.get("srv", False))
-        tls = bool(config.get("tls", False))
-        tlsCAFile = config.get("tlsCAFile", None)
-
-        endpoint = "{}:{}".format(host, port)
 
-        self.mongo_client = mongo_client_factory.create_client(host, port, username, password, srv, tls, tlsCAFile)
+        self.mongo_client = mongo_client_factory.create_client(uri, username, password)
         self.database = self.mongo_client.request_store
         self.store = self.database[request_collection]
 
@@ -56,11 +49,11 @@ def __init__(self, config=None, metric_store_config=None):
             self.metric_store = metric_store.create_metric_store(metric_store_config)
 
         self.storage_metric_collector = MongoStorageMetricCollector(
-            endpoint, self.mongo_client, "request_store", request_collection
+            uri, self.mongo_client, "request_store", request_collection
         )
         self.request_store_metric_collector = MongoRequestStoreMetricCollector()
 
-        logging.info("MongoClient configured to open at {}".format(endpoint))
+        logging.info("MongoClient configured to open at {}".format(uri))
 
     def get_type(self):
         return "mongodb"

diff --git a/tests/unit/test_mongo_client_factory.py b/tests/unit/test_mongo_client_factory.py
@@ -1,59 +1,44 @@
+import typing
 from unittest import mock
 
 from polytope_server.common import mongo_client_factory
 
 
 @mock.patch("polytope_server.common.mongo_client_factory.pymongo.MongoClient", autospec=True)
 def test_create_without_credentials(mock_mongo: mock.Mock):
-    mongo_client_factory.create_client("host", "123", username=None, password=None, tls=False)
+    mongo_client_factory.create_client("mongodb://host:123")
 
-    _verify(mock_mongo, "host:123", False, False)
+    _verify(mock_mongo, "mongodb://host:123", None, None)
 
 
 @mock.patch("polytope_server.common.mongo_client_factory.pymongo.MongoClient", autospec=True)
-def test_create_with_srv(mock_mongo: mock.Mock):
-    mongo_client_factory.create_client("host", "123", username=None, password=None, srv=True, tls=False)
+def test_create_without_password_credentials(mock_mongo: mock.Mock):
+    mongo_client_factory.create_client("mongodb+srv://host:123", username="admin")
 
-    _verify(mock_mongo, "host:123", True, False)
+    _verify(mock_mongo, "mongodb+srv://host:123", None, None)
 
 
 @mock.patch("polytope_server.common.mongo_client_factory.pymongo.MongoClient", autospec=True)
-def test_create_with_credentials(mock_mongo: mock.Mock):
-    mongo_client_factory.create_client("host", "123", username="admin", password="admin", tls=False)
-
-    _verify(mock_mongo, "admin:admin@host:123", False, False)
-
+def test_create_without_username_credentials(mock_mongo: mock.Mock):
+    mongo_client_factory.create_client("host:123", password="password")
 
-@mock.patch("polytope_server.common.mongo_client_factory.pymongo.MongoClient", autospec=True)
-def test_create_without_credentials_tls(mock_mongo: mock.Mock):
-    mongo_client_factory.create_client("host", "123", username=None, password=None, tls=True)
-
-    _verify(mock_mongo, "host:123", False, True)
+    _verify(mock_mongo, "host:123", None, None)
 
 
 @mock.patch("polytope_server.common.mongo_client_factory.pymongo.MongoClient", autospec=True)
-def test_create_with_credentials_tls(mock_mongo: mock.Mock):
-    mongo_client_factory.create_client("host", "123", username="admin", password="admin", tls=True)
-
-    _verify(mock_mongo, "admin:admin@host:123", False, True)
-
-
-@mock.patch("polytope_server.common.mongo_client_factory.pymongo.MongoClient", autospec=True)
-def test_create_with_tlsCAfile(mock_mongo: mock.Mock):
-    mongo_client_factory.create_client(
-        "host", "123", username="admin", password="admin", tls=True, tlsCAFile="/test/ca.pem"
-    )
+def test_create_with_credentials(mock_mongo: mock.Mock):
+    mongo_client_factory.create_client("mongodb+srv://host", username="admin", password="est123123")
 
-    _verify(mock_mongo, "admin:admin@host:123", False, True, "/test/ca.pem")
+    _verify(mock_mongo, "mongodb+srv://host", "admin", "est123123")
 
 
-def _verify(mock_mongo: mock.Mock, endpoint: str, srv: bool, tls: bool, tlsCAFile=None):
+def _verify(
+    mock_mongo: mock.Mock, endpoint: str, username: typing.Optional[str] = None, password: typing.Optional[str] = None
+):
     mock_mongo.assert_called_once()
     args, kwargs = mock_mongo.call_args
-    if srv:
-        assert args[0] == f"mongodb+srv://{endpoint}"
-    else:
-        assert args[0] == f"mongodb://{endpoint}"
-
-    assert kwargs["tls"] == tls
-    assert kwargs["tlsCAFile"] == tlsCAFile
+    assert args[0] == endpoint
+    if username:
+        assert kwargs["username"] == username
+    if password:
+        assert kwargs["password"] == password