debuggingdebugging

5 files changed, 132 insertions, 84 deletions

diff --git a/pypaste/server/__init__.py b/pypaste/server/__init__.py
index 5d2fd6d..722abc1 100644
--- a/pypaste/server/__init__.py
+++ b/pypaste/server/__init__.py
@@ -131,6 +131,10 @@ class Storage:
     async def vacuum(self, size: int) -> None:
         pass
 
+    @abstractmethod
+    async def get_used(self) -> Optional[int]:
+        pass
+
     async def read_paste_info(self, key: Key) -> Optional[PasteInfo]:
         async with self.connection.execute(
             "select pastes.datetime,pastes.size,pastes.syntax from pastes where pastes.key=? limit 1",
@@ -254,6 +258,16 @@ class App:
 
         return web.HTTPOk(text=url)
 
+    async def vacuum(self):
+        while True:
+            try:
+                await self.storage.vacuum(self.config.storage_max_bytes)
+            except Exception as e:
+                log_error(f"failed to vacuum: {e}")
+                raise e
+
+            await asyncio.sleep(60)
+
 
 def test_humanize_dehumanize_roundtrip() -> None:
     key = keygen(6)
diff --git a/pypaste/server/__main__.py b/pypaste/server/__main__.py
index be83ba2..23efb44 100644
--- a/pypaste/server/__main__.py
+++ b/pypaste/server/__main__.py
@@ -155,6 +155,8 @@ async def main() -> int:
 
     log_info("starting pypaste")
 
+    asyncio.create_task(pypaste.vacuum())
+
     try:
         await asyncio.Event().wait()
     except asyncio.exceptions.CancelledError:
diff --git a/pypaste/server/s3/__init__.py b/pypaste/server/s3/__init__.py
index 6a30358..78d9731 100644
--- a/pypaste/server/s3/__init__.py
+++ b/pypaste/server/s3/__init__.py
@@ -100,17 +100,11 @@ class S3(Storage):
 
     async def vacuum(self, max: int) -> None:
         while True:
-            async with self.connection.execute(
-                (
-                    "select sum(pastes.size) from pastes "
-                    "inner join s3 on s3.key "
-                    "where s3.key=pastes.key"
-                )
-            ) as cursor:
-                if (row := await cursor.fetchone()) is None:
+            match await self.get_use():
+                case int(use):
+                    pass
+                case None:
                     return
-                else:
-                    use = row[0]
 
             async with self.connection.execute(
                 (
@@ -130,3 +124,22 @@ class S3(Storage):
                 await self.delete(oldest)
             else:
                 return
+
+    async def get_use(self) -> Optional[int]:
+        async with self.connection.execute(
+            (
+                (
+                    "select sum(pastes.size) from pastes "
+                    "inner join sqlite on s3.key "
+                    "where pastes.key=s3.key "
+                    "order by pastes.datetime "
+                )
+            )
+        ) as cursor:
+            match await cursor.fetchone():
+                case [int(use)]:
+                    return use
+                case None:
+                    return None
+                case _:
+                    raise Exception("unreachable")
diff --git a/pypaste/server/sqlite/__init__.py b/pypaste/server/sqlite/__init__.py
index be07db6..db62ea1 100644
--- a/pypaste/server/sqlite/__init__.py
+++ b/pypaste/server/sqlite/__init__.py
@@ -85,33 +85,47 @@ class Sqlite(Storage):
 
     async def vacuum(self, max: int) -> None:
         while True:
-            async with self.connection.execute(
-                (
-                    "select sum(pastes.size) from pastes "
-                    "inner join sqlite on sqlite.key "
-                    "where pastes.key=sqlite.key"
-                )
-            ) as cursor:
-                if (row := await cursor.fetchone()) is None:
-                    return
-                else:
-                    use = row[0]
+            if (use := await self.get_use()) is None:
+                return
+
+            if (oldest := await self.oldest()) is None:
+                return
+
+            if use >= max:
+                await self.delete(oldest)
+            else:
+                return
 
-            async with self.connection.execute(
+    async def get_use(self) -> Optional[int]:
+        async with self.connection.execute(
+            (
                 (
-                    "select pastes.key, pastes.key_length from pastes "
+                    "select sum(pastes.size) from pastes "
                     "inner join sqlite on sqlite.key "
                     "where pastes.key=sqlite.key "
-                    "order by pastes.datetime "
-                    "limit 1"
                 )
-            ) as cursor:
-                if (row := await cursor.fetchone()) is None:
-                    return
-                else:
-                    oldest = Key(row[0], row[1])
+            )
+        ) as cursor:
+            match await cursor.fetchone():
+                case [int(use)]:
+                    return use
+                case None:
+                    return None
+                case _:
+                    raise Exception("unreachable")
 
-            if use > max:
-                await self.delete(oldest)
-            else:
-                return
+    async def oldest(self) -> Optional[Key]:
+        async with self.connection.execute(
+            (
+                "select pastes.key,pastes.key_length from pastes "
+                "inner join sqlite on sqlite.key "
+                "where pastes.key=sqlite.key "
+            )
+        ) as cursor:
+            match await cursor.fetchone():
+                case [bytes(data), int(size)]:
+                    return Key(data, size)
+                case None:
+                    return None
+                case _:
+                    raise Exception("unreachable")
diff --git a/tests/test_storage.py b/tests/test_storage.py
index dde2f40..029e413 100755
--- a/tests/test_storage.py
+++ b/tests/test_storage.py
@@ -5,12 +5,14 @@ import os
 import asyncio
 import tempfile
 import aiosqlite
+import zstandard
+import string
 from pypaste.server import Paste, Storage, keygen
 from pypaste.server.sqlite import Sqlite
 from pypaste.server.s3 import S3
 from datetime import datetime
 from pathlib import Path
-from typing import List
+from secrets import choice
 
 
 def truncate(path: Path) -> None:
@@ -18,6 +20,23 @@ def truncate(path: Path) -> None:
         f.truncate(0)
 
 
+def random_string(length: int) -> str:
+    return "".join(choice(string.ascii_letters) for _ in range(length))
+
+
+async def test_vacuum(storage: Storage):
+    # generate random string to avoid it getting compressed
+    content = random_string(32)
+    size_compressed = len(zstandard.compress(content.encode()))
+
+    for x in range(1024 // size_compressed + 1):
+        await storage.insert(Paste(datetime.now(), None, content), keygen(6))
+
+    await storage.vacuum(256)
+
+    assert (use := await storage.get_used()) is not None and use <= 256
+
+
 async def test_exists_but_not_in_our_table(storage: Storage) -> None:
     key = keygen(6)
 
@@ -67,57 +86,43 @@ async def test_insert_retrieve(storage: Storage) -> None:
     assert paste.text == "hello world"
 
 
-async def main() -> int:
-    stores: List[Storage] = []
-
-    with tempfile.TemporaryDirectory() as tmpdir:
-        f = Path(tmpdir) / "database"
-        truncate(f)
-        async with aiosqlite.connect(f) as connection:
-            await connection.execute(
-                (
-                    "create table pastes("
-                    "key blob,"
-                    "key_length int,"
-                    "datetime text,"
-                    "size int,"
-                    "syntax text"
-                    ")"
-                )
-            )
-
-            sqlite_storage = Sqlite(connection)
-            await sqlite_storage.setup()
-            stores.append(sqlite_storage)
-
-            try:
-                os.environ["PYPASTE_TEST_S3"]
-                test_s3 = True
-            except KeyError:
-                test_s3 = False
-
-            if test_s3:
-                s3_storage = S3(
-                    connection,
-                    os.environ["PYPASTE_TEST_ENDPOINT"],
-                    os.environ["PYPASTE_TEST_REGION"],
-                    os.environ["PYPASTE_TEST_BUCKET"],
-                    os.environ["PYPASTE_TEST_ACCESS_KEY"],
-                    os.environ["PYPASTE_TEST_SECRET_KEY"],
-                )
-                await s3_storage.setup()
-                stores.append(s3_storage)
-
-            for store in stores:
-                await asyncio.gather(
-                    test_insert_retrieve(store),
-                    test_insert_retrieve(store),
-                    test_delete(store),
-                    test_delete(store),
-                    test_exists_but_not_in_our_table(store),
-                    test_exists_but_not_in_our_table(store),
+async def test_sqlite(tests):
+    for test in tests:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            f = Path(tmpdir) / "pastes.sqlite"
+            truncate(f)
+
+            async with aiosqlite.connect(f) as connection:
+                await connection.execute(
+                    (
+                        "create table pastes("
+                        "key blob,"
+                        "key_length int,"
+                        "datetime text,"
+                        "size int,"
+                        "syntax text"
+                        ")"
+                    )
                 )
 
+                storage = Sqlite(connection)
+
+                await storage.setup()
+
+                await test(storage)
+
+
+async def main() -> int:
+    tests = [
+        test_insert_retrieve,
+        test_delete,
+        test_exists,
+        test_exists_but_not_in_our_table,
+        test_vacuum,
+    ]
+
+    await test_sqlite(tests)
+
     return 0