diff options
| -rw-r--r-- | pypaste/__init__.py | 4 | ||||
| -rw-r--r-- | pypaste/__main__.py | 14 | ||||
| -rw-r--r-- | pypaste/sqlite/__init__.py | 103 |
3 files changed, 119 insertions, 2 deletions
diff --git a/pypaste/__init__.py b/pypaste/__init__.py index 0206e9e..4932f76 100644 --- a/pypaste/__init__.py +++ b/pypaste/__init__.py @@ -92,6 +92,10 @@ class Storage: connection: aiosqlite.Connection @abstractmethod + async def setup(self) -> None: + pass + + @abstractmethod async def insert(self, paste: Paste) -> None: pass diff --git a/pypaste/__main__.py b/pypaste/__main__.py index c7704d5..5215d15 100644 --- a/pypaste/__main__.py +++ b/pypaste/__main__.py @@ -17,8 +17,9 @@ import sys import os import asyncio import aiosqlite -from pypaste import App, AppConfig, log_error, log_info +from pypaste import App, AppConfig, Storage, log_error, log_info from pypaste.s3 import S3 +from pypaste.sqlite import Sqlite from socket import socket, AF_UNIX, SOCK_STREAM from argparse import ArgumentParser from aiohttp import web @@ -48,6 +49,8 @@ async def main() -> int: s3parser.add_argument("--access-key", required=True) s3parser.add_argument("--secret-key", type=Path, required=True) + subparsers.add_parser("sqlite") + args = parser.parse_args() try: @@ -99,7 +102,7 @@ async def main() -> int: log_error(f"failed to load secret key from {args.secret_key}: {e}") return 1 - storage = S3( + storage: Storage = S3( connection, args.endpoint, args.region, @@ -109,6 +112,13 @@ async def main() -> int: ) await storage.setup() + case "sqlite": + storage = Sqlite(connection) + + await storage.setup() + + case _: + raise Exception(f"unsupported storage backend: {args.command}") pypaste = App(config, storage) diff --git a/pypaste/sqlite/__init__.py b/pypaste/sqlite/__init__.py new file mode 100644 index 0000000..5700356 --- /dev/null +++ b/pypaste/sqlite/__init__.py @@ -0,0 +1,103 @@ +import asyncio +import zstandard +import aiosqlite +from pypaste import Storage, Paste +from dataclasses import dataclass +from typing import Optional + + +@dataclass +class Sqlite(Storage): + connection: aiosqlite.Connection + + async def setup(self) -> None: + await self.connection.execute( + "create table if not exists sqlite(key text, data blob)" + ) + + await self.connection.commit() + + async def insert(self, paste: Paste) -> None: + def compress(): + return zstandard.compress(paste.text.encode()) + + data = await asyncio.to_thread(compress) + + await self.connection.execute( + "insert into pastes values(?, ?, ?, ?)", + (paste.key, paste.dt.isoformat(), len(data), paste.syntax), + ) + + await self.connection.execute( + "insert into sqlite values(?, ?)", + ( + paste.key, + data, + ), + ) + + await self.connection.commit() + + async def retrieve(self, key: str) -> Optional[Paste]: + if not await self.exists(key): + return None + + async with self.connection.execute( + "select sqlite.data from sqlite where key=? limit 1", (key,) + ) as cursor: + match await cursor.fetchone(): + case [bytes(data)]: + pass + case _: + raise Exception("unreachable") + + row = await self.read_row(key) + + assert row is not None + + (dt, size, syntax) = row + + def decompress(): + return zstandard.decompress(data).decode() + + text = await asyncio.to_thread(decompress) + + return Paste(key, dt, syntax, text) + + async def delete(self, key: str) -> None: + await self.connection.execute("delete from pastes where key=?", (key,)) + + await self.connection.execute("delete from sqlite where key=?", (key,)) + + 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 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 = row[0] + + if use > max: + await self.delete(oldest) + else: + return |