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 pprint import pformat
from typing import Any, Dict, Optional from typing import Any, Dict, Optional
from urllib.parse import urlparse from urllib.parse import urlparse
from aiohttp.client_exceptions import ClientConnectionError from aiohttp.client_exceptions import ClientConnectionError
from jsonschema import Draft4Validator, FormatChecker, validators from jsonschema import Draft4Validator, FormatChecker, validators
from playhouse.sqliteq import SqliteQueueDatabase from playhouse.sqliteq import SqliteQueueDatabase
@ -50,10 +49,9 @@ from nio import (
) )
from nio.crypto import Sas from nio.crypto import Sas
from nio.store import SqliteStore from nio.store import SqliteStore
from pantalaimon.index import INDEXING_ENABLED from pantalaimon.index import INDEXING_ENABLED
from pantalaimon.log import logger from pantalaimon.log import logger
from pantalaimon.store import FetchTask, MediaInfo from pantalaimon.store import FetchTask, MediaInfo, PanSqliteStore
from pantalaimon.thread_messages import ( from pantalaimon.thread_messages import (
DaemonResponse, DaemonResponse,
InviteSasSignal, InviteSasSignal,
@ -162,7 +160,7 @@ class PanClient(AsyncClient):
media_info=None, media_info=None,
): ):
config = config or AsyncClientConfig( 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) super().__init__(homeserver, user_id, device_id, store_path, config, ssl, proxy)

View File

@ -121,6 +121,8 @@ class ServerConfig:
the room history. the room history.
history_fetch_delay (int): The delay between room history fetching history_fetch_delay (int): The delay between room history fetching
requests in seconds. requests in seconds.
store_forgetful (bool): Enable or disable discarding of previous sessions
from the store.
""" """
name = attr.ib(type=str) name = attr.ib(type=str)
@ -137,7 +139,7 @@ class ServerConfig:
index_encrypted_only = attr.ib(type=bool, default=True) index_encrypted_only = attr.ib(type=bool, default=True)
indexing_batch_size = attr.ib(type=int, default=100) indexing_batch_size = attr.ib(type=int, default=100)
history_fetch_delay = attr.ib(type=int, default=3) history_fetch_delay = attr.ib(type=int, default=3)
store_forgetful = attr.ib(type=bool, default=False)
@attr.s @attr.s
class PanConfig: class PanConfig:
@ -201,7 +203,7 @@ class PanConfig:
proxy = section.geturl("Proxy") proxy = section.geturl("Proxy")
search_requests = section.getboolean("SearchRequests") search_requests = section.getboolean("SearchRequests")
index_encrypted_only = section.getboolean("IndexEncryptedOnly") index_encrypted_only = section.getboolean("IndexEncryptedOnly")
store_forgetful = config["Default"].getboolean("StoreForgetful")
indexing_batch_size = section.getint("IndexingBatchSize") indexing_batch_size = section.getint("IndexingBatchSize")
if not 1 < indexing_batch_size <= 1000: if not 1 < indexing_batch_size <= 1000:
@ -243,6 +245,7 @@ class PanConfig:
index_encrypted_only, index_encrypted_only,
indexing_batch_size, indexing_batch_size,
history_fetch_delay / 1000, history_fetch_delay / 1000,
store_forgetful,
) )
self.servers[section_name] = server_conf self.servers[section_name] = server_conf

View File

@ -150,6 +150,7 @@ class ProxyDaemon:
pan_client.user_id = user_id pan_client.user_id = user_id
pan_client.access_token = token pan_client.access_token = token
pan_client.load_store() pan_client.load_store()
pan_client.store.forgetful = self.conf.store_forgetful
self.pan_clients[user_id] = pan_client self.pan_clients[user_id] = pan_client
loop.create_task( loop.create_task(

View File

@ -25,6 +25,9 @@ from nio.store import (
DeviceTrustState, DeviceTrustState,
use_database, use_database,
use_database_atomic, use_database_atomic,
SqliteStore,
MegolmInboundSessions,
ForwardedChains
) )
from peewee import SQL, DoesNotExist, ForeignKeyField, Model, SqliteDatabase, TextField from peewee import SQL, DoesNotExist, ForeignKeyField, Model, SqliteDatabase, TextField
from cachetools import LRUCache from cachetools import LRUCache
@ -372,3 +375,20 @@ class PanStore:
store[account.user_id] = device_store store[account.user_id] = device_store
return 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)