Merge pull request #297 from alphagov/refactor-data-layer-step-1

Refactor data layer step 1

backdrop/core/storage/mongo.py

@@ -0,0 +1,55 @@
+import os
+import logging
+import datetime
+from itertools import imap
+
+import pymongo
+from pymongo.errors import AutoReconnect
+from bson import Code
+
+from .. import timeutils
+from ... import statsd
+
+
+logger = logging.getLogger(__name__)
+
+__all__ = ['MongoStorageEngine']
+
+
+def get_mongo_client(hosts, port):
+    """Return an appropriate mongo client
+    """
+    client_list = ','.join(':'.join([host, str(port)]) for host in hosts)
+
+    # We can't always guarantee we'll be on 'production'
+    # so we allow jenkins to add the set as a variable
+    # Some test environments / other envs have their own sets e.g. 'gds-ci'
+    replica_set = os.getenv('MONGO_REPLICA_SET', 'production')
+
+    if replica_set == '':
+        return pymongo.MongoClient(client_list)
+    else:
+        return pymongo.MongoReplicaSetClient(
+            client_list, replicaSet=replica_set)
+
+
+class MongoStorageEngine(object):
+    @classmethod
+    def create(cls, hosts, port, database):
+        return cls(get_mongo_client(hosts, port), database)
+
+    def __init__(self, mongo, database):
+        self._mongo = mongo
+        self._db = mongo[database]
+
+    def alive(self):
+        return self._mongo.alive()
+
+    def dataset_exists(self, dataset_id):
+        return dataset_id in self._db.collection_names()
+
+    def create_dataset(self, dataset_id, size):
+        if size > 0:
+            self._db.create_collection(dataset_id, capped=True, size=size)
+        else:
+            self._db.create_collection(dataset_id, capped=False)

backdrop/read/api.py

@@ -14,6 +14,8 @@
 from ..core.repository import DataSetConfigRepository
 from ..core.timeutils import as_utc
 
+from ..core.storage.mongo import MongoStorageEngine
+
 
 GOVUK_ENV = getenv("GOVUK_ENV", "development")
 
@@ -30,6 +32,10 @@
     app.config['MONGO_PORT'],
     app.config['DATABASE_NAME']
 )
+storage = MongoStorageEngine.create(
+    app.config['MONGO_HOSTS'],
+    app.config['MONGO_PORT'],
+    app.config['DATABASE_NAME'])
 
 data_set_repository = DataSetConfigRepository(
     app.config['STAGECRAFT_URL'],
@@ -78,7 +84,7 @@ def http_error_handler(e):
 @cache_control.nocache
 def health_check():
 
-    if not db.alive():
+    if not storage.alive():
         return jsonify(status='error',
                        message='cannot connect to database'), 500
 

backdrop/write/api.py

@@ -11,6 +11,8 @@
 from ..core.errors import ParseError, ValidationError
 from ..core import database, log_handler, cache_control
 
+from ..core.storage.mongo import MongoStorageEngine
+
 from .validation import auth_header_is_valid, extract_bearer_token
 
 
@@ -29,6 +31,10 @@
     app.config['MONGO_PORT'],
     app.config['DATABASE_NAME']
 )
+storage = MongoStorageEngine.create(
+    app.config['MONGO_HOSTS'],
+    app.config['MONGO_PORT'],
+    app.config['DATABASE_NAME'])
 
 data_set_repository = DataSetConfigRepository(
     app.config['STAGECRAFT_URL'],
@@ -88,7 +94,7 @@ def http_error_handler(e):
 @app.route('/_status', methods=['GET'])
 @cache_control.nocache
 def health_check():
-    if db.alive():
+    if storage.alive():
         return jsonify(status='ok', message='database seems fine')
     else:
         abort(500, 'cannot connect to database')
@@ -172,7 +178,7 @@ def create_collection_for_dataset(dataset_name):
     if not _allow_create_collection(request.headers.get('Authorization')):
         abort(401, 'Unauthorized: invalid or no token given.')
 
-    if db.collection_exists(dataset_name):
+    if storage.dataset_exists(dataset_name):
         abort(400, 'Collection exists with that name.')
 
     try:
@@ -185,10 +191,7 @@ def create_collection_for_dataset(dataset_name):
     if capped_size is None or not isinstance(capped_size, int):
         abort(400, 'You must specify an int capped_size of 0 or more')
 
-    if capped_size == 0:
-        db.create_uncapped_collection(dataset_name)
-    else:
-        db.create_capped_collection(dataset_name, capped_size)
+    storage.create_dataset(dataset_name, capped_size)
 
     return jsonify(status='ok', message='Created "{}"'.format(dataset_name))
 

tests/core/storage/test_mongo.py

@@ -0,0 +1,56 @@
+"""
+Integration tests for the mongo storage engine.
+Unit tests are in doctest format in the module itself.
+"""
+from hamcrest import assert_that, is_, has_item, has_entries
+from nose.tools import assert_raises
+
+from pymongo.errors import CollectionInvalid
+
+from backdrop.core.storage.mongo import MongoStorageEngine
+
+
+class TestMongoStorageEngine(object):
+    def setup(self):
+        self.engine = MongoStorageEngine.create(
+            ['localhost'], 27017, 'backdrop_test')
+        self.mongo = self.engine._mongo
+        self.db = self.mongo['backdrop_test']
+
+    def teardown(self):
+        self.mongo.drop_database('backdrop_test')
+
+    # ALIVE
+    def test_alive_returns_true_if_mongo_is_up(self):
+        assert_that(self.engine.alive(), is_(True))
+
+    # EXISTS
+    def test_exists_returns_false_if_data_set_does_not_exist(self):
+        assert_that(self.engine.dataset_exists('foo_bar'), is_(False))
+
+    def test_exists_returns_true_if_data_set_exists(self):
+        self.db.create_collection('foo_bar')
+
+        assert_that(self.engine.dataset_exists('foo_bar'), is_(True))
+
+    # CREATE
+    def test_create_a_non_capped_collection(self):
+        self.engine.create_dataset('foo_bar', 0)
+
+        assert_that(self.db.collection_names(), has_item('foo_bar'))
+        assert_that(self.db['foo_bar'].options(), has_entries(
+            {'capped': False}))
+
+    def test_create_a_capped_collection(self):
+        self.engine.create_dataset('foo_bar', 100)
+
+        assert_that(self.db.collection_names(), has_item('foo_bar'))
+        assert_that(self.db['foo_bar'].options(), has_entries(
+            {'capped': True, 'size': 100}))
+
+    def test_create_fails_if_collection_already_exists(self):
+        # TODO: reraise a backdrop exception
+        self.engine.create_dataset('foo_bar', 0)
+
+        assert_raises(CollectionInvalid,
+                      self.engine.create_dataset, 'foo_bar', 0)

tests/write/test_flask_integration.py

@@ -191,9 +191,9 @@ def test_api_exposes_a_healthcheck(self):
         assert_that(entity["status"], is_("ok"))
 
     @patch("backdrop.write.api.statsd")
-    @patch("backdrop.write.api.db")
-    def test_exception_handling(self, db, statsd):
-        db.alive.side_effect = ValueError("BOOM")
+    @patch("backdrop.write.api.storage")
+    def test_exception_handling(self, storage, statsd):
+        storage.alive.side_effect = ValueError("BOOM")
 
         response = self.app.get("/_status")