User authentication with tests.

gns3server/db/database.py

@@ -1,26 +0,0 @@
-#!/usr/bin/env python
-#
-# Copyright (C) 2020 GNS3 Technologies Inc.
-#
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU General Public License as published by
-# the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
-#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-# GNU General Public License for more details.
-#
-# You should have received a copy of the GNU General Public License
-# along with this program.  If not, see <http://www.gnu.org/licenses/>.
-
-import os
-
-from sqlalchemy.ext.asyncio import create_async_engine
-from sqlalchemy.orm import declarative_base
-
-SQLALCHEMY_DATABASE_URL = os.environ.get("DATABASE_URI", "sqlite:///./sql_app.db")
-
-engine = create_async_engine(SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False})
-Base = declarative_base()

gns3server/db/models.py

@@ -21,7 +21,10 @@ from sqlalchemy import Boolean, Column, ForeignKey, Integer, String, DateTime, f
 from sqlalchemy.orm import relationship
 from sqlalchemy.types import TypeDecorator, CHAR
 from sqlalchemy.dialects.postgresql import UUID
-from .database import Base
+
+from sqlalchemy.orm import declarative_base
+
+Base = declarative_base()
 
 
 class GUID(TypeDecorator):
@@ -68,11 +71,14 @@ class BaseTable(Base):
                         onupdate=func.current_timestamp())
 
 
+def generate_uuid():
+    return str(uuid.uuid4())
+
 class User(BaseTable):
 
     __tablename__ = "users"
 
-    user_id = Column(GUID, primary_key=True, default=str(uuid.uuid4()))
+    user_id = Column(GUID, primary_key=True, default=generate_uuid)
     username = Column(String, unique=True, index=True)
     email = Column(String, unique=True, index=True)
     full_name = Column(String)

gns3server/db/repositories/base.py

@@ -16,14 +16,10 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 from sqlalchemy.ext.asyncio import AsyncSession
-from ..database import engine
 
 
 class BaseRepository:
 
-    async def db(self):
-        session = AsyncSession(engine)
-        try:
-            yield session
-        finally:
-            await session.close()
+    def __init__(self, db_session: AsyncSession) -> None:
+
+        self._db_session = db_session

gns3server/db/repositories/users.py

@@ -20,7 +20,6 @@ from typing import Optional, List
 from sqlalchemy import select, update, delete
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from ..database import engine
 from .base import BaseRepository
 
 import gns3server.db.models as models
@@ -30,73 +29,66 @@ from gns3server.services import auth_service
 
 class UsersRepository(BaseRepository):
 
-    def __init__(self) -> None:
+    def __init__(self, db_session: AsyncSession) -> None:
 
-        super().__init__()
+        super().__init__(db_session)
         self._auth_service = auth_service
 
     async def get_user(self, user_id: UUID) -> Optional[models.User]:
 
-        async with AsyncSession(engine) as session:
-            result = await session.execute(select(models.User).where(models.User.user_id == user_id))
-            return result.scalars().first()
+        result = await self._db_session.execute(select(models.User).where(models.User.user_id == user_id))
+        return result.scalars().first()
 
     async def get_user_by_username(self, username: str) -> Optional[models.User]:
 
-        async with AsyncSession(engine) as session:
-            result = await session.execute(select(models.User).where(models.User.username == username))
-            return result.scalars().first()
+        result = await self._db_session.execute(select(models.User).where(models.User.username == username))
+        return result.scalars().first()
 
     async def get_user_by_email(self, email: str) -> Optional[models.User]:
 
-        async with AsyncSession(engine) as session:
-            result = await session.execute(select(models.User).where(models.User.email == email))
-            return result.scalars().first()
+        result = await self._db_session.execute(select(models.User).where(models.User.email == email))
+        return result.scalars().first()
 
     async def get_users(self) -> List[models.User]:
 
-        async with AsyncSession(engine) as session:
-            result = await session.execute(select(models.User))
-            return result.scalars().all()
+        result = await self._db_session.execute(select(models.User))
+        return result.scalars().all()
 
     async def create_user(self, user: schemas.UserCreate) -> models.User:
 
-        async with AsyncSession(engine) as session:
-            hashed_password = self._auth_service.hash_password(user.password)
-            db_user = models.User(username=user.username,
-                                  email=user.email,
-                                  full_name=user.full_name,
-                                  hashed_password=hashed_password)
-            session.add(db_user)
-            await session.commit()
-            await session.refresh(db_user)
-            return db_user
+        hashed_password = self._auth_service.hash_password(user.password)
+        db_user = models.User(
+            username=user.username,
+            email=user.email,
+            full_name=user.full_name,
+            hashed_password=hashed_password
+        )
+        self._db_session.add(db_user)
+        await self._db_session.commit()
+        await self._db_session.refresh(db_user)
+        return db_user
 
     async def update_user(self, user_id: UUID, user_update: schemas.UserUpdate) -> Optional[models.User]:
 
-        async with AsyncSession(engine) as session:
+        update_values = user_update.dict(exclude_unset=True)
+        password = update_values.pop("password", None)
+        if password:
+            update_values["hashed_password"] = self._auth_service.hash_password(password=password)
 
-            update_values = user_update.dict(exclude_unset=True)
-            password = update_values.pop("password", None)
-            if password:
-                update_values["hashed_password"] = self._auth_service.hash_password(password=password)
+        query = update(models.User) \
+            .where(models.User.user_id == user_id) \
+            .values(update_values)
 
-            print(update_values)
-            query = update(models.User) \
-                .where(models.User.user_id == user_id) \
-                .values(update_values)
-
-            await session.execute(query)
-            await session.commit()
-            return await self.get_user(user_id)
+        await self._db_session.execute(query)
+        await self._db_session.commit()
+        return await self.get_user(user_id)
 
     async def delete_user(self, user_id: UUID) -> bool:
 
-        async with AsyncSession(engine) as session:
-            query = delete(models.User).where(models.User.user_id == user_id)
-            result = await session.execute(query)
-            await session.commit()
-            return result.rowcount > 0
+        query = delete(models.User).where(models.User.user_id == user_id)
+        result = await self._db_session.execute(query)
+        await self._db_session.commit()
+        return result.rowcount > 0
             #except:
             #    await session.rollback()
 

gns3server/db/tasks.py

@@ -15,20 +15,26 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
-from sqlalchemy.exc import SQLAlchemyError
+import os
+
+from fastapi import FastAPI
+from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.ext.asyncio import create_async_engine
 
-from .database import engine
 from .models import Base
 
 import logging
 log = logging.getLogger(__name__)
 
 
-async def connect_to_db() -> None:
+async def connect_to_db(app: FastAPI) -> None:
 
+    db_url = os.environ.get("GNS3_DATABASE_URI", "sqlite:///./sql_app.db")
+    engine = create_async_engine(db_url, connect_args={"check_same_thread": False}, future=True)
     try:
         async with engine.begin() as conn:
             await conn.run_sync(Base.metadata.create_all)
             log.info("Successfully connected to the database")
+        app.state._db_engine = engine
     except SQLAlchemyError as e:
         log.error(f"Error while connecting to the database: {e}")