Support PEP-567 context variables

Invoke the given function within a copy of the current PEP-567
context in `greenlet_spawn()`, so that subsequent sync methods could
retrieve the correct context variable.

Adjusted the greenlet integration, which provides support for Python asyncio
in SQLAlchemy, to accommodate for the handling of Python ``contextvars``
(introduced in Python 3.7) for ``greenlet`` versions greater than 0.4.17.
Greenlet version 0.4.17 added automatic handling of contextvars in a
backwards-incompatible way; we've coordinated with the greenlet authors to
add a preferred API for this in versions subsequent to 0.4.17 which is now
supported by SQLAlchemy's greenlet integration.  For greenlet versions prior
to 0.4.17 no behavioral change is needed, version 0.4.17 itself is blocked
from the dependencies.

Fixes: #5615
Closes: #5616
Pull-request: https://github.com/sqlalchemy/sqlalchemy/pull/5616
Pull-request-sha: dcf42f7b78ef3e983fda17f7190cda7b4511e3b4

Change-Id: I378953ee69e1ef2535ba51b97f28cbbbf9de3161

diff --git a/doc/build/changelog/unreleased_14/5615.rst b/doc/build/changelog/unreleased_14/5615.rst
new file mode 100644 (file)
index 0000000..df45ea1
--- /dev/null
+++ b/doc/build/changelog/unreleased_14/5615.rst
@@ -0,0 +1,13 @@
+.. change::
+  :tags: bug, asyncio
+  :tickets: 5615
+
+  Adjusted the greenlet integration, which provides support for Python asyncio
+  in SQLAlchemy, to accommodate for the handling of Python ``contextvars``
+  (introduced in Python 3.7) for ``greenlet`` versions greater than 0.4.17.
+  Greenlet version 0.4.17 added automatic handling of contextvars in a
+  backwards-incompatible way; we've coordinated with the greenlet authors to
+  add a preferred API for this in versions subsequent to 0.4.17 which is now
+  supported by SQLAlchemy's greenlet integration.  For greenlet versions prior
+  to 0.4.17 no behavioral change is needed, version 0.4.17 itself is blocked
+  from the dependencies.
\ No newline at end of file

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..ba53ea63522fa9c7819a1b92f0416458731b8093 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,34 @@ class TestAsyncioCompat(fixtures.TestBase):
             await greenlet_spawn(go)
 
         await to_await
+
+    @async_test
+    @testing.requires.python37
+    async def test_contextvars(self):
+        import asyncio
+        import contextvars
+
+        var = contextvars.ContextVar("var")
+        concurrency = 5
+
+        async def async_inner(val):
+            eq_(val, var.get())
+            return var.get()
+
+        def inner(val):
+            retval = await_only(async_inner(val))
+            eq_(val, var.get())
+            eq_(retval, val)
+            return retval
+
+        async def task(val):
+            var.set(val)
+            return await greenlet_spawn(inner, val)
+
+        values = {
+            await coro
+            for coro in asyncio.as_completed(
+                [task(i) for i in range(concurrency)]
+            )
+        }
+        eq_(values, set(range(concurrency)))

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]