text: #!/usr/bin/env python
  # --coding:utf-8--
  
  # Copyright (c) 2020 vesoft inc. All rights reserved.
  #
  # This source code is licensed under Apache 2.0 License.
  
  
  import logging
  import socket
  
  from _thread import RLock
  from nebula2.Exception import (
      InValidHostname,
      PartNotFoundException,
      SpaceNotFoundException,
      TagNotFoundException,
      EdgeNotFoundException
  )
  from nebula2.common.ttypes import HostAddr, ErrorCode
  from nebula2.meta.ttypes import (
      HostStatus,
      ListTagsReq,
      ListEdgesReq,
      ListSpacesReq,
      GetPartsAllocReq,
      ListHostsReq,
      HostRole
  )
  from nebula2.meta import (
      ttypes,
      MetaService
  )
  
  from nebula2.fbthrift.transport import TSocket, TTransport
  from nebula2.fbthrift.protocol import TBinaryProtocol
  
  
  class MetaClient(object):
      def __init__(self, addresses, timeout):
          if len(addresses) == 0:
              raise RuntimeError('Input empty addresses')
          self._timeout = timeout
          self._connection = None
          self._retry_count = 3
          self._addresses = addresses
          for address in addresses:
              try:
                  socket.gethostbyname(address[0])
              except Exception:
                  raise InValidHostname(str(address[0]))
          self._leader = self._addresses[0]
          self._lock = RLock()
  
      def open(self):
          """open the connection to connect meta service
  
          :eturn: void
          """
          try:
              self.close()
              s = TSocket.TSocket(self._leader[0], self._leader[1])
              if self._timeout > 0:
                  s.setTimeout(self._timeout)
              transport = TTransport.TBufferedTransport(s)
              protocol = TBinaryProtocol.TBinaryProtocol(transport)
              transport.open()
              self._connection = MetaService.Client(protocol)
          except Exception:
              raise
  
      def list_tags(self, space_id):
          """get all version tags
  
          :param space_id: the specified space id
          :eturn: list<TagItem>
          """
          with self._lock:
              if self._connection is None:
                  raise RuntimeError('The connection is no open')
              req = ListTagsReq()
              req.space_id = space_id
              count = 0
              while count < self._retry_count:
                  resp = self._connection.listTags(req)
                  if resp.code != ErrorCode.SUCCEEDED:
                      if resp.code == ErrorCode.E_LEADER_CHANGED:
                          self.update_leader(resp.leader)
                          count = count + 1
                          continue
                      raise RuntimeError("List tags from space id:{} failed, error code: {}"
                                         .format(space_id, ErrorCode._VALUES_TO_NAMES.get(resp.code)))
                  return resp.tags
              raise RuntimeError("List tags from space id:{} failed, error code: {}"
                                 .format(space_id, ErrorCode._VALUES_TO_NAMES.get(resp.code)))
  
      def list_edges(self, space_id):
          """get all version edge
  
         :param space_id: the specified space id
         :eturn: list<EdgeItem>
         """
          with self._lock:
              if self._connection is None:
                  raise RuntimeError('The connection is no open')
              req = ListEdgesReq()
              req.space_id = space_id
              count = 0
              while count < self._retry_count:
                  resp = self._connection.listEdges(req)
                  if resp.code != ErrorCode.SUCCEEDED:
                      if resp.code == ErrorCode.E_LEADER_CHANGED:
                          self.update_leader(resp.leader)
                          count = count + 1
                          continue
                      raise RuntimeError("List edges from space id:{} failed, error code: {}"
                                         .format(space_id, ErrorCode._VALUES_TO_NAMES.get(resp.code)))
                  return resp.edges
              raise RuntimeError("List edges from space id:{} failed, error code: {}"
                                 .format(space_id, ErrorCode._VALUES_TO_NAMES.get(resp.code)))
  
      def list_spaces(self):
          """get all spaces info
  
          :eturn: list<IdName>
          """
          with self._lock:
              if self._connection is None:
                  raise RuntimeError('The connection is no open')
              req = ListSpacesReq()
              count = 0
              while count < self._retry_count:
                  resp = self._connection.listSpaces(req)
                  if resp.code != ErrorCode.SUCCEEDED:
                      if resp.code == ErrorCode.E_LEADER_CHANGED:
                          self.update_leader(resp.leader)
                          count = count + 1
                          continue
                      raise RuntimeError("List spaces failed, error code: {}"
                                         .format(ErrorCode._VALUES_TO_NAMES.get(resp.code)))
                  return resp.spaces
              raise RuntimeError("List spaces failed, error code: {}"
                                 .format(ErrorCode._VALUES_TO_NAMES.get(resp.code)))
  
      def list_hosts(self):
          """get all online hosts info
  
          :eturn: list<HostItem>
          """
          with self._lock:
              if self._connection is None:
                  raise RuntimeError('The connection is no open')
              req = ListHostsReq()
              req.role = HostRole.STORAGE
              count = 0
              while count < self._retry_count:
                  resp = self._connection.listHosts(req)
                  if resp.code != ErrorCode.SUCCEEDED:
                      if resp.code == ErrorCode.E_LEADER_CHANGED:
                          self.update_leader(resp.leader)
                          count = count + 1
                          continue
                      raise RuntimeError("List spaces failed, error code: {}"
                                         .format(ErrorCode._VALUES_TO_NAMES.get(resp.code)))
                  valid_hosts = []
                  for host in resp.hosts:
                      if host.status == HostStatus.ONLINE:
                          valid_hosts.append(host)
                  return valid_hosts
              raise RuntimeError("List spaces failed, error code: {}"
                                 .format(ErrorCode._VALUES_TO_NAMES.get(resp.code)))
  
      def get_parts_alloc(self, space_id):
          """get all parts info of the specified space id
  
          :param space_id:
          :eturn: map<PartitionID, list<HostAddr>>
          """
          with self._lock:
              if self._connection is None:
                  raise RuntimeError('The connection is no open')
              req = GetPartsAllocReq()
              req.space_id = space_id
              count = 0
              while count < self._retry_count:
                  resp = self._connection.getPartsAlloc(req)
                  if resp.code != ErrorCode.SUCCEEDED:
                      if resp.code == ErrorCode.E_LEADER_CHANGED:
                          self.update_leader(resp.leader)
                          count = count + 1
                          continue
                      raise RuntimeError("List parts from space id:{} failed, error code: {}"
                                         .format(space_id, ErrorCode._VALUES_TO_NAMES.get(resp.code)))
                  return resp.parts
              raise RuntimeError("List parts from space id:{} failed, error code: {}"
                                 .format(space_id, ErrorCode._VALUES_TO_NAMES.get(resp.code)))
  
      def close(self):
          """close the connection
  
          :eturn: void
          """
          try:
              if self._connection is not None:
                  self._connection._iprot.trans.close()
          except Exception:
              raise
  
      def update_leader(self, leader):
          """update the leader meta info when hanppen leader change
  
          :param leader: the address of meta leader
          :eturn:
          """
          try:
              self._leader = (leader.host, leader.port)
              self.open()
          except Exception as e:
              logging.error(e)
  
      def __del__(self):
          self.close()
  
  
  class MetaCache(object):
      class SpaceCache:
          def __init__(self):
              self.space_id = 0
              self.space_name = ''
              self.tag_items = {}
              self.edge_items = {}
              self.parts_alloc = {}
  
          def __repr__(self):
              return 'space_id: {}, space_name: {}, tag_items: {}, edge_items: {}, parts_alloc: {}'\
                  .format(self.space_id,
                          self.space_name,
                          self.tag_items,
                          self.edge_items,
                          self.parts_alloc)
  
      def __init__(self, meta_addrs, timeout=2000, load_period=10, decode_type='utf-8'):
          self._decode_type = decode_type
          self._load_period = load_period
          self._lock = RLock()
          self._space_caches = {}
          self._space_id_names = {}
          self._storage_addrs = []
          self._storage_leader = {}
          self._close = False
          self._meta_client = MetaClient(meta_addrs, timeout)
          self._meta_client.open()
  
          # load meta data
          self._load_all()
  
      def close(self):
          """close the metaClient
  
          :eturn: void
          """
          if self._close:
              return
          self._close = True
          if self._meta_client is not None:
              self._meta_client.close()
  
      def __del__(self):
          self.close()
  
      def _load_all(self):
          """load all space info and schema info from meta services
  
          :eturn: void
          """
          try:
              spaces = self._meta_client.list_spaces()
              space_caches = {}
              space_id_names = {}
              for space in spaces:
                  space_id = space.id.get_space_id()
                  space_cache = MetaCache.SpaceCache()
                  space_cache.space_id = space_id
                  space_cache.space_name = space.name.decode('utf-8')
                  space_id_names[space_id] = space_cache.space_name
                  tags = self._meta_client.list_tags(space_id)
                  edges = self._meta_client.list_edges(space_id)
                  parts_alloc = self._meta_client.get_parts_alloc(space_id)
                  for tag in tags:
                      tag_name = tag.tag_name.decode(self._decode_type)
                      if tag_name not in space_cache.tag_items.keys():
                          space_cache.tag_items[tag_name] = tag
                      else:
                          if space_cache.tag_items[tag_name].version < tag.version:
                              space_cache.tag_items[tag_name] = tag
                  for edge in edges:
                      edge_name = edge.edge_name.decode(self._decode_type)
                      if edge_name not in space_cache.edge_items.keys():
                          space_cache.edge_items[edge_name] = edge
                      else:
                          if space_cache.edge_items[edge_name].version < edge.version:
                              space_cache.edge_items[edge_name] = edge
                      space_cache.edge_items[edge.edge_name.decode(self._decode_type)] = edge
                  space_cache.parts_alloc = parts_alloc
                  space_caches[space.name.decode(self._decode_type)] = space_cache
  
              hosts = self._meta_client.list_hosts()
              storage_addrs = []
              for host_item in hosts:
                  storage_addrs.append(host_item.hostAddr)
  
              with self._lock:
                  self._storage_addrs = storage_addrs
                  self._space_caches = space_caches
                  self._space_id_names = space_id_names
                  for space_name in self._space_caches.keys():
                      if space_name in self._storage_leader.keys():
                          continue
                      parts_alloc = self._space_caches[space_name].parts_alloc
                      self._storage_leader[space_name] = {}
                      for part_id in parts_alloc:
                          self._storage_leader[space_name][part_id] = parts_alloc[part_id][0]
          except Exception as x:
              logging.error('Update meta data failed: {}'.format(x))
              import traceback
              logging.error(traceback.format_exc())
  
      def get_all_storage_addrs(self):
          """get all storage address
  
          :return: list[HostAddr]
          """
          return self._storage_addrs
  
      def get_tag_id(self, space_name, tag_name):
          """get tag id
  
          :param space_name:
          :param tag_name:
          :return: tag_id
          """
          with self._lock:
              tag_item = self._get_tag_item(space_name, tag_name)
              return tag_item.tag_id
  
      def get_edge_type(self, space_name, edge_name):
          """get edge type
  
          :param space_name:
          :param edge_name:
          :return: edge_type
          """
          with self._lock:
              edge_item = self._get_edge_item(space_name, edge_name)
              return edge_item.edge_type
  
      def get_space_id(self, space_name):
          """get space id
  
          :param space_name:
          :return: space_id
          """
          with self._lock:
              if space_name not in self._space_caches.keys():
                  self._load_all()
                  if space_name not in self._space_caches.keys():
                      raise SpaceNotFoundException(space_name)
              return self._space_caches[space_name].space_id
  
      def get_tag_schema(self, space_name, tag_name):
          """get tag schema
  
          :param space_name:
          :param tag_name:
          :return: schema
          """
          tag_item = self._get_tag_item(space_name, tag_name)
          return tag_item.schema
  
      def get_edge_schema(self, space_name, edge_name):
          """get edge schema
  
          :param space_name:
          :param edge_name:
          :return: schema
          """
          edge_item = self._get_edge_item(space_name, edge_name)
          return edge_item.schema
  
      def get_part_leader(self, space_name, part_id):
          """
  
          :param space_name:
          :param part_id:
          :return: storage ip port: HostAddr
          """
          part_leaders = self.get_part_leaders(space_name)
          if part_id not in part_leaders.keys():
              raise PartNotFoundException(part_id)
          return part_leaders[part_id]
  
      def get_part_leaders(self, space_name):
          """get all part leader info of the space
  
          :param space_name: space name
          :eturn: map<PartitionID, HostAddr>
          """
          with self._lock:
              if space_name not in self._storage_leader.keys():
                  self._load_all()
                  if space_name not in self._storage_leader.keys():
                      raise SpaceNotFoundException(space_name)
              return self._storage_leader[space_name]
  
      def get_part_alloc(self, space_name):
          """get all part info of the space
  
          :param space_name: space name
          :eturn: map<PartitionID, list<HostAddr>>
          """
          with self._lock:
              if space_name not in self._space_caches.keys():
                  self._load_all()
                  if space_name not in self._space_caches.keys():
                      raise SpaceNotFoundException(space_name)
              return self._space_caches[space_name].parts_alloc
  
      def _get_tag_item(self, space_name, tag_name):
          with self._lock:
              if space_name not in self._space_caches.keys():
                  self._load_all()
                  if space_name not in self._space_caches.keys():
                      raise SpaceNotFoundException(space_name)
              space_info = self._space_caches[space_name]
              if tag_name not in space_info.tag_items.keys():
                  self._load_all()
                  if tag_name not in space_info.tag_items.keys():
                      raise TagNotFoundException(tag_name)
              return space_info.tag_items[tag_name]
  
      def _get_edge_item(self, space_name, edge_name):
          with self._lock:
              if space_name not in self._space_caches.keys():
                  self._load_all()
                  if space_name not in self._space_caches.keys():
                      raise SpaceNotFoundException(space_name)
              space_info = self._space_caches[space_name]
              if edge_name not in space_info.edge_items.keys():
                  self._load_all()
                  if edge_name not in space_info.edge_items.keys():
                      raise EdgeNotFoundException(edge_name)
              return space_info.edge_items[edge_name]
  
      def update_storage_leader(self, space_id, part_id, address: HostAddr):
          """if the storage leader change, storage client need to call this function
  
          :param space_id:
          :param part_id:
          :param address: HostAddr, if the address is None, it means the leader can't connect,
          choose the peer as leader
          :return: coid
          """
          with self._lock:
              if space_id not in self._space_id_names.keys():
                  logging.error("Space name:{} is not found".format(space_id))
                  return
              space_name = self._space_id_names.get(space_id)
              if part_id not in self._storage_leader[space_name].keys():
                  logging.error("part_id:{} is not found".format(space_name))
                  return
              if address is not None:
                  self._storage_leader[space_name][part_id] = address
                  return
              part_addresses = self._space_caches[space_name].parts_alloc.get(part_id)
              for part_addr in part_addresses:
                  if part_addr == address:
                      continue
                  self._storage_leader[space_name][part_id] = part_addr
  
  
  