11from typing import Any
22
3+ import discord
34from discord .ext import commands , tasks
45from influxdb_client .client .influxdb_client import InfluxDBClient
56from influxdb_client .client .write .point import Point
1213
1314
1415class InfluxLogger (commands .Cog ):
15- def __init__ (self , bot : Tux ):
16+ def __init__ (self , bot : Tux ) -> None :
1617 self .bot = bot
1718 self .db = DatabaseController ()
1819 self .influx_write_api : Any | None = None
1920 self .influx_org : str = ""
2021
2122 if self .init_influx ():
23+ self ._log_guild_stats .start ()
2224 self .logger .start ()
2325 else :
2426 logger .warning ("InfluxDB logger failed to init. Check .env configuration if you want to use it." )
@@ -42,7 +44,46 @@ def init_influx(self) -> bool:
4244 return True
4345 return False
4446
45- @tasks .loop (seconds = 60 )
47+ @tasks .loop (seconds = 60 , name = "influx_guild_stats" )
48+ async def _log_guild_stats (self ) -> None :
49+ """Logs guild statistics to InfluxDB."""
50+ if not self .bot .is_ready () or not self .influx_write_api :
51+ logger .debug ("Bot not ready or InfluxDB writer not initialized, skipping InfluxDB logging." )
52+ return
53+
54+ for guild in self .bot .guilds :
55+ online_members = sum (m .status != discord .Status .offline for m in guild .members )
56+
57+ tags = {"guild" : guild .name }
58+ fields = {
59+ "members" : guild .member_count ,
60+ "online" : online_members ,
61+ }
62+
63+ point = {"measurement" : "guild_stats" , "tags" : tags , "fields" : fields }
64+
65+ self .influx_write_api .write (bucket = "tux_stats" , org = self .influx_org , record = point )
66+
67+ @_log_guild_stats .before_loop
68+ async def before_log_guild_stats (self ) -> None :
69+ """Wait until the bot is ready."""
70+ await self .bot .wait_until_ready ()
71+
72+ @_log_guild_stats .error
73+ async def on_log_guild_stats_error (self , error : BaseException ) -> None :
74+ """Handles errors in the guild stats logging loop."""
75+ logger .error (f"Error in InfluxDB guild stats logger loop: { error } " )
76+ if isinstance (error , Exception ):
77+ self .bot .sentry_manager .capture_exception (error )
78+ else :
79+ raise error
80+
81+ async def cog_unload (self ) -> None :
82+ if self .influx_write_api :
83+ self ._log_guild_stats .cancel ()
84+ self .logger .cancel ()
85+
86+ @tasks .loop (seconds = 60 , name = "influx_db_logger" )
4687 async def logger (self ) -> None :
4788 """Log statistics to InfluxDB at regular intervals.
4889
@@ -55,40 +96,54 @@ async def logger(self) -> None:
5596 influx_bucket = "tux stats"
5697
5798 # Collect the guild list from the database
58- try :
59- guild_list = await self .db .guild .find_many (where = {})
99+ guild_list = await self .db .guild .find_many (where = {})
100+
101+ # Iterate through each guild and collect metrics
102+ for guild in guild_list :
103+ if not guild .guild_id :
104+ continue
60105
61- # Iterate through each guild and collect metrics
62- for guild in guild_list :
63- if not guild .guild_id :
64- continue
106+ guild_id = int (guild .guild_id )
65107
66- guild_id = int (guild .guild_id )
108+ # Collect data by querying controllers
109+ starboard_stats = await self .db .starboard_message .find_many (where = {"message_guild_id" : guild_id })
67110
68- # Collect data by querying controllers
69- starboard_stats = await self .db .starboard_message .find_many (where = {"message_guild_id" : guild_id })
111+ snippet_stats = await self .db .snippet .find_many (where = {"guild_id" : guild_id })
70112
71- snippet_stats = await self .db .snippet .find_many (where = {"guild_id" : guild_id })
113+ afk_stats = await self .db .afk .find_many (where = {"guild_id" : guild_id })
72114
73- afk_stats = await self .db .afk .find_many (where = {"guild_id" : guild_id })
115+ case_stats = await self .db .case .find_many (where = {"guild_id" : guild_id })
74116
75- case_stats = await self .db .case .find_many (where = {"guild_id" : guild_id })
117+ # Create data points with type ignores for InfluxDB methods
118+ # The InfluxDB client's type hints are incomplete
119+ points : list [Point ] = [
120+ Point ("guild stats" ).tag ("guild" , guild_id ).field ("starboard count" , len (starboard_stats )), # type: ignore
121+ Point ("guild stats" ).tag ("guild" , guild_id ).field ("snippet count" , len (snippet_stats )), # type: ignore
122+ Point ("guild stats" ).tag ("guild" , guild_id ).field ("afk count" , len (afk_stats )), # type: ignore
123+ Point ("guild stats" ).tag ("guild" , guild_id ).field ("case count" , len (case_stats )), # type: ignore
124+ ]
76125
77- # Create data points with type ignores for InfluxDB methods
78- # The InfluxDB client's type hints are incomplete
79- points : list [Point ] = [
80- Point ("guild stats" ).tag ("guild" , guild_id ).field ("starboard count" , len (starboard_stats )), # type: ignore
81- Point ("guild stats" ).tag ("guild" , guild_id ).field ("snippet count" , len (snippet_stats )), # type: ignore
82- Point ("guild stats" ).tag ("guild" , guild_id ).field ("afk count" , len (afk_stats )), # type: ignore
83- Point ("guild stats" ).tag ("guild" , guild_id ).field ("case count" , len (case_stats )), # type: ignore
84- ]
126+ # Write to InfluxDB
127+ self .influx_write_api .write (bucket = influx_bucket , org = self .influx_org , record = points )
85128
86- # Write to InfluxDB
87- self .influx_write_api .write (bucket = influx_bucket , org = self .influx_org , record = points )
129+ @logger .before_loop
130+ async def before_logger (self ) -> None :
131+ """Wait until the bot is ready."""
132+ await self .bot .wait_until_ready ()
88133
89- except Exception as e :
90- logger .error (f"Error collecting metrics for InfluxDB: { e } " )
134+ @logger .error
135+ async def on_logger_error (self , error : BaseException ) -> None :
136+ """Handles errors in the logger loop."""
137+ logger .error (f"Error in InfluxDB logger loop: { error } " )
138+ if isinstance (error , Exception ):
139+ self .bot .sentry_manager .capture_exception (error )
140+ else :
141+ raise error
91142
92143
93144async def setup (bot : Tux ) -> None :
94- await bot .add_cog (InfluxLogger (bot ))
145+ # Only load the cog if InfluxDB configuration is available
146+ if all ([CONFIG .INFLUXDB_TOKEN , CONFIG .INFLUXDB_URL , CONFIG .INFLUXDB_ORG ]):
147+ await bot .add_cog (InfluxLogger (bot ))
148+ else :
149+ logger .warning ("InfluxDB configuration incomplete, skipping InfluxLogger cog" )
0 commit comments