From 9b26d6b2ae82caf00a4745c8166dcf35d2f03bbb Mon Sep 17 00:00:00 2001 From: HgO Date: Fri, 8 Jul 2022 23:04:04 +0200 Subject: [PATCH] fix context manager --- matrix_alertbot/alertmanager.py | 1 - matrix_alertbot/main.py | 127 ++++++++++++++++---------------- matrix_alertbot/webhook.py | 1 - 3 files changed, 64 insertions(+), 65 deletions(-) diff --git a/matrix_alertbot/alertmanager.py b/matrix_alertbot/alertmanager.py index fcadf46..20c9d04 100644 --- a/matrix_alertbot/alertmanager.py +++ b/matrix_alertbot/alertmanager.py @@ -6,7 +6,6 @@ from typing import Any, Dict, List import aiohttp import pytimeparse from aiohttp import ClientError -from aiotools import AsyncContextManager from diskcache import Cache from matrix_alertbot.errors import ( diff --git a/matrix_alertbot/main.py b/matrix_alertbot/main.py index a35c2a3..5042b7f 100644 --- a/matrix_alertbot/main.py +++ b/matrix_alertbot/main.py @@ -5,7 +5,6 @@ import sys from asyncio import TimeoutError from time import sleep -import aiotools from aiohttp import ClientConnectionError, ServerDisconnectedError from diskcache import Cache from nio import ( @@ -52,72 +51,53 @@ def create_matrix_client(config: Config) -> AsyncClient: return client -async def start_matrix_client(cache: Cache, config: Config) -> bool: - async with aiotools.closing_async(create_matrix_client(config)) as client: - # Configure Alertmanager client - async with AlertmanagerClient(config.alertmanager_url, cache) as alertmanager: - # Set up event callbacks - callbacks = Callbacks(client, alertmanager, cache, config) - client.add_event_callback(callbacks.message, (RoomMessageText,)) - client.add_event_callback( - callbacks.invite_event_filtered_callback, (InviteMemberEvent,) - ) - client.add_event_callback(callbacks.decryption_failure, (MegolmEvent,)) - client.add_event_callback(callbacks.unknown, (UnknownEvent,)) +async def start_matrix_client( + client: AsyncClient, cache: Cache, config: Config +) -> bool: + # Keep trying to reconnect on failure (with some time in-between) + while True: + try: + if config.user_token: + # Use token to log in + client.load_store() - # Keep trying to reconnect on failure (with some time in-between) - while True: + # Sync encryption keys with the server + if client.should_upload_keys: + await client.keys_upload() + else: + # Try to login with the configured username/password try: - if config.user_token: - # Use token to log in - client.load_store() - - # Sync encryption keys with the server - if client.should_upload_keys: - await client.keys_upload() - else: - # Try to login with the configured username/password - try: - login_response = await client.login( - password=config.user_password, - device_name=config.device_name, - ) - - # Check if login failed - if type(login_response) == LoginError: - logger.error( - "Failed to login: %s", login_response.message - ) - return False - except LocalProtocolError as e: - # There's an edge case here where the user hasn't installed the correct C - # dependencies. In that case, a LocalProtocolError is raised on login. - logger.fatal( - "Failed to login. Have you installed the correct dependencies? " - "https://github.com/poljar/matrix-nio#installation " - "Error: %s", - e, - ) - return False - - # Login succeeded! - - logger.info(f"Logged in as {config.user_id}") - await client.sync_forever(timeout=30000, full_state=True) - - except (ClientConnectionError, ServerDisconnectedError, TimeoutError): - logger.warning( - "Unable to connect to homeserver, retrying in 15s..." + login_response = await client.login( + password=config.user_password, + device_name=config.device_name, ) - # Sleep so we don't bombard the server with login requests - sleep(15) + # Check if login failed + if type(login_response) == LoginError: + logger.error("Failed to login: %s", login_response.message) + return False + except LocalProtocolError as e: + # There's an edge case here where the user hasn't installed the correct C + # dependencies. In that case, a LocalProtocolError is raised on login. + logger.fatal( + "Failed to login. Have you installed the correct dependencies? " + "https://github.com/poljar/matrix-nio#installation " + "Error: %s", + e, + ) + return False + # Login succeeded! -async def start_webhook_server(cache: Cache, config: Config) -> None: - async with aiotools.closing_async(create_matrix_client(config)) as client: - async with Webhook(client, cache, config) as webhook_server: - await webhook_server.start() + logger.info(f"Logged in as {config.user_id}") + await client.sync_forever(timeout=30000, full_state=True) + except (ClientConnectionError, ServerDisconnectedError, TimeoutError): + logger.warning("Unable to connect to homeserver, retrying in 15s...") + + # Sleep so we don't bombard the server with login requests + sleep(15) + finally: + client.close() def main() -> None: @@ -133,14 +113,35 @@ def main() -> None: # Read the parsed config file and create a Config object config = Config(config_path) + client = create_matrix_client(config) + # Configure the cache cache = Cache(config.cache_dir) + # Configure Alertmanager client + alertmanager = AlertmanagerClient(config.alertmanager_url, cache) + + # Set up event callbacks + callbacks = Callbacks(client, alertmanager, cache, config) + client.add_event_callback(callbacks.message, (RoomMessageText,)) + client.add_event_callback( + callbacks.invite_event_filtered_callback, (InviteMemberEvent,) + ) + client.add_event_callback(callbacks.decryption_failure, (MegolmEvent,)) + client.add_event_callback(callbacks.unknown, (UnknownEvent,)) + + # Configure webhook server + webhook_server = Webhook(client, cache, config) + loop = asyncio.get_event_loop() - loop.create_task(start_webhook_server(cache, config)) - loop.create_task(start_matrix_client(cache, config)) + loop.create_task(webhook_server.start()) + loop.create_task(start_matrix_client(client, cache, config)) try: loop.run_forever() except Exception as e: logger.error(e) + finally: + loop.run_until_complete(webhook_server.close()) + loop.run_until_complete(alertmanager.close()) + loop.run_until_complete(client.close()) diff --git a/matrix_alertbot/webhook.py b/matrix_alertbot/webhook.py index d098b3d..56b2752 100644 --- a/matrix_alertbot/webhook.py +++ b/matrix_alertbot/webhook.py @@ -4,7 +4,6 @@ import logging from typing import Any from aiohttp import web, web_request -from aiotools import AsyncContextManager from diskcache import Cache from nio import AsyncClient, SendRetryError