diff --git a/garven/calculate_dev_cluster.py b/garven/calculate_dev_cluster.py index 0ee8b03..e63281e 100644 --- a/garven/calculate_dev_cluster.py +++ b/garven/calculate_dev_cluster.py @@ -1,10 +1,9 @@ import os -def get_dev_cluster() -> str: - guild_id = 601219766258106399 +def get_cluster_and_shard_for_guild_id(guild_id: int) -> tuple[str, str]: shard_id = (guild_id >> 22) % int(os.environ["TOTAL_SHARDS"]) - number_of_shards_per_cluster = 5 + number_of_shards_per_cluster = int(os.environ["SHARDS_PER_CLUSTER"]) clusters = { cid: [ i @@ -14,11 +13,16 @@ def get_dev_cluster() -> str: + number_of_shards_per_cluster, ) ] - for cid in range(1, 20) + for cid in range(1, 25) # Arb number of clusters } for cluster, shards in clusters.items(): if shard_id in shards: - return str(cluster) + return str(cluster), str(shard_id) raise ValueError("Maths went wrong") + + +def get_dev_cluster() -> str: + cluster, _ = get_cluster_and_shard_for_guild_id(601219766258106399) + return cluster diff --git a/garven/routers/aggregate.py b/garven/routers/aggregate.py index 5fae810..58b7ea0 100644 --- a/garven/routers/aggregate.py +++ b/garven/routers/aggregate.py @@ -9,8 +9,9 @@ from fastapi import Request from zonis import RequestFailed +from garven.calculate_dev_cluster import get_cluster_and_shard_for_guild_id from garven.dependencies import get_auth_header -from garven.schema import Statistic, CachedItemsStatistic +from garven.schema import Statistic, CachedItemsStatistic, ShardInfo, Message if TYPE_CHECKING: from zonis.server import Server @@ -21,11 +22,28 @@ prefix="/aggregate", dependencies=[Depends(get_auth_header)], tags=["Aggregate"], - responses={200: {"model": Statistic}}, ) -@aggregate_router.get("/guilds/count", description="Fetch an up to date guild count.") +@aggregate_router.get( + "/guilds/{guild_id}/shard_info", + responses={200: {"model": ShardInfo}, 400: {"model": Message}}, + description="Fetch the shard and cluster for a guild", +) +async def get_guild_shard_info(request: Request, guild_id: int): + try: + cluster_id, shard_id = get_cluster_and_shard_for_guild_id(guild_id) + except ValueError: + return Message(message="guild_id was incorrect / out of range") + + return ShardInfo(shard_id=shard_id, cluster_id=cluster_id) + + +@aggregate_router.get( + "/guilds/count", + description="Fetch an up to date guild count.", + responses={200: {"model": Statistic}}, +) async def guild_count(request: Request): z: Server = request.app.zonis statistic = Statistic(statistic=0) diff --git a/garven/schema/__init__.py b/garven/schema/__init__.py index 2ed62af..bd2aa4a 100644 --- a/garven/schema/__init__.py +++ b/garven/schema/__init__.py @@ -1,3 +1,3 @@ -from .message import Message +from .message import Message, ShardInfo from .statistic import Statistic, CachedItemsStatistic from .ratelimited import RateLimited diff --git a/garven/schema/message.py b/garven/schema/message.py index 167e583..fae8657 100644 --- a/garven/schema/message.py +++ b/garven/schema/message.py @@ -10,3 +10,11 @@ class Config: "message": "Invalid API key.", } } + + +class ShardInfo(BaseModel): + shard_id: str + cluster_id: str + + class Config: + schema_extra = {"example": {"shard_id": "9", "cluster_id": "1"}} diff --git a/garven/schema/premium/__init__.py b/garven/schema/premium/__init__.py deleted file mode 100644 index 7b742f2..0000000 --- a/garven/schema/premium/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .shared_guilds import SharedGuildsRequest, SharedGuildsResponse diff --git a/garven/schema/premium/shared_guilds.py b/garven/schema/premium/shared_guilds.py deleted file mode 100644 index e1f7075..0000000 --- a/garven/schema/premium/shared_guilds.py +++ /dev/null @@ -1,18 +0,0 @@ -from pydantic import BaseModel, Field - - -class SharedGuildsRequest(BaseModel): - guild_ids: list[int] = Field( - description="A list of guild ids to check if the bot is in." - ) - - -class SharedGuildsResponse(BaseModel): - shared_guilds: list[int] = Field( - description="A list of guild id's that the bot is also in." - ) - partial_response: bool = Field( - False, - description="This will be true when the returned data " - "does not accurately represent the entire expected dataset.", - )