From f0a304b909f54a7bc215f292290aa3be56eb0d9a Mon Sep 17 00:00:00 2001
From: "Gavin M. Roy" <gavinr@aweber.com>
Date: Fri, 13 Mar 2015 15:05:52 -0400
Subject: [PATCH] Handle connection errors better

---
 consulate/cli.py | 50 +++++++++++++++++++++++++++++++++++-------------
 1 file changed, 37 insertions(+), 13 deletions(-)

diff --git a/consulate/cli.py b/consulate/cli.py
index bf99672..3777ed5 100644
--- a/consulate/cli.py
+++ b/consulate/cli.py
@@ -3,9 +3,16 @@
 import json
 import sys
 
+from requests import exceptions
+
 from consulate import api
 
 
+def connection_error():
+    sys.stderr.write('ERROR: Could not connect to consul\n')
+    sys.exit(1)
+
+
 def parse_cli_args():
     parser = argparse.ArgumentParser(description='CLI utilities for Consul')
 
@@ -73,41 +80,58 @@ def parse_cli_args():
 def main():
     args = parse_cli_args()
     session = api.Session(args.api_host, args.api_port, args.datacenter,
-                            args.token)
+                          args.token)
 
     if args.command == 'register':
-
         check = args.path if args.ctype == 'check' else None
         interval = '%ss' % args.interval if args.ctype == 'check' else None
         ttl = '%ss' % args.duration if args.ctype == 'ttl' else None
         tags = args.tags.split(',') if args.tags else None
-        session.agent.service.register(args.name,
-                                       args.service_id,
-                                       args.port,
-                                       tags, check, interval, ttl)
+        try:
+            session.agent.service.register(args.name,
+                                           args.service_id,
+                                           args.port,
+                                           tags, check, interval, ttl)
+        except exceptions.ConnectionError:
+                    connection_error()
 
     elif args.command == 'kv':
 
         if args.action == 'backup':
             with open(args.backup_file, 'wb') as handle:
-                handle.write(json.dumps(session.kv.records(),
-                                        ensure_ascii=False,
-                                        indent=2))
+                try:
+                    handle.write(json.dumps(session.kv.records(),
+                                            ensure_ascii=False,
+                                            indent=2))
+                except exceptions.ConnectionError:
+                            connection_error()
 
         elif args.action == 'restore':
             with open(args.restore_file, 'rb') as handle:
                 data = json.load(handle)
                 for row in data:
-                    session.kv.set_record(row[0], row[1], row[2])
+                    try:
+                        session.kv.set_record(row[0], row[1], row[2])
+                    except exceptions.ConnectionError:
+                        connection_error()
 
         elif args.action == 'del':
-            del session.kv[args.key]
+            try:
+                del session.kv[args.key]
+            except exceptions.ConnectionError:
+                connection_error()
 
         elif args.action == 'get':
-            sys.stdout.write("%s\n" % session.kv.get(args.key))
+            try:
+                sys.stdout.write("%s\n" % session.kv.get(args.key))
+            except exceptions.ConnectionError:
+                connection_error()
 
         elif args.action == 'set':
-            session.kv[args.key] = args.value
+            try:
+                session.kv[args.key] = args.value
+            except exceptions.ConnectionError:
+                connection_error()
 
 
 if __name__ == '__main__':