Support PEP-567 context variables

Fixes: #5615

diff --git a/lib/sqlalchemy/util/_concurrency_py3k.py b/lib/sqlalchemy/util/_concurrency_py3k.py
index 5d11bf92c16e1ad5806b715cf72a9637d73404b1..dcee057134e9cacca68b717b4f08694c46395e7c 100644 (file)
--- a/lib/sqlalchemy/util/_concurrency_py3k.py
+++ b/lib/sqlalchemy/util/_concurrency_py3k.py
@@ -8,6 +8,16 @@ import greenlet
 
 from .. import exc
 
+try:
+    from contextvars import copy_context as _copy_context
+
+    # If greenlet.gr_context is present in current version of greenlet,
+    # it will be set with a copy of the current context on creation.
+    # Refs: https://github.com/python-greenlet/greenlet/pull/198
+    getattr(greenlet.greenlet, "gr_context")
+except (ImportError, AttributeError):
+    _copy_context = None
+
 
 # implementation based on snaury gist at
 # https://gist.github.com/snaury/202bf4f22c41ca34e56297bae5f33fef
@@ -18,6 +28,8 @@ class _AsyncIoGreenlet(greenlet.greenlet):
     def __init__(self, fn, driver):
         greenlet.greenlet.__init__(self, fn, driver)
         self.driver = driver
+        if _copy_context is not None:
+            self.gr_context = _copy_context()
 
 
 def await_only(awaitable: Coroutine) -> Any:

diff --git a/setup.cfg b/setup.cfg
index e8cd0dfb4beca3e97a9c915911c3d682aba64cc2..1912fd3cd6fe4b86a6de23486477696d66a60a9d 100644 (file)
--- a/setup.cfg
+++ b/setup.cfg
@@ -39,7 +39,7 @@ package_dir =
     =lib
 install_requires =
     importlib-metadata;python_version<"3.8"
-    greenlet;python_version>="3"
+    greenlet != 0.4.17;python_version>="3"
 
 [options.extras_require]
 asyncio =

diff --git a/test/base/test_concurrency_py3k.py b/test/base/test_concurrency_py3k.py
index 10b89291e0b709f05148d8ab2afce4b01bdddd88..44a9337d63d160af7cd52e74f3c7a66751f4c5a1 100644 (file)
--- a/test/base/test_concurrency_py3k.py
+++ b/test/base/test_concurrency_py3k.py
@@ -1,4 +1,5 @@
 from sqlalchemy import exc
+from sqlalchemy import testing
 from sqlalchemy.testing import async_test
 from sqlalchemy.testing import eq_
 from sqlalchemy.testing import expect_raises_message
@@ -7,6 +8,11 @@ from sqlalchemy.util import await_fallback
 from sqlalchemy.util import await_only
 from sqlalchemy.util import greenlet_spawn
 
+try:
+    from greenlet import greenlet
+except ImportError:
+    greenlet = None
+
 
 async def run1():
     return 1
@@ -101,3 +107,36 @@ class TestAsyncioCompat(fixtures.TestBase):
             await greenlet_spawn(go)
 
         await to_await
+
+    @async_test
+    @testing.requires.python37
+    @testing.skip_if(lambda: not hasattr(greenlet, "gr_context"))
+    async def test_contextvars(self):
+        import asyncio
+        import contextvars
+
+        var = contextvars.ContextVar("var")
+        event = asyncio.Event()
+        counter = [0]
+        concurrency = 5
+
+        async def async_inner(val):
+            eq_(val, var.get())
+
+        def inner(val):
+            await_only(async_inner(val))
+            eq_(val, var.get())
+
+        async def task(val):
+            var.set(val)
+            counter[0] += 1
+            if counter[0] == concurrency:
+                event.set()
+            await event.wait()
+            await greenlet_spawn(inner, val)
+
+        done, _ = await asyncio.wait(
+            [asyncio.ensure_future(task(i)) for i in range(concurrency)]
+        )
+        for fut in done:
+            await fut

diff --git a/tox.ini b/tox.ini
index 2167d3bb1204acbf9243a3fae2a4e348eabaf365..6cfcf62efc25acd3287a7625d3a849d2d51580de 100644 (file)
--- a/tox.ini
+++ b/tox.ini
@@ -17,7 +17,7 @@ usedevelop=
 
 deps=pytest>=4.6.11 # this can be 6.x once we are on python 3 only
      pytest-xdist
-     greenlet
+     greenlet != 0.4.17
      mock; python_version < '3.3'
      importlib_metadata; python_version < '3.8'
      postgresql: .[postgresql]