1 files changed, 227 insertions, 0 deletions
diff --git a/clustering/znode.py b/clustering/znode.py
new file mode 100644
index 00000000..3b39865b
--- /dev/null
+++ b/clustering/znode.py
@@ -0,0 +1,227 @@
+#!/usr/bin/python
+# Copyright 2015 WP Engine, Inc. All rights reserved.
+#
+# 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: znode
+version_added: "2.0"
+short_description: Create, delete, retrieve, and update znodes using ZooKeeper.
+options:
+    hosts:
+        description:
+            - A list of ZooKeeper servers (format '[server]:[port]').
+        required: true
+    path:
+        description:
+            - The path of the znode.
+        required: true
+    value:
+        description:
+            - The value assigned to the znode.
+        default: None
+        required: false
+    op:
+        description:
+            - An operation to perform. Mutually exclusive with state.
+        default: None
+        required: false
+    state:
+        description:
+            - The state to enforce. Mutually exclusive with op.
+        default: None
+        required: false
+    timeout:
+        description:
+            - The amount of time to wait for a node to appear.
+        default: 300
+        required: false
+requirements:
+    - kazoo >= 2.1
+author: "Trey Perry (@treyperry)"
+"""
+
+EXAMPLES = """
+# Creating or updating a znode with a given value
+- action: znode hosts=localhost:2181 name=/mypath value=myvalue state=present
+
+# Getting the value and stat structure for a znode
+- action: znode hosts=localhost:2181 name=/mypath op=get
+
+# Listing a particular znode's children
+- action: znode hosts=localhost:2181 name=/zookeeper op=list
+
+# Waiting 20 seconds for a znode to appear at path /mypath
+- action: znode hosts=localhost:2181 name=/mypath op=wait timeout=20
+
+# Deleting a znode at path /mypath
+- action: znode hosts=localhost:2181 name=/mypath state=absent
+"""
+
+try:
+    from kazoo.client import KazooClient
+    from kazoo.exceptions import NoNodeError, ZookeeperError
+    from kazoo.handlers.threading import KazooTimeoutError
+    KAZOO_INSTALLED = True
+except ImportError:
+    KAZOO_INSTALLED = False
+
+
+def main():
+    module = AnsibleModule(
+        argument_spec=dict(
+            hosts=dict(required=True, type='str'),
+            name=dict(required=True, type='str'),
+            value=dict(required=False, default=None, type='str'),
+            op=dict(required=False, default=None, choices=['get', 'wait', 'list']),
+            state=dict(choices=['present', 'absent']),
+            timeout=dict(required=False, default=300, type='int')
+        ),
+        supports_check_mode=False
+    )
+
+    if not KAZOO_INSTALLED:
+        module.fail_json(msg='kazoo >= 2.1 is required to use this module. Use pip to install it.')
+
+    check = check_params(module.params)
+    if not check['success']:
+        module.fail_json(msg=check['msg'])
+
+    zoo = KazooCommandProxy(module)
+    try:
+        zoo.start()
+    except KazooTimeoutError:
+        module.fail_json(msg='The connection to the ZooKeeper ensemble timed out.')
+
+    command_dict = {
+        'op': {
+            'get': zoo.get,
+            'list': zoo.list,
+            'wait': zoo.wait
+        },
+        'state': {
+            'present': zoo.present,
+            'absent': zoo.absent
+        }
+    }
+
+    command_type = 'op' if 'op' in module.params and module.params['op'] is not None else 'state'
+    method = module.params[command_type]
+    result, result_dict = command_dict[command_type][method]
+    zoo.shutdown()
+
+    if result:
+        module.exit_json(**result_dict)
+    else:
+        module.fail_json(**result_dict)
+
+
+def check_params(params):
+    if not params['state'] and not params['op']:
+        return {'success': False, 'msg': 'Please define an operation (op) or a state.'}
+
+    if params['state'] and params['op']:
+        return {'success': False, 'msg': 'Please choose an operation (op) or a state, but not both.'}
+
+    return {'success': True}
+
+
+class KazooCommandProxy():
+    def __init__(self, module):
+        self.module = module
+        self.zk = KazooClient(module.params['hosts'])
+
+    def absent(self):
+        return self._absent(self.module.params['name'])
+
+    def exists(self, znode):
+        return self.zk.exists(znode)
+
+    def list(self):
+        children = self.zk.get_children(self.module.params['name'])
+        return True, {'count': len(children), 'items': children, 'msg': 'Retrieved znodes in path.',
+                      'znode': self.module.params['name']}
+
+    def present(self):
+        return self._present(self.module.params['name'], self.module.params['value'])
+
+    def get(self):
+        return self._get(self.module.params['name'])
+
+    def shutdown(self):
+        self.zk.stop()
+        self.zk.close()
+
+    def start(self):
+        self.zk.start()
+
+    def wait(self):
+        return self._wait(self.module.params['name'], self.module.params['timeout'])
+
+    def _absent(self, znode):
+        if self.exists(znode):
+            self.zk.delete(znode)
+            return True, {'changed': True, 'msg': 'The znode was deleted.'}
+        else:
+            return True, {'changed': False, 'msg': 'The znode does not exist.'}
+
+    def _get(self, path):
+        if self.exists(path):
+            value, zstat = self.zk.get(path)
+            stat_dict = {}
+            for i in dir(zstat):
+                if not i.startswith('_'):
+                    attr = getattr(zstat, i)
+                    if type(attr) in (int, str):
+                        stat_dict[i] = attr
+            result = True, {'msg': 'The node was retrieved.', 'znode': path, 'value': value,
+                            'stat': stat_dict}
+        else:
+            result = False, {'msg': 'The requested node does not exist.'}
+
+        return result
+
+    def _present(self, path, value):
+        if self.exists(path):
+            (current_value, zstat) = self.zk.get(path)
+            if value != current_value:
+                self.zk.set(path, value)
+                return True, {'changed': True, 'msg': 'Updated the znode value.', 'znode': path,
+                              'value': value}
+            else:
+                return True, {'changed': False, 'msg': 'No changes were necessary.', 'znode': path, 'value': value}
+        else:
+            self.zk.create(path, value, makepath=True)
+            return True, {'changed': True, 'msg': 'Created a new znode.', 'znode': path, 'value': value}
+
+    def _wait(self, path, timeout, interval=5):
+        lim = time.time() + timeout
+
+        while time.time() < lim:
+            if self.exists(path):
+                return True, {'msg': 'The node appeared before the configured timeout.',
+                              'znode': path, 'timeout': timeout}
+            else:
+                time.sleep(interval)
+
+        return False, {'msg': 'The node did not appear before the operation timed out.', 'timeout': timeout,
+                       'znode': path}
+
+from ansible.module_utils.basic import *
+
+main()
+