fix context manager

This commit is contained in:
HgO 2022-07-08 23:04:04 +02:00
parent 9b3ef85e76
commit 9b26d6b2ae
3 changed files with 64 additions and 65 deletions

View file

@ -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 (

View file

@ -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())

View file

@ -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