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

5 files changed, 64 insertions, 2 deletions

diff --git a/doc/build/changelog/unreleased_14/5615.rst b/doc/build/changelog/unreleased_14/5615.rst
new file mode 100644
index 000000000..df45ea109
--- /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 5d11bf92c..dcee05713 100644
--- 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 e8cd0dfb4..1912fd3cd 100644
--- 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 10b89291e..ba53ea635 100644
--- 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 2167d3bb1..6cfcf62ef 100644
--- 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]