From 3bbb01ed3a8f459f137eff5131006c2b15b85ce1 Mon Sep 17 00:00:00 2001 From: Hayley Date: Sat, 11 Oct 2025 20:25:36 +1100 Subject: [PATCH 1/3] Add animals plugin and cats command --- plugins/animals.py | 301 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 301 insertions(+) create mode 100644 plugins/animals.py diff --git a/plugins/animals.py b/plugins/animals.py new file mode 100644 index 0000000..5b86e69 --- /dev/null +++ b/plugins/animals.py @@ -0,0 +1,301 @@ + + +import abc +from collections import defaultdict +from dataclasses import dataclass +import logging +from typing import Any, Literal, Optional + +from bot.acl import privileged +from bot.config import plugin_config_command +from discord.ext.commands import group +import discord +from bot.commands import Context, plugin_command +from discord.ext.commands import command +import aiohttp +import datetime as dt + +from util.discord import UserError + +CAT_API_ROOT = 'https://api.thecatapi.com/v1/images/search' +DOG_API_ROOT = 'https://api.thedogapi.com/v1/images/search' +TIMEOUT = 10 # seconds + +def with_suffix(obj: dict[str, Any], *path: str, suffix: str) -> Optional[str]: + val = obj + try: + for p in path: + val = val[p] + except (KeyError, TypeError): + return None + if val is not None: + return f'{val}{suffix}' + return None + +class RateLimiter(abc.ABC): + @abc.abstractmethod + def is_allowed(self, ctx: Context) -> tuple[bool, str]: + pass + +class PerPersonRateLimiter(RateLimiter): + def __init__(self, calls: int, period: dt.timedelta): + self.calls = calls + self.period = period + self.users: dict[int, list[dt.datetime]] = defaultdict(list) + + def __repr__(self) -> str: + return f'{self.__class__.__name__}<{self.calls=}, {self.period=}>' + + def __str__(self) -> str: + return f'{self.__class__.__name__}<{self.calls} calls per {self.period}>' + + def _is_allowed(self, user_id: int) -> tuple[bool, str]: + now = dt.datetime.now() + timestamps = self.users[user_id] + # Remove timestamps outside the period + while timestamps and now - timestamps[0] > self.period: + timestamps.pop(0) + if len(timestamps) < self.calls: + timestamps.append(now) + return True, '' + else: + next_time = timestamps[0] + self.period + return False, f'Rate limit exceeded: {self.calls} calls per {self.period}. Please try again .' + + def is_allowed(self, ctx: Context): + user_id = ctx.author.id + return self._is_allowed(user_id) + +class GlobalRateLimiter(PerPersonRateLimiter): + def is_allowed(self, ctx: Context): + return self._is_allowed(0) # Use a single user ID for global rate limiting + +@dataclass +class CatRequest: + limit: Optional[int] = None + page: Optional[int] = None + order: Literal['ASC', 'DESC', 'RAND', None] = None + has_breeds: Optional[bool] = None + breed_ids: Optional[list[str]] = None + #sub_id: Optional[str] + + def __post_init__(self): + try: + assert self.limit is None or 1 <= self.limit <= 100 + assert self.page is None or 0 <= self.page + except AssertionError: + raise ValueError() + + def to_dict(self) -> dict[str, Any]: + result = {} + if self.limit is not None: + result['limit'] = self.limit + if self.page is not None: + result['page'] = self.page + if self.order is not None: + result['order'] = self.order + if self.has_breeds is not None: + result['has_breeds'] = int(self.has_breeds) + if self.breed_ids is not None: + result['breed_ids'] = ','.join(self.breed_ids) + return result + +@dataclass +class DogRequest: + size: Literal['full', 'med', 'small', 'thumb'] = 'small' + mime_types: Optional[list[Literal['jpg', 'png', 'gif']]] = None + format: Literal['json', 'src'] = 'json' + order: Literal['ASC', 'DESC', 'RAND', None] = None + limit: Optional[int] = None + page: Optional[int] = None + has_breeds: Optional[bool] = None + + def to_dict(self) -> dict[str, Any]: + result = {} + result['size'] = self.size + if self.mime_types is not None: + result['mime_types'] = ','.join(self.mime_types) + result['format'] = self.format + result['order'] = self.order + result['limit'] = self.limit + result['page'] = self.page + if self.has_breeds is not None: + result['has_breeds'] = int(self.has_breeds) + return {k: v for k, v in result.items() if v is not None} + +@dataclass +class AnimalResponse: + id: str + url: str + width: int + height: int + categories: Optional[list[Any]] = None + breeds: Optional[list[dict[str, Any]]] = None + + def get_weight(self) -> Optional[str]: + if self.breeds: + breed = self.breeds[0] + return ( + with_suffix(breed, 'weight', 'metric', suffix=' kg') or + with_suffix(breed, 'weight', 'imperial', suffix=' lbs') or + with_suffix(breed, 'weight', suffix='') + ) + return None + + def get_life_span(self) -> Optional[str]: + if self.breeds: + breed = self.breeds[0] + lifespan = breed.get('life_span', None) + if not lifespan: + return None + if 'years' in lifespan: + return lifespan + else: + return f'{lifespan} years' + return None + + + def get_description(self) -> str: + if self.breeds: + breed = self.breeds[0] + out = {} + out['Name'] = breed.get('name', None) + out['Temperament'] = breed.get('temperament', None) + out['Origin'] = breed.get('origin', None) + out['Description'] = breed.get('description', None) + out['Weight'] = self.get_weight() + out['Life span'] = self.get_life_span() + if len(self.breeds) > 1: + other_breeds = [b.get('name', None) for b in self.breeds[1:]] + out['Other breeds'] = ', '.join(b for b in other_breeds if b) + description = '\n'.join(f'{key}: {value}' for key, value in out.items() if value) + return description + else: + return 'No breed information available.' + +class AnimalApi: + def __init__(self, api_root: str, api_key: Optional[str], rate_limiters: list[RateLimiter]) -> None: + self.api_root = api_root + self.api_key = api_key + self.rate_limiters = rate_limiters + + def configure(self, api_key: str) -> None: + self.api_key = api_key + + def get_configuration(self) -> str: + out = '' + out += f'API Root: {self.api_root}\n' + out += f'API Key: {self.api_key}\n' + for limiter in self.rate_limiters: + out += f'Rate Limiter: {limiter}\n' + return out + + async def fetch_random_animal(self, ctx: Context, params: dict[str, Any]) -> AnimalResponse: + logging.debug(f'Fetching random animal with params: {params}') + if not self.api_key: + raise UserError('Animal API is not configured.') + headers = {'x-api-key': self.api_key} + for limiter in self.rate_limiters: + allowed, message = limiter.is_allowed(ctx) + if not allowed: + raise UserError(message) + async with aiohttp.ClientSession() as session: + async with session.get(self.api_root, params=params, headers=headers, timeout=TIMEOUT) as resp: + data = await resp.json() + if not data: + raise UserError('No animal found!') + return AnimalResponse(**data[0]) + +class CatApi(AnimalApi): + def __init__(self, api_key: Optional[str], rate_limiters: list[RateLimiter]) -> None: + super().__init__(CAT_API_ROOT, api_key, rate_limiters) + + async def fetch_random_cat(self, ctx: Context, req: CatRequest): + return await self.fetch_random_animal(ctx, req.to_dict()) + +# cat api key should be set via the config command +cat_api = CatApi('', [ + # reasonable rate limits, I don't see a good way to configure these dynamically + PerPersonRateLimiter(calls=3, period=dt.timedelta(minutes=1)), + PerPersonRateLimiter(calls=20, period=dt.timedelta(hours=1)), + PerPersonRateLimiter(calls=100, period=dt.timedelta(days=1)), + + GlobalRateLimiter(calls=10, period=dt.timedelta(minutes=1)), +]) + +class DogApi(AnimalApi): + def __init__(self, api_key: Optional[str], rate_limiters: list[RateLimiter]) -> None: + super().__init__(DOG_API_ROOT, api_key, rate_limiters) + + async def fetch_random_dog(self, ctx: Context, req: DogRequest): + return await self.fetch_random_animal(ctx, req.to_dict()) + +# dog api key should be set via the config command +dog_api = DogApi('', [ + # reasonable rate limits, I don't see a good way to configure these dynamically + PerPersonRateLimiter(calls=3, period=dt.timedelta(minutes=1)), + PerPersonRateLimiter(calls=20, period=dt.timedelta(hours=1)), + PerPersonRateLimiter(calls=100, period=dt.timedelta(days=1)), + + GlobalRateLimiter(calls=10, period=dt.timedelta(minutes=1)), +]) + +@plugin_command +@privileged +@command('cat') +async def random_cat(ctx: Context) -> None: + ''' Fetches and displays a random cat image ''' + + async with ctx.typing(): + # Prepare request params + req = CatRequest() + + cat = await cat_api.fetch_random_cat(ctx, req) + embed = discord.Embed( + title='Here is your random cat! 🐱', + url=cat.url, + color=discord.Color.random() + ) + embed.set_image(url=cat.url) + footer = cat.get_description() + if footer: + footer += '\n' + footer += 'thecatapi.com' + embed.set_footer(text=footer) + await ctx.send(embed=embed) + +@plugin_command +@privileged +@command('dog') +async def random_dog(ctx: Context) -> None: + ''' Fetches and displays a random dog image ''' + + async with ctx.typing(): + # Prepare request params + req = DogRequest() + + dog = await dog_api.fetch_random_dog(ctx, req) + embed = discord.Embed( + title='Here is your random dog! 🐶', + url=dog.url, + color=discord.Color.random() + ) + embed.set_image(url=dog.url) + footer = dog.get_description() + if footer: + footer += '\n' + footer += 'thedogapi.com' + embed.set_footer(text=footer) + await ctx.send(embed=embed) + +@plugin_config_command +@privileged +@group("animals", invoke_without_command=True) +async def animals_config(ctx: Context, cat_api_key: Optional[str], dog_api_key: Optional[str]) -> None: + ''' Configures the animals plugin ''' + if cat_api_key: + cat_api.configure(api_key=cat_api_key) + if dog_api_key: + dog_api.configure(api_key=dog_api_key) + + await ctx.send(f'Animals plugin configuration:\n**Cat**\n{cat_api.get_configuration()}\n**Dog**\n{dog_api.get_configuration()}') \ No newline at end of file From 3a7555647f7a6008e9b6c8e14d77b2e18a2847fb Mon Sep 17 00:00:00 2001 From: Hayley Date: Mon, 13 Oct 2025 03:55:24 +1100 Subject: [PATCH 2/3] removed rate limiter from animals command, added slightly more useful config, still broken --- plugins/animals.py | 333 ++++++++++++++++++++------------------------- 1 file changed, 151 insertions(+), 182 deletions(-) diff --git a/plugins/animals.py b/plugins/animals.py index 5b86e69..55b9ef8 100644 --- a/plugins/animals.py +++ b/plugins/animals.py @@ -1,7 +1,6 @@ import abc -from collections import defaultdict from dataclasses import dataclass import logging from typing import Any, Literal, Optional @@ -13,85 +12,60 @@ from bot.commands import Context, plugin_command from discord.ext.commands import command import aiohttp -import datetime as dt +from sqlalchemy import TEXT +from sqlalchemy.orm import Mapped, mapped_column +import sqlalchemy.orm +from sqlalchemy.ext.asyncio import async_sessionmaker + +import plugins from util.discord import UserError +import util.db CAT_API_ROOT = 'https://api.thecatapi.com/v1/images/search' DOG_API_ROOT = 'https://api.thedogapi.com/v1/images/search' TIMEOUT = 10 # seconds -def with_suffix(obj: dict[str, Any], *path: str, suffix: str) -> Optional[str]: - val = obj - try: - for p in path: - val = val[p] - except (KeyError, TypeError): - return None - if val is not None: - return f'{val}{suffix}' - return None +logger = logging.getLogger(__name__) -class RateLimiter(abc.ABC): - @abc.abstractmethod - def is_allowed(self, ctx: Context) -> tuple[bool, str]: - pass +cat_api: Optional['AnimalApi'] = None +dog_api: Optional['AnimalApi'] = None -class PerPersonRateLimiter(RateLimiter): - def __init__(self, calls: int, period: dt.timedelta): - self.calls = calls - self.period = period - self.users: dict[int, list[dt.datetime]] = defaultdict(list) +http = aiohttp.ClientSession() +plugins.finalizer(http.close) - def __repr__(self) -> str: - return f'{self.__class__.__name__}<{self.calls=}, {self.period=}>' - - def __str__(self) -> str: - return f'{self.__class__.__name__}<{self.calls} calls per {self.period}>' - - def _is_allowed(self, user_id: int) -> tuple[bool, str]: - now = dt.datetime.now() - timestamps = self.users[user_id] - # Remove timestamps outside the period - while timestamps and now - timestamps[0] > self.period: - timestamps.pop(0) - if len(timestamps) < self.calls: - timestamps.append(now) - return True, '' - else: - next_time = timestamps[0] + self.period - return False, f'Rate limit exceeded: {self.calls} calls per {self.period}. Please try again .' +registry = sqlalchemy.orm.registry() +sessionmaker = async_sessionmaker(util.db.engine, expire_on_commit=False) - def is_allowed(self, ctx: Context): - user_id = ctx.author.id - return self._is_allowed(user_id) - -class GlobalRateLimiter(PerPersonRateLimiter): - def is_allowed(self, ctx: Context): - return self._is_allowed(0) # Use a single user ID for global rate limiting +@registry.mapped +class GlobalConfig: + __tablename__ = "config" + __table_args__ = {"schema": "animals"} + + id: Mapped[int] = mapped_column(primary_key=True, default=0) + cat_api_key: Mapped[str] = mapped_column(TEXT) + dog_api_key: Mapped[str] = mapped_column(TEXT) +class AnimalRequest(abc.ABC): + @abc.abstractmethod + def to_dict(self) -> dict[str, Any]: + ''' Converts the request to a dictionary of query parameters ''' + pass + @dataclass -class CatRequest: - limit: Optional[int] = None - page: Optional[int] = None - order: Literal['ASC', 'DESC', 'RAND', None] = None - has_breeds: Optional[bool] = None - breed_ids: Optional[list[str]] = None - #sub_id: Optional[str] - - def __post_init__(self): - try: - assert self.limit is None or 1 <= self.limit <= 100 - assert self.page is None or 0 <= self.page - except AssertionError: - raise ValueError() +class CatRequest(AnimalRequest): + # https://developers.thecatapi.com/view-account/ylX4blBYT9FaoVd6OhvR + + limit: Optional[int] = None # API defaults to 1 + order: Literal['ASC', 'DESC', 'RAND', None] = None # API defaults to RAND + #page: Optional[int] = None # only relevant for ASC/DESC search + has_breeds: Optional[bool] = None # API defaults to all + breed_ids: Optional[list[str]] = None # API defaults to all def to_dict(self) -> dict[str, Any]: result = {} if self.limit is not None: result['limit'] = self.limit - if self.page is not None: - result['page'] = self.page if self.order is not None: result['order'] = self.order if self.has_breeds is not None: @@ -101,14 +75,16 @@ def to_dict(self) -> dict[str, Any]: return result @dataclass -class DogRequest: +class DogRequest(AnimalRequest): + # https://docs.thedogapi.com/docs/examples/images + size: Literal['full', 'med', 'small', 'thumb'] = 'small' - mime_types: Optional[list[Literal['jpg', 'png', 'gif']]] = None + mime_types: Optional[list[Literal['jpg', 'png', 'gif']]] = None # API defaults to all format: Literal['json', 'src'] = 'json' - order: Literal['ASC', 'DESC', 'RAND', None] = None - limit: Optional[int] = None - page: Optional[int] = None - has_breeds: Optional[bool] = None + order: Literal['ASC', 'DESC', 'RAND', None] = None # API defaults to RAND + limit: Optional[int] = None # API defaults to 1 + #page: Optional[int] = None # only relevant for ASC/DESC search + has_breeds: Optional[bool] = None # API defaults to all def to_dict(self) -> dict[str, Any]: result = {} @@ -118,13 +94,15 @@ def to_dict(self) -> dict[str, Any]: result['format'] = self.format result['order'] = self.order result['limit'] = self.limit - result['page'] = self.page if self.has_breeds is not None: result['has_breeds'] = int(self.has_breeds) return {k: v for k, v in result.items() if v is not None} @dataclass class AnimalResponse: + # We could define separate CatResponse and DogResponse classes, + # but they're similar enough that it's simpler to just have one. + id: str url: str width: int @@ -133,28 +111,37 @@ class AnimalResponse: breeds: Optional[list[dict[str, Any]]] = None def get_weight(self) -> Optional[str]: + # The weight comes back in different formats depending on the API and the breed if self.breeds: breed = self.breeds[0] - return ( - with_suffix(breed, 'weight', 'metric', suffix=' kg') or - with_suffix(breed, 'weight', 'imperial', suffix=' lbs') or - with_suffix(breed, 'weight', suffix='') - ) + try: + return breed['weight']['metric'] + ' kg' + except KeyError: + pass + + try: + return breed['weight']['imperial'] + ' lbs' + except KeyError: + pass + + try: + return breed['weight'] + '' + except KeyError: + pass return None def get_life_span(self) -> Optional[str]: if self.breeds: breed = self.breeds[0] - lifespan = breed.get('life_span', None) - if not lifespan: - return None - if 'years' in lifespan: + try: + lifespan = breed['life_span'] + if 'years' not in lifespan: + lifespan += ' years' return lifespan - else: - return f'{lifespan} years' + except KeyError: + pass return None - def get_description(self) -> str: if self.breeds: breed = self.breeds[0] @@ -174,71 +161,38 @@ def get_description(self) -> str: return 'No breed information available.' class AnimalApi: - def __init__(self, api_root: str, api_key: Optional[str], rate_limiters: list[RateLimiter]) -> None: + # The actual API structure is identical, just with different parameters and API keys. + + def __init__(self, api_root: str, api_key: str) -> None: self.api_root = api_root self.api_key = api_key - self.rate_limiters = rate_limiters - - def configure(self, api_key: str) -> None: - self.api_key = api_key - - def get_configuration(self) -> str: - out = '' - out += f'API Root: {self.api_root}\n' - out += f'API Key: {self.api_key}\n' - for limiter in self.rate_limiters: - out += f'Rate Limiter: {limiter}\n' - return out - - async def fetch_random_animal(self, ctx: Context, params: dict[str, Any]) -> AnimalResponse: - logging.debug(f'Fetching random animal with params: {params}') - if not self.api_key: - raise UserError('Animal API is not configured.') + + async def fetch_random_animal(self, req: Optional[AnimalRequest] = None) -> AnimalResponse: + params = req.to_dict() if req else {} + logger.debug(f'Fetching random animal with params: {params}') headers = {'x-api-key': self.api_key} - for limiter in self.rate_limiters: - allowed, message = limiter.is_allowed(ctx) - if not allowed: - raise UserError(message) - async with aiohttp.ClientSession() as session: - async with session.get(self.api_root, params=params, headers=headers, timeout=TIMEOUT) as resp: - data = await resp.json() - if not data: - raise UserError('No animal found!') - return AnimalResponse(**data[0]) - -class CatApi(AnimalApi): - def __init__(self, api_key: Optional[str], rate_limiters: list[RateLimiter]) -> None: - super().__init__(CAT_API_ROOT, api_key, rate_limiters) - - async def fetch_random_cat(self, ctx: Context, req: CatRequest): - return await self.fetch_random_animal(ctx, req.to_dict()) - -# cat api key should be set via the config command -cat_api = CatApi('', [ - # reasonable rate limits, I don't see a good way to configure these dynamically - PerPersonRateLimiter(calls=3, period=dt.timedelta(minutes=1)), - PerPersonRateLimiter(calls=20, period=dt.timedelta(hours=1)), - PerPersonRateLimiter(calls=100, period=dt.timedelta(days=1)), - - GlobalRateLimiter(calls=10, period=dt.timedelta(minutes=1)), -]) - -class DogApi(AnimalApi): - def __init__(self, api_key: Optional[str], rate_limiters: list[RateLimiter]) -> None: - super().__init__(DOG_API_ROOT, api_key, rate_limiters) - - async def fetch_random_dog(self, ctx: Context, req: DogRequest): - return await self.fetch_random_animal(ctx, req.to_dict()) - -# dog api key should be set via the config command -dog_api = DogApi('', [ - # reasonable rate limits, I don't see a good way to configure these dynamically - PerPersonRateLimiter(calls=3, period=dt.timedelta(minutes=1)), - PerPersonRateLimiter(calls=20, period=dt.timedelta(hours=1)), - PerPersonRateLimiter(calls=100, period=dt.timedelta(days=1)), - - GlobalRateLimiter(calls=10, period=dt.timedelta(minutes=1)), -]) + + async with http.get( + self.api_root, + params=params, + headers=headers, + timeout=TIMEOUT, + ) as resp: + data = await resp.json() + if not data: + raise UserError('No animal found!') + return AnimalResponse(**data[0]) + +def animal_to_embed(animal: AnimalResponse, title: str) -> discord.Embed: + embed = discord.Embed( + title=title, + url=animal.url, + color=discord.Color.random() + ) + embed.set_image(url=animal.url) + footer = animal.get_description() + embed.set_footer(text=footer) + return embed @plugin_command @privileged @@ -247,21 +201,11 @@ async def random_cat(ctx: Context) -> None: ''' Fetches and displays a random cat image ''' async with ctx.typing(): - # Prepare request params - req = CatRequest() - - cat = await cat_api.fetch_random_cat(ctx, req) - embed = discord.Embed( - title='Here is your random cat! 🐱', - url=cat.url, - color=discord.Color.random() - ) - embed.set_image(url=cat.url) - footer = cat.get_description() - if footer: - footer += '\n' - footer += 'thecatapi.com' - embed.set_footer(text=footer) + if not cat_api: + raise UserError('Cat API is not configured.') + + cat = await cat_api.fetch_random_animal() + embed = animal_to_embed(cat, title='Here is your random cat! 🐱') await ctx.send(embed=embed) @plugin_command @@ -271,31 +215,56 @@ async def random_dog(ctx: Context) -> None: ''' Fetches and displays a random dog image ''' async with ctx.typing(): - # Prepare request params - req = DogRequest() - - dog = await dog_api.fetch_random_dog(ctx, req) - embed = discord.Embed( - title='Here is your random dog! 🐶', - url=dog.url, - color=discord.Color.random() - ) - embed.set_image(url=dog.url) - footer = dog.get_description() - if footer: - footer += '\n' - footer += 'thedogapi.com' - embed.set_footer(text=footer) + if not dog_api: + raise UserError('Dog API is not configured.') + + dog = await dog_api.fetch_random_animal() + embed = animal_to_embed(dog, title='Here is your random dog! 🐶') await ctx.send(embed=embed) +@plugins.init +async def init() -> None: + global cat_api, dog_api + await util.db.init(util.db.get_ddl(registry.metadata.create_all)) + + async with sessionmaker() as session: + conf = await session.get(GlobalConfig, 0) + if conf: + cat_api = AnimalApi( + CAT_API_ROOT, + conf.cat_api_key + ) + dog_api = AnimalApi( + DOG_API_ROOT, + conf.dog_api_key + ) + else: + raise UserError('No configuration found for animals plugin.') + @plugin_config_command +@group("animals") @privileged -@group("animals", invoke_without_command=True) -async def animals_config(ctx: Context, cat_api_key: Optional[str], dog_api_key: Optional[str]) -> None: - ''' Configures the animals plugin ''' - if cat_api_key: - cat_api.configure(api_key=cat_api_key) - if dog_api_key: - dog_api.configure(api_key=dog_api_key) - - await ctx.send(f'Animals plugin configuration:\n**Cat**\n{cat_api.get_configuration()}\n**Dog**\n{dog_api.get_configuration()}') \ No newline at end of file +async def config(ctx: Context) -> None: + pass + +@config.command("cat_api_key") +async def set_cat_api_key(ctx: Context, api_key: str) -> None: + global cat_api + cat_api = AnimalApi(CAT_API_ROOT, api_key) + async with sessionmaker() as session: + conf = await session.get(GlobalConfig, 0) + assert conf + conf.cat_api_key = api_key + await session.commit() + await ctx.send(f'\u2705') + +@config.command("dog_api_key") +async def set_dog_api_key(ctx: Context, api_key: str) -> None: + global dog_api + dog_api = AnimalApi(DOG_API_ROOT, api_key) + async with sessionmaker() as session: + conf = await session.get(GlobalConfig, 0) + assert conf + conf.dog_api_key = api_key + await session.commit() + await ctx.send(f'\u2705') \ No newline at end of file From 53803aa6cef9740e72a10b2160c900b9dac6bdbd Mon Sep 17 00:00:00 2001 From: Hayley Date: Mon, 13 Oct 2025 13:20:00 +1100 Subject: [PATCH 3/3] possibly doing the right thing with config?? --- ...9174adad82bc86607c7ed6a5a1fde54f8bd3dc.sql | 9 ++++ plugins/animals.py | 45 +++++++++++-------- 2 files changed, 36 insertions(+), 18 deletions(-) create mode 100644 migrations/plugins.animals-037e8d94bf0df85816f7fb3f37f34267400bab98-219174adad82bc86607c7ed6a5a1fde54f8bd3dc.sql diff --git a/migrations/plugins.animals-037e8d94bf0df85816f7fb3f37f34267400bab98-219174adad82bc86607c7ed6a5a1fde54f8bd3dc.sql b/migrations/plugins.animals-037e8d94bf0df85816f7fb3f37f34267400bab98-219174adad82bc86607c7ed6a5a1fde54f8bd3dc.sql new file mode 100644 index 0000000..2569010 --- /dev/null +++ b/migrations/plugins.animals-037e8d94bf0df85816f7fb3f37f34267400bab98-219174adad82bc86607c7ed6a5a1fde54f8bd3dc.sql @@ -0,0 +1,9 @@ +create schema if not exists animals; +create table if not exists animals.config ( + id integer primary key, + cat_api_key text not null, + dog_api_key text not null +); +insert into animals.config (id, cat_api_key, dog_api_key) + values (0, '', '') + on conflict (id) do nothing; \ No newline at end of file diff --git a/plugins/animals.py b/plugins/animals.py index 55b9ef8..de16c14 100644 --- a/plugins/animals.py +++ b/plugins/animals.py @@ -13,7 +13,7 @@ from discord.ext.commands import command import aiohttp -from sqlalchemy import TEXT +from sqlalchemy import TEXT, BigInteger, Computed from sqlalchemy.orm import Mapped, mapped_column import sqlalchemy.orm from sqlalchemy.ext.asyncio import async_sessionmaker @@ -26,10 +26,7 @@ DOG_API_ROOT = 'https://api.thedogapi.com/v1/images/search' TIMEOUT = 10 # seconds -logger = logging.getLogger(__name__) -cat_api: Optional['AnimalApi'] = None -dog_api: Optional['AnimalApi'] = None http = aiohttp.ClientSession() plugins.finalizer(http.close) @@ -42,10 +39,22 @@ class GlobalConfig: __tablename__ = "config" __table_args__ = {"schema": "animals"} - id: Mapped[int] = mapped_column(primary_key=True, default=0) + id: Mapped[int] = mapped_column(BigInteger, Computed("0"), primary_key=True) cat_api_key: Mapped[str] = mapped_column(TEXT) dog_api_key: Mapped[str] = mapped_column(TEXT) + def __str__(self) -> str: + out = 'Animal plugin configuration:\n' + out += f'Cat API key: {"set" if self.cat_api_key else "not set"}\n' + out += f'Dog API key: {"set" if self.dog_api_key else "not set"}\n' + return out + +conf: GlobalConfig +logger = logging.getLogger(__name__) + +cat_api: Optional['AnimalApi'] = None +dog_api: Optional['AnimalApi'] = None + class AnimalRequest(abc.ABC): @abc.abstractmethod def to_dict(self) -> dict[str, Any]: @@ -224,22 +233,22 @@ async def random_dog(ctx: Context) -> None: @plugins.init async def init() -> None: - global cat_api, dog_api + global cat_api, dog_api, conf await util.db.init(util.db.get_ddl(registry.metadata.create_all)) async with sessionmaker() as session: - conf = await session.get(GlobalConfig, 0) - if conf: - cat_api = AnimalApi( - CAT_API_ROOT, - conf.cat_api_key - ) - dog_api = AnimalApi( - DOG_API_ROOT, - conf.dog_api_key - ) - else: - raise UserError('No configuration found for animals plugin.') + c = await session.get(GlobalConfig, 0) + assert c, 'No configuration found for animals plugin.' + conf = c + logger.info(f'Loaded animal plugin configuration: {conf.cat_api_key=}, {conf.dog_api_key=}') + cat_api = AnimalApi( + CAT_API_ROOT, + conf.cat_api_key + ) + dog_api = AnimalApi( + DOG_API_ROOT, + conf.dog_api_key + ) @plugin_config_command @group("animals")