pvc/add_zk_node_to_vm.py

141 lines
3.8 KiB
Python
Executable File

#!/usr/bin/env python3
# add_zk_node_to_vm.py - Debugging tool to add a new ZK node to all existing
# cluster entries.
# Part of the Parallel Virtual Cluster (PVC) system
#
# Copyright (C) 2018 Joshua M. Boniface <joshua@boniface.me>
#
# This program 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.
#
# This program 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 this program. If not, see <https://www.gnu.org/licenses/>.
#
###############################################################################
import socket
import click
import kazoo.client
#
# Connect and disconnect from Zookeeper
#
def startZKConnection(zk_host):
zk_conn = kazoo.client.KazooClient(hosts=zk_host)
zk_conn.start()
return zk_conn
def stopZKConnection(zk_conn):
zk_conn.stop()
zk_conn.close()
return 0
########################
########################
## ##
## CLICK COMPONENTS ##
## ##
########################
########################
zk_host = ''
myhostname = socket.gethostname()
CONTEXT_SETTINGS = dict(help_option_names=['-h', '--help'], max_content_width=120)
###############################################################################
# node
###############################################################################
@click.command(name='node', short_help='Add field to all node configs.')
@click.argument(
'field', default=None
)
@click.argument(
'data', default='', required=False
)
def node(field, data):
"""
Add FIELD to all node configs with optional DATA.
For example, 'foo [bar]` creates Zookeeper node '/domains/*/foo' with contents 'bar'.
"""
# Open a Zookeeper connection
zk_conn = startZKConnection(zk_host)
full_node_list = zk_conn.get_children('/nodes')
transaction = zk_conn.transaction()
for node in full_node_list:
transaction.create('/nodes/{}/{}'.format(node, field), data.encode('ascii'))
results = transaction.commit()
# Close the Zookeeper connection
stopZKConnection(zk_conn)
###############################################################################
# vm
###############################################################################
@click.command(name='vm', short_help='Add field to all node configs.')
@click.argument(
'field', default=None
)
@click.argument(
'data', default='', required=False
)
def vm(field, data):
"""
Add FIELD to all VM configs with optional DATA.
For example, 'foo [bar]` creates Zookeeper node '/nodes/*/foo' with contents 'bar'.
"""
# Open a Zookeeper connection
zk_conn = startZKConnection(zk_host)
full_vm_list = zk_conn.get_children('/domains')
transaction = zk_conn.transaction()
for vm in full_vm_list:
transaction.create('/domains/{}/{}'.format(vm, field), data.encode('ascii'))
results = transaction.commit()
# Close the Zookeeper connection
stopZKConnection(zk_conn)
@click.group(context_settings=CONTEXT_SETTINGS)
@click.option(
'-z', '--zookeeper', '_zk_host', envvar='PVC_ZOOKEEPER', default='{}:2181'.format(myhostname), show_default=True,
help='Zookeeper connection string.'
)
def cli(_zk_host):
"""
Parallel Virtual Cluster CLI management tool
"""
global zk_host
zk_host = _zk_host
cli.add_command(node)
cli.add_command(vm)
#
# Main entry point
#
def main():
return cli(obj={})
if __name__ == '__main__':
main()