elstat/elstat/manager.py

79 lines
2.0 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 _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
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