Support async pghistory.context (#220)

* Support async pghistory.context

* Fix test

Co-authored-by: qqii <[email protected]>

Author: wesleykendall

docs/requirements.txt

@@ -1,6 +1,7 @@
 arrow==1.3.0 ; python_full_version >= "3.9.0" and python_version < "4"
 asgiref==3.7.2 ; python_full_version >= "3.9.0" and python_version < "4"
 babel==2.13.0 ; python_full_version >= "3.9.0" and python_version < "4"
+backports-asyncio-runner==1.2.0 ; python_version >= "3.9" and python_version < "3.11"
 backrefs==5.8 ; python_version >= "3.9" and python_version < "4"
 beautifulsoup4==4.11.1 ; python_full_version >= "3.9.0" and python_version < "4"
 binaryornot==0.4.4 ; python_full_version >= "3.9.0" and python_version < "4"
@@ -81,6 +82,7 @@ pymdown-extensions==10.3 ; python_version >= "3.9" and python_version < "4"
 pyproject-api==1.8.0 ; python_full_version >= "3.9.0" and python_version < "4"
 pyproject-hooks==1.2.0 ; python_full_version >= "3.9.0" and python_version < "4.0"
 pyright==1.1.399 ; python_full_version >= "3.9.0" and python_version < "4"
+pytest-asyncio==1.1.0 ; python_version >= "3.9" and python_version < "4"
 pytest-cov==6.1.1 ; python_version >= "3.9" and python_version < "4"
 pytest-django==4.11.1 ; python_full_version >= "3.9.0" and python_version < "4"
 pytest-dotenv==0.5.2 ; python_full_version >= "3.9.0" and python_version < "4"

pghistory/runtime.py

@@ -5,6 +5,7 @@
 import uuid
 from typing import Any, Dict, Tuple, Union
 
+from asgiref.sync import sync_to_async
 from django.db import connection
 
 from pghistory import config, utils
@@ -184,7 +185,13 @@ def __enter__(self):
 
         return _tracker.value
 
+    async def __aenter__(self):
+        return await sync_to_async(self.__enter__)()
+
     def __exit__(self, *exc):
         if self._pre_execute_hook:
             delattr(_tracker, "value")
             self._pre_execute_hook.__exit__(*exc)
+
+    async def __aexit__(self, *exc):
+        return await sync_to_async(self.__exit__)(*exc)

pghistory/tests/test_tracking.py

@@ -6,6 +6,25 @@
 import pghistory.tests.models as test_models
 
 
+@pytest.mark.asyncio
+@pytest.mark.django_db
+async def test_async_context():
+    async with pghistory.context(example="context"):
+        dummy = await test_models.BigAutoFieldModel.objects.acreate()
+
+    event = await test_models.BigAutoFieldModel.pgh_event_model.objects.select_related(
+        "pgh_context"
+    ).aget()
+    assert event.pgh_obj_id == dummy.id
+    assert event.pgh_context is not None
+    assert event.pgh_context.metadata == {"example": "context"}
+
+    # async pytest django does not clean up the database. Clean up manually
+    await test_models.BigAutoFieldModel.pgh_event_model.objects.all().adelete()
+    await test_models.BigAutoFieldModel.objects.all().adelete()
+    await pghistory.models.Context.objects.all().adelete()
+
+
 @pytest.mark.django_db(transaction=True)
 def test_concurrent_index_creation():
     """

poetry.lock

@@ -87,6 +87,7 @@ dj-database-url = "2.3.0"
 psycopg2-binary = "2.9.10"
 pytest-django = "4.11.1"
 django-dynamic-fixture = "4.0.1"
+pytest-asyncio = "1.1.0"
 
 [tool.pytest.ini_options]
 xfail_strict = true