Refactor heartbeat to delegate to the modulestores and sql

LMS-2801

common/djangoapps/contentserver/tests/test.py

@@ -15,7 +15,6 @@ from django.test.utils import override_settings
 from student.models import CourseEnrollment
 
 from xmodule.contentstore.django import contentstore, _CONTENTSTORE
-from xmodule.contentstore.content import StaticContent
 from xmodule.modulestore.django import modulestore
 from opaque_keys.edx.locations import SlashSeparatedCourseKey
 from xmodule.modulestore.tests.django_utils import (studio_store_config,

common/djangoapps/heartbeat/tests/test_heartbeat.py

@@ -0,0 +1,45 @@
+"""
+Test the heartbeat
+"""
+from django.test.client import Client
+from django.core.urlresolvers import reverse
+import json
+from django.db.utils import DatabaseError
+import mock
+from django.test.utils import override_settings
+from django.conf import settings
+from django.test.testcases import TestCase
+from xmodule.modulestore.tests.django_utils import mongo_store_config
+
+TEST_MODULESTORE = mongo_store_config(settings.TEST_ROOT / "data")
+
+@override_settings(MODULESTORE=TEST_MODULESTORE)
+class HeartbeatTestCase(TestCase):
+    """
+    Test the heartbeat
+    """
+
+    def setUp(self):
+        self.client = Client()
+        self.heartbeat_url = reverse('heartbeat')
+        return super(HeartbeatTestCase, self).setUp()
+
+    def tearDown(self):
+        return super(HeartbeatTestCase, self).tearDown()
+
+    def test_success(self):
+        response = self.client.get(self.heartbeat_url)
+        self.assertEqual(response.status_code, 200)
+
+    def test_sql_fail(self):
+        with mock.patch('heartbeat.views.connection') as mock_connection:
+            mock_connection.cursor.return_value.execute.side_effect = DatabaseError
+            response = self.client.get(self.heartbeat_url)
+            self.assertEqual(response.status_code, 503)
+            response_dict = json.loads(response.content)
+            self.assertIn('SQL', response_dict)
+
+    def test_mongo_fail(self):
+        with mock.patch('pymongo.MongoClient.alive', return_value=False):
+            response = self.client.get(self.heartbeat_url)
+            self.assertEqual(response.status_code, 503)

common/djangoapps/heartbeat/views.py

@@ -1,18 +1,33 @@
-import json
-from datetime import datetime
-from pytz import UTC
-from django.http import HttpResponse
 from xmodule.modulestore.django import modulestore
 from dogapi import dog_stats_api
+from util.json_request import JsonResponse
+from django.db import connection
+from django.db.utils import DatabaseError
+from xmodule.exceptions import HeartbeatFailure
 
 
 @dog_stats_api.timed('edxapp.heartbeat')
 def heartbeat(request):
     """
-    Simple view that a loadbalancer can check to verify that the app is up
+    Simple view that a loadbalancer can check to verify that the app is up. Returns a json doc
+    of service id: status or message. If the status for any service is anything other than True,
+    it returns HTTP code 503 (Service Unavailable); otherwise, it returns 200.
     """
-    output = {
-        'date': datetime.now(UTC).isoformat(),
-        'courses': [course.location.to_deprecated_string() for course in modulestore().get_courses()],
-    }
-    return HttpResponse(json.dumps(output, indent=4))
+    # This refactoring merely delegates to the default modulestore (which if it's mixed modulestore will
+    # delegate to all configured modulestores) and a quick test of sql. A later refactoring may allow
+    # any service to register itself as participating in the heartbeat. It's important that all implementation
+    # do as little as possible but give a sound determination that they are ready.
+    try:
+        output = modulestore().heartbeat()
+    except HeartbeatFailure as fail:
+        return JsonResponse({fail.service: unicode(fail)}, status=503)
+
+    cursor = connection.cursor()
+    try:
+        cursor.execute("SELECT CURRENT_DATE")
+        cursor.fetchone()
+        output['SQL'] = True
+    except DatabaseError as fail:
+        return JsonResponse({'SQL': unicode(fail)}, status=503)
+
+    return JsonResponse(output)

common/lib/xmodule/xmodule/exceptions.py

@@ -38,3 +38,20 @@ class UndefinedContext(Exception):
     Tried to access an xmodule field which needs a different context (runtime) to have a value.
     """
     pass
+
+
+class HeartbeatFailure(Exception):
+    """
+    Raised when heartbeat fails.
+    """
+
+    def __unicode__(self, *args, **kwargs):
+        return self.message
+
+
+    def __init__(self, msg, service):
+        """
+        In addition to a msg, provide the name of the service.
+        """
+        self.service = service
+        return super(HeartbeatFailure, self).__init__(msg)

common/lib/xmodule/xmodule/modulestore/__init__.py

@@ -362,6 +362,12 @@ class ModuleStoreReadBase(ModuleStoreRead):
         else:
             return any(c.id == course_id for c in self.get_courses())
 
+    def heartbeat(self):
+        """
+        Is this modulestore ready?
+        """
+        # default is to say yes by not raising an exception
+        return {'default_impl': True}
 
 class ModuleStoreWriteBase(ModuleStoreReadBase, ModuleStoreWrite):
     '''

common/lib/xmodule/xmodule/modulestore/mixed.py

@@ -19,6 +19,7 @@ from opaque_keys.edx.keys import CourseKey, UsageKey
 from xmodule.modulestore.mongo.base import MongoModuleStore
 from xmodule.modulestore.split_mongo.split import SplitMongoModuleStore
 from opaque_keys.edx.locations import SlashSeparatedCourseKey
+import itertools
 
 log = logging.getLogger(__name__)
 
@@ -332,6 +333,18 @@ class MixedModuleStore(ModuleStoreWriteBase):
             courses.extend(modulestore.get_courses_for_wiki(wiki_slug))
         return courses
 
+    def heartbeat(self):
+        """
+        Delegate to each modulestore and package the results for the caller.
+        """
+        # could be done in parallel threads if needed
+        return dict(
+            itertools.chain.from_iterable(
+                store.heartbeat().iteritems()
+                for store in self.modulestores.itervalues()
+            )
+        )
+
 
 def _compare_stores(left, right):
     """

common/lib/xmodule/xmodule/modulestore/mongo/base.py

@@ -38,6 +38,7 @@ from xmodule.modulestore.inheritance import own_metadata, InheritanceMixin, inhe
 from xmodule.tabs import StaticTab, CourseTabList
 from xblock.core import XBlock
 from opaque_keys.edx.locations import SlashSeparatedCourseKey
+from xmodule.exceptions import HeartbeatFailure
 
 log = logging.getLogger(__name__)
 
@@ -1034,3 +1035,12 @@ class MongoModuleStore(ModuleStoreWriteBase):
 
         field_data = KvsFieldData(kvs)
         return field_data
+
+    def heartbeat(self):
+        """
+        Check that the db is reachable.
+        """
+        if self.database.connection.alive():
+            return {MONGO_MODULESTORE_TYPE: True}
+        else:
+            raise HeartbeatFailure("Can't connect to {}".format(self.database.name), 'mongo')

common/lib/xmodule/xmodule/modulestore/split_mongo/mongo_connection.py

@@ -4,6 +4,7 @@ Segregation of pymongo functions from the data modeling mechanisms for split mod
 import re
 import pymongo
 from bson import son
+from xmodule.exceptions import HeartbeatFailure
 
 class MongoConnection(object):
     """
@@ -41,6 +42,15 @@ class MongoConnection(object):
         self.structures.write_concern = {'w': 1}
         self.definitions.write_concern = {'w': 1}
 
+    def heartbeat(self):
+        """
+        Check that the db is reachable.
+        """
+        if self.database.connection.alive():
+            return True
+        else:
+            raise HeartbeatFailure("Can't connect to {}".format(self.database.name))
+
     def get_structure(self, key):
         """
         Get the structure from the persistence mechanism whose id is the given key

common/lib/xmodule/xmodule/modulestore/split_mongo/split.py

@@ -1769,3 +1769,9 @@ class SplitMongoModuleStore(ModuleStoreWriteBase):
         """
         courses = []
         return courses
+
+    def heartbeat(self):
+        """
+        Check that the db is reachable.
+        """
+        return {SPLIT_MONGO_MODULESTORE_TYPE: self.db_connection.heartbeat()}

common/lib/xmodule/xmodule/modulestore/xml.py

@@ -815,3 +815,13 @@ class XMLModuleStore(ModuleStoreReadBase):
         """
         courses = self.get_courses()
         return [course.location for course in courses if (course.wiki_slug == wiki_slug)]
+
+    def heartbeat(self):
+        """
+        Ensure that every known course is loaded and ready to go. Really, just return b/c
+        if this gets called the __init__ finished which means the courses are loaded.
+
+        Returns the course count
+        """
+        return {XML_MODULESTORE_TYPE: True}
+