Merge pull request gmr#47 from amorken/master

Add support for HTTP health checks and CLI support for deregistering services

consulate/api/agent.py

@@ -60,13 +60,15 @@ def register(self, name,
                      check_id=None,
                      interval=None,
                      ttl=None,
-                     notes=None):
+                     notes=None,
+                     http=None):
             """Add a new check to the local agent. Checks are either a script
             or TTL type. The agent is responsible for managing the status of
             the check and keeping the Catalog in sync.
 
             The ``name`` field is mandatory, as is either ``script`` and
-            ``interval`` or ``ttl``. Only one of ``script`` and ``interval``
+            ``interval``, ``http`` and ``interval`` or ``ttl``.
+            Only one of ``script`` and ``interval``, ``http`` and  ``interval``
             or ``ttl`` should be provided.  If an ``check_id`` is not
             provided, it is set to ``name``. You cannot have  duplicate
             ``check_id`` entries per agent, so it may be necessary to provide
@@ -75,14 +77,19 @@ def register(self, name,
 
             If a ``script`` is provided, the check type is a script, and Consul
             will evaluate the script every ``interval`` to update the status.
+            If a ``http`` URL is provided, Consul will poll the URL every
+            ``interval`` to update the status - only 2xx results are considered
+            healthy.
             If a ``ttl`` type is used, then the ``ttl`` update APIs must be
             used to periodically update the state of the check.
 
             :param str name: The check name
+            :param str http: The URL to poll for health checks
             :param str script: The path to the script to run
             :param str check_id: The optional check id
             :param int interval: The interval to run the check
             :param int ttl: The ttl to specify for the check
+            :param str notes: Administrative notes.
             :rtype: bool
             :raises: ValueError
 
@@ -93,12 +100,22 @@ def register(self, name,
             elif script and ttl:
                 raise ValueError('Can not specify script and ttl together')
 
+            if http and not interval:
+                raise ValueError('Must specify interval when using http')
+            elif http and ttl:
+                raise ValueError('Can not specify http and ttl together')
+
+            if http and script:
+                raise ValueError('Can not specify script and http together')
+
+
             # Register the check
             return self._put_no_response_body(['register'], None, {
                 'ID': check_id,
                 'Name': name,
                 'Notes': notes,
                 'Script': script,
+                'HTTP': http,
                 'Interval': interval,
                 'TTL': ttl
             })
@@ -161,17 +178,19 @@ def register(self, name,
                      tags=None,
                      check=None,
                      interval=None,
-                     ttl=None):
+                     ttl=None,
+                     httpcheck=None):
             """Add a new service to the local agent.
 
             :param str name: The name of the service
             :param str service_id: The id for the service (optional)
             :param str address: The service IP address
             :param int port: The service port
             :param list tags: A list of tags for the service
-            :param str check: The path to the check to run
-            :param str interval: The script execution interval
+            :param str check: The path to the check script to run
+            :param str interval: The check execution interval
             :param str ttl: The TTL for external script check pings
+            :param str httpcheck: An URL to check every interval
             :rtype: bool
             :raises: ValueError
 
@@ -181,21 +200,34 @@ def register(self, name,
                 raise ValueError('port must be an integer')
             elif tags and not isinstance(tags, list):
                 raise ValueError('tags must be a list of strings')
-            elif check and ttl:
+            elif (check or httpcheck) and ttl:
                 raise ValueError('Can not specify both a check and ttl')
 
+            if (check or httpcheck) and not interval:
+                raise ValueError('An interval is required for check scripts and http checks.')
+
+            check_spec = None
+            if check:
+                check_spec = {'script': check,
+                              'interval': interval}
+            elif httpcheck:
+                check_spec = {'HTTP': httpcheck,
+                              'interval': interval}
+            elif ttl:
+                check_spec = {'TTL': ttl}
+
             # Build the payload to send to consul
             payload = {
                 'id': service_id,
                 'name': name,
                 'port': port,
                 'address': address,
-                'tags': tags,
-                'check': {'script': check,
-                          'interval': interval,
-                          'ttl': ttl}
+                'tags': tags
             }
 
+            if check_spec:
+                payload['check'] = check_spec
+
             for key in list(payload.keys()):
                 if payload[key] is None:
                     del payload[key]

consulate/cli.py

@@ -78,7 +78,7 @@ def add_register_args(parser):
     registerp.add_argument('name', help='The service name')
     registerp.add_argument('-a', '--address', default=None,
                            help='Specify an address')
-    registerp.add_argument('-p', '--port', default=None, help='Specify a port')
+    registerp.add_argument('-p', '--port', default=None, type=int, help='Specify a port')
     registerp.add_argument('-s', '--service-id', default=None,
                            help='Specify a service ID')
     registerp.add_argument('-t', '--tags', default=[],
@@ -91,11 +91,27 @@ def add_register_args(parser):
                        help='How often to run the check script')
     check.add_argument('path', default=None,
                        help='Path to the script invoked by Consul')
+    httpcheck = rsparsers.add_parser('httpcheck',
+                                 help='Define an HTTP-based check')
+    httpcheck.add_argument('interval', default=10, type=int,
+                       help='How often to run the check script')
+    httpcheck.add_argument('url', default=None,
+                       help='HTTP URL to be polled by Consul')
     rsparsers.add_parser('no-check', help='Do not enable service monitoring')
     ttl = rsparsers.add_parser('ttl', help='Define a duration based TTL check')
     ttl.add_argument('duration', type=int, default=10,
                      help='TTL duration for a service with missing check data')
 
+def add_deregister_args(parser):
+    """Add the deregister command and arguments.
+
+    :param argparse.Subparser parser: parser
+
+    """
+    # Service registration
+    registerp = parser.add_parser('deregister',
+                                  help='Deregister a service for this node')
+    registerp.add_argument('service_id', help='The service registration id')
 
 def parse_cli_args():
     """Create the argument parser and add the arguments"""
@@ -118,6 +134,7 @@ def parse_cli_args():
 
     sparser = parser.add_subparsers(title='Commands', dest='command')
     add_register_args(sparser)
+    add_deregister_args(sparser)
     add_kv_args(sparser)
     return parser.parse_args()
 
@@ -247,12 +264,25 @@ def register(consul, args):
 
     """
     check = args.path if args.ctype == 'check' else None
-    interval = '%ss' % args.interval if args.ctype == 'check' else None
+    httpcheck = args.url if args.ctype == 'httpcheck' else None
+    interval = '%ss' % args.interval if 'check' in args.ctype else None
     ttl = '%ss' % args.duration if args.ctype == 'ttl' else None
     tags = args.tags.split(',') if args.tags else None
     try:
         consul.agent.service.register(args.name, args.service_id, args.address,
-                                      args.port, tags, check, interval, ttl)
+                                      args.port, tags, check, interval, ttl, httpcheck)
+    except exceptions.ConnectionError:
+        connection_error()
+
+def deregister(consul, args):
+    """Handle service deregistration.
+
+    :param consulate.api_old.Consul consul: The Consul instance
+    :param argparser.namespace args: The cli args
+
+    """
+    try:
+        consul.agent.service.deregister(args.service_id)
     except exceptions.ConnectionError:
         connection_error()
 
@@ -281,5 +311,7 @@ def main():
                               args.token, args.api_scheme, adapter)
     if args.command == 'register':
         register(consul, args)
+    elif args.command == 'deregister':
+        deregister(consul, args)
     elif args.command == 'kv':
         KV_ACTIONS[args.action](consul, args)