allow psycopg query types from plain migrations

tunnelworks · Dec 5, 2024 · f9c3a66 · f9c3a66
1 parent 0d0d655
commit f9c3a66
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/postcar/db/migrations/classes.py b/src/postcar/db/migrations/classes.py
@@ -3,6 +3,7 @@
 
 
 if t.TYPE_CHECKING:
+    from psycopg.abc import Query
     from postcar._types import Connection, Module
 
 
@@ -34,8 +35,8 @@ async def migrate(self) -> None:
     async def rollback(self) -> None:
         raise NotImplementedError()
 
-    def get_forward(self) -> t.Optional[str]:
+    def get_forward(self) -> t.Optional["Query"]:
         return None
 
-    def get_rollback(self) -> t.Optional[str]:
+    def get_rollback(self) -> t.Optional["Query"]:
         return None