fix context manager
This commit is contained in:
parent
9b3ef85e76
commit
9b26d6b2ae
3 changed files with 64 additions and 65 deletions
|
@ -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 (
|
||||
|
|
|
@ -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())
|
||||
|
|
|
@ -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
|
||||
|
||||
|
|
Loading…
Reference in a new issue