From 1819a276cb255eb6de51febb9e3eec9423ad757c Mon Sep 17 00:00:00 2001
From: holger krekel <holger@merlinux.eu>
Date: Wed, 20 Mar 2024 17:02:04 +0100
Subject: [PATCH] implement persistence via marshal

---
 chatmaild/src/chatmaild/metadata.py           | 30 ++++++++++++++++---
 .../src/chatmaild/tests/test_metadata.py      | 17 +++++------
 2 files changed, 34 insertions(+), 13 deletions(-)

diff --git a/chatmaild/src/chatmaild/metadata.py b/chatmaild/src/chatmaild/metadata.py
index cd965eda..ace92336 100644
--- a/chatmaild/src/chatmaild/metadata.py
+++ b/chatmaild/src/chatmaild/metadata.py
@@ -12,6 +12,7 @@ import sys
 import logging
 import os
 import requests
+import marshal
 
 
 DICTPROXY_LOOKUP_CHAR = "L"
@@ -25,14 +26,35 @@ DICTPROXY_TRANSACTION_CHARS = "SBC"
 class Notifier:
     def __init__(self, metadata_dir):
         self.metadata_dir = metadata_dir
-        self.guid2token = {}
         self.to_notify_queue = Queue()
 
+    def get_metadata(self, guid):
+        guid_path = self.metadata_dir.joinpath(guid)
+        if guid_path.exists():
+            with guid_path.open("rb") as f:
+                return marshal.load(f)
+        return {}
+
+    def set_metadata(self, guid, guid_data):
+        guid_path = self.metadata_dir.joinpath(guid)
+        write_path = guid_path.with_suffix(".tmp")
+        with write_path.open("wb") as f:
+            marshal.dump(guid_data, f)
+        os.rename(write_path, guid_path)
+
     def set_token(self, guid, token):
-        self.guid2token[guid] = token
+        guid_data = self.get_metadata(guid)
+        guid_data["token"] = token
+        self.set_metadata(guid, guid_data)
+
+    def del_token(self, guid):
+        guid_data = self.get_metadata(guid)
+        if "token" in guid_data:
+            del guid_data["token"]
+            self.set_metadata(guid, guid_data)
 
     def get_token(self, guid):
-        return self.guid2token.get(guid)
+        return self.get_metadata(guid).get("token")
 
     def new_message_for_guid(self, guid):
         self.to_notify_queue.put(guid)
@@ -54,7 +76,7 @@ class Notifier:
             if response.status_code == 410:
                 # 410 Gone status code
                 # means the token is no longer valid.
-                del self.guid2token[guid]
+                self.del_token(guid)
 
 
 def handle_dovecot_protocol(rfile, wfile, notifier):
diff --git a/chatmaild/src/chatmaild/tests/test_metadata.py b/chatmaild/src/chatmaild/tests/test_metadata.py
index 3148e81e..2a3c122b 100644
--- a/chatmaild/src/chatmaild/tests/test_metadata.py
+++ b/chatmaild/src/chatmaild/tests/test_metadata.py
@@ -27,27 +27,26 @@ def test_handle_dovecot_request_happy_path(notifier):
     # lookups return the same NOTFOUND result
     res = handle_dovecot_request("Lpriv/123/chatmail", transactions, notifier)
     assert res == "N\n"
-    assert not notifier.guid2token and not transactions
+    assert notifier.get_token("guid00") is None and not transactions
 
     # set device token in a transaction
     tx = "1111"
     msg = f"B{tx}\tuser"
     res = handle_dovecot_request(msg, transactions, notifier)
-    assert not res and not notifier.guid2token
+    assert not res and notifier.get_token("guid00") is None
     assert transactions == {tx: "O\n"}
 
     msg = f"S{tx}\tpriv/guid00/devicetoken\t01234"
     res = handle_dovecot_request(msg, transactions, notifier)
     assert not res
     assert len(transactions) == 1
-    assert len(notifier.guid2token) == 1
-    assert notifier.guid2token["guid00"] == "01234"
+    assert notifier.get_token("guid00") == "01234"
 
     msg = f"C{tx}"
     res = handle_dovecot_request(msg, transactions, notifier)
     assert res == "O\n"
     assert len(transactions) == 0
-    assert notifier.guid2token["guid00"] == "01234"
+    assert notifier.get_token("guid00") == "01234"
 
     # trigger notification for incoming message
     assert handle_dovecot_request(f"B{tx}\tuser", transactions, notifier) is None
@@ -72,7 +71,7 @@ def test_handle_dovecot_protocol_set_devicetoken(notifier):
     )
     wfile = io.BytesIO()
     handle_dovecot_protocol(rfile, wfile, notifier)
-    assert notifier.guid2token["guid00"] == "01234"
+    assert notifier.get_token("guid00") == "01234"
     assert wfile.getvalue() == b"O\n"
 
 
@@ -125,7 +124,7 @@ def test_notifier_thread_run(notifier):
     notifier.thread_run_one(ReqMock())
     url, data, timeout = requests[0]
     assert data == "01234"
-    assert len(notifier.guid2token) == 1
+    assert notifier.get_token("guid00") == "01234"
 
 
 def test_notifier_thread_run_gone_removes_token(notifier):
@@ -142,8 +141,8 @@ def test_notifier_thread_run_gone_removes_token(notifier):
 
     notifier.set_token("guid00", "01234")
     notifier.new_message_for_guid("guid00")
-    assert notifier.guid2token["guid00"] == "01234"
+    assert notifier.get_token("guid00") == "01234"
     notifier.thread_run_one(ReqMock())
     url, data, timeout = requests[0]
     assert data == "01234"
-    assert len(notifier.guid2token) == 0
+    assert notifier.get_token("guid00") is None