Add option to forget sessions when inserting a new session

This commit is contained in:
Will Hunt 2020-09-22 15:17:41 +01:00
parent 908c1a16be
commit d3b535f070
4 changed files with 28 additions and 6 deletions

View File

@ -18,7 +18,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
@ -50,10 +49,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,
@ -162,7 +160,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)

View File

@ -121,6 +121,8 @@ class ServerConfig:
the room history.
history_fetch_delay (int): The delay between room history fetching
requests in seconds.
store_forgetful (bool): Enable or disable discarding of previous sessions
from the store.
"""
name = attr.ib(type=str)
@ -137,7 +139,7 @@ class ServerConfig:
index_encrypted_only = attr.ib(type=bool, default=True)
indexing_batch_size = attr.ib(type=int, default=100)
history_fetch_delay = attr.ib(type=int, default=3)
store_forgetful = attr.ib(type=bool, default=False)
@attr.s
class PanConfig:
@ -201,7 +203,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")
if not 1 < indexing_batch_size <= 1000:
@ -243,6 +245,7 @@ class PanConfig:
index_encrypted_only,
indexing_batch_size,
history_fetch_delay / 1000,
store_forgetful,
)
self.servers[section_name] = server_conf

View File

@ -150,6 +150,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(

View File

@ -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
@ -372,3 +375,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)