Skip to content

Commit

Permalink
feat(services): split out create bot to function
Browse files Browse the repository at this point in the history
  • Loading branch information
@RobertRosca
RobertRosca committed Sep 24, 2024
1 parent 596bd58 commit d1f4372
Showing 1 changed file with 88 additions and 76 deletions.
164 changes: 88 additions & 76 deletions src/zulip_write_only_proxy/services.py
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import datetime
from typing import TYPE_CHECKING, Annotated

import fastapi
Expand All @@ -10,8 +11,8 @@
if TYPE_CHECKING:
from .settings import Settings

CLIENT_REPO: repositories.BaseRepository = None # type: ignore[assignment]
ZULIPRC_REPO: repositories.BaseRepository = None # type: ignore[assignment]
CLIENT_REPO: repositories.BaseRepository[models.ScopedClient] = None # type: ignore[assignment]
ZULIPRC_REPO: repositories.BaseRepository[models.BotConfig] = None # type: ignore[assignment]


async def configure(settings: "Settings", _: fastapi.FastAPI | None):
Expand All @@ -20,11 +21,13 @@ async def configure(settings: "Settings", _: fastapi.FastAPI | None):
global CLIENT_REPO, ZULIPRC_REPO

ZULIPRC_REPO = repositories.BaseRepository(
file=settings.config_dir / "zuliprc.json", model=models.BotConfig
file=settings.config_dir / "zuliprc.json",
model=models.BotConfig, # type: ignore[assignment]
)

CLIENT_REPO = repositories.BaseRepository(
file=settings.config_dir / "clients.json", model=models.ScopedClient
file=settings.config_dir / "clients.json",
model=models.ScopedClient, # type: ignore[assignment]
)

logger.info(
Expand All @@ -35,6 +38,61 @@ async def configure(settings: "Settings", _: fastapi.FastAPI | None):
await ZULIPRC_REPO.load()


async def get_or_create_bot(
proposal_no: int,
bot_email: str | None = None,
bot_key: str | None = None,
bot_site: str = "https://mylog.connect.xfel.eu/",
bot_id: int | None = None,
):
created_at = None

if bot_site and bot_id and (bot := await ZULIPRC_REPO.get(f"{bot_site}/{bot_id}")):
return bot

if not bot_email or not bot_key:
res = await mymdc.CLIENT.get_zulip_bot_credentials(proposal_no)
bot_email = res.get("bot_email")
bot_key = res.get("bot_key")

if not bot_email or not bot_key:
raise fastapi.HTTPException(
status_code=422,
detail=(
f"A bot could not be found for proposal '{proposal_no}' via "
"MyMdC. To add a client with a new bot provide both bot_email bot_key."
),
)

if not bot_id:
profile = zulip.Client(
email=bot_email, api_key=bot_key, site=bot_site
).get_profile()
bot_id = profile.get("user_id")
created_at = profile.get("date_joined")

if not bot_id:
raise fastapi.HTTPException(
status_code=422,
detail=(
f"could not get bot id from zulip for bot '{bot_email=}, {bot_site=}'."
),
)

bot = models.BotConfig(
id=bot_id,
key=SecretStr(bot_key),
email=bot_email,
site=Url(bot_site),
created_at=created_at or datetime.datetime.utcnow(),
proposal_no=proposal_no,
)

await ZULIPRC_REPO.insert(bot)

return bot


async def create_client(
new_client: Annotated[models.ScopedClientCreate, fastapi.Depends()],
created_by: str,
Expand All @@ -47,64 +105,19 @@ async def create_client(
)
logger.debug("Stream name from MyMdC", stream=new_client.stream)

name = new_client.bot_name or new_client.proposal_no
key, email, site = (new_client.bot_key, new_client.bot_email, new_client.bot_site)
_id = None

if name not in await ZULIPRC_REPO.list():
logger.debug("Bot zuliprc not present")
if not all((key, email, site)):
_id, key, email, site = await mymdc.CLIENT.get_zulip_bot_credentials(
new_client.proposal_no
)
logger.debug(
"Bot credentials from MyMdC",
bot_name=name,
bot_email=email,
bot_key=key,
bot_id=_id,
)

if not key or not email:
raise fastapi.HTTPException(
status_code=422,
detail=(
f"bot '{name}' does not exist, and a bot could not "
f"be found for proposal '{new_client.proposal_no}' via MyMdC. To "
"add a client with a new bot provide both bot_email bot_key."
),
)

if not _id:
profile = zulip.Client(email=email, api_key=key, site=site).get_profile()
_id = profile.get("user_id")
if not _id:
raise fastapi.HTTPException(
status_code=422,
detail=(
f"could not get bot id from zulip for bot '{name=}, {email=}'."
),
)

bot = models.BotConfig(
id=_id,
key=SecretStr(key),
email=email,
site=Url(site),
created_at=new_client.created_at,
proposal_no=new_client.proposal_no,
)

await ZULIPRC_REPO.insert(bot)
else:
bot = await ZULIPRC_REPO.get(str(name))
bot = await get_or_create_bot(
new_client.proposal_no, bot_id=new_client.bot_id, bot_site=new_client.bot_site
)

client = models.ScopedClient.model_validate({
**new_client.model_dump(),
"created_by": created_by,
"bot_id": bot.id,
"bot_site": bot.site,
})
client = models.ScopedClient(
proposal_no=new_client.proposal_no,
stream=new_client.stream,
bot_id=bot.id,
bot_site=bot.site,
token=new_client.token,
created_at=new_client.created_at,
created_by=created_by,
)

await CLIENT_REPO.insert(client)

Expand All @@ -117,33 +130,32 @@ async def get_client(key: str | None) -> models.ScopedClient:
if key is None:
raise fastapi.HTTPException(status_code=403, detail="Not authenticated")

try:
client = await CLIENT_REPO.get(key, by="token")
except KeyError as e:
client = await CLIENT_REPO.get(key, by="token")

if client is None:
raise fastapi.HTTPException(
status_code=401, detail="Unauthorised", headers={"HX-Location": "/"}
) from e
)

_bot_key = f"{client.bot_site.host}/{client.bot_id}"
try:
bot_config = await ZULIPRC_REPO.get(_bot_key)
except KeyError as e:
bot = await ZULIPRC_REPO.get(client._bot_key)

if bot is None:
raise fastapi.HTTPException(
status_code=401, detail=f"Bot configuration not found for {_bot_key}"
) from e
status_code=401, detail=f"Bot configuration not found for {client._bot_key}"
)

client._client = zulip.Client(
email=bot_config.email,
api_key=bot_config.key.get_secret_value(),
site=str(bot_config.site),
email=bot.email,
api_key=bot.key.get_secret_value(),
site=str(bot.site),
)

return client


async def get_bot(bot_name: str) -> models.BotConfig:
async def get_bot(bot_name: str) -> models.BotConfig | None:
return await ZULIPRC_REPO.get(bot_name)


async def list_clients() -> list[models.ScopedClientWithToken]:
async def list_clients() -> list[models.ScopedClient]:
return await CLIENT_REPO.list()

0 comments on commit d1f4372

Please sign in to comment.