From 383f197e58e5715844c5edcdbb7c7bdd172f9bb3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Damir=20Jeli=C4=87?= <poljar@termina.org.uk>
Date: Tue, 11 Jun 2019 16:06:35 +0200
Subject: [PATCH] client: Store and load sync tokens.

---
 pantalaimon/client.py | 12 +++++++++++-
 pantalaimon/store.py  | 42 ++++++++++++++++++++++++++++++++++++++++--
 tests/store_test.py   | 10 ++++++++++
 3 files changed, 61 insertions(+), 3 deletions(-)

diff --git a/pantalaimon/client.py b/pantalaimon/client.py
index fbfb5a5..7180e81 100644
--- a/pantalaimon/client.py
+++ b/pantalaimon/client.py
@@ -276,6 +276,12 @@ class PanClient(AsyncClient):
 
         self.index.commit()
 
+        self.pan_store.save_token(
+            self.server_name,
+            self.user_id,
+            self.next_batch
+        )
+
         for room_id, room in response.rooms.join.items():
             if room.timeline.limited:
                 logger.info("Room {} had a limited timeline, queueing "
@@ -389,6 +395,11 @@ class PanClient(AsyncClient):
             }
         }
 
+        self.next_batch = self.pan_store.load_token(
+            self.server_name,
+            self.user_id
+        )
+
         # We don't store any room state so initial sync needs to be with the
         # full_state parameter. Subsequent ones are normal.
         while True:
@@ -415,7 +426,6 @@ class PanClient(AsyncClient):
 
         await self.sync_forever(timeout, sync_filter)
 
-
     async def start_sas(self, message, device):
         try:
             await self.start_key_verification(device)
diff --git a/pantalaimon/store.py b/pantalaimon/store.py
index 4d4eba3..7cb536d 100644
--- a/pantalaimon/store.py
+++ b/pantalaimon/store.py
@@ -16,9 +16,10 @@ import datetime
 import json
 import os
 from collections import defaultdict
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple
 
 import attr
+from nio import RoomMessage
 from nio.store import (Accounts, DeviceKeys, DeviceTrustState, TrustState,
                        use_database)
 from peewee import (SQL, DateTimeField, DoesNotExist, ForeignKeyField, Model,
@@ -98,6 +99,16 @@ class UserMessages(Model):
         column_name="event_id")
 
 
+class PanSyncTokens(Model):
+    token = TextField()
+    user = ForeignKeyField(
+        model=ServerUsers,
+        column_name="user_id")
+
+    class Meta:
+        constraints = [SQL("UNIQUE(user_id)")]
+
+
 @attr.s
 class ClientInfo:
     user_id = attr.ib(type=str)
@@ -120,6 +131,7 @@ class PanStore:
         Profile,
         Event,
         UserMessages,
+        PanSyncTokens
     ]
 
     def __attrs_post_init__(self):
@@ -153,10 +165,36 @@ class PanStore:
         except DoesNotExist:
             return None
 
+    @use_database
+    def save_token(self, server, pan_user, token):
+        # type: (str, str, str) -> None
+        """Save a sync token for a pan user."""
+        server = Servers.get(name=server)
+        user = ServerUsers.get(server=server, user_id=pan_user)
+
+        PanSyncTokens.replace(user=user, token=token).execute()
+
+    @use_database
+    def load_token(self, server, pan_user):
+        # type: (str, str) -> Optional[str]
+        """Load a sync token for a pan user.
+
+        Returns the sync token if one is found.
+        """
+        server = Servers.get(name=server)
+        user = ServerUsers.get(server=server, user_id=pan_user)
+
+        token = PanSyncTokens.get_or_none(user=user)
+
+        if token:
+            return token.token
+
+        return None
+
     @use_database
     def save_event(self, server, pan_user, event, room_id, display_name,
                    avatar_url):
-        # type (str, str, str, RoomMessage, str, str, str) -> Optional[int]
+        # type: (str, str, str, RoomMessage, str, str, str) -> Optional[int]
         """Save an event to the store.
 
         Returns the database id of the event.
diff --git a/tests/store_test.py b/tests/store_test.py
index 02ff7e8..ee9bbcb 100644
--- a/tests/store_test.py
+++ b/tests/store_test.py
@@ -116,3 +116,13 @@ class TestClass(object):
         event_dict = panstore.load_event_by_columns("example", user, found_id)
 
         assert event_dict["result"] == event.source
+
+    def test_token_storing(self, panstore_with_users):
+        panstore = panstore_with_users
+        accounts = panstore.load_all_users()
+        user, _ = accounts[0]
+
+        assert not panstore.load_token("example", user)
+        panstore.save_token("example", user, "abc123")
+
+        assert "abc123" == panstore.load_token("example", user)