Merge pull request alphagov#44 from alphagov/fix-mongo-status

rjw1 · rjw1 · commit c5eaf4bf5a8a · 2014-04-11T14:37:37.000+01:00
Fix and improve Mongo commands
diff --git a/mongo.py b/mongo.py
@@ -9,16 +9,26 @@
 def node_name(node_name):
     return node_name.split('.production').pop(0)
 
-def strip_isodate(raw_output):
-    return re.sub(r'ISODate\((.*?)\)', r'\1', raw_output)
+def strip_dates(raw_output):
+    stripped_isodates = re.sub(r'ISODate\((.*?)\)', r'\1', raw_output)
+    return re.sub(r'Timestamp\((.*?)\)', r'"\1"', stripped_isodates)
 
 def mongo_command(command):
     return "mongo --quiet --eval 'printjson(%s)'" % command
 
-def run_mongo_command(command):
-    return json.loads(
-            strip_isodate(
-                run(mongo_command(command))))
+def run_mongo_command(command, command_warn_only=False):
+    if command_warn_only:
+        with settings(warn_only=True):
+            response = run(mongo_command(command))
+    else:
+        response = run(mongo_command(command))
+
+    if response.return_code == 252:
+        dict_response = {"return_code": 252}
+    else:
+        dict_response = json.loads(strip_dates(response))
+
+    return dict_response
 
 @task(default=True)
 @roles('class-mongo')
@@ -39,20 +49,32 @@ def status():
     with hide('output'):
         status = run_mongo_command("rs.status()")
         print("Status at %s" % status['date'])
+        parsed_statuses = []
+        parsed_statuses.append(['Name', 'Health', 'State', 'Heartbeat'])
+
         for member_status in status['members']:
             name = node_name(member_status['name'])
             health = "OK" if member_status['health'] == 1 else "ERROR"
             state = member_status['stateStr']
             heartbeat = member_status.get('lastHeartbeat', '')
-            print("%s %s % -10s %s" % (name, health, state, heartbeat))
+            parsed_statuses.append([name, health, state, heartbeat])
+
+        for status in parsed_statuses:
+            print("| {0:<22} | {1:<8} | {2:<10} | {3:<22} |".format(status[0], status[1], status[2], status[3]))
 
 @task
 def step_down_primary(seconds='100'):
     """Step down as primary for a given number of seconds (default: 100)"""
     with hide('output'):
-        result = run_mongo_command("rs.stepDown(%s)" % seconds)
-        if result['ok'] == 1:
-            print("Stepped down for %s seconds" % seconds)
-        else:
+        # Mongo returns an exit code of 252 when the primary steps down, as well as disconnecting
+        # the current console session. This means that we have to run the command with Fabric's
+        # warn_only enabled, or Fabric will error.
+        result = run_mongo_command("rs.stepDown(%s)" % seconds, command_warn_only=True)
+
+        if 'ok' in result and result['ok'] == 0:
             print("Failed to step down: %s\nTry running mongo.find_primary" % result['errmsg'])
+            return 0
 
+        if 'return_code' in result and result['return_code'] == 252:
+            print "Received a 252 exit code from Mongo. There may have been disconnection warnings"
+            print "during the step down from primary that caused this. Verify the primary with mongo.find_primary."
