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

     1  #!/usr/bin/env python3
     2  # -*- coding: utf-8 -*-
     3  import os
     4  import sys
     5  import yaml
     6  def write_file(config, filename, overwrite):
     7      if not overwrite and os.path.exists(filename):
     8          pass
     9      else:
    10          with open(filename, 'w') as f:
    11              f.write(config)
    12  def read_file_lines(file):
    13      ret = []
    14      for line in file.readlines():
    15          line = line.strip()
    16          if line and not line.startswith('#'):
    17              ret.append(line)
    18      return ret
    19  def postgresql_conf_to_dict(file_path):
    20      with open(file_path, 'r') as f:
    21          content = f.read()
    22      lines = content.splitlines()
    23      result = {}
    24      for line in lines:
    25          if line.startswith('#'):
    26              continue
    27          if '=' not in line:
    28              continue
    29          key, value = line.split('=', 1)
    30          result[key.strip()] = value.strip().strip("'")
    31      return result
    32  def main(filename):
    33      restore_dir = os.environ.get('RESTORE_DATA_DIR', '')
    34      local_config = yaml.safe_load(
    35          os.environ.get('SPILO_CONFIGURATION', os.environ.get('PATRONI_CONFIGURATION', ''))) or {}
    36      if not 'postgresql' in local_config:
    37          local_config['postgresql'] = {}
    38      postgresql = local_config['postgresql']
    39      postgresql['config_dir'] = '/home/postgres/pgdata/conf'
    40      postgresql['custom_conf'] = '/home/postgres/conf/postgresql.conf'
    41      # add pg_hba.conf
    42      with open('/home/postgres/conf/pg_hba.conf', 'r') as f:
    43          lines = read_file_lines(f)
    44          if lines:
    45              postgresql['pg_hba'] = lines
    46      if restore_dir and os.path.isfile(os.path.join(restore_dir, 'kb_restore.signal')):
    47          if not 'bootstrap' in local_config:
    48              local_config['bootstrap'] = {}
    49          with open('/home/postgres/conf/kb_restore.conf', 'r') as f:
    50              local_config['bootstrap'].update(yaml.safe_load(f))
    51      # point in time recovery(PITR)
    52      data_dir = os.environ.get('PGDATA', '')
    53      if os.path.isfile("/home/postgres/pgdata/conf/recovery.conf"):
    54          with open('/home/postgres/conf/kb_pitr.conf', 'r') as f:
    55              pitr_config = yaml.safe_load(f)
    56              re_config = postgresql_conf_to_dict("/home/postgres/pgdata/conf/recovery.conf")
    57              pitr_config[pitr_config['method']]['recovery_conf'].update(re_config)
    58              local_config['bootstrap'].update(pitr_config)
    59      write_file(yaml.dump(local_config, default_flow_style=False), filename, True)
    60  if __name__ == '__main__':
    61      main(sys.argv[1])