From 944f0093a919a0daff40b32cdbdc247761aa42e9 Mon Sep 17 00:00:00 2001 From: Ske Date: Thu, 2 Aug 2018 00:36:50 +0200 Subject: [PATCH] Add basic HTTP API --- docker-compose.yml | 17 ++- src/{bot.Dockerfile => Dockerfile} | 1 - src/api_main.py | 122 ++++++++++++++++++ src/pluralkit/__init__.py | 24 +++- src/pluralkit/api/__init__.py | 0 src/pluralkit/bot/commands/misc_commands.py | 3 +- src/pluralkit/bot/commands/switch_commands.py | 7 +- src/pluralkit/bot/commands/system_commands.py | 5 +- src/pluralkit/bot/utils.py | 41 +----- src/pluralkit/db.py | 12 ++ src/pluralkit/stats.py | 12 ++ src/pluralkit/utils.py | 44 +++++++ src/requirements.txt | 1 + 13 files changed, 238 insertions(+), 51 deletions(-) rename src/{bot.Dockerfile => Dockerfile} (82%) create mode 100644 src/api_main.py create mode 100644 src/pluralkit/api/__init__.py create mode 100644 src/pluralkit/utils.py diff --git a/docker-compose.yml b/docker-compose.yml index 93d94be2..6da0b3eb 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -1,9 +1,10 @@ version: '3' services: bot: - build: - context: src/ - dockerfile: bot.Dockerfile + build: src/ + entrypoint: + - python + - bot_main.py depends_on: - db - influx @@ -11,6 +12,16 @@ services: - CLIENT_ID - TOKEN restart: always + api: + build: src/ + entrypoint: + - python + - api_main.py + depends_on: + - db + restart: always + ports: + - "2939:8080" db: image: postgres:alpine volumes: diff --git a/src/bot.Dockerfile b/src/Dockerfile similarity index 82% rename from src/bot.Dockerfile rename to src/Dockerfile index 4d6ced1f..204cad22 100644 --- a/src/bot.Dockerfile +++ b/src/Dockerfile @@ -7,5 +7,4 @@ ADD requirements.txt /app RUN pip install --trusted-host pypi.python.org -r requirements.txt ADD . /app -ENTRYPOINT ["python", "bot_main.py"] diff --git a/src/api_main.py b/src/api_main.py new file mode 100644 index 00000000..35a89338 --- /dev/null +++ b/src/api_main.py @@ -0,0 +1,122 @@ +import logging + +from aiohttp import web + +from pluralkit import db, utils + +logging.basicConfig(level=logging.INFO, format="[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s") +logger = logging.getLogger("pluralkit.api") + + +def db_handler(f): + async def inner(request): + async with request.app["pool"].acquire() as conn: + return await f(request, conn) + + return inner + + +@db_handler +async def get_system(request: web.Request, conn): + system = await db.get_system_by_hid(conn, request.match_info["id"]) + + if not system: + raise web.HTTPNotFound() + + members = await db.get_all_members(conn, system.id) + + system_json = system.to_json() + system_json["members"] = [member.to_json() for member in members] + return web.json_response(system_json) + + +@db_handler +async def get_member(request: web.Request, conn): + member = await db.get_member_by_hid(conn, request.match_info["id"]) + + if not member: + raise web.HTTPNotFound() + + return web.json_response(member.to_json()) + + +@db_handler +async def get_switches(request: web.Request, conn): + system = await db.get_system_by_hid(conn, request.match_info["id"]) + + if not system: + raise web.HTTPNotFound() + + switches = await utils.get_front_history(conn, system.id, 99999) + + data = [{ + "timestamp": stamp.isoformat(), + "members": [member.hid for member in members] + } for stamp, members in switches] + + return web.json_response(data) + +@db_handler +async def get_message(request: web.Request, conn): + message = await db.get_message(conn, request.match_info["id"]) + if not message: + raise web.HTTPNotFound() + + return web.json_response(message.to_json()) + +@db_handler +async def get_switch(request: web.Request, conn): + system = await db.get_system_by_hid(conn, request.match_info["id"]) + + if not system: + raise web.HTTPNotFound() + + members, stamp = await utils.get_fronters(conn, system.id) + if not stamp: + # No switch has been registered at all + raise web.HTTPNotFound() + + data = { + "timestamp": stamp.isoformat(), + "members": [member.to_json() for member in members] + } + return web.json_response(data) + +@db_handler +async def get_switch_name(request: web.Request, conn): + system = await db.get_system_by_hid(conn, request.match_info["id"]) + + if not system: + raise web.HTTPNotFound() + + members, stamp = await utils.get_fronters(conn, system.id) + return web.Response(text=members[0].name if members else "(nobody)") + +@db_handler +async def get_switch_color(request: web.Request, conn): + system = await db.get_system_by_hid(conn, request.match_info["id"]) + + if not system: + raise web.HTTPNotFound() + + members, stamp = await utils.get_fronters(conn, system.id) + return web.Response(text=members[0].color if members else "#ffffff") + +app = web.Application() +app.add_routes([ + web.get("/systems/{id}", get_system), + web.get("/systems/{id}/switches", get_switches), + web.get("/systems/{id}/switch", get_switch), + web.get("/systems/{id}/switch/name", get_switch_name), + web.get("/systems/{id}/switch/color", get_switch_color), + web.get("/members/{id}", get_member), + web.get("/messages/{id}", get_message) +]) + + +async def run(): + app["pool"] = await db.connect() + return app + + +web.run_app(run()) diff --git a/src/pluralkit/__init__.py b/src/pluralkit/__init__.py index d933762f..0b12fe6d 100644 --- a/src/pluralkit/__init__.py +++ b/src/pluralkit/__init__.py @@ -11,6 +11,15 @@ class System(namedtuple("System", ["id", "hid", "name", "description", "tag", "a avatar_url: str created: datetime + def to_json(self): + return { + "id": self.hid, + "name": self.name, + "description": self.description, + "tag": self.tag, + "avatar_url": self.avatar_url + } + class Member(namedtuple("Member", ["id", "hid", "system", "color", "avatar_url", "name", "birthday", "pronouns", "description", "prefix", "suffix", "created"])): id: int hid: str @@ -23,4 +32,17 @@ class Member(namedtuple("Member", ["id", "hid", "system", "color", "avatar_url", description: str prefix: str suffix: str - created: datetime \ No newline at end of file + created: datetime + + def to_json(self): + return { + "id": self.hid, + "name": self.name, + "color": self.color, + "avatar_url": self.avatar_url, + "birthday": self.birthday.isoformat() if self.birthday else None, + "pronouns": self.pronouns, + "description": self.description, + "prefix": self.prefix, + "suffix": self.suffix + } \ No newline at end of file diff --git a/src/pluralkit/api/__init__.py b/src/pluralkit/api/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/pluralkit/bot/commands/misc_commands.py b/src/pluralkit/bot/commands/misc_commands.py index 053238db..c8b443fe 100644 --- a/src/pluralkit/bot/commands/misc_commands.py +++ b/src/pluralkit/bot/commands/misc_commands.py @@ -6,6 +6,7 @@ from typing import List from discord.utils import oauth_url +import pluralkit.utils from pluralkit.bot import utils from pluralkit.bot.commands import * @@ -52,7 +53,7 @@ async def invite_link(ctx: CommandContext, args: List[str]): async def export(ctx: CommandContext, args: List[str]): members = await db.get_all_members(ctx.conn, ctx.system.id) accounts = await db.get_linked_accounts(ctx.conn, ctx.system.id) - switches = await utils.get_front_history(ctx.conn, ctx.system.id, 999999) + switches = await pluralkit.utils.get_front_history(ctx.conn, ctx.system.id, 999999) system = ctx.system data = { diff --git a/src/pluralkit/bot/commands/switch_commands.py b/src/pluralkit/bot/commands/switch_commands.py index 3ead2a6c..e4ee781c 100644 --- a/src/pluralkit/bot/commands/switch_commands.py +++ b/src/pluralkit/bot/commands/switch_commands.py @@ -5,6 +5,7 @@ from typing import List import dateparser import humanize +import pluralkit.utils from pluralkit import Member from pluralkit.bot import utils from pluralkit.bot.commands import * @@ -27,7 +28,7 @@ async def switch_member(ctx: MemberCommandContext, args: List[str]): # Compare requested switch IDs and existing fronter IDs to check for existing switches # Lists, because order matters, it makes sense to just swap fronters member_ids = [member.id for member in members] - fronter_ids = (await utils.get_fronter_ids(ctx.conn, ctx.system.id))[0] + fronter_ids = (await pluralkit.utils.get_fronter_ids(ctx.conn, ctx.system.id))[0] if member_ids == fronter_ids: if len(members) == 1: raise CommandError("{} is already fronting.".format(members[0].name)) @@ -51,7 +52,7 @@ async def switch_member(ctx: MemberCommandContext, args: List[str]): @command(cmd="switch out", description="Registers a switch with no one in front.", category="Switching commands") async def switch_out(ctx: MemberCommandContext, args: List[str]): # Get current fronters - fronters, _ = await utils.get_fronter_ids(ctx.conn, system_id=ctx.system.id) + fronters, _ = await pluralkit.utils.get_fronter_ids(ctx.conn, system_id=ctx.system.id) if not fronters: raise CommandError("There's already no one in front.") @@ -79,7 +80,7 @@ async def switch_move(ctx: MemberCommandContext, args: List[str]): # Make sure it all runs in a big transaction for atomicity async with ctx.conn.transaction(): # Get the last two switches to make sure the switch to move isn't before the second-last switch - last_two_switches = await utils.get_front_history(ctx.conn, ctx.system.id, count=2) + last_two_switches = await pluralkit.utils.get_front_history(ctx.conn, ctx.system.id, count=2) if len(last_two_switches) == 0: raise CommandError("There are no registered switches for this system.") diff --git a/src/pluralkit/bot/commands/system_commands.py b/src/pluralkit/bot/commands/system_commands.py index b2e0412e..91d5b563 100644 --- a/src/pluralkit/bot/commands/system_commands.py +++ b/src/pluralkit/bot/commands/system_commands.py @@ -4,6 +4,7 @@ from urllib.parse import urlparse import humanize +import pluralkit.utils from pluralkit.bot import utils from pluralkit.bot.commands import * @@ -153,7 +154,7 @@ async def system_fronter(ctx: CommandContext, args: List[str]): if system is None: raise CommandError("Can't find system \"{}\".".format(args[0])) - fronters, timestamp = await utils.get_fronters(ctx.conn, system_id=system.id) + fronters, timestamp = await pluralkit.utils.get_fronters(ctx.conn, system_id=system.id) fronter_names = [member.name for member in fronters] embed = utils.make_default_embed(None) @@ -182,7 +183,7 @@ async def system_fronthistory(ctx: CommandContext, args: List[str]): raise CommandError("Can't find system \"{}\".".format(args[0])) lines = [] - front_history = await utils.get_front_history(ctx.conn, system.id, count=10) + front_history = await pluralkit.utils.get_front_history(ctx.conn, system.id, count=10) for i, (timestamp, members) in enumerate(front_history): # Special case when no one's fronting if len(members) == 0: diff --git a/src/pluralkit/bot/utils.py b/src/pluralkit/bot/utils.py index e75b8075..9d4cf626 100644 --- a/src/pluralkit/bot/utils.py +++ b/src/pluralkit/bot/utils.py @@ -1,16 +1,13 @@ -from datetime import datetime import logging import random import re -from typing import List, Tuple import string -import asyncio -import asyncpg import discord import humanize from pluralkit import System, Member, db +from pluralkit.utils import get_fronters logger = logging.getLogger("pluralkit.utils") @@ -53,42 +50,6 @@ def parse_channel_mention(mention: str, server: discord.Server) -> discord.Chann except ValueError: return None -async def get_fronter_ids(conn, system_id) -> (List[int], datetime): - switches = await db.front_history(conn, system_id=system_id, count=1) - if not switches: - return [], None - - if not switches[0]["members"]: - return [], switches[0]["timestamp"] - - return switches[0]["members"], switches[0]["timestamp"] - -async def get_fronters(conn, system_id) -> (List[Member], datetime): - member_ids, timestamp = await get_fronter_ids(conn, system_id) - - # Collect in dict and then look up as list, to preserve return order - members = {member.id: member for member in await db.get_members(conn, member_ids)} - return [members[member_id] for member_id in member_ids], timestamp - -async def get_front_history(conn, system_id, count) -> List[Tuple[datetime, List[Member]]]: - # Get history from DB - switches = await db.front_history(conn, system_id=system_id, count=count) - if not switches: - return [] - - # Get all unique IDs referenced - all_member_ids = {id for switch in switches for id in switch["members"]} - - # And look them up in the database into a dict - all_members = {member.id: member for member in await db.get_members(conn, list(all_member_ids))} - - # Collect in array and return - out = [] - for switch in switches: - timestamp = switch["timestamp"] - members = [all_members[id] for id in switch["members"]] - out.append((timestamp, members)) - return out async def get_system_fuzzy(conn, client: discord.Client, key) -> System: if isinstance(key, discord.User): diff --git a/src/pluralkit/db.py b/src/pluralkit/db.py index 9f5c8f6b..91d4ea49 100644 --- a/src/pluralkit/db.py +++ b/src/pluralkit/db.py @@ -6,6 +6,7 @@ import time import asyncpg import asyncpg.exceptions +from discord.utils import snowflake_time from pluralkit import System, Member, stats @@ -210,6 +211,17 @@ class MessageInfo(namedtuple("MemberInfo", ["mid", "channel", "member", "content system_name: str system_hid: str + def to_json(self): + return { + "id": str(self.mid), + "channel": str(self.channel), + "member": self.hid, + "system": self.system_hid, + "message_sender": str(self.sender), + "content": self.content, + "timestamp": snowflake_time(self.mid).isoformat() + } + @db_wrap async def get_message_by_sender_and_id(conn, message_id: str, sender_id: str) -> MessageInfo: row = await conn.fetchrow("""select diff --git a/src/pluralkit/stats.py b/src/pluralkit/stats.py index b832529d..63aa873b 100644 --- a/src/pluralkit/stats.py +++ b/src/pluralkit/stats.py @@ -7,6 +7,9 @@ async def connect(): await client.create_database(db="pluralkit") async def report_db_query(query_name, time, success): + if not client: + return + await client.write({ "measurement": "database_query", "tags": {"query": query_name}, @@ -14,6 +17,9 @@ async def report_db_query(query_name, time, success): }) async def report_command(command_name, execution_time, response_time): + if not client: + return + await client.write({ "measurement": "command", "tags": {"command": command_name}, @@ -21,12 +27,18 @@ async def report_command(command_name, execution_time, response_time): }) async def report_webhook(time, success): + if not client: + return + await client.write({ "measurement": "webhook", "fields": {"response_time": time, "success": int(success)} }) async def report_periodical_stats(conn): + if not client: + return + from pluralkit import db systems = await db.system_count(conn) diff --git a/src/pluralkit/utils.py b/src/pluralkit/utils.py new file mode 100644 index 00000000..cfa2da9e --- /dev/null +++ b/src/pluralkit/utils.py @@ -0,0 +1,44 @@ +from datetime import datetime +from typing import List, Tuple + +from pluralkit import db, Member + + +async def get_fronter_ids(conn, system_id) -> (List[int], datetime): + switches = await db.front_history(conn, system_id=system_id, count=1) + if not switches: + return [], None + + if not switches[0]["members"]: + return [], switches[0]["timestamp"] + + return switches[0]["members"], switches[0]["timestamp"] + + +async def get_fronters(conn, system_id) -> (List[Member], datetime): + member_ids, timestamp = await get_fronter_ids(conn, system_id) + + # Collect in dict and then look up as list, to preserve return order + members = {member.id: member for member in await db.get_members(conn, member_ids)} + return [members[member_id] for member_id in member_ids], timestamp + + +async def get_front_history(conn, system_id, count) -> List[Tuple[datetime, List[Member]]]: + # Get history from DB + switches = await db.front_history(conn, system_id=system_id, count=count) + if not switches: + return [] + + # Get all unique IDs referenced + all_member_ids = {id for switch in switches for id in switch["members"]} + + # And look them up in the database into a dict + all_members = {member.id: member for member in await db.get_members(conn, list(all_member_ids))} + + # Collect in array and return + out = [] + for switch in switches: + timestamp = switch["timestamp"] + members = [all_members[id] for id in switch["members"]] + out.append((timestamp, members)) + return out \ No newline at end of file diff --git a/src/requirements.txt b/src/requirements.txt index f271b53c..dd796da9 100644 --- a/src/requirements.txt +++ b/src/requirements.txt @@ -1,3 +1,4 @@ +aiodns aiohttp aioinflux asyncpg