consul.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
# @Version     : Python 3.11.4
# @Software    : Sublime Text 4
# @Author      : StudentCWZ
# @Email       : StudentCWZ@outlook.com
# @Date        : 2023/10/28 15:06
# @File        : consul.py
# @Description :
"""

from typing import Any

import consul
import yaml


class ConsulConfig:
    def __init__(self, host='localhost', port=8500, token=None, dc='dc1'):
        """
        Initialize a ConsulConfig instance.

        :param host: The host of the Consul server. Default is 'localhost'.
        :param port: The port of the Consul server. Default is 8500.
        :param token: The Consul token. Default is None.
        :param dc: The datacenter to use. Default is 'dc1'.
        """
        self.token = token
        self.dc = dc
        # Create a Consul client instance
        self.client = consul.Consul(host=host, port=port, token=token, dc=dc)

    def get(self, key: str) -> Any:
        """
        Retrieve the value of a key from the Consul KV store.

        :param key: The key to retrieve.
        :return: The value of the key.
        :raises KeyError: If the key is not found in Consul.
        """
        # Retrieve the key-value pair from Consul
        _, data = self.client.kv.get(key=key, token=self.token, dc=self.dc)
        # If the key-value pair is not found, raise an exception
        if data is None:
            raise KeyError(f'Key {key} not found in Consul.')
        # Load the value from the key-value pair as YAML
        return yaml.load(data['Value'], Loader=yaml.FullLoader)