import time import asyncio import logging from .consts import ADAPTERS class ServiceWorker: def __init__(self, manager, name, service): self.name = name self.manager = manager self.service = service self.adapter = ADAPTERS[service['adapter']] self.log = logging.getLogger(f'elstat.service.{name}') self.last_poll = None self._start() async def work(self): return await self.adapter.query(self, self.service['adapter_args']) async def process_work(self, result: tuple): """Process given adapter result and insert into the database.""" columns = self.adapter.spec['db'] conn = self.manager.conn timestamp = int(time.time() * 1000) args_str = ','.join(['?'] * (len(result) + 1)) query = f""" INSERT INTO {self.name} ({','.join(columns)}) VALUES ({args_str}) """ conn.execute(query, (timestamp,) + result) conn.commit() async def _work_loop(self): try: while True: self.log.info(f'polling {self.name}') self.last_poll = time.monotonic() res = await self.work() self.manager.state[self.name] = res await self.process_work(res) await asyncio.sleep(self.service['poll']) except asyncio.CancelledError: self.log.info('cancelled, stopping') except Exception: self.log.exception('fail on work loop, retrying') await self._work_loop() def stop(self): """Stop the current worker.""" self._work_task.cancel() def _start(self): self.log.info(f'starting work loop for {self.name}') self._work_task = self.manager.loop.create_task(self._work_loop())