In my Django project, I must interact with several server, through websockets, and it needs constant connection with them. How can I define a management command to connect to all servers, and then in my API, I call each one of them that I want with their URL?
I wrote an async manager and management command to create connections to all of them:
# servers/websocket_manager.py
import asyncio
class WebSocketManager:
def __init__(self):
self.clients = {}
async def add_client(self, url):
client = WebSocketClient(url)
await client.connect()
self.clients[url] = client
def get_client(self, url):
return self.clients.get(url)
async def close_all_clients(self):
await asyncio.gather(*[client.close() for client in self.clients.values()])
self.clients.clear()
class WebSocketClient:
def __init__(self, url):
self.url = url
self.connection = None
async def connect(self):
# Logic to establish the WebSocket connection
pass
async def close(self):
# Logic to close the WebSocket connection
if self.connection:
await self.connection.close()
# management/commands/init_websockets.py
from django.core.management.base import BaseCommand
from servers.websocket_manager import WebSocketManager
from servers.models import Server
import asyncio
from multiprocessing import Process
class Command(BaseCommand):
help = 'Initializes WebSocket connections to specified URLs'
def handle(self, *args, **options):
manager = WebSocketManager()
servers = Server.objects.all()
urls = [server.url for server in servers]
process = Process(target=self.run_event_loop, args=(manager, urls))
process.start()
def run_event_loop(self, manager, urls):
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
loop.run_until_complete(self.initialize_connections(manager, urls))
async def initialize_connections(self, manager, urls):
tasks = [asyncio.create_task(manager.add_client(url)) for url in urls]
await asyncio.gather(*tasks)
print('WebSocket connections initialized.')
But as it runs in a different process, I don’t have access to the object or its connection. How can I handle it? Do I need to run the process?