@@ -36,6 +36,17 @@ async def _ensure_async(conn: asyncpg.Connection) -> None:
3636 used_credits INTEGER NOT NULL DEFAULT 0,
3737 updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
3838 );
39+ CREATE TABLE IF NOT EXISTS bot_credit_user_limits (
40+ user_id BIGINT PRIMARY KEY,
41+ limit INTEGER NOT NULL,
42+ updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
43+ );
44+ CREATE TABLE IF NOT EXISTS bot_credit_unlimited_roles (
45+ role_id BIGINT PRIMARY KEY,
46+ role_name TEXT,
47+ guild_id BIGINT,
48+ created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
49+ );
3950 """
4051 )
4152
@@ -185,3 +196,158 @@ def resolve_user_limit_from_roles(*, member_roles: list[tuple[int, str]]) -> int
185196 if isinstance (lim , int ) and lim > max_limit :
186197 max_limit = lim
187198 return max_limit
199+
200+
201+ def has_unlimited_from_roles (* , member_roles : list [tuple [int , str ]]) -> bool :
202+ # From settings
203+ names = set ((settings .credit_unlimited_role_names or []))
204+ ids = set (int (x ) for x in (settings .credit_unlimited_role_ids or []))
205+ for rid , name in member_roles :
206+ if rid in ids :
207+ return True
208+ if name and name in names :
209+ return True
210+ # From DB
211+ async def run () -> bool :
212+ conn = await asyncpg .connect (_dsn ())
213+ try :
214+ await _ensure_async (conn )
215+ role_ids = [int (rid ) for rid , _ in member_roles if rid ]
216+ if not role_ids :
217+ return False
218+ rows = await conn .fetch (
219+ "SELECT role_id FROM bot_credit_unlimited_roles WHERE role_id = ANY($1::BIGINT[])" ,
220+ role_ids ,
221+ )
222+ return bool (rows )
223+ finally :
224+ await conn .close ()
225+
226+ return asyncio .run (run ())
227+
228+
229+ def get_user_limit_override (user_id : int ) -> Optional [int ]:
230+ async def run () -> Optional [int ]:
231+ conn = await asyncpg .connect (_dsn ())
232+ try :
233+ await _ensure_async (conn )
234+ row = await conn .fetchrow ("SELECT limit FROM bot_credit_user_limits WHERE user_id=$1" , int (user_id ))
235+ return int (row [0 ]) if row else None
236+ finally :
237+ await conn .close ()
238+
239+ return asyncio .run (run ())
240+
241+
242+ def set_user_limit (user_id : int , limit : int ) -> None :
243+ async def run ():
244+ conn = await asyncpg .connect (_dsn ())
245+ try :
246+ await _ensure_async (conn )
247+ await conn .execute (
248+ """
249+ INSERT INTO bot_credit_user_limits(user_id, limit)
250+ VALUES ($1, $2)
251+ ON CONFLICT (user_id) DO UPDATE SET limit=EXCLUDED.limit, updated_at=NOW()
252+ """ ,
253+ int (user_id ), int (limit )
254+ )
255+ finally :
256+ await conn .close ()
257+
258+ asyncio .run (run ())
259+
260+
261+ def clear_user_limit (user_id : int ) -> None :
262+ async def run ():
263+ conn = await asyncpg .connect (_dsn ())
264+ try :
265+ await _ensure_async (conn )
266+ await conn .execute ("DELETE FROM bot_credit_user_limits WHERE user_id=$1" , int (user_id ))
267+ finally :
268+ await conn .close ()
269+
270+ asyncio .run (run ())
271+
272+
273+ def add_unlimited_role (role_id : int , role_name : Optional [str ], guild_id : Optional [int ]) -> None :
274+ async def run ():
275+ conn = await asyncpg .connect (_dsn ())
276+ try :
277+ await _ensure_async (conn )
278+ await conn .execute (
279+ """
280+ INSERT INTO bot_credit_unlimited_roles(role_id, role_name, guild_id)
281+ VALUES ($1, $2, $3)
282+ ON CONFLICT (role_id) DO UPDATE SET role_name=EXCLUDED.role_name, guild_id=EXCLUDED.guild_id
283+ """ ,
284+ int (role_id ), role_name , int (guild_id ) if guild_id else None
285+ )
286+ finally :
287+ await conn .close ()
288+
289+ asyncio .run (run ())
290+
291+
292+ def remove_unlimited_role (role_id : int ) -> None :
293+ async def run ():
294+ conn = await asyncpg .connect (_dsn ())
295+ try :
296+ await _ensure_async (conn )
297+ await conn .execute ("DELETE FROM bot_credit_unlimited_roles WHERE role_id=$1" , int (role_id ))
298+ finally :
299+ await conn .close ()
300+
301+ asyncio .run (run ())
302+
303+
304+ def list_unlimited_roles () -> list [tuple [int , Optional [str ], Optional [int ]]]:
305+ async def run () -> list [tuple [int , Optional [str ], Optional [int ]]]:
306+ conn = await asyncpg .connect (_dsn ())
307+ try :
308+ await _ensure_async (conn )
309+ rows = await conn .fetch ("SELECT role_id, role_name, guild_id FROM bot_credit_unlimited_roles ORDER BY created_at" )
310+ out : list [tuple [int , Optional [str ], Optional [int ]]] = []
311+ for r in rows :
312+ out .append ((int (r ["role_id" ]), r ["role_name" ], int (r ["guild_id" ]) if r ["guild_id" ] else None ))
313+ return out
314+ finally :
315+ await conn .close ()
316+
317+ return asyncio .run (run ())
318+
319+
320+ def get_usage (user_id : int ) -> tuple [int , int ]:
321+ """Return (user_used, global_used) for current period."""
322+ period = _period_start ()
323+
324+ async def run () -> tuple [int , int ]:
325+ conn = await asyncpg .connect (_dsn ())
326+ try :
327+ await _ensure_async (conn )
328+ row_u = await conn .fetchrow ("SELECT used_credits FROM bot_credits_user WHERE user_id=$1 AND period_start=$2" , int (user_id ), period )
329+ row_g = await conn .fetchrow ("SELECT used_credits FROM bot_credits_global WHERE period_start=$1" , period )
330+ return (int (row_u [0 ]) if row_u else 0 , int (row_g [0 ]) if row_g else 0 )
331+ finally :
332+ await conn .close ()
333+
334+ return asyncio .run (run ())
335+
336+
337+ def compute_user_policy (* , user_id : int , member_roles : list [tuple [int , str ]], is_admin : bool ) -> tuple [bool , int ]:
338+ """Return (unlimited, per_user_limit).
339+
340+ unlimited ignores per-user limit but still respects global cap.
341+ """
342+ # Admins unlimited by default
343+ if is_admin :
344+ return True , 10 ** 9
345+ if has_unlimited_from_roles (member_roles = member_roles ):
346+ return True , 10 ** 9
347+ # Per-user override
348+ ul = get_user_limit_override (user_id )
349+ if isinstance (ul , int ):
350+ return False , max (1 , int (ul ))
351+ # Rank-based
352+ limit = resolve_user_limit_from_roles (member_roles = member_roles )
353+ return False , max (1 , int (limit ))
0 commit comments