zkcli.py

#! /usr/bin/env python2

import argparse

from kazoo.client import KazooClient
from kazoo.exceptions import NoNodeError

ZK_DEFAULT_HOST = 'localhost:2181'


def get_command_handler(zk, args):
    try:
        data, _ = zk.get(args.path)
    except NoNodeError:
        data = args.or_default
    return data.decode('utf-8')


def set_command_handler(zk, args):
    if zk.exists(args.path):
        zk.set(args.path, args.new_value)
    else:
        zk.create(args.path, args.new_value)


def delete_command_handler(zk, args):
    try:
        zk.delete(args.path, version=-1, recursive=args.recursive)
    except NoNodeError:
        print '%s not exists' % args.path


arg_parser = argparse.ArgumentParser()
arg_parser.add_argument('-s', '--server', help='ZooKeeper server address', default=ZK_DEFAULT_HOST)

arg_subparsers = arg_parser.add_subparsers(help='sub-command help')
arg_subparsers.required = True

get_command_parser = arg_subparsers.add_parser('get', help='fetches value along of the specified path')
get_command_parser.add_argument('path', type=str)
get_command_parser.add_argument('or_default')
get_command_parser.set_defaults(func=get_command_handler)

set_command_parser = arg_subparsers.add_parser('set', help='creates or updates value along of the specified path')
set_command_parser.add_argument('path', type=str)
set_command_parser.add_argument('new_value')
set_command_parser.set_defaults(func=set_command_handler)

delete_command_parser = arg_subparsers.add_parser('delete', help='delete value along of the specified path')
delete_command_parser.add_argument('--recursive', action='store_true')
delete_command_parser.add_argument('path', type=str)
delete_command_parser.set_defaults(func=delete_command_handler)

args = arg_parser.parse_args()

zk = KazooClient(hosts=args.server)
zk.start()

print args.func(zk, args)