Fix Sanic support

Author: fantix

gino/ext/sanic.py

@@ -1,10 +1,14 @@
 # noinspection PyPackageRequirements
 from sanic.exceptions import NotFound
+from sqlalchemy.engine.url import URL
+try:
+    # noinspection PyPackageRequirements
+    from aiocontextvars import enable_inherit, disable_inherit
+except ImportError:
+    enable_inherit = disable_inherit = lambda: None
 
 from ..api import Gino as _Gino, GinoExecutor as _Executor
-from ..local import enable_task_local, disable_task_local
-from ..connection import GinoConnection as _Connection
-from ..pool import GinoPool as _Pool
+from ..engine import GinoConnection as _Connection, GinoEngine as _Engine
 
 
 class SanicModelMixin:
@@ -36,7 +40,9 @@ async def first_or_404(self, *args, **kwargs):
 
 
 # noinspection PyClassHasNoInit
-class GinoPool(_Pool):
+class GinoEngine(_Engine):
+    connection_cls = GinoConnection
+
     async def first_or_404(self, *args, **kwargs):
         rv = await self.first(*args, **kwargs)
         if rv is None:
@@ -65,51 +71,59 @@ class Gino(_Gino):
     """
     model_base_classes = _Gino.model_base_classes + (SanicModelMixin,)
     query_executor = GinoExecutor
-    connection_cls = GinoConnection
-    pool_cls = GinoPool
+
+    def __init__(self, app=None, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        if app is not None:
+            self.init_app(app)
 
     def init_app(self, app):
-        task_local_enabled = [False]
+        inherit_enabled = [False]
 
         if app.config.setdefault('DB_USE_CONNECTION_FOR_REQUEST', True):
             @app.middleware('request')
             async def on_request(request):
-                request['connection_ctx'] = ctx = self.acquire(lazy=True)
-                request['connection'] = await ctx.__aenter__()
+                request['connection'] = await self.acquire(lazy=True)
 
             @app.middleware('response')
             async def on_response(request, _):
-                ctx = request.pop('connection_ctx', None)
-                request.pop('connection', None)
-                if ctx is not None:
-                    await ctx.__aexit__(None, None, None)
+                conn = request.pop('connection', None)
+                if conn is not None:
+                    await self.release(conn)
 
         @app.listener('before_server_start')
         async def before_server_start(_, loop):
             if app.config.setdefault('DB_USE_CONNECTION_FOR_REQUEST', True):
-                enable_task_local(loop)
-                task_local_enabled[0] = True
-
-            await self.create_pool(
-                host=app.config.setdefault('DB_HOST', 'localhost'),
-                port=app.config.setdefault('DB_PORT', 5432),
-                user=app.config.setdefault('DB_USER', 'postgres'),
-                password=app.config.setdefault('DB_PASSWORD', ''),
-                database=app.config.setdefault('DB_DATABASE', 'postgres'),
+                enable_inherit(loop)
+                inherit_enabled[0] = True
+
+            await self.set_bind(
+                URL(
+                    drivername=app.config.setdefault('DB_DRIVER', 'asyncpg'),
+                    host=app.config.setdefault('DB_HOST', 'localhost'),
+                    port=app.config.setdefault('DB_PORT', 5432),
+                    username=app.config.setdefault('DB_USER', 'postgres'),
+                    password=app.config.setdefault('DB_PASSWORD', ''),
+                    database=app.config.setdefault('DB_DATABASE', 'postgres'),
+                ),
                 min_size=app.config.setdefault('DB_POOL_MIN_SIZE', 5),
                 max_size=app.config.setdefault('DB_POOL_MAX_SIZE', 10),
                 loop=loop,
             )
 
         @app.listener('after_server_stop')
         async def after_server_stop(_, loop):
-            await self._bind.close()
-            if task_local_enabled[0]:
-                disable_task_local(loop)
-                task_local_enabled[0] = False
+            await self.pop_bind().close()
+            if inherit_enabled[0]:
+                disable_inherit(loop)
+                inherit_enabled[0] = False
 
     async def first_or_404(self, *args, **kwargs):
         rv = await self.first(*args, **kwargs)
         if rv is None:
             raise NotFound('No such data')
         return rv
+
+    async def set_bind(self, bind, loop=None, **kwargs):
+        kwargs.setdefault('engine_cls', GinoEngine)
+        return await super().set_bind(bind, loop=loop, **kwargs)

tests/test_sanic.py

@@ -0,0 +1,73 @@
+import gino
+import sanic
+import pytest
+from sanic.response import text, json
+from gino.ext.sanic import Gino
+
+from .models import DB_ARGS, PG_URL
+
+
+# noinspection PyShadowingNames
+@pytest.fixture
+async def app():
+    app = sanic.Sanic()
+    app.config['DB_HOST'] = DB_ARGS['host']
+    app.config['DB_PORT'] = DB_ARGS['port']
+    app.config['DB_USER'] = DB_ARGS['user']
+    app.config['DB_PASSWORD'] = DB_ARGS['password']
+    app.config['DB_DATABASE'] = DB_ARGS['database']
+
+    db = Gino(app)
+
+    class User(db.Model):
+        __tablename__ = 'gino_users'
+
+        id = db.Column(db.BigInteger(), primary_key=True)
+        nickname = db.Column(db.Unicode(), default='noname')
+
+    @app.route('/')
+    async def root(request):
+        return text('Hello, world!')
+
+    @app.route('/users/<uid:int>')
+    async def get_user(request, uid):
+        return json((await User.get_or_404(uid)).to_dict())
+
+    @app.route('/users', methods=['POST'])
+    async def add_user(request):
+        u = await User.create(nickname=request.form.get('name'))
+        await u.query.gino.first_or_404()
+        await db.first_or_404(u.query)
+        await db.bind.first_or_404(u.query)
+        await request['connection'].first_or_404(u.query)
+        return json(u.to_dict())
+
+    e = await gino.create_engine(PG_URL)
+    try:
+        try:
+            await db.gino.create_all(e)
+            yield app
+        finally:
+            await db.gino.drop_all(e)
+    finally:
+        await e.close()
+
+
+def test_index_returns_200(app):
+    request, response = app.test_client.get('/')
+    assert response.status == 200
+    assert response.text == 'Hello, world!'
+
+
+def test(app):
+    request, response = app.test_client.get('/users/1')
+    assert response.status == 404
+
+    request, response = app.test_client.post('/users',
+                                             data=dict(name='fantix'))
+    assert response.status == 200
+    assert response.json == dict(id=1, nickname='fantix')
+
+    request, response = app.test_client.get('/users/1')
+    assert response.status == 200
+    assert response.json == dict(id=1, nickname='fantix')