improve error handling by narrowing down exception types and using transactionsHEADmaster

5 files changed, 42 insertions, 23 deletions

diff --git a/pypaste/server/__init__.py b/pypaste/server/__init__.py
index 6f79b65..185a400 100644
--- a/pypaste/server/__init__.py
+++ b/pypaste/server/__init__.py
@@ -21,6 +21,7 @@ from dataclasses import dataclass
 from typing import Optional
 from pypaste import log_error, log_warning, log_info
 from pygments import highlight
+from pygments.util import ClassNotFound
 from pygments.lexers import guess_lexer, get_lexer_by_name, get_lexer_for_mimetype
 from pygments.formatters import HtmlFormatter
 from pygments.styles import get_style_by_name
@@ -48,12 +49,12 @@ def pygmentize(
                 lexer = get_lexer_for_mimetype(mime)
             case _:
                 lexer = get_lexer_by_name("text")
-    except Exception:
+    except ClassNotFound:
         lexer = get_lexer_by_name("text")
 
     try:
         s = get_style_by_name(style)
-    except Exception as e:
+    except ClassNotFound as e:
         log_warning(f"failed to find style: {style}: {e}")
         s = get_style_by_name("default")
 
@@ -103,6 +104,10 @@ class PasteInfo:
     syntax: Optional[str]
 
 
+class StorageError(Exception):
+    pass
+
+
 @dataclass
 class Storage:
     connection: aiosqlite.Connection
@@ -177,13 +182,13 @@ class App:
 
         try:
             dehumanized_key = dehumanize(key)
-        except Exception as e:
+        except ValueError as e:
             log_warning(f"invalid key parameter: {key}: {e}")
             return web.HTTPBadRequest(text="invalid key")
 
         try:
             paste = await self.storage.retrieve(dehumanized_key)
-        except Exception as e:
+        except StorageError as e:
             log_error(f"failed to retrieve paste {key}: {e}")
             return web.HTTPInternalServerError()
 
@@ -234,7 +239,7 @@ class App:
 
         try:
             data = await request.read()
-        except Exception as e:
+        except OSError as e:
             log_error(f"failed to read data: {e}")
             return web.HTTPInternalServerError(text="failed to read data")
 
@@ -250,7 +255,7 @@ class App:
         try:
             paste = Paste(datetime.now(), syntax, text)
             await self.storage.insert(paste, key)
-        except Exception as e:
+        except StorageError as e:
             log_error(f"failed to insert paste {key} to storage: {e}")
             return web.HTTPInternalServerError()
 
@@ -277,7 +282,7 @@ class App:
                     case [None, None]:
                         pass
 
-            except Exception as e:
+            except StorageError as e:
                 log_error(f"failed to vacuum: {e}")
             finally:
                 await asyncio.sleep(60 * 5)
diff --git a/pypaste/server/__main__.py b/pypaste/server/__main__.py
index 54bd963..daa8b26 100644
--- a/pypaste/server/__main__.py
+++ b/pypaste/server/__main__.py
@@ -17,6 +17,7 @@ import sys
 import os
 import asyncio
 import aiosqlite
+from aiosqlite import DatabaseError
 from pypaste import log_error, log_info
 from pypaste.server import App, AppConfig, Storage
 from pypaste.server.s3 import S3
@@ -134,7 +135,7 @@ async def main() -> int:
 
     try:
         connection = await aiosqlite.connect(args.database)
-    except Exception as e:
+    except DatabaseError as e:
         log_error(f"failed to connect to database {args.database}: {e}")
         return 1
 
@@ -151,7 +152,7 @@ async def main() -> int:
             )
         )
         await connection.commit()
-    except Exception as e:
+    except DatabaseError as e:
         log_error(f"failed to initialize database: {e}")
         return 1
 
diff --git a/pypaste/server/s3/__init__.py b/pypaste/server/s3/__init__.py
index eeded21..09e8194 100644
--- a/pypaste/server/s3/__init__.py
+++ b/pypaste/server/s3/__init__.py
@@ -16,8 +16,9 @@
 import asyncio
 import zstandard
 import aiosqlite
-from pypaste.server import Storage, Paste, Key
-from pypaste.server.s3.bucket import Bucket
+from aiosqlite import DatabaseError
+from pypaste.server import Storage, Paste, Key, StorageError
+from pypaste.server.s3.bucket import Bucket, BucketError
 from dataclasses import dataclass
 from typing import Optional
 
@@ -58,9 +59,12 @@ class S3(Storage):
         try:
             await self.bucket.put(key.data.hex(), compressed)
             await self.connection.commit()
-        except Exception as e:
+        except BucketError as e:
             await self.connection.rollback()
-            raise e
+            raise StorageError(str(e))
+        except DatabaseError as e:
+            await self.connection.rollback()
+            raise StorageError(str(e))
 
     async def retrieve(self, key: Key) -> Optional[Paste]:
         if not await self.exists(key):
@@ -82,15 +86,19 @@ class S3(Storage):
         return Paste(info.dt, info.syntax, text)
 
     async def delete(self, key: Key) -> None:
+        await self.connection.execute("begin")
         await self.connection.execute("delete from pastes where key=?", (key.data,))
         await self.connection.execute("delete from s3 where key=?", (key.data,))
 
         try:
             await self.bucket.delete(key.data.hex())
             await self.connection.commit()
-        except Exception as e:
+        except BucketError as e:
+            await self.connection.rollback()
+            raise StorageError(str(e))
+        except DatabaseError as e:
             await self.connection.rollback()
-            raise e
+            raise StorageError(str(e))
 
     async def exists(self, key: Key) -> bool:
         async with self.connection.execute(
diff --git a/pypaste/server/s3/bucket.py b/pypaste/server/s3/bucket.py
index 484faa6..a5d1f29 100644
--- a/pypaste/server/s3/bucket.py
+++ b/pypaste/server/s3/bucket.py
@@ -22,6 +22,10 @@ from hashlib import sha256
 from typing import Optional
 
 
+class BucketError(Exception):
+    pass
+
+
 @dataclass
 class Request:
     url: str
@@ -71,7 +75,7 @@ class Bucket:
                     case 404:
                         return None
                     case _:
-                        raise Exception(
+                        raise BucketError(
                             f"failed to get {self.endpoint}/{self.bucket}/{key} with status {get.status}"
                         )
 
@@ -109,7 +113,7 @@ class Bucket:
         async with aiohttp.ClientSession() as client:
             async with client.put(url, headers=headers, data=data) as put:
                 if put.status != 200:
-                    raise Exception(
+                    raise BucketError(
                         f"failed put {self.endpoint}/{self.bucket}/{key} with {put.status}"
                     )
 
@@ -144,6 +148,6 @@ class Bucket:
         async with aiohttp.ClientSession() as client:
             async with client.delete(url, headers=headers) as delete:
                 if delete.status != 204:
-                    raise Exception(
+                    raise BucketError(
                         f"failed to delete {self.endpoint}/{self.bucket}/{key} with {delete.status}"
                     )
diff --git a/pypaste/server/sqlite/__init__.py b/pypaste/server/sqlite/__init__.py
index 5cd00a6..bbb1467 100644
--- a/pypaste/server/sqlite/__init__.py
+++ b/pypaste/server/sqlite/__init__.py
@@ -1,7 +1,8 @@
 import asyncio
 import zstandard
 import aiosqlite
-from pypaste.server import Storage, Paste, Key
+from aiosqlite import DatabaseError
+from pypaste.server import Storage, Paste, Key, StorageError
 from dataclasses import dataclass
 from typing import Optional
 from uuid import uuid4, UUID
@@ -58,9 +59,9 @@ class Sqlite(Storage):
             )
 
             await self.connection.commit()
-        except Exception as e:
+        except DatabaseError as e:
             await self.connection.rollback()
-            raise e
+            raise StorageError(str(e))
 
     async def retrieve(self, key: Key) -> Optional[Paste]:
         async with self.connection.execute(
@@ -115,10 +116,10 @@ class Sqlite(Storage):
             )
 
             await self.connection.commit()
-        except Exception as e:
+        except DatabaseError as e:
             await self.connection.rollback()
 
-            raise e
+            raise StorageError(str(e))
 
     async def exists(self, key: Key) -> bool:
         async with self.connection.execute(