diff --git a/pantalaimon/client.py b/pantalaimon/client.py index 4e0ea85..dad2936 100644 --- a/pantalaimon/client.py +++ b/pantalaimon/client.py @@ -19,7 +19,6 @@ from collections import defaultdict from pprint import pformat from typing import Any, Dict, Optional from urllib.parse import urlparse - from aiohttp.client_exceptions import ClientConnectionError from jsonschema import Draft4Validator, FormatChecker, validators from playhouse.sqliteq import SqliteQueueDatabase @@ -51,10 +50,9 @@ from nio import ( ) from nio.crypto import Sas from nio.store import SqliteStore - from pantalaimon.index import INDEXING_ENABLED from pantalaimon.log import logger -from pantalaimon.store import FetchTask, MediaInfo +from pantalaimon.store import FetchTask, MediaInfo, PanSqliteStore from pantalaimon.thread_messages import ( DaemonResponse, InviteSasSignal, @@ -163,7 +161,7 @@ class PanClient(AsyncClient): media_info=None, ): config = config or AsyncClientConfig( - store=store_class or SqliteStore, store_name="pan.db" + store=store_class or PanSqliteStore, store_name="pan.db" ) super().__init__(homeserver, user_id, device_id, store_path, config, ssl, proxy) diff --git a/pantalaimon/config.py b/pantalaimon/config.py index 0d09d5c..5421cf3 100644 --- a/pantalaimon/config.py +++ b/pantalaimon/config.py @@ -127,6 +127,8 @@ class ServerConfig: pantalaimon on startup. sync_stop_after (int): The number of seconds to wait since the client has requested a /sync, before stopping a sync. + store_forgetful (bool): Enable or disable discarding of previous sessions + from the store. """ name = attr.ib(type=str) @@ -145,7 +147,7 @@ class ServerConfig: history_fetch_delay = attr.ib(type=int, default=3) sync_on_startup = attr.ib(type=bool, default=False) sync_stop_after = attr.ib(type=int, default=600) - + store_forgetful = attr.ib(type=bool, default=False) @attr.s class PanConfig: @@ -209,7 +211,7 @@ class PanConfig: proxy = section.geturl("Proxy") search_requests = section.getboolean("SearchRequests") index_encrypted_only = section.getboolean("IndexEncryptedOnly") - + store_forgetful = config["Default"].getboolean("StoreForgetful") indexing_batch_size = section.getint("IndexingBatchSize") sync_on_startup = False #section.getboolean("SyncOnStartup") @@ -256,6 +258,7 @@ class PanConfig: history_fetch_delay / 1000, sync_on_startup, sync_stop_after, + store_forgetful, ) self.servers[section_name] = server_conf diff --git a/pantalaimon/daemon.py b/pantalaimon/daemon.py index 49c99b9..68941a3 100755 --- a/pantalaimon/daemon.py +++ b/pantalaimon/daemon.py @@ -160,6 +160,7 @@ class ProxyDaemon: pan_client.user_id = user_id pan_client.access_token = token pan_client.load_store() + pan_client.store.forgetful = self.conf.store_forgetful self.pan_clients[user_id] = pan_client loop.create_task( diff --git a/pantalaimon/store.py b/pantalaimon/store.py index c1b4691..cb5b853 100644 --- a/pantalaimon/store.py +++ b/pantalaimon/store.py @@ -25,6 +25,9 @@ from nio.store import ( DeviceTrustState, use_database, use_database_atomic, + SqliteStore, + MegolmInboundSessions, + ForwardedChains ) from peewee import SQL, DoesNotExist, ForeignKeyField, Model, SqliteDatabase, TextField from cachetools import LRUCache @@ -443,3 +446,20 @@ class PanStore: store[account.user_id] = device_store return store + +class PanSqliteStore(SqliteStore): + forgetful = False + + @use_database + def save_inbound_group_session(self, session): + """Save the provided Megolm inbound group session to the database. + Args: + session (InboundGroupSession): The session to save. + """ + # Delete previous sessions + if self.forgetful: + MegolmInboundSessions.delete().where( + (MegolmInboundSessions.sender_key == session.sender_key) | + (MegolmInboundSessions.room_id == session.room_id) + ).execute() + super().save_inbound_group_session(session) \ No newline at end of file