#24 omg

Author: HEKUCHAN

api/api/db/dependencies.py

@@ -1,17 +1,19 @@
 from typing import AsyncGenerator
 
 from sqlalchemy.ext.asyncio import AsyncSession
-from starlette.requests import Request
+from starlette.requests import Request, HTTPConnection
 
 
-async def get_db_session(request: Request) -> AsyncGenerator[AsyncSession, None]:
+async def get_db_session(
+    connection: HTTPConnection,
+) -> AsyncGenerator[AsyncSession, None]:
     """
     Create and get database session.
 
-    :param request: current request.
+    :param connection: current request.
     :yield: database session.
     """
-    session: AsyncSession = request.app.state.db_session_factory()
+    session: AsyncSession = connection.app.state.db_session_factory()
 
     try:  # noqa: WPS501
         yield session

api/api/db/migrations/versions/2024-02-02-07-12_1522526dbc9c.py

@@ -1,15 +1,15 @@
-"""Add Rooms, User, RoomLinkUser Table
+"""add tables
 
-Revision ID: 25615d348393
+Revision ID: a52660341927
 Revises: 
-Create Date: 2024-02-02 06:45:34.408905
+Create Date: 2024-02-03 17:07:40.830802
 
 """
 import sqlalchemy as sa
 from alembic import op
 
 # revision identifiers, used by Alembic.
-revision = "25615d348393"
+revision = "a52660341927"
 down_revision = None
 branch_labels = None
 depends_on = None

api/api/db/migrations/versions/2024-02-02-06-45_25615d348393.py api/api/db/migrations/versions/2024-02-03-17-07_a52660341927.py

@@ -1,16 +1,16 @@
-"""Add Rooms, User, RoomLinkUser Table
+"""empty message
 
-Revision ID: 668743a5e60c
-Revises: 25615d348393
-Create Date: 2024-02-02 07:05:50.008640
+Revision ID: 9b57dbe84d37
+Revises: a52660341927
+Create Date: 2024-02-03 17:52:34.833449
 
 """
 import sqlalchemy as sa
 from alembic import op
 
 # revision identifiers, used by Alembic.
-revision = "668743a5e60c"
-down_revision = "25615d348393"
+revision = "9b57dbe84d37"
+down_revision = "a52660341927"
 branch_labels = None
 depends_on = None
 

api/api/db/migrations/versions/2024-02-02-07-05_668743a5e60c.py api/api/db/migrations/versions/2024-02-03-17-52_9b57dbe84d37.py

@@ -21,8 +21,12 @@ class RoomLinkUser(Base):
         primary_key=True,
     )
     room: Mapped["RoomModel"] = relationship(
-        "RoomModel", back_populates="user_associations"
+        "RoomModel",
+        back_populates="user_associations",
+        viewonly=True,
     )
     user: Mapped["UserModel"] = relationship(
-        "UserModel", back_populates="room_associations"
+        "UserModel",
+        back_populates="room_associations",
+        viewonly=True,
     )

api/api/db/models/associations/room_user.py

@@ -23,9 +23,10 @@ class RoomModel(Base):
     owner_id: Mapped[UUID] = mapped_column(UUID(as_uuid=True), ForeignKey("users.id"))
 
     users: Mapped[List["UserModel"]] = relationship(
-        secondary="room_link_user",
-        back_populates="rooms",
+        secondary="room_link_user", back_populates="rooms", overlaps="room"
     )
     user_associations: Mapped[List[RoomLinkUser]] = relationship(
         back_populates="room",
+        overlaps="rooms",
+        viewonly=True,
     )

api/api/db/models/room_model.py

@@ -22,9 +22,10 @@ class UserModel(Base):
     )
 
     rooms: Mapped[List["RoomModel"]] = relationship(
-        secondary="room_link_user",
-        back_populates="users",
+        secondary="room_link_user", back_populates="users", overlaps="user"
     )
     room_associations: Mapped[List[RoomLinkUser]] = relationship(
         back_populates="user",
+        overlaps="users",
+        viewonly=True,
     )

api/api/db/models/user_model.py

@@ -1,19 +1,22 @@
 from typing import Optional
+from api.db.dependencies import get_db_session
 from api.db.models.user_model import UserModel
-from api.schemas.user import UserInfo
 
 from fastapi import Cookie, Depends, Header, HTTPException, status
+from starlette.requests import HTTPConnection
 from jose import ExpiredSignatureError, JWTError, jwt
+from sqlalchemy.ext.asyncio import AsyncSession
 
 from api.db.dao.user_dao import UserDAO
 from api.settings import settings
 
 
 async def with_authenticate(
     user_dao: UserDAO = Depends(),
+    access_token: str = Cookie(default=None),
     authorization: Optional[str] = Header(default=None),
 ) -> UserModel:
-    if authorization is None:
+    if authorization is None and access_token is None:
         HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,
             detail="Authorization credentials is missing.",
@@ -22,6 +25,8 @@ async def with_authenticate(
 
     if authorization is not None:
         jwt_token = authorization.rsplit(maxsplit=1)[-1]
+    elif access_token is not None:
+        jwt_token = access_token
     else:
         raise HTTPException(
             status_code=401,
@@ -56,3 +61,28 @@ async def with_authenticate(
         detail="Not found user.",
         headers={"WWW-Authenticate": 'Bearer error="not_found_user"'},
     )
+
+
+async def ws_with_autheticate(
+    connection: HTTPConnection,
+    db_session: AsyncSession = Depends(get_db_session),
+):
+    jwt_token = connection.cookies.get("access_token", None)
+
+    if jwt_token is None:
+        return None
+
+    try:
+        payload = jwt.decode(
+            jwt_token, settings.token_secret_key, algorithms=[settings.token_algorithm]
+        )
+    except:
+        return None
+
+    user_dao = UserDAO(db_session)
+    user = await user_dao.get_user(payload["user_id"])
+
+    if user is not None:
+        return user
+
+    return None

api/api/dependencies/auth.py

@@ -0,0 +1,2 @@
+class NotFoundException(Exception):
+    pass

api/api/libs/errors.py

@@ -1,7 +1,7 @@
 import enum
 from pathlib import Path
 from tempfile import gettempdir
-from typing import Literal
+from typing import List, Literal
 
 from pydantic_settings import BaseSettings, SettingsConfigDict
 from yarl import URL
@@ -36,6 +36,12 @@ class Settings(BaseSettings):
     port: int = 8000
     web_uri: str = "http://localhost:3000/"
 
+    origins: List[str] = [
+        "http://localhost:3000",
+        "http://127.0.0.1:3000",
+        "http://web:3000",
+    ]
+
     # quantity of workers for uvicorn
     workers_count: int = 1
     # Enable uvicorn reloading

api/api/settings.py

@@ -18,7 +18,6 @@ async def generate_token(
     refresh_token: str = Cookie(default=None),
     user_dao: UserDAO = Depends(),
 ) -> Response:
-    print(refresh_token)
     if refresh_token is None and token_dto.refresh_token is None:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,

api/api/web/api/token/views.py

@@ -1,9 +1,11 @@
 from importlib import metadata
 from pathlib import Path
+from api.settings import settings
 
 from fastapi import FastAPI
 from fastapi.responses import UJSONResponse
 from fastapi.staticfiles import StaticFiles
+from fastapi.middleware.cors import CORSMiddleware
 
 from api.logging import configure_logging
 from api.web.api.router import api_router
@@ -30,6 +32,15 @@ def get_app() -> FastAPI:
         openapi_url="/api/openapi.json",
         default_response_class=UJSONResponse,
     )
+
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=settings.origins,
+        allow_credentials=True,
+        allow_methods=["*"],
+        allow_headers=["*"],
+    )
+
     # Adds startup and shutdown events.
     register_startup_event(app)
     register_shutdown_event(app)

api/api/web/application.py

@@ -0,0 +1,104 @@
+from typing import List, Set
+import uuid
+from api.db.models.room_model import RoomModel
+from api.db.models.user_model import UserModel
+from api.libs.errors import NotFoundException
+from api.web.ws.rooms.room import Room
+from api.web.ws.rooms.send.schemas import SendContent, SendJoinRoomBody
+from api.web.ws.rooms.user import User
+from fastapi import WebSocket
+
+
+class ConnectionManager:
+    def __init__(self):
+        self.rooms: List[Room] = []
+
+    async def connect(
+        self,
+        username: str,
+        websocket: WebSocket,
+        user_model: UserModel,
+        room_model: RoomModel,
+    ):
+        room = self.get_room_or_none(room_model.id)
+        if room is None:
+            room = self.create_room(room_model.id, room_model)
+
+        user = room.get_user_or_none(user_model.id)
+        if user is None:
+            user = self.create_user(
+                username,
+                websocket,
+                user_model,
+                add_room=room,
+            )
+        else:
+            user.websocket = websocket
+
+        try:
+            message_content = SendContent(
+                type="join_room",
+                body=SendJoinRoomBody(
+                    user_id=user_model.id,
+                    username=user.username,
+                ).model_dump(),
+            )
+            await room.send_json(message_content.model_dump(), exclude_users=[user])
+        except:
+            pass
+
+    async def disconnect(self, room_model: RoomModel, user_model: UserModel):
+        room = self.get_room_or_none(room_model.id)
+        if room is not None:
+            user = room.get_user_or_none(user_model.id)
+            if user is not None:
+                room.leave_user(user)
+
+    def create_room(self, room_id: uuid.UUID, room_model: RoomModel) -> Room:
+        room = Room(room_id, room_model)
+        self.add_room(room)
+        return room
+
+    def create_user(self, username: str, websocket: WebSocket, user_model: UserModel, add_room: Room | None = None) -> User:
+        user = User(
+            id=user_model.id,
+            username=username,
+            websocket=websocket,
+            user_model=user_model,
+        )
+
+        if add_room is not None:
+            add_room.add_user(user)
+
+        return user
+
+    def get_room(self, room_id: uuid.UUID) -> Room:
+        for room in self.rooms:
+            if room.id == room_id:
+                return room
+        raise NotFoundException("Not found room from room_id.")
+
+    def get_room_or_none(
+        self, room_uuid: uuid.UUID
+    ) -> Room | None:
+        try:
+            return self.get_room(room_uuid)
+        except NotFoundException:
+            return None
+
+    def add_room(self, room: Room) -> None:
+        self.rooms.append(room)
+
+    def delete_room(self, room: Room) -> None:
+        if room in self.rooms:
+            self.rooms.discard(room)
+
+    def delete_room_by_id(self, room_id: uuid.UUID) -> None:
+        room = self.get_room_by_idI(room_id)
+        self.delete_room(room)
+
+    def get_room_by_id(self, uuid: uuid.UUID) -> Room | None:
+        for room in self.rooms:
+            if room.id == uuid:
+                return room
+        return None

api/api/web/ws/rooms/connection_manager.py

@@ -0,0 +1,34 @@
+import uuid
+from datetime import datetime
+from typing import Any, Literal, Union
+
+from pydantic import BaseModel
+
+
+ReceiveMessageType = Literal[
+    "chat",
+    "join_room",
+    "leave_room",
+]
+
+
+class ReceiveChatBody(BaseModel):
+    text: str
+    created_at: datetime
+
+
+class ReceiveJoinRoomBody(BaseModel):
+    username: str
+
+
+class ReceiveChangeNameBody(BaseModel):
+    username: str
+
+
+class ReceiveLeaveRoomBody(BaseModel):
+    pass
+
+
+class ReceiveContent(BaseModel):
+    type: ReceiveMessageType
+    body: Any