98 lines
2.6 KiB
Python
98 lines
2.6 KiB
Python
import logging
|
|
|
|
from typing import List
|
|
|
|
from .consts import ADAPTERS
|
|
from .worker import ServiceWorker
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
_COLUMNS = {
|
|
'timestamp': 'timestamp bigint',
|
|
'status': 'status bool',
|
|
'latency': 'latency bigint',
|
|
}
|
|
|
|
|
|
class ServiceManager:
|
|
def __init__(self, app):
|
|
self.app = app
|
|
self.cfg = app.cfg
|
|
self.conn = app.conn
|
|
self.loop = app.loop
|
|
|
|
self.workers = {}
|
|
self.state = {}
|
|
self.subscribers = {}
|
|
|
|
self._start()
|
|
|
|
def _make_db_table(self, name: str, service: dict):
|
|
adapter = ADAPTERS[service['adapter']]
|
|
|
|
columnstr = map(_COLUMNS.get, adapter.spec['db'])
|
|
columnstr = ',\n'.join(columnstr)
|
|
|
|
log.info(f'Making table for {name}')
|
|
self.conn.executescript(f"""
|
|
CREATE TABLE IF NOT EXISTS {name} (
|
|
{columnstr}
|
|
);
|
|
""")
|
|
|
|
def _check(self, columns: tuple, field: str, chan_name: str = None):
|
|
if chan_name is None:
|
|
chan_name = field
|
|
|
|
if field in columns and chan_name not in self.subscribers:
|
|
self.subscribers[chan_name] = []
|
|
log.info(f'Created channel {chan_name}')
|
|
|
|
def _create_channels(self, worker):
|
|
columns = worker.adapter.spec['db']
|
|
|
|
self._check(columns, 'status')
|
|
self._check(columns, 'latency', f'latency:{worker.name}')
|
|
|
|
def _start(self):
|
|
for name, service in self.cfg.SERVICES.items():
|
|
self._make_db_table(name, service)
|
|
|
|
# spawn a service worker
|
|
serv_worker = ServiceWorker(self, name, service)
|
|
self.workers[name] = serv_worker
|
|
self.state[name] = None
|
|
|
|
self._create_channels(serv_worker)
|
|
|
|
def close(self):
|
|
for worker in self.workers.values():
|
|
worker.stop()
|
|
|
|
def subscribe(self, channels: List[str], websocket):
|
|
"""Subscribe to a list of channels."""
|
|
subscribed = []
|
|
|
|
for chan in channels:
|
|
try:
|
|
self.subscribers[chan].append(websocket.client_id)
|
|
except KeyError:
|
|
pass
|
|
|
|
return subscribed
|
|
|
|
def unsub_all(self, websocket):
|
|
"""Unsubscribe a websocket from all known channels."""
|
|
unsub = []
|
|
|
|
for chan, subs in self.subscribers.items():
|
|
log.info(f'Unsubscribing {websocket.client_id} from {chan}')
|
|
try:
|
|
subs.remove(websocket.client_id)
|
|
unsub.append(chan)
|
|
except ValueError:
|
|
pass
|
|
|
|
log.info(f'unsubscribed {websocket.client_id} from {unsub}')
|
|
return unsub
|