Upgrade demo to SQLAlchemy 2 ORM (#639)

Author: Dreamsorcerer

.github/workflows/ci.yaml

@@ -22,7 +22,7 @@ jobs:
     - name: Setup Python
       uses: actions/setup-python@v4
       with:
-        python-version: 3.9
+        python-version: '3.10'
         cache: 'pip'
         cache-dependency-path: '**/requirements*.txt'
     - name: Install dependencies

demo/database_auth/__main__.py

@@ -0,0 +1,6 @@
+from aiohttp.web import run_app
+
+from .main import init_app
+
+if __name__ == "__main__":
+    run_app(init_app())

demo/database_auth/db.py

@@ -1,34 +1,37 @@
 import sqlalchemy as sa
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
 
 
-metadata = sa.MetaData()
+class Base(DeclarativeBase):
+    metadata = sa.MetaData(naming_convention={
+        "ix": "ix_%(column_0_label)s",
+        "uq": "uq_%(table_name)s_%(column_0_name)s",
+        "ck": "ck_%(table_name)s_%(column_0_name)s",
+        "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
+        "pk": "pk_%(table_name)s"
+    })
 
 
-users = sa.Table(
-    'users', metadata,
-    sa.Column('id', sa.Integer, nullable=False),
-    sa.Column('login', sa.String(256), nullable=False),
-    sa.Column('passwd', sa.String(256), nullable=False),
-    sa.Column('is_superuser', sa.Boolean, nullable=False,
-              server_default='FALSE'),
-    sa.Column('disabled', sa.Boolean, nullable=False,
-              server_default='FALSE'),
+class User(Base):
+    """A user and their credentials."""
 
-    # indices
-    sa.PrimaryKeyConstraint('id', name='user_pkey'),
-    sa.UniqueConstraint('login', name='user_login_key'),
-)
+    __tablename__ = "users"
 
+    id: Mapped[int] = mapped_column(primary_key=True)
+    username: Mapped[str] = mapped_column(sa.String(256), unique=True, index=True)
+    password: Mapped[str] = mapped_column(sa.String(256))
+    is_superuser: Mapped[bool] = mapped_column(
+        default=False, server_default=sa.sql.expression.false())
+    disabled: Mapped[bool] = mapped_column(
+        default=False, server_default=sa.sql.expression.false())
+    permissions = relationship("Permission", cascade="all, delete")
 
-permissions = sa.Table(
-    'permissions', metadata,
-    sa.Column('id', sa.Integer, nullable=False),
-    sa.Column('user_id', sa.Integer, nullable=False),
-    sa.Column('perm_name', sa.String(64), nullable=False),
 
-    # indices
-    sa.PrimaryKeyConstraint('id', name='permission_pkey'),
-    sa.ForeignKeyConstraint(['user_id'], [users.c.id],
-                            name='user_permission_fkey',
-                            ondelete='CASCADE'),
-)
+class Permission(Base):
+    """A permission that grants a user access to something."""
+
+    __tablename__ = "permissions"
+
+    user_id: Mapped[int] = mapped_column(
+        sa.ForeignKey(User.id, ondelete="CASCADE"), primary_key=True)
+    name: Mapped[str] = mapped_column(sa.String(64), primary_key=True)

demo/database_auth/db_auth.py

@@ -1,62 +1,52 @@
 from enum import Enum
-from typing import Any, Optional, Union
 
 import sqlalchemy as sa
 from passlib.hash import sha256_crypt
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+from sqlalchemy.orm import selectinload
 
 from aiohttp_security.abc import AbstractAuthorizationPolicy
-from . import db
+from .db import User
+
+
+def _where_authorized(identity: str) -> tuple[sa.sql.ColumnElement[bool], ...]:
+    return (User.username == identity, ~User.disabled)
 
 
 class DBAuthorizationPolicy(AbstractAuthorizationPolicy):
-    def __init__(self, dbengine: Any):
-        self.dbengine = dbengine
-
-    async def authorized_userid(self, identity: str) -> Optional[str]:
-        async with self.dbengine.acquire() as conn:
-            where = sa.and_(db.users.c.login == identity,
-                            sa.not_(db.users.c.disabled))  # type: ignore[no-untyped-call]
-            query = db.users.count().where(where)
-            ret = await conn.scalar(query)
-            if ret:
-                return identity
-            else:
-                return None
-
-    async def permits(self, identity: Optional[str], permission: Union[str, Enum],
-                      context: None = None) -> bool:
-        async with self.dbengine.acquire() as conn:
-            where = sa.and_(db.users.c.login == identity,
-                            sa.not_(db.users.c.disabled))  # type: ignore[no-untyped-call]
-            query = db.users.select().where(where)
-            ret = await conn.execute(query)
-            user = await ret.fetchone()
-            if user is not None:
-                user_id = user[0]
-                is_superuser = user[3]
-                if is_superuser:
-                    return True
-
-                where = db.permissions.c.user_id == user_id
-                query = db.permissions.select().where(where)
-                ret = await conn.execute(query)
-                result = await ret.fetchall()
-                if ret is not None:
-                    for record in result:
-                        if record.perm_name == permission:
-                            return True
+    def __init__(self, dbsession: async_sessionmaker[AsyncSession]):
+        self.dbsession = dbsession
+
+    async def authorized_userid(self, identity: str) -> str | None:
+        where = _where_authorized(identity)
+        async with self.dbsession() as sess:
+            user_id = await sess.scalar(sa.select(User.id).where(*where))
+        return str(user_id) if user_id else None
+
+    async def permits(self, identity: str | None, permission: str | Enum,
+                      context: dict[str, object] | None = None) -> bool:
+        if identity is None:
+            return False
 
+        where = _where_authorized(identity)
+        stmt = sa.select(User).options(selectinload(User.permissions)).where(*where)
+        async with self.dbsession() as sess:
+            user = await sess.scalar(stmt)
+
+        if user is None:
             return False
+        if user.is_superuser:
+            return True
+        return any(p.name == permission for p in user.permissions)
+
+
+async def check_credentials(db_session: async_sessionmaker[AsyncSession],
+                            username: str, password: str) -> bool:
+    where = _where_authorized(username)
+    async with db_session() as sess:
+        hashed_pw = await sess.scalar(sa.select(User.password).where(*where))
 
+    if hashed_pw is None:
+        return False
 
-async def check_credentials(db_engine: Any, username: str, password: str) -> bool:
-    async with db_engine.acquire() as conn:
-        where = sa.and_(db.users.c.login == username,
-                        sa.not_(db.users.c.disabled))  # type: ignore[no-untyped-call]
-        query = db.users.select().where(where)
-        ret = await conn.execute(query)
-        user = await ret.fetchone()
-        if user is not None:
-            hashed = user[2]
-            return sha256_crypt.verify(password, hashed)
-    return False
+    return sha256_crypt.verify(password, hashed_pw)

demo/database_auth/handlers.py

@@ -8,7 +8,7 @@
 from .db_auth import check_credentials
 
 
-class Web(object):
+class Web:
     index_template = dedent("""
         <!doctype html>
             <head></head>
@@ -32,20 +32,18 @@ async def index(self, request: web.Request) -> web.Response:
                 message='Hello, {username}!'.format(username=username))
         else:
             template = self.index_template.format(message='You need to login')
-        response = web.Response(body=template.encode())
-        return response
+        return web.Response(text=template, content_type="text/html")
 
     async def login(self, request: web.Request) -> NoReturn:
         invalid_resp = web.HTTPUnauthorized(body=b"Invalid username/password combination")
         form = await request.post()
         login = form.get('login')
         password = form.get('password')
-        db_engine = request.app["db_engine"]
 
         if not (isinstance(login, str) and isinstance(password, str)):
             raise invalid_resp
 
-        if await check_credentials(db_engine, login, password):
+        if await check_credentials(request.app["db_session"], login, password):
             response = web.HTTPFound("/")
             await remember(request, response, login)
             raise response
@@ -54,26 +52,22 @@ async def login(self, request: web.Request) -> NoReturn:
 
     async def logout(self, request: web.Request) -> web.Response:
         await check_authorized(request)
-        response = web.Response(body=b'You have been logged out')
+        response = web.Response(text="You have been logged out")
         await forget(request, response)
         return response
 
     async def internal_page(self, request: web.Request) -> web.Response:
         await check_permission(request, 'public')
-        response = web.Response(
-            body=b'This page is visible for all registered users')
-        return response
+        return web.Response(text="This page is visible for all registered users")
 
     async def protected_page(self, request: web.Request) -> web.Response:
         await check_permission(request, 'protected')
-        response = web.Response(body=b'You are on protected page')
-        return response
+        return web.Response(text="You are on protected page")
 
     def configure(self, app: web.Application) -> None:
         router = app.router
         router.add_route('GET', '/', self.index, name='index')
         router.add_route('POST', '/login', self.login, name='login')
         router.add_route('GET', '/logout', self.logout, name='logout')
         router.add_route('GET', '/public', self.internal_page, name='public')
-        router.add_route('GET', '/protected', self.protected_page,
-                         name='protected')
+        router.add_route('GET', '/protected', self.protected_page, name='protected')

demo/database_auth/main.py

@@ -1,59 +1,48 @@
-import asyncio
-from typing import Tuple
-
 from aiohttp import web
-from aiohttp_session import setup as setup_session
-from aiohttp_session.redis_storage import RedisStorage
-from aiopg.sa import create_engine
-from aioredis import create_pool  # type: ignore[attr-defined]
+from aiohttp_session import SimpleCookieStorage, setup as setup_session
+from sqlalchemy.ext.asyncio import (AsyncEngine, AsyncSession, async_sessionmaker,
+                                    create_async_engine)
 
 from aiohttp_security import SessionIdentityPolicy
 from aiohttp_security import setup as setup_security
+from .db import Base, Permission, User
 from .db_auth import DBAuthorizationPolicy
 from .handlers import Web
 
 
-async def init(loop: asyncio.AbstractEventLoop) -> Tuple[asyncio.Server, web.Application,
-                                                         web.Server]:
-    redis_pool = await create_pool(('localhost', 6379))
-    db_engine = await create_engine(  # type: ignore[no-untyped-call]  # noqa: S106
-        user="aiohttp_security", password="aiohttp_security",
-        database="aiohttp_security", host="127.0.0.1")
+async def init_db(db_engine: AsyncEngine, db_session: async_sessionmaker[AsyncSession]) -> None:
+    """Initialise DB with sample data."""
+    async with db_engine.begin() as conn:
+        await conn.run_sync(Base.metadata.create_all)
+    async with db_session.begin() as sess:
+        pw = "$5$rounds=535000$2kqN9fxCY6Xt5/pi$tVnh0xX87g/IsnOSuorZG608CZDFbWIWBr58ay6S4pD"
+        sess.add(User(username="admin", password=pw, is_superuser=True))
+        moderator = User(username="moderator", password=pw)
+        user = User(username="user", password=pw)
+        sess.add(moderator)
+        sess.add(user)
+    async with db_session.begin() as sess:
+        sess.add(Permission(user_id=moderator.id, name="protected"))
+        sess.add(Permission(user_id=moderator.id, name="public"))
+        sess.add(Permission(user_id=user.id, name="public"))
+
+
+async def init_app() -> web.Application:
     app = web.Application()
-    app["db_engine"] = db_engine
-    setup_session(app, RedisStorage(redis_pool))
-    setup_security(app,
-                   SessionIdentityPolicy(),
-                   DBAuthorizationPolicy(db_engine))
-
-    web_handlers = Web()
-    web_handlers.configure(app)
 
-    handler = app.make_handler()
-    srv = await loop.create_server(handler, '127.0.0.1', 8080)
-    print('Server started at http://127.0.0.1:8080')
-    return srv, app, handler
+    db_engine = create_async_engine("sqlite+aiosqlite:///:memory:")
+    app["db_session"] = async_sessionmaker(db_engine, expire_on_commit=False)
 
+    await init_db(db_engine, app["db_session"])
 
-async def finalize(srv: asyncio.Server, app: web.Application, handler: web.Server) -> None:
-    sock = srv.sockets[0]
-    app.loop.remove_reader(sock.fileno())
-    sock.close()
-
-    await handler.shutdown(1.0)
-    srv.close()
-    await srv.wait_closed()
-    await app.cleanup()
+    setup_session(app, SimpleCookieStorage())
+    setup_security(app, SessionIdentityPolicy(), DBAuthorizationPolicy(app["db_session"]))
 
+    web_handlers = Web()
+    web_handlers.configure(app)
 
-def main() -> None:
-    loop = asyncio.get_event_loop()
-    srv, app, handler = loop.run_until_complete(init(loop))
-    try:
-        loop.run_forever()
-    except KeyboardInterrupt:
-        loop.run_until_complete((finalize(srv, app, handler)))
+    return app
 
 
-if __name__ == '__main__':
-    main()
+if __name__ == "__main__":
+    web.run_app(init_app())