Add basic HTTP API

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:

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"]
 

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())

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
@@ -24,3 +33,16 @@ class Member(namedtuple("Member", ["id", "hid", "system", "color", "avatar_url",
     prefix: str
     suffix: str
     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
+        }

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 = {

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.")
 

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:

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):

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

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)

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

src/requirements.txt

@@ -1,3 +1,4 @@
+aiodns
 aiohttp
 aioinflux
 asyncpg