📝 Add docstring everywhere

Signed-off-by: Muhammed Hussein Karimi <[email protected]>

Author: mhkarimi1383

haproxy_redis_sentinel/cli.py

@@ -71,6 +71,9 @@ def run(
             rich_help_panel="HAProxy Info"
         )] = "current_master",
 ):
+    """
+    Run the HAProxy and Redis Sentinel handler.
+    """
     handler = Handler(
         sentinel_host=sentinel_host,
         sentinel_port=sentinel_port,

haproxy_redis_sentinel/handler.py

@@ -1,7 +1,7 @@
 from time import sleep
 from typing import Any
 from enum import StrEnum
-from utils import send_command
+from utils import send_command, is_empty
 from redis import Redis
 from redis.retry import Retry
 from redis.backoff import FullJitterBackoff
@@ -11,13 +11,20 @@
 
 
 class HAProxyOutput(StrEnum):
+    """
+    Enum for HAProxy output messages.
+    """
     SERVER_REGISTERED = "New server registered."
     SERVER_DELETED = "Server deleted."
     SERVER_NOT_FOUND = "No such server."
     BACKEND_NOT_FOUND = 'No such backend.'
 
 
 class Handler(object):
+    """
+    Handler class for managing HAProxy and Redis Sentinel.
+    """
+
     def __init__(
         self,
         sentinel_host: str = "127.0.0.1",
@@ -28,6 +35,16 @@ def __init__(
         haproxy_backend: str = "redis_master",
         haproxy_server_name: str = "current_master",
     ) -> None:
+        """
+        Initialize the Handler class.
+            :param sentinel_host: Redis Sentinel host
+            :param sentinel_port: Redis Sentinel port
+            :param sentinel_password: Redis Sentinel password
+            :param master_name: Redis master name
+            :param haproxy_socket: HAProxy socket path
+            :param haproxy_backend: HAProxy backend name
+            :param haproxy_server_name: HAProxy server name
+        """
         self.conn = Redis(
             host=sentinel_host,
             port=sentinel_port,
@@ -41,6 +58,10 @@ def __init__(
         self.haproxy_server_name = haproxy_server_name
 
     def get_master_address(self) -> str:
+        """
+        Get the address of the current master from Redis Sentinel.
+            :return: Address of the current master
+        """
         address = None
         sentinel_info: dict[str, Any] = self.conn.info()  # type: ignore
         try:
@@ -56,12 +77,22 @@ def get_master_address(self) -> str:
         return address
 
     def send_command(self, commands: str | list[str], log_data=True) -> str:
+        """
+        Send a command to the HAProxy socket and return the response.
+            :param commands: Command to send (string or list of strings)
+            :param log_data: Whether to log the command and output
+            :return: Response from the HAProxy socket
+        """
         out = send_command(self.haproxy_socket, commands)
         if log_data:
             info(f"HAProxy command: {commands}, Output: {out}")
         return out
 
     def shutdown_current_server(self) -> str:
+        """
+        Shuts down the current sesions of the server and sets it to maintenance mode.
+            :return: Response from the HAProxy socket
+        """
         return self.send_command(
             [
                 f"set server {self.haproxy_backend}/{self.haproxy_server_name} state maint",  # noqa: E501
@@ -74,6 +105,11 @@ def remove_current_server(
             ignore_notfound: bool = True,
             shutdown: bool = True,
     ) -> None:
+        """
+        Remove the current server from HAProxy.
+            :param ignore_notfound: Whether to ignore "not found" errors
+            :param shutdown: Whether to shut down the current server
+        """
         out = ""
         if shutdown:
             out = self.shutdown_current_server()
@@ -91,6 +127,10 @@ def remove_current_server(
             raise Exception(f"Error while removing old server: {out}")
 
     def add_server(self, address: str):
+        """
+        Add a new server to HAProxy (Used for initial sets).
+            :param address: Address of the new server
+        """
         out = self.send_command(
             f"add server {self.haproxy_backend}/{self.haproxy_server_name} {address}"  # noqa: E501
         )
@@ -101,6 +141,11 @@ def add_server(self, address: str):
         )
 
     def set_server_address(self, host: str, port: int):
+        """
+        Set the address of the current server in HAProxy.
+            :param host: Host of the new server
+            :param port: Port of the new server
+        """
         self.send_command(
             [
                 f"set server {self.haproxy_backend}/{self.haproxy_server_name} addr {host}",  # noqa: E501
@@ -110,6 +155,10 @@ def set_server_address(self, host: str, port: int):
         )
 
     def subscriber(self):
+        """
+        Subscribe to Redis Sentinel events and handle master failover (Blocking).
+            :return: None
+        """
         pubsub = self.conn.pubsub()
         pubsub.subscribe("+switch-master")
         for message in pubsub.listen():
@@ -123,38 +172,58 @@ def subscriber(self):
             self.set_server_address(master_info[3], int(master_info[4]))
 
     def set_initial_server(self):
+        """
+        Set the initial server in HAProxy.
+            :return: None
+        """
+        info("Setting initial server...")
         self.remove_current_server()
         return self.add_server(self.get_master_address())
 
     def haproxy_server_checker(self):
+        """
+        Check the HAProxy server status and set the initial server if needed.
+            :return: None
+        """
         stats: list[list[dict | None] | None] | None = orjson.loads(
             self.send_command(
                 f"show stat {self.haproxy_backend} 4 -1 json",
                 log_data=False,
             )
         )
-        if stats in (None, [], {}):
+        if is_empty(stats):
+            info("Empty data in stats, Setting initial server...")
             return self.set_initial_server()
-        for group in stats:
-            if group in (None, [], {}):
+        for group in stats:  # type: ignore
+            if is_empty(group):
+                info("Empty group in stats, Setting initial server...")
                 return self.set_initial_server()
-            for item in group:
-                if (item is None) or (not isinstance(item, dict)):
+            for item in group:  # type: ignore
+                if is_empty(item) or (not isinstance(item, dict)):
                     continue
                 field = item.get("field", {})
-                if (field is None) or (not isinstance(field, dict)):
+                if is_empty(field) or (not isinstance(field, dict)):
                     continue
                 if field.get("name") == "addr":
                     addr_value = item.get("value", {}).get("value", "")
-                    if not addr_value or len(addr_value) < 0:
+                    if is_empty(addr_value):
+                        info("Empty addr value in stats, Setting initial server...")  # noqa: E501
                         return self.set_initial_server()
 
     def haproxy_server_checker_worker(self):
+        """
+        Worker for checking the HAProxy server status (Blocking).
+            :return: None
+        """
         while True:
             self.haproxy_server_checker()
             sleep(2)
 
     def start_worker(self):
+        """
+        Start the worker processes for subscriber and HAProxy server checker (Blocking).
+            :return: None
+        """
         subscriber_process = Process(target=self.subscriber, name="subscriber")
         haproxy_server_checker_process = Process(
             target=self.haproxy_server_checker_worker,

haproxy_redis_sentinel/logger.py

@@ -3,10 +3,17 @@
 
 
 def log_prefix() -> str:
-    return f"{datetime.now().strftime("%c")} "
+    """
+    Returns a formatted string with the current date and time.
+    """
+    return f"{datetime.now().strftime('%c')} "
 
 
-def info(msg):
+def info(msg: str):
+    """
+    Prints an info message with a green color.
+        :param msg: Message to print
+    """
     typer.echo(
         log_prefix() + typer.style(
             msg,
@@ -18,6 +25,10 @@ def info(msg):
 
 
 def error(msg: str):
+    """
+    Prints an error message with a red color.
+        :param msg: Message to print
+    """
     typer.echo(
         log_prefix() + typer.style(
             msg,

haproxy_redis_sentinel/utils.py

@@ -1,13 +1,23 @@
 import socket
 
-__all__ = ["send_command"]
+__all__ = ["send_command", "is_empty"]
 
 
 def encode_command(command: str) -> bytes:
+    """
+    Encode a command string into bytes.
+        :param command: Command string to encode
+        :return: Encoded command as bytes
+    """
     return f"{command};\n".encode("utf-8")
 
 
-def recvall(sock: socket.socket):
+def recvall(sock: socket.socket) -> bytes:
+    """
+    Receive all data from a socket until EOF.
+        :param sock: Socket to receive data from
+        :return: Received data as bytes
+    """
     BUFF_SIZE = 1024
     data = b''
     while True:
@@ -20,6 +30,12 @@ def recvall(sock: socket.socket):
 
 
 def send_command(addr: str, command: str | list[str]) -> str:
+    """
+    Send a command to a socket address and return the response.
+        :param addr: Socket address (IP:port or path)
+        :param command: Command to send (string or list of strings)
+        :return: Response from the socket
+    """
     result = ""
     if len(addr.split(":")) == 2:
         addr_parts = addr.split(":")
@@ -42,3 +58,12 @@ def send_command(addr: str, command: str | list[str]) -> str:
     finally:
         unix_socket.close()
     return result
+
+
+def is_empty(value):
+    """
+    Check if a value is empty.
+        :param value: Value to check
+        :return: True if empty, False otherwise
+    """
+    return value in (None, "", [], {})