github.com/1aal/kubeblocks@v0.0.0-20231107070852-e1c03e598921/deploy/pulsar/scripts/zookeeper.py (about)

     1  #!/usr/bin/python3
     2  from kazoo.client import KazooClient
     3  from kazoo.exceptions import NodeExistsError, NoNodeError
     4  import os
     5  import sys
     6  import logging
     7  
     8  logging.basicConfig(level=logging.INFO,
     9                      format='%(asctime)s %(filename)s[line:%(lineno)d] %(levelname)s %(message)s',
    10                      datefmt='%a, %d %b %Y %H:%M:%S', )
    11  logger = logging.getLogger()
    12  
    13  
    14  class ZKClient(object):
    15      def __init__(self, zk_url):
    16          self.client = KazooClient(hosts=zk_url)
    17          self.client.start()
    18  
    19      def create_node(self, path, data=None):
    20          try:
    21              self.client.create(path, data.encode() if data else None)
    22              logger.info(f"Node {path} created successfully.")
    23          except NodeExistsError:
    24              self.stop_with_error(f"Node {path} already exists.")
    25          except Exception as e:
    26              self.stop_with_error(f"Failed to create node {path}: {e}")
    27  
    28      def get_node(self, path):
    29          try:
    30              data, _ = self.client.get(path)
    31              logger.info(f"Data at {path}: {data.decode()}")
    32          except NoNodeError:
    33              self.stop_with_error(f"Node {path} does not exist.")
    34          except Exception as e:
    35              self.stop_with_error(f"Failed to read node {path}: {e}")
    36  
    37      def update_node(self, path, data):
    38          try:
    39              self.client.set(path, data.encode())
    40              logger.info(f"Node {path} updated successfully.")
    41          except NoNodeError:
    42              self.stop_with_error(f"Node {path} does not exist.")
    43          except Exception as e:
    44              self.stop_with_error(f"Failed to update node {path}: {e}")
    45  
    46      def delete_node(self, path):
    47          try:
    48              self.client.delete(path)
    49              logger.info(f"Node {path} deleted successfully.")
    50          except NoNodeError:
    51              self.stop_with_error(f"Node {path} does not exist.")
    52          except Exception as e:
    53              self.stop_with_error(f"Failed to delete node {path}: {e}")
    54  
    55      def stop(self):
    56          self.client.stop()
    57  
    58      def stop_with_error(self, message):
    59          logger.error(message)
    60          self.client.stop()
    61          exit(1)
    62  
    63  
    64  if __name__ == '__main__':
    65      args = sys.argv[1:]
    66      if len(args) != 2:
    67          raise Exception("requires 2 arguments.")
    68      op = args[0]
    69      path = args[1]
    70      zk_url = os.environ.get('zkURL')
    71      zk_client = ZKClient(zk_url)
    72      if op == "delete":
    73          zk_client.delete_node(path)
    74      elif op == "create":
    75          zk_client.create_node(path)
    76      elif op == "get":
    77          zk_client.get_node(path)
    78      else:
    79          zk_client.stop_with_error(f"Unknown operation: {op}")
    80      zk_client.stop()