Add a zk-shell debug script
This is like the real zk-shell except: 1) It supports SSL 2) It automatically decompresses znodes if necessary 3) It handles sharded data 4) It doesn't have many commands Change-Id: I2124f216f08f68ca63b763cd6f6adce452c07f79
This commit is contained in:
James E. Blair
parent
1f3f724bbb
commit
a396505490
|
|
@ -0,0 +1,179 @@
|
|||
# Copyright 2022 Acme Gating, LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License"); you may
|
||||
# not use this file except in compliance with the License. You may obtain
|
||||
# a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
|
||||
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
|
||||
# License for the specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
# Inspect ZK contents like zk-shell; handles compressed and sharded
|
||||
# data.
|
||||
|
||||
import argparse
|
||||
import pathlib
|
||||
import cmd
|
||||
import sys
|
||||
import textwrap
|
||||
import zlib
|
||||
|
||||
import kazoo.client
|
||||
from kazoo.exceptions import NoNodeError
|
||||
|
||||
|
||||
def resolve_path(path, rest):
|
||||
newpath = path / rest
|
||||
newparts = []
|
||||
for part in newpath.parts:
|
||||
if part == '.':
|
||||
continue
|
||||
elif part == '..':
|
||||
newparts.pop()
|
||||
else:
|
||||
newparts.append(part)
|
||||
return pathlib.PurePosixPath(*newparts)
|
||||
|
||||
|
||||
class REPL(cmd.Cmd):
|
||||
def __init__(self, args):
|
||||
self.path = pathlib.PurePosixPath('/')
|
||||
super().__init__()
|
||||
kwargs = {}
|
||||
if args.cert:
|
||||
kwargs['use_ssl'] = True
|
||||
kwargs['keyfile'] = args.key
|
||||
kwargs['certfile'] = args.cert
|
||||
kwargs['ca'] = args.ca
|
||||
self.client = kazoo.client.KazooClient(args.host, **kwargs)
|
||||
self.client.start()
|
||||
|
||||
@property
|
||||
def prompt(self):
|
||||
return f'{self.path}> '
|
||||
|
||||
def do_EOF(self, path):
|
||||
sys.exit(0)
|
||||
|
||||
def do_ls(self, path):
|
||||
'List znodes: ls [PATH]'
|
||||
if path:
|
||||
mypath = self.path / path
|
||||
else:
|
||||
mypath = self.path
|
||||
try:
|
||||
for child in self.client.get_children(str(mypath)):
|
||||
print(child)
|
||||
except NoNodeError:
|
||||
print(f'No such node: {mypath}')
|
||||
|
||||
def do_cd(self, path):
|
||||
'Change the working path: cd PATH'
|
||||
if path:
|
||||
newpath = resolve_path(self.path, path)
|
||||
if self.client.exists(str(newpath)):
|
||||
self.path = newpath
|
||||
else:
|
||||
print(f'No such node: {newpath}')
|
||||
|
||||
def do_pwd(self):
|
||||
'Print the working path'
|
||||
print(self.path)
|
||||
|
||||
def help_get(self):
|
||||
print(textwrap.dedent(self.do_get.__doc__))
|
||||
|
||||
def do_get(self, args):
|
||||
"""\
|
||||
Get znode value: get PATH [-v]
|
||||
|
||||
-v: output metadata about the path
|
||||
"""
|
||||
args = args.split(' ')
|
||||
path = args[0]
|
||||
args = args[1:]
|
||||
path = resolve_path(self.path, path)
|
||||
try:
|
||||
compressed_data, zstat = self.client.get(str(path))
|
||||
except NoNodeError:
|
||||
print(f'No such node: {path}')
|
||||
return
|
||||
was_compressed = False
|
||||
try:
|
||||
data = zlib.decompress(compressed_data)
|
||||
was_compressed = True
|
||||
except zlib.error:
|
||||
data = compressed_data
|
||||
if '-v' in args:
|
||||
print(f'Compressed: {was_compressed}')
|
||||
print(f'Size: {len(data)}')
|
||||
print(f'Compressed size: {len(compressed_data)}')
|
||||
print(f'Zstat: {zstat}')
|
||||
print(data)
|
||||
|
||||
def help_unshard(self):
|
||||
print(textwrap.dedent(self.do_unshard.__doc__))
|
||||
|
||||
def do_unshard(self, args):
|
||||
"""\
|
||||
Get the unsharded value: get PATH [-v]
|
||||
|
||||
-v: output metadata about the path
|
||||
"""
|
||||
args = args.split(' ')
|
||||
path = args[0]
|
||||
args = args[1:]
|
||||
path = resolve_path(self.path, path)
|
||||
|
||||
try:
|
||||
shards = sorted(self.client.get_children(str(path)))
|
||||
except NoNodeError:
|
||||
print(f'No such node: {path}')
|
||||
return
|
||||
compressed_data = b''
|
||||
data = b''
|
||||
for shard in shards:
|
||||
d, _ = self.client.get(str(path / shard))
|
||||
compressed_data += d
|
||||
if compressed_data:
|
||||
data = zlib.decompress(compressed_data)
|
||||
|
||||
if '-v' in args:
|
||||
print(f'Size: {len(data)}')
|
||||
print(f'Compressed size: {len(compressed_data)}')
|
||||
print(data)
|
||||
|
||||
def do_rm(self, args):
|
||||
'Delete znode: rm PATH [-r]'
|
||||
args = args.split(' ')
|
||||
path = args[0]
|
||||
args = args[1:]
|
||||
path = resolve_path(self.path, path)
|
||||
if '-r' in args:
|
||||
recursive = True
|
||||
else:
|
||||
recursive = False
|
||||
try:
|
||||
self.client.delete(str(path), recursive=recursive)
|
||||
except NoNodeError:
|
||||
print(f'No such node: {path}')
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('host', help='ZK host string')
|
||||
parser.add_argument('--cert', help='Path to TLS certificate')
|
||||
parser.add_argument('--key', help='Path to TLS key')
|
||||
parser.add_argument('--ca', help='Path to TLS CA cert')
|
||||
args = parser.parse_args()
|
||||
|
||||
repl = REPL(args)
|
||||
repl.cmdloop()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Loading…
Reference in New Issue