Merge pull request #376 from BracketJohn/add_async_apply_schema

procrastinate-org · Jan 7, 2021 · 4e96713 · 4e96713
2 parents 0895281 + 1cf5fdb
commit 4e96713
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/procrastinate/schema.py b/procrastinate/schema.py
@@ -23,3 +23,7 @@ def get_migrations_path() -> str:
     def apply_schema(self) -> None:
         queries = self.get_schema()
         self.connector.execute_query(query=queries)
+
+    async def apply_schema_async(self) -> None:
+        queries = self.get_schema()
+        await self.connector.execute_query_async(query=queries)
diff --git a/tests/unit/test_schema.py b/tests/unit/test_schema.py
@@ -1,5 +1,7 @@
 from collections import defaultdict
 
+import pytest
+
 
 def test_get_schema(app):
     assert app.schema_manager.get_schema().startswith("-- Procrastinate Schema")
@@ -15,3 +17,12 @@ def test_apply_schema(app, connector):
     app.schema_manager.apply_schema()
 
     assert connector.queries == [("apply_schema", {})]
+
+
+@pytest.mark.asyncio
+async def test_apply_schema_async(app, connector):
+    connector.reverse_queries = defaultdict(lambda: "apply_schema")
+    connector.set_schema_version_run = lambda *a, **kw: None
+    await app.schema_manager.apply_schema_async()
+
+    assert connector.queries == [("apply_schema", {})]