1+ import logging
2+ from typing import Literal
3+
4+ from asgiref .sync import async_to_sync , sync_to_async
5+ from channels .db import database_sync_to_async
6+ from channels .layers import get_channel_layer
7+ from django .template .loader import render_to_string
8+
9+ from director .apps .sites .models import Site , Operation
10+ from director .apps .users .models import User
11+ from redis import asyncio as aioredis
12+
13+ from channels .generic .websocket import AsyncWebsocketConsumer
14+ import json
15+
16+ from director .settings import REDIS_URL
17+
18+ logger = logging .getLogger (__name__ )
19+
20+ _redis = None
21+
22+ async def _get_redis () -> aioredis .Redis :
23+ """
24+ Returns:
25+ An :class:`aioredis.Redis` instance, and avoids making a new connection if one already exists
26+ This helps save resources
27+ """
28+ global _redis
29+ if _redis is None :
30+ _redis = await aioredis .from_url (REDIS_URL )
31+ return _redis
32+
33+
34+ class SiteStatusConsumer (AsyncWebsocketConsumer ):
35+ async def connect (self ):
36+ self .site_id = self .scope ["url_route" ]["kwargs" ]["site_id" ]
37+ self .group_name = f"site_{ self .site_id }
38+ self .redis = await _get_redis ()
39+
40+ self .user = self .scope ["user" ]
41+ if not self .user .is_authenticated :
42+ await self .close ()
43+ return
44+
45+ try :
46+ if self .user .is_superuser :
47+ self .site = await Site .objects .aget (id = self .site_id )
48+ else :
49+ self .site = await Site .objects .aget (id = self .site_id , users = self .user )
50+ except Site .DoesNotExist :
51+ await self .close ()
52+ return
53+
54+ await self .channel_layer .group_add (self .group_name , self .channel_name )
55+ await self .accept ()
56+
57+ # Add this user to a redis set of active users
58+ await self .redis .sadd (self .group_name , self .user .id )
59+ # Let all connected user's know this user is now connected
60+ await self .channel_layer .group_send (
61+ self .group_name ,
62+ {"type" : "send_active_users" }
63+ )
64+
65+ async def disconnect (self , close_code ):
66+ await self .channel_layer .group_discard (self .group_name , self .channel_name )
67+
68+ # Remove this user from the redis set of active users
69+ await self .redis .srem (self .group_name , self .user .id )
70+ # Make sure other users know this user's connection is going away
71+ await self .channel_layer .group_send (
72+ self .group_name ,
73+ {"type" : "send_active_users" }
74+ )
75+
76+ async def send_power_event (self , event ):
77+ power_event_template = render_to_string ("base_with_site_header_and_nav.html.partials/status_indicator.html" , {"site" : self .site , "status" : event ["status" ]})
78+ await self .send (text_data = power_event_template .strip ())
79+
80+ async def send_operation_updated_event (self , event ):
81+ operation = await Operation .objects .filter (site = self .site ).afirst ()
82+ power_event_template = await sync_to_async (render_to_string )("site_dashboard.html.partials/operations_and_process_log.html" , {"operation" : operation })
83+ await self .send (text_data = power_event_template .strip ())
84+
85+ async def send_process_log_data_event (self , event ):
86+ pass
87+
88+ async def send_active_users (self , event ):
89+ members = await self .redis .smembers (self .group_name )
90+ users = []
91+ for member_id in members :
92+ try :
93+ user = await User .objects .aget (id = member_id .decode ("utf-8" ))
94+ users .append (user )
95+ except User .DoesNotExist :
96+ logger .exception ("Non critical error: a user does not exist when parsing the redis active users set"
97+ f"User id was { member_id } { self .group_name } )
98+
99+ active_users_template = render_to_string ("base_with_site_header_and_nav.html.partials/active_users.html" , {"active_users" : users , "user" : self .user })
100+ await self .send (text_data = active_users_template .strip ())
101+
102+ def trigger_power_event (site : Site , status : Literal ["Online" , "Offline" , "Restarting" , "Starting" , "Stopping" ]):
103+ async_to_sync (get_channel_layer ().group_send )(
104+ f"site_{ site .id } ,
105+ {
106+ "type" : "send_power_event" ,
107+ "status" : status ,
108+ }
109+ )
110+
111+ def trigger_operation_updated_event (site : Site ):
112+ async_to_sync (get_channel_layer ().group_send )(
113+ f"site_{ site .id } ,
114+ {
115+ "type" : "send_operation_updated_event" ,
116+ }
117+ )
118+
119+ def trigger_process_log_data_event (site : Site ):
120+ async_to_sync (get_channel_layer ().group_send )(
121+ f"site_{ site .id } ,
122+ {
123+ "type" : "send_operation_updated_event" ,
124+ }
125+ )
0 commit comments