github.com/projecteru2/core@v0.0.0-20240321043226-06bcc1c23f58/scripts/redis_as_broker.py (about)

     1  #!/usr/bin/env python
     2  # -*- coding: utf-8 -*-
     3  
     4  import argparse
     5  import sys
     6  
     7  import etcd3
     8  import redis
     9  
    10  
    11  def migrate(root, etcd_client, redis_client, dry_run):
    12      range_start = root
    13      range_end = etcd3.utils.increment_last_byte(
    14          etcd3.utils.to_bytes(range_start)
    15      )
    16  
    17      while True:
    18          range_request = etcd3.etcdrpc.RangeRequest()
    19          range_request.key = etcd3.utils.to_bytes(range_start)
    20          range_request.keys_only = False
    21          range_request.range_end = etcd3.utils.to_bytes(range_end)
    22          range_request.sort_order = etcd3.etcdrpc.RangeRequest.ASCEND
    23          range_request.sort_target = etcd3.etcdrpc.RangeRequest.KEY
    24          range_request.serializable = True
    25          range_request.limit = 1000
    26  
    27          resp = etcd_client.kvstub.Range(
    28              range_request,
    29              etcd_client.timeout,
    30              credentials=etcd_client.call_credentials,
    31              metadata=etcd_client.metadata,
    32          )
    33  
    34          for kv in resp.kvs:
    35              key = kv.key.decode('utf-8')
    36              if not key.startswith(root + '/'):
    37                  continue
    38  
    39              key = key.replace(root, '')
    40              value = kv.value.decode('utf-8')
    41  
    42              if dry_run:
    43                  print('will set %s to %s' % (key, value))
    44                  continue
    45  
    46              redis_client.set(key, value)
    47              print('key %s is set to %s' % (key, value))
    48  
    49          if not resp.more:
    50              break
    51  
    52          range_start = etcd3.utils.increment_last_byte(kv.key)
    53  
    54  
    55  def getargs():
    56      ap = argparse.ArgumentParser()
    57      ap.add_argument('--source', dest='source_root_prefix', help='source root prefix', default='/eru')
    58      ap.add_argument('--etcd-host', default='127.0.0.1')
    59      ap.add_argument('--etcd-port', type=int, default=2379)
    60      ap.add_argument('--redis-host', default='127.0.0.1')
    61      ap.add_argument('--redis-port', type=int, default=6379)
    62      ap.add_argument('--redis-db', type=int, default=0)
    63      ap.add_argument('--dry-run', dest='dry_run', action='store_true', help='dry run, will not actually migrate')
    64      return ap.parse_args()
    65  
    66  
    67  def main():
    68      args = getargs()
    69  
    70      etcd_client = etcd3.client(host=args.etcd_host, port=args.etcd_port)
    71      redis_client = redis.StrictRedis(host=args.redis_host, port=args.redis_port, db=args.redis_db)
    72      migrate(args.source_root_prefix, etcd_client, redis_client, args.dry_run)
    73      return 0
    74  
    75  
    76  if __name__ == '__main__':
    77      sys.exit(main())