From 5c3753a6e1b87b799152aad451d14c92f923ad9d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Damir=20Jeli=C4=87?= <poljar@termina.org.uk>
Date: Mon, 22 Jul 2019 16:21:54 +0200
Subject: [PATCH] client: Add support for the Sqlite Queue database from
 peewee.

---
 pantalaimon/client.py    | 19 ++++++++++++++++++-
 pantalaimon/daemon.py    |  4 ++++
 tests/conftest.py        |  3 ++-
 tests/pan_client_test.py |  4 +++-
 4 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/pantalaimon/client.py b/pantalaimon/client.py
index c185ef3..210b6cf 100644
--- a/pantalaimon/client.py
+++ b/pantalaimon/client.py
@@ -20,6 +20,7 @@ from typing import Any, Dict, Optional
 
 from aiohttp.client_exceptions import ClientConnectionError
 from jsonschema import Draft4Validator, FormatChecker, validators
+from playhouse.sqliteq import SqliteQueueDatabase
 from nio import (
     AsyncClient,
     ClientConfig,
@@ -121,6 +122,18 @@ class InvalidLimit(Exception):
     pass
 
 
+class SqliteQStore(SqliteStore):
+    def _create_database(self):
+        return SqliteQueueDatabase(
+            self.database_path,
+            pragmas=(("foregign_keys", 1), ("secure_delete", 1))
+        )
+
+    def close(self):
+        self.database.stop()
+
+
+
 class PanClient(AsyncClient):
     """A wrapper class around a nio AsyncClient extending its functionality."""
 
@@ -137,8 +150,9 @@ class PanClient(AsyncClient):
         config=None,
         ssl=None,
         proxy=None,
+        store_class=None,
     ):
-        config = config or ClientConfig(store=SqliteStore, store_name="pan.db")
+        config = config or ClientConfig(store=store_class or SqliteQStore, store_name="pan.db")
         super().__init__(homeserver, user_id, device_id, store_path, config, ssl, proxy)
 
         index_dir = os.path.join(store_path, server_name, user_id)
@@ -619,6 +633,9 @@ class PanClient(AsyncClient):
 
             self.history_fetcher_task = None
 
+        if isinstance(self.store, SqliteQueueDatabase):
+            self.store.close()
+
         self.history_fetch_queue = asyncio.Queue()
 
     def pan_decrypt_event(self, event_dict, room_id=None, ignore_failures=True):
diff --git a/pantalaimon/daemon.py b/pantalaimon/daemon.py
index f30f379..dc2e4b1 100755
--- a/pantalaimon/daemon.py
+++ b/pantalaimon/daemon.py
@@ -34,6 +34,7 @@ from pantalaimon.client import (
     PanClient,
     UnknownRoomError,
     validate_json,
+    SqliteQStore
 )
 from pantalaimon.index import INDEXING_ENABLED, InvalidQueryError
 from pantalaimon.log import logger
@@ -77,6 +78,7 @@ class ProxyDaemon:
     recv_queue = attr.ib()
     proxy = attr.ib(default=None)
     ssl = attr.ib(default=None)
+    client_store_class = attr.ib(default=SqliteQStore)
 
     decryption_timeout = 10
     unverified_send_timeout = 60
@@ -128,6 +130,7 @@ class ProxyDaemon:
                 store_path=self.data_dir,
                 ssl=self.ssl,
                 proxy=self.proxy,
+                store_class=self.client_store_class
             )
             pan_client.user_id = user_id
             pan_client.access_token = token
@@ -546,6 +549,7 @@ class ProxyDaemon:
             store_path=self.data_dir,
             ssl=self.ssl,
             proxy=self.proxy,
+            store_class=self.client_store_class
         )
         response = await pan_client.login(password, "pantalaimon")
 
diff --git a/tests/conftest.py b/tests/conftest.py
index d536fd5..f473ad8 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -129,7 +129,8 @@ async def pan_proxy_server(tempdir, aiohttp_server):
         send_queue=pan_queue.async_q,
         recv_queue=ui_queue.async_q,
         proxy=None,
-        ssl=False
+        ssl=False,
+        client_store_class=SqliteStore
     )
 
     app.add_routes([
diff --git a/tests/pan_client_test.py b/tests/pan_client_test.py
index c6b6041..74a4478 100644
--- a/tests/pan_client_test.py
+++ b/tests/pan_client_test.py
@@ -16,6 +16,7 @@ from nio import (
     SyncResponse,
     Timeline,
 )
+from nio.store import SqliteStore
 
 from pantalaimon.client import PanClient
 from pantalaimon.config import ServerConfig
@@ -45,7 +46,8 @@ async def client(tmpdir, loop):
         queue.async_q,
         "@example:example.org",
         "DEVICEID",
-        tmpdir
+        tmpdir,
+        store_class=SqliteStore
     )
 
     yield pan_client