#!/usr/bin/env python3 import aiohttp import aiohttp.web import aiomas import asyncio import mimeparse import sqlalchemy import sys import json import datetime import pytz import os import common.rpc import common.postgres from common.config import config class Poison: pass class Server(common.rpc.Server): router = aiomas.rpc.Service() def __init__(self): super().__init__() self.engine, self.metadata = common.postgres.get_engine_and_metadata() self.queues = [] async def negotiate(self, request): request.headers.getall('Accept', "*/*") mime_type = mimeparse.best_match(['application/json', 'text/event-stream'], ",".join(request.headers.getall('Accept', "*/*"))) if mime_type == 'text/event-stream': return await self.event_stream(request) elif mime_type == 'application/json': return await self.json(request) else: raise NotImplementedError(mime_type) def get_last_events(self, request): try: last_event_id = int(request.headers.get('Last-Event-Id', request.query.get('last-event-id'))) except (ValueError, TypeError): last_event_id = None interval = request.query.get('interval') if interval is not None and last_event_id is None: last_event_id = 0 if last_event_id is not None: events = self.metadata.tables['events'] query = sqlalchemy.select([ events.c.id, events.c.event, events.c.data, events.c.time ]) query = query.where(events.c.id > last_event_id) if interval is not None: query = query.where(events.c.time > sqlalchemy.func.current_timestamp() - sqlalchemy.cast(interval, sqlalchemy.Interval)) query = query.order_by(events.c.id) try: with self.engine.begin() as conn: return [ {'id': id, 'event': event, 'data': dict(data, time=time.isoformat())} for id, event, data, time in conn.execute(query) ] except sqlalchemy.exc.DataError as e: raise aiohttp.web.HTTPBadRequest from e return [] async def event_stream(self, request): queue = asyncio.Queue() for event in self.get_last_events(request): await queue.put(event) self.queues.append(queue) response = aiohttp.web.StreamResponse() response.enable_chunked_encoding() response.headers['Access-Control-Allow-Origin'] = '*' response.headers['Content-Type'] = 'text/event-stream; charset=utf-8' response.headers['Vary'] = "Accept" await response.prepare(request) while True: try: try: event = await asyncio.wait_for(queue.get(), 15) if event['event'] is Poison: break await response.write(b"id:%d\n" % event['id']) await response.write(b"event:%s\n" % event['event'].encode('utf-8')) await response.write(b"data:%s\n" % json.dumps(event['data']).encode('utf-8')) await response.write(b"\n") queue.task_done() except asyncio.TimeoutError: await response.write(b":keep-alive\n\n") except IOError: break self.queues.remove(queue) return response async def json(self, request): return aiohttp.web.json_response({ 'events': self.get_last_events(request), }, headers={"Vary": "Accept", 'Access-Control-Allow-Origin': request.headers.get('Origin', '*')}) async def cors_preflight(self, request): return aiohttp.web.Response(headers={ 'Access-Control-Allow-Origin': request.headers.get('Origin', '*'), }) @aiomas.expose async def event(self, event, data, time=None): if time is None: time = datetime.datetime.now(tz=pytz.utc) events = self.metadata.tables['events'] with self.engine.begin() as conn: id, = conn.execute(events.insert().returning(events.c.id), event=event, data=data, time=time, ).first() event = { 'id': id, 'event': event, 'data': dict(data, time=time.isoformat()), } for queue in self.queues: await queue.put(event) async def on_shutdown(self, app): for queue in self.queues: await queue.put({'event': Poison}) server = None srv = None handler = None app = None async def main(loop): global server, srv, app, handler try: os.unlink(config['eventsocket']) except FileNotFoundError: pass server = Server() await server.start(config['eventsocket'], config['event_port']) app = aiohttp.web.Application() app.router.add_route('GET', '/api/v2/events', server.negotiate) app.router.add_route('OPTIONS', '/api/v2/events', server.cors_preflight) app.on_shutdown.append(server.on_shutdown) handler = app.make_handler() srv = await loop.create_server(handler, 'localhost', 8080) if sys.platform == "win32": # On Windows Ctrl+C doesn't interrupt `select()`. def windows_is_butts(): asyncio.get_event_loop().call_later(5, windows_is_butts) windows_is_butts() async def cleanup(): global server, srv, app, handler srv.close() await server.close() await srv.wait_closed() await app.shutdown() await handler.finish_connections(60.0) await app.cleanup() loop = asyncio.get_event_loop() loop.run_until_complete(main(loop)) try: loop.run_forever() except KeyboardInterrupt: pass finally: loop.run_until_complete(cleanup()) loop.close()