Add Reference class

src/media_observer/storage.py

@@ -11,6 +11,7 @@ from media_observer.article import (
 )
 from media_observer.storage_abstraction import (
     Table,
+    Reference,
     Column,
     UniqueIndex,
     View,
@@ -35,7 +36,7 @@ table_frontpages = Table(
         Column(name="id", primary_key=True),
         Column(
             name="site_id",
-            references="sites (id) ON DELETE CASCADE",
+            references=Reference("sites", "id", on_delete="cascade"),
         ),
         Column(name="timestamp", type_="timestamp with time zone"),
         Column(name="timestamp_virtual", type_="timestamp with time zone"),
@@ -64,15 +65,15 @@ table_main_articles = Table(
         Column(name="url", type_="TEXT"),
         Column(
             name="frontpage_id",
-            references="frontpages (id) ON DELETE CASCADE",
+            references=Reference("frontpages", "id", on_delete="cascade"),
         ),
         Column(
             name="article_id",
-            references="articles (id) ON DELETE CASCADE",
+            references=Reference("articles", "id", on_delete="cascade"),
         ),
         Column(
             name="title_id",
-            references="titles (id) ON DELETE CASCADE",
+            references=Reference("titles", "id", on_delete="cascade"),
         ),
     ],
 )
@@ -84,15 +85,15 @@ table_top_articles = Table(
         Column(name="rank", type_="INTEGER"),
         Column(
             name="frontpage_id",
-            references="frontpages (id) ON DELETE CASCADE",
+            references=Reference("frontpages", "id", on_delete="cascade"),
         ),
         Column(
             name="article_id",
-            references="articles (id) ON DELETE CASCADE",
+            references=Reference("articles", "id", on_delete="cascade"),
         ),
         Column(
             name="title_id",
-            references="titles (id) ON DELETE CASCADE",
+            references=Reference("titles", "id", on_delete="cascade"),
         ),
     ],
 )
@@ -100,7 +101,9 @@ table_embeddings = Table(
     name="embeddings",
     columns=[
         Column(name="id", primary_key=True),
-        Column(name="title_id", references="titles (id) ON DELETE CASCADE"),
+        Column(
+            name="title_id", references=Reference("titles", "id", on_delete="cascade")
+        ),
         Column(name="vector", type_="bytea"),
     ],
 )

src/media_observer/storage_abstraction.py

@@ -21,19 +21,31 @@ class UniqueIndex:
         await conn.execute(stmt)
 
 
+@frozen
+class Reference:
+    table_name: str
+    column_name: str
+    on_delete: str | None = None
+
+    @property
+    def as_sql(self):
+        on_delete = "ON DELETE CASCADE" if self.on_delete == "cascade" else None
+        return f"{self.table_name} ({self.column_name}) {on_delete}"
+
+
 @frozen
 class Column:
     name: str
     type_: str | None = None
     primary_key: bool = False
-    references: str | None = None
+    references: Reference | None = None
 
     @property
     def attrs(self):
         if self.primary_key:
             return "SERIAL PRIMARY KEY"
         elif self.references is not None:
-            return f"INTEGER REFERENCES {self.references}"
+            return f"INTEGER REFERENCES {self.references.as_sql}"
         elif self.type_ is not None:
             return self.type_
         else: