diff options
| -rw-r--r-- | pypaste/server/__init__.py | 8 | ||||
| -rw-r--r-- | pypaste/server/s3/__init__.py | 61 | ||||
| -rw-r--r-- | pypaste/server/sqlite/__init__.py | 61 | ||||
| -rwxr-xr-x | tests/test_storage.py | 26 |
4 files changed, 106 insertions, 50 deletions
diff --git a/pypaste/server/__init__.py b/pypaste/server/__init__.py index 5d2fd6d..da5e538 100644 --- a/pypaste/server/__init__.py +++ b/pypaste/server/__init__.py @@ -131,6 +131,14 @@ class Storage: async def vacuum(self, size: int) -> None: pass + @abstractmethod + async def storage_use(self) -> Optional[int]: + pass + + @abstractmethod + async def oldest(self) -> Optional[Key]: + 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", diff --git a/pypaste/server/s3/__init__.py b/pypaste/server/s3/__init__.py index 6a30358..eeded21 100644 --- a/pypaste/server/s3/__init__.py +++ b/pypaste/server/s3/__init__.py @@ -98,33 +98,44 @@ class S3(Storage): ) as cursor: return await cursor.fetchone() is not None + async def storage_use(self) -> Optional[int]: + async with self.connection.execute( + ( + "select sum(pastes.size) from pastes " + "inner join s3 on s3.key=pastes.key" + ) + ) as cursor: + match await cursor.fetchone(): + case [int(use)]: + return use + case None: + return None + case _: + raise Exception("unreachable") + + async def oldest(self) -> Optional[Key]: + async with self.connection.execute( + ( + "select pastes.key,pastes.key_length from pastes " + "inner join s3 on s3.key=pastes.key " + "order by pastes.datetime" + ) + ) as cursor: + match await cursor.fetchone(): + case [bytes(data), int(length)]: + return Key(data, length) + case None: + return None + case _: + raise Exception("unreachable") + 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: - return - else: - use = row[0] - - async with self.connection.execute( - ( - "select pastes.key from pastes " - "inner join s3 on s3.key " - "where s3.key=pastes.key " - "order by pastes.datetime " - "limit 1" - ) - ) as cursor: - if (row := await cursor.fetchone()) is None: - return - else: - oldest = row[0] + if (use := await self.storage_use()) is None: + return + + if (oldest := await self.oldest()) is None: + return if use > max: await self.delete(oldest) diff --git a/pypaste/server/sqlite/__init__.py b/pypaste/server/sqlite/__init__.py index be07db6..9ae49ad 100644 --- a/pypaste/server/sqlite/__init__.py +++ b/pypaste/server/sqlite/__init__.py @@ -83,33 +83,44 @@ class Sqlite(Storage): ) as cursor: return await cursor.fetchone() is not None + async def storage_use(self) -> Optional[int]: + async with self.connection.execute( + ( + "select sum(pastes.size) from pastes " + "inner join sqlite on sqlite.key=pastes.key" + ) + ) as cursor: + match await cursor.fetchone(): + case [int(use)]: + return use + case None: + return None + case _: + raise Exception("unreachable") + + 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=pastes.key " + "order by pastes.datetime" + ) + ) as cursor: + match await cursor.fetchone(): + case [data, length]: + return Key(data, length) + case None: + return None + case _: + raise Exception("unreachable") + 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] - - async with self.connection.execute( - ( - "select pastes.key, pastes.key_length 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]) + if (use := await self.storage_use()) is None: + return + + if (oldest := await self.oldest()) is None: + return if use > max: await self.delete(oldest) diff --git a/tests/test_storage.py b/tests/test_storage.py index f2b8912..62dd7b2 100755 --- a/tests/test_storage.py +++ b/tests/test_storage.py @@ -5,6 +5,7 @@ import os import asyncio import tempfile import aiosqlite +import zstandard from pypaste.server import Paste, Storage, keygen from pypaste.server.sqlite import Sqlite from pypaste.server.s3 import S3 @@ -17,6 +18,29 @@ def truncate(path: Path) -> None: f.truncate(0) +async def test_storage_use(storage: Storage) -> None: + content = "hello pypaste" + length = len(zstandard.compress(content.encode())) + + for x in range(5): + paste = Paste(datetime.now(), None, content) + key = keygen(6) + await storage.insert(paste, key) + + assert (use := await storage.storage_use()) is not None and use == length * 5 + + +async def test_oldest(storage: Storage) -> None: + pastes = [] + for x in range(5): + paste = Paste(datetime.now(), None, "hello pypaste") + key = keygen(6) + pastes.append(key) + await storage.insert(paste, key) + + assert (oldest := await storage.oldest()) is not None and oldest == pastes[0] + + async def test_exists_but_not_in_our_table(storage: Storage) -> None: key = keygen(6) @@ -125,6 +149,8 @@ async def main() -> int: test_delete, test_exists, test_exists_but_not_in_our_table, + test_oldest, + test_storage_use, ] await test_sqlite(tests) |