Merge pull request #10 from ecmwf/MeteoSwiss-APN-feature/merge_back_t…

…o_ecmwf

Meteo swiss apn feature/merge back to ecmwf

polytope_server/common/authentication/mongoapikey_authentication.py

@@ -20,8 +20,7 @@
 
 from datetime import datetime
 
-import pymongo
-
+from .. import mongo_client_factory
 from ..auth import User
 from ..exceptions import ForbiddenRequest
 from ..metric_collector import MongoStorageMetricCollector
@@ -40,19 +39,18 @@ 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")
 
-        endpoint = "{}:{}".format(host, port)
-        self.mongo_client = pymongo.MongoClient(endpoint, journal=True, connect=False)
+        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)
 
@@ -63,7 +61,6 @@ def authentication_info(self):
         return "Authenticate with Polytope API Key from ../auth/keys"
 
     def authenticate(self, credentials: str) -> User:
-
         # credentials should be of the form '<ApiKey>'
         res = self.keys.find_one({"key.key": credentials})
         if res is None:

polytope_server/common/authentication/mongodb_authentication.py

@@ -22,8 +22,7 @@
 import binascii
 import hashlib
 
-import pymongo
-
+from .. import mongo_client_factory
 from ..auth import User
 from ..exceptions import ForbiddenRequest
 from ..metric_collector import MongoStorageMetricCollector
@@ -32,19 +31,18 @@
 
 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")
 
-        endpoint = "{}:{}".format(host, port)
-        self.mongo_client = pymongo.MongoClient(endpoint, journal=True, connect=False)
+        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)
@@ -59,7 +57,6 @@ def authentication_info(self):
         return "Authenticate with username and password"
 
     def authenticate(self, credentials: str) -> User:
-
         # credentials should be of the form 'base64(<username>:<API_key>)'
         try:
             decoded = base64.b64decode(credentials).decode("utf-8")

polytope_server/common/authorization/mongodb_authorization.py

@@ -18,8 +18,7 @@
 # does it submit to any jurisdiction.
 #
 
-import pymongo
-
+from .. import mongo_client_factory
 from ..auth import User
 from ..metric_collector import MongoStorageMetricCollector
 from . import authorization
@@ -29,23 +28,22 @@ 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")
 
-        endpoint = "{}:{}".format(self.host, self.port)
-        self.mongo_client = pymongo.MongoClient(endpoint, journal=True, connect=False)
+        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)
 
     def get_roles(self, user: User) -> list:
-
         if user.realm != self.realm():
             raise ValueError(
                 "Trying to authorize a user in the wrong realm, expected {}, got {}".format(self.realm(), user.realm)

polytope_server/common/caching/caching.py

@@ -29,9 +29,9 @@
 from typing import Dict, Union
 
 import pymemcache
-import pymongo
 import redis
 
+from .. import mongo_client_factory
 from ..metric import MetricType
 from ..metric_collector import (
     DictStorageMetricCollector,
@@ -195,17 +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)
-        endpoint = "{}:{}".format(host, port)
+        uri = cache_config.get("uri", "mongodb://localhost:27017")
+
+        username = cache_config.get("username")
+        password = cache_config.get("password")
+
         collection = cache_config.get("collection", "cache")
-        self.client = pymongo.MongoClient(host + ":" + str(port), journal=False, connect=False)
+        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):
@@ -220,7 +223,6 @@ def get(self, key):
         return obj["data"]
 
     def set(self, key, object, lifetime):
-
         if lifetime == 0 or lifetime is None:
             expiry = datetime.datetime.max
         else:
@@ -324,7 +326,6 @@ def __call__(self, f):
 
         @functools.wraps(f)
         def wrapper(*args, **kwargs):
-
             cache.cancelled = False
 
             if self.cache is None:

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

polytope_server/common/identity/mongodb_identity.py

@@ -18,8 +18,7 @@
 # does it submit to any jurisdiction.
 #
 
-import pymongo
-
+from .. import mongo_client_factory
 from ..authentication.mongodb_authentication import MongoAuthentication
 from ..exceptions import Conflict, NotFound
 from ..metric_collector import MetricCollector, MongoStorageMetricCollector
@@ -29,12 +28,17 @@
 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")
 
-        endpoint = "{}:{}".format(self.host, self.port)
-        self.mongo_client = pymongo.MongoClient(endpoint, journal=True, connect=False)
+        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.realm = config.get("realm")
@@ -47,12 +51,11 @@ 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()
 
     def add_user(self, username: str, password: str, roles: list) -> bool:
-
         if self.users.find_one({"username": username}) is not None:
             raise Conflict("Username already registered")
 
@@ -70,7 +73,6 @@ def add_user(self, username: str, password: str, roles: list) -> bool:
         return True
 
     def remove_user(self, username: str) -> bool:
-
         result = self.users.delete_one({"username": username})
         if result.deleted_count > 0:
             return True

polytope_server/common/keygenerator/mongodb_keygenerator.py

@@ -22,8 +22,7 @@
 import uuid
 from datetime import datetime, timedelta
 
-import pymongo
-
+from .. import mongo_client_factory
 from ..auth import User
 from ..exceptions import ForbiddenRequest
 from ..metric_collector import MongoStorageMetricCollector
@@ -34,19 +33,19 @@ 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")
-        endpoint = "{}:{}".format(host, port)
-        self.mongo_client = pymongo.MongoClient(endpoint, journal=True, connect=False)
+        username = config.get("username")
+        password = config.get("password")
+
+        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:
             raise ForbiddenRequest("Not allowed to create an API Key for users in realm {}".format(user.realm))
 

polytope_server/common/metric_collector/queue_metric_collector.py

@@ -46,3 +46,21 @@ def total_queued(self):
         channel = connection.channel()
         q = channel.queue_declare(queue=self.queue_name, durable=True, passive=True)
         return q.method.message_count
+
+
+class SQSQueueMetricCollector(QueueMetricCollector):
+    def __init__(self, host, client):
+        self.host = host
+        self.client = client
+
+    def total_queued(self):
+        response = self.client.get_queue_attributes(
+            QueueUrl=self.host,
+            AttributeNames=[
+                "ApproximateNumberOfMessages",
+                "ApproximateNumberOfMessagesDelayed",
+                "ApproximateNumberOfMessagesNotVisible",
+            ],
+        )
+        values = response.get("Attributes", {}).values()
+        return sum(map(int, values))

polytope_server/common/metric_collector/storage_metric_collector.py

@@ -112,7 +112,7 @@ def storage_space_used(self):
         return space_used
 
     def total_entries(self):
-        return self.store.count()
+        return self.store.count_documents({})
 
     def db_name(self):
         return self.database

polytope_server/common/metric_store/mongodb_metric_store.py

@@ -22,6 +22,7 @@
 
 import pymongo
 
+from .. import mongo_client_factory
 from ..metric import (
     CacheInfo,
     Metric,
@@ -38,13 +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")
 
-        endpoint = "{}:{}".format(host, port)
+        username = config.get("username")
+        password = config.get("password")
 
-        self.mongo_client = pymongo.MongoClient(endpoint, journal=True, connect=False)
+        self.mongo_client = mongo_client_factory.create_client(uri, username, password)
         self.database = self.mongo_client.metric_store
         self.store = self.database[metric_collection]
 
@@ -58,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"
@@ -85,7 +86,6 @@ def get_metric(self, uuid):
             return None
 
     def get_metrics(self, ascending=None, descending=None, limit=None, **kwargs):
-
         all_slots = []
 
         found_type = None

polytope_server/common/mongo_client_factory.py

@@ -0,0 +1,14 @@
+import typing
+
+import pymongo
+
+
+def create_client(
+    uri: str,
+    username: typing.Optional[str] = None,
+    password: typing.Optional[str] = None,
+) -> pymongo.MongoClient:
+    if username and password:
+        return pymongo.MongoClient(host=uri, journal=True, connect=False, username=username, password=password)
+    else:
+        return pymongo.MongoClient(host=uri, journal=True, connect=False)

polytope_server/common/queue/queue.py

@@ -80,7 +80,7 @@ def collect_metric_info(
         """Collect dictionary of metrics"""
 
 
-queue_dict = {"rabbitmq": "RabbitmqQueue"}
+queue_dict = {"rabbitmq": "RabbitmqQueue", "sqs": "SQSQueue"}
 
 
 def create_queue(queue_config):