socket_client.py

import asyncio
import logging
import sys
import threading
import time
from typing import Dict, List, Optional

import socketio
from chromatrace import LoggingConfig, LoggingSettings
from chromatrace.tracer import trace_id_ctx
from pydantic import BaseModel, Field

logging_config = LoggingConfig(
    LoggingSettings(
        enable_tracing=True,
        file_path="./logs/app-client.log",
        enable_file_logging=True,
        show_process_id=True,
    )
)
trace_id = "123456"
trace_id_ctx.set(f"C-{trace_id}")


class SocketClientConfig(BaseModel):
    server_url: str = "http://0.0.0.0:8001"
    namespaces: List[str] = ["/"]
    headers: Optional[Dict] = Field(default=None)
    socketio_path: str = "/socket.io"
    auth: Optional[Dict] = Field(default=None)
    logger: bool = False
    engineio_logger: bool = False
    reconnection: bool = True
    reconnection_delay: int = 3
    reconnection_attempts: int = 10


class SocketClient:
    def __init__(
        self,
        config: SocketClientConfig,
    ):
        self.logger = logging_config.get_logger(self.__class__.__name__)
        self.config = config
        self.logger.info("Initializing Socket Client to connect to the server.")
        self.logger.info(
            "Server URL: %s%s", self.config.server_url, self.config.socketio_path
        )
        self.sio = socketio.AsyncClient(
            handle_sigint=True,
            logger=self.logger if self.config.logger else False,
            engineio_logger=self.logger if self.config.engineio_logger else False,
            reconnection=self.config.reconnection,
            reconnection_delay=self.config.reconnection_delay,
            reconnection_attempts=self.config.reconnection_attempts,
        )

        self._client_loop = asyncio.new_event_loop()
        self.logger.info("Socket Client is initialized.")

    def start_background_loop(self, loop: asyncio.AbstractEventLoop) -> None:
        asyncio.set_event_loop(loop)
        loop.run_forever()

    @property
    def client_loop(self):
        return self._client_loop

    async def connect_to_server(self):
        try:
            self.logger.info("Connecting to the server...")
            await self.sio.connect(
                self.config.server_url,
                headers=self.config.headers,
                socketio_path=self.config.socketio_path,
                auth=self.config.auth,
                namespaces=self.config.namespaces,
            )
            self.logger.info("Connected to the server successfully.")
        except ConnectionError:
            self.logger.error("Connection failed.")
        await self.sio.wait()

    def call_backs(self):
        @self.sio.on("connect", namespace="/")
        async def connect():
            self.logger.info("Connection established to the server with handshake.")

        @self.sio.on("disconnect", namespace="/")
        def disconnect():
            self.logger.info("Server disconnected.")
            self.client_loop.stop()

        @self.sio.on("message", namespace="/")
        async def message(data):
            self.logger.info("Message from the server in `message` event: %s", data)

    async def run(self):
        self.call_backs()
        await self.connect_to_server()


def handle_user_input(base_client: SocketClient, logger: logging.Logger):
    while True:
        if not base_client.sio.connected:
            logger.info(
                "Client connection with server has been lost, exiting the client..."
            )
            sys.exit(1)

        time.sleep(0.5)  # to enhance showcase
        user_input = input('Enter a message (or "exit" to quit): ')
        try:
            if user_input.lower() == "exit":
                break
            else:
                asyncio.run_coroutine_threadsafe(
                    base_client.sio.emit(
                        event="message",
                        data={"message": user_input},
                        namespace="/",
                    ),
                    base_client.client_loop,
                )
        except KeyboardInterrupt:
            break


def create_client():
    token = "Bearer XXXXXXXXXXX"
    config = SocketClientConfig(
        headers={
            "authorization": token,
            "x-trace-id": trace_id,
        },
        auth={"Authorization": token},
    )
    base_client = SocketClient(config=config)

    # starting the client background loop
    th = threading.Thread(
        target=base_client.start_background_loop,
        args=(base_client.client_loop,),
        daemon=True,
    )
    th.start()

    return base_client


def main():
    logger = logging_config.get_logger(__name__)
    logger.info("Starting client...")
    base_client = create_client()

    # running the client in the background loop
    asyncio.run_coroutine_threadsafe(base_client.run(), base_client.client_loop)

    time.sleep(0.5)
    logger.info("Ready to check the connection status and sending messages...")
    handle_user_input(base_client, logger)

    logger.info("Connection is closing...")
    base_client.sio.disconnect()


if __name__ == "__main__":
    main()