initial commit.

library/riak

@@ -0,0 +1,285 @@
+#!/usr/bin/python
+# -*- coding: utf-8 -*-
+
+# (c) 2013, James Martin <jmartin@basho.com>, Drew Kerrigan <dkerrigan@basho.com>
+#
+# This file is part of Ansible
+#
+# Ansible is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# Ansible is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with Ansible.  If not, see <http://www.gnu.org/licenses/>.
+#
+DOCUMENTATION = '''
+---
+module: Riak
+short_description: This module handles some common Riak operations
+description:
+     - This module can be used to join nodes to a cluster, check
+       the status of the cluster.
+version_added: "1.2"
+options:
+  command:
+    description:
+      - The command you would like to perform against the cluster.
+    required: false
+    default: null
+    aliases: []
+    choices: ['ping', 'kv_test', 'join', 'plan', 'commit']
+  config_dir:
+    description:
+      - The path to the riak configuration directory
+    required: false
+    default: /etc/riak
+    aliases: []
+  http_conn:
+    description:
+      - The ip address and port that is listening for Riak HTTP queries
+    required: false
+    default: 127.0.0.1:8098
+    aliases: []
+  ring_ready:
+    description:
+      - Returns true if all nodes agree on the ring.
+    required: false
+    default: false
+    aliases: []
+    type: 'bool'
+  target_node:
+    description:
+      - The target node for certain operations (join, ping)
+    required: false
+    default: riak@127.0.0.1
+    aliases: []
+  wait_for_handoffs:
+    description:
+      - Waits for handoffs to complete before continuing.  This can take awhile and should generally be used with async mode.
+    required: false
+    default: null
+    aliases: []
+    type: 'bool'
+  wait_for_service:
+    description:
+      - Waits for a riak service to come online before continuing.
+    required: false
+    default: kv
+    aliases: []
+    choices: ['kv']
+examples:
+   - code: "riak: command=join target_node=riak@10.1.1.1"
+     description: "Join's a Riak node to another node"
+   - code: "riak: wait_for_handoffs=true"
+     description: "Wait for handoffs to finish.  Use with async and poll."
+   - code: "riak: wait_for_service=kv"
+     description: "Wait for riak_kv service to startup"
+'''
+
+
+
+import re
+import os.path
+import urllib2
+import json
+import time
+
+
+def _run(cmd):
+    # returns (rc, stdout, stderr) from shell command
+    process = subprocess.Popen(cmd, stdout=subprocess.PIPE,
+                               stderr=subprocess.PIPE, shell=True)
+    stdout, stderr = process.communicate()
+    return (process.returncode, stdout, stderr)
+
+
+def is_number(s):
+    try:
+        float(s)
+        return True
+    except ValueError:
+        return False
+
+
+def status_to_json():
+    # remove all unnecessary symbols and whitespace
+    rc, out, err = _run("riak-admin status 2> /dev/null")
+    if rc == 0:
+        stats = out
+    else:
+        module.fail_json(msg="Could not properly gather stats")
+
+    for line in raw_stats.splitlines():
+        stats += line.strip() + '\n'
+
+    stats = stats.replace('<<', '').replace('>>', '')
+    stats = stats.replace('\\n', '').replace(',\n', ',')
+    stats = stats.replace(": '", ': ').replace("'\n", "\n")
+    stats = stats.replace(': "[', ': [').replace(']"\n', "]\n")
+
+    stats = stats.replace('"', "'")
+
+    matchObj = re.compile(r"^(.*) : (.*)", re.M | re.I)
+
+    json_stats = '{'
+
+    for match in matchObj.finditer(stats):
+        key, value = match.groups()
+
+        if (value[0] == "'"):
+            value = value[1:-1]
+
+        if not is_number(value):
+            value = '"' + value + '"'
+        json_stats += '"' + key + '":' + value + ','
+
+    json_stats = json_stats[0:-1]
+    json_stats += '}'
+
+    return json_stats
+
+
+def main():
+
+    ansible_facts = {}
+    arg_spec = dict(
+        command=dict(required=False, default=None, choices=[
+                    'ping', 'kv_test', 'join', 'plan', 'commit']),
+        config_dir=dict(default='/etc/riak'),
+        http_conn=dict(required=False, default='127.0.0.1:8098'),
+        target_node=dict(default='riak@127.0.0.1', required=False),
+        wait_for_handoffs=dict(default=False, type='bool'),
+        wait_for_service=dict(
+            required=False, default=None, choices=['kv'])
+    )
+    global module
+    module = AnsibleModule(argument_spec=arg_spec)
+
+
+    command = module.params.get('command')
+    config_dir = module.params.get('config_dir')
+    http_conn = module.params.get('http_conn')
+    target_node = module.params.get('target_node')
+    wait_for_handoffs = module.params.get('wait_for_handoffs')
+    wait_for_service = module.params.get('wait_for_service')
+
+    rc = 0
+    err = ''
+    out = ''
+
+    #make sure riak commands are on the path
+    for item in ['riak', 'riak-admin']:
+        rc, out, err = _run('which %s' % item)
+        if rc == 1:
+            module.fail_json(msg='Could not find path to %s executable' % item)
+
+    rc, out, err = _run(
+        "riak version 2> /dev/null |grep ^riak|cut -f2 -d' '|tr -d '('")
+    if rc == 0:
+        version = out.strip()
+    else:
+        module.fail_json(msg='Could not determine Riak version')
+
+    try:
+        stats_raw = urllib2.urlopen(
+            'http://%s/stats' % (http_conn), None, 5).read()
+    except urllib2.HTTPError, e:
+        stats_raw = status_to_json()
+    except urllib2.URLError, e:
+        stats_raw = status_to_json()
+
+    stats = json.loads(stats_raw)
+
+    node_name = stats['nodename']
+    nodes = stats['ring_members']
+    ring_size = stats['ring_creation_size']
+
+
+
+    result = {'node_name': node_name,
+              'nodes': nodes,
+              'ring_size': ring_size,
+              'version': version}
+
+    if command == 'ping':
+        rc, out, err = _run('riak ping %s' % target_node)
+        if rc == 0:
+            result['ping'] = out
+        else:
+            module.fail_json(msg=out)
+
+    elif command == 'kv_test':
+        rc, out, err = _run('riak-admin test')
+        if rc == 0:
+            result['kv_test'] = out
+        else:
+            module.fail_json(msg=out)
+
+    elif command == 'join':
+        if nodes.count(node_name) == 1 and len(nodes) > 1:
+            result['join'] = 'Node is already in cluster or staged to be in cluster.'
+        else:
+            rc, out, err = _run('riak-admin cluster join %s' % target_node)
+            if rc == 0:
+                result['join'] = out
+                result['changed'] = True
+            else:
+                module.fail_json(msg=out)
+
+    elif command == 'plan':
+        rc, out, err = _run('riak-admin cluster plan %s' % target_node)
+        if rc == 0:
+            result['plan'] = out
+            if out.find('Staged Changes') != -1:
+                result['changed'] = True
+        else:
+            module.fail_json(msg=out)
+
+    elif command == 'commit':
+
+        rc, out, err = _run('riak-admin cluster commit %s' % target_node)
+        if rc == 0:
+            result['commit'] = out
+            changed = True
+        else:
+            module.fail_json(msg=out)
+
+    rc = 0
+    err = ''
+    out = ''
+    wait = 0
+
+# this could take a while, recommend to run in async mode
+    if wait_for_handoffs:
+        while wait == 0:
+            time.sleep(10)
+            rc, out, err = _run('riak-admin transfers 2> /dev/null')
+            if out.find('No transfers active') != -1:
+                result['handoffs'] = 'No transfers active.'
+                break
+
+# this could take a while, recommend to run in async mode
+    if wait_for_service:
+        rc, out, err = _run('riak-admin wait_for_service riak_%s %s' % (
+            wait_for_service, node_name))
+        result['service'] = out
+
+
+    rc, out, err = _run('riak-admin ringready 2> /dev/null')
+    if rc == 0 and out.find('TRUE All nodes agree on the ring') != -1:
+        result['ring_ready'] = True
+    else:
+        result['ring_ready'] = False
+
+    module.exit_json(**result)
+
+# this is magic, see lib/ansible/module_common.py
+#<<INCLUDE_ANSIBLE_MODULE_COMMON>>
+
+main()