From 8ad1c0c5948ba5175ed4f4f93448c2d50d0c727c Mon Sep 17 00:00:00 2001
From: Fantix King <fantix.king@gmail.com>
Date: Sun, 27 Sep 2020 15:08:18 -0500
Subject: [PATCH] Support PEP-567 context variables

Fixes: #5615
---
 lib/sqlalchemy/util/_concurrency_py3k.py | 12 ++++++++
 setup.cfg                                |  2 +-
 test/base/test_concurrency_py3k.py       | 39 ++++++++++++++++++++++++
 tox.ini                                  |  2 +-
 4 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/lib/sqlalchemy/util/_concurrency_py3k.py b/lib/sqlalchemy/util/_concurrency_py3k.py
index 5d11bf92c1..dcee057134 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 e8cd0dfb4b..1912fd3cd6 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 10b89291e0..44a9337d63 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,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 2167d3bb12..6cfcf62efc 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]
-- 
2.47.3