Support for aiohttp (tortoise#473)

Rodrigo Martins de Oliveira · web-flow · commit d501981d6bc9 · 2020-08-01T18:45:57.000+02:00
diff --git a/.coveragerc b/.coveragerc
@@ -9,5 +9,6 @@ omit =
     tortoise/contrib/sanic/*
     tortoise/contrib/starlette/*
     tortoise/contrib/fastapi/*
+    tortoise/contrib/aiohttp/*
 [report]
 show_missing = True
diff --git a/CONTRIBUTORS.rst b/CONTRIBUTORS.rst
@@ -35,6 +35,7 @@ Contributors
 * Seo Jiyeon ``jiyeonseo``
 * Blake Watters ``blakewatters``
 * Alexey Tylindus ``@mirrorrim``
+* Rodrigo Oliveira ``@allrod5``
 
 Special Thanks
 ==============
diff --git a/docs/contrib.rst b/docs/contrib.rst
@@ -11,6 +11,7 @@ Contrib
    contrib/quart
    contrib/sanic
    contrib/starlette
+   contrib/aiohttp
 
 
 
diff --git a/docs/contrib/aiohttp.rst b/docs/contrib/aiohttp.rst
@@ -0,0 +1,17 @@
+.. _contrib_aiohttp:
+
+================================
+Tortoise-ORM aiohttp integration
+================================
+
+We have a lightweight integration util ``tortoise.contrib.aiohttp`` which has a single function ``register_tortoise`` which sets up Tortoise-ORM on startup and cleans up on teardown.
+
+See the :ref:`example_aiohttp`
+
+Reference
+=========
+
+.. automodule:: tortoise.contrib.aiohttp
+    :members:
+    :undoc-members:
+    :show-inheritance:
diff --git a/docs/examples.rst b/docs/examples.rst
@@ -13,3 +13,4 @@ Examples
    examples/quart
    examples/sanic
    examples/starlette
+   examples/aiohttp
diff --git a/docs/examples/aiohttp.rst b/docs/examples/aiohttp.rst
@@ -0,0 +1,24 @@
+.. _example_aiohttp:
+
+===============
+AIOHTTP Example
+===============
+
+This is an example of the  :ref:`contrib_aiohttp`
+
+**Usage:**
+
+.. code-block:: sh
+
+    python3 main.py
+
+
+models.py
+=========
+.. literalinclude::  ../../examples/aiohttp/models.py
+
+main.py
+=======
+.. literalinclude::  ../../examples/aiohttp/main.py
+
+
diff --git a/examples/aiohttp/README.rst b/examples/aiohttp/README.rst
@@ -0,0 +1,11 @@
+Tortoise-ORM aiohttp example
+============================
+
+We have a lightweight integration util ``tortoise.contrib.aiohttp`` which has a single function ``register_tortoise`` which sets up Tortoise-ORM on startup and cleans up on teardown.
+
+Usage
+-----
+
+.. code-block:: sh
+
+    python3 main.py
diff --git a/examples/aiohttp/__init__.py b/examples/aiohttp/__init__.py
diff --git a/examples/aiohttp/main.py b/examples/aiohttp/main.py
@@ -0,0 +1,30 @@
+# pylint: disable=E0401,E0611
+import logging
+
+from aiohttp import web
+
+from models import Users
+from tortoise.contrib.aiohttp import register_tortoise
+
+logging.basicConfig(level=logging.DEBUG)
+
+
+async def list_all(request):
+    users = await Users.all()
+    return web.json_response({"users": [str(user) for user in users]})
+
+
+async def add_user(request):
+    user = await Users.create(name="New User")
+    return web.json_response({"user": str(user)})
+
+
+app = web.Application()
+app.add_routes([web.get("/", list_all), web.post("/user", add_user)])
+register_tortoise(
+    app, db_url="sqlite://:memory:", modules={"models": ["models"]}, generate_schemas=True
+)
+
+
+if __name__ == "__main__":
+    web.run_app(app, port=5000)
diff --git a/examples/aiohttp/models.py b/examples/aiohttp/models.py
@@ -0,0 +1,9 @@
+from tortoise import Model, fields
+
+
+class Users(Model):
+    id = fields.IntField(pk=True)
+    name = fields.CharField(50)
+
+    def __str__(self):
+        return f"User {self.id}: {self.name}"
diff --git a/tests/requirements-pypy.txt b/tests/requirements-pypy.txt
@@ -10,7 +10,7 @@ apipkg==1.5               # via execnet
 asynctest==0.13.0         # via -r requirements-pypy.in
 attrs==19.3.0             # via pytest
 certifi==2020.6.20        # via requests
-cffi==1.14.0              # via cryptography
+cffi==1.14.1              # via cryptography
 chardet==3.0.4            # via requests
 coverage==5.2.1           # via coveralls, pytest-cov
 coveralls==2.1.1          # via -r requirements-pypy.in
@@ -23,15 +23,15 @@ iso8601==0.1.12           # via -r ../requirements.txt
 more-itertools==8.4.0     # via pytest
 packaging==20.4           # via pytest
 pluggy==0.13.1            # via pytest
-py==1.9.0                 # via pytest
+py==1.9.0                 # via pytest, pytest-forked
 pycparser==2.20           # via cffi
 pydantic==1.6.1           # via -r requirements-pypy.in
 pymysql==0.9.2            # via aiomysql
 pyparsing==2.4.7          # via packaging
 pypika==0.37.16           # via -r ../requirements.txt
 pytest-cov==2.10.0        # via -r requirements-pypy.in
-pytest-forked==1.2.0      # via pytest-xdist
-pytest-xdist==1.33.0      # via -r requirements-pypy.in
+pytest-forked==1.3.0      # via pytest-xdist
+pytest-xdist==1.34.0      # via -r requirements-pypy.in
 pytest==5.3.5             # via -r requirements-pypy.in, pytest-cov, pytest-forked, pytest-xdist
 requests==2.24.0          # via coveralls
 six==1.15.0               # via cryptography, packaging, pytest-xdist
diff --git a/tests/requirements.in b/tests/requirements.in
@@ -52,3 +52,6 @@ pydantic
 
 # FastAPI support
 fastapi
+
+# Aiohttp support
+aiohttp
diff --git a/tests/requirements.txt b/tests/requirements.txt
@@ -6,6 +6,7 @@
 #
 -e .
 aiofiles==0.5.0           # via quart, sanic
+aiohttp==3.6.2            # via -r requirements.in
 aiomysql==0.0.20          # via -r requirements.in
 aiosqlite==0.15.0         # via -r ../docs/../requirements.txt, tortoise-orm
 alabaster==0.7.12         # via sphinx
diff --git a/tortoise/contrib/aiohttp/__init__.py b/tortoise/contrib/aiohttp/__init__.py
@@ -0,0 +1,91 @@
+import logging
+from typing import Dict, List, Optional
+
+from aiohttp import web  # pylint: disable=E0401
+
+from tortoise import Tortoise
+
+
+def register_tortoise(
+    app: web.Application,
+    config: Optional[dict] = None,
+    config_file: Optional[str] = None,
+    db_url: Optional[str] = None,
+    modules: Optional[Dict[str, List[str]]] = None,
+    generate_schemas: bool = False,
+) -> None:
+    """
+    Registers ``on_startup`` and ``on_shutdown`` hooks to set-up and tear-down
+    Tortoise-ORM inside a Aiohttp webserver.
+
+    You can configure using only one of ``config``, ``config_file``
+    and ``(db_url, modules)``.
+
+    Parameters
+    ----------
+    app:
+        Aiohttp app.
+    config:
+        Dict containing config:
+
+        Example
+        -------
+
+        .. code-block:: python3
+
+            {
+                'connections': {
+                    # Dict format for connection
+                    'default': {
+                        'engine': 'tortoise.backends.asyncpg',
+                        'credentials': {
+                            'host': 'localhost',
+                            'port': '5432',
+                            'user': 'tortoise',
+                            'password': 'qwerty123',
+                            'database': 'test',
+                        }
+                    },
+                    # Using a DB_URL string
+                    'default': 'postgres://postgres:qwerty123@localhost:5432/events'
+                },
+                'apps': {
+                    'models': {
+                        'models': ['__main__'],
+                        # If no default_connection specified, defaults to 'default'
+                        'default_connection': 'default',
+                    }
+                }
+            }
+
+    config_file:
+        Path to .json or .yml (if PyYAML installed) file containing config with
+        same format as above.
+    db_url:
+        Use a DB_URL string. See :ref:`db_url`
+    modules:
+        Dictionary of ``key``: [``list_of_modules``] that defined "apps" and modules that
+        should be discovered for models.
+    generate_schemas:
+        True to generate schema immediately. Only useful for dev environments
+        or SQLite ``:memory:`` databases
+
+    Raises
+    ------
+    ConfigurationError
+        For any configuration error
+    """
+
+    async def init_orm(app):  # pylint: disable=W0612
+        await Tortoise.init(config=config, config_file=config_file, db_url=db_url, modules=modules)
+        logging.info(f"Tortoise-ORM started, {Tortoise._connections}, {Tortoise.apps}")
+        if generate_schemas:
+            logging.info("Tortoise-ORM generating schema")
+            await Tortoise.generate_schemas()
+
+    async def close_orm(app):  # pylint: disable=W0612
+        await Tortoise.close_connections()
+        logging.info("Tortoise-ORM shutdown")
+
+    app.on_startup.append(init_orm)
+    app.on_cleanup.append(close_orm)
