diff --git a/.pyflyby b/.pyflyby index 048eee1..db34f23 100644 --- a/.pyflyby +++ b/.pyflyby @@ -12,6 +12,7 @@ from learn_sql_model.models.hero import Hero from learn_sql_model.models.hero import HeroCreate from learn_sql_model.models.hero import HeroDelete from learn_sql_model.models.hero import HeroRead +from learn_sql_model.models.hero import Heros from learn_sql_model.models.hero import HeroUpdate from learn_sql_model.models.new import new from learn_sql_model.models.new import newCreate diff --git a/learn_sql_model/api/hero.py b/learn_sql_model/api/hero.py index 0f14b4d..bc17008 100644 --- a/learn_sql_model/api/hero.py +++ b/learn_sql_model/api/hero.py @@ -82,4 +82,4 @@ async def get_heros( session: Session = Depends(get_session), ) -> list[Hero]: "get all heros" - return HeroRead.list(session=session) + return Heros.list(session=session) diff --git a/learn_sql_model/api/websocket.py b/learn_sql_model/api/websocket.py index 670f90f..b4a5e0f 100644 --- a/learn_sql_model/api/websocket.py +++ b/learn_sql_model/api/websocket.py @@ -1,7 +1,10 @@ -from fastapi import APIRouter, WebSocket, WebSocketDisconnect +from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect from fastapi.responses import HTMLResponse +from sqlmodel import Session from learn_sql_model.api.websocket_connection_manager import manager +from learn_sql_model.config import get_session +from learn_sql_model.models.hero import Heros web_socket_router = APIRouter() @@ -70,3 +73,14 @@ async def websocket_endpoint(websocket: WebSocket): except WebSocketDisconnect: manager.disconnect(websocket, id) await manager.broadcast(f"Client #{id} left the chat", id) + + +@web_socket_router.websocket("/wsecho") +async def websocket_endpoint( + websocket: WebSocket, + session: Session = Depends(get_session), +): + await websocket.accept() + while True: + heros = Heros.list(session=session) + await websocket.send_text(heros.json()) diff --git a/learn_sql_model/game/game.py b/learn_sql_model/game/game.py index b4db071..6bef22e 100644 --- a/learn_sql_model/game/game.py +++ b/learn_sql_model/game/game.py @@ -12,13 +12,7 @@ import pygame import typer from typer import Typer -from learn_sql_model.models.hero import ( - Hero, - HeroCreate, - HeroDelete, - HeroRead, - HeroUpdate, -) +from learn_sql_model.models.hero import Hero, HeroCreate, HeroDelete, HeroUpdate speed = 10 @@ -42,7 +36,7 @@ class Client: self.moving_left = False self.moving_right = False - self.others = [hero for hero in HeroRead.list() if hero.id != self.hero.id] + self.others = [hero for hero in Heros.list() if hero.id != self.hero.id] self.ticks = 0 atexit.register(self.quit) @@ -73,7 +67,7 @@ class Client: HeroUpdate( **{k: v for k, v in self.hero.dict().items() if v is not None} ).update() - self.others = [hero for hero in HeroRead.list() if hero.id != self.hero.id] + self.others = [hero for hero in Heros.list() if hero.id != self.hero.id] def render(self): # Console().print(self.hero) diff --git a/learn_sql_model/models/hero.py b/learn_sql_model/models/hero.py index 718fb33..8c2455b 100644 --- a/learn_sql_model/models/hero.py +++ b/learn_sql_model/models/hero.py @@ -54,6 +54,10 @@ class HeroRead(HeroBase): raise HTTPException(status_code=404, detail="Hero not found") return hero + +class Heros(BaseModel): + heros: list[Hero] + @classmethod def list( self, @@ -64,17 +68,23 @@ class HeroRead(HeroBase): ) -> Hero: # with config.database.session as session: - if session is None: - engine = get_config().database.engine - with Session(engine) as session: + def get_heros(session, where, offset, limit): statement = select(Hero) if where != "None" and where is not None: from sqlmodel import text statement = statement.where(text(where)) statement = statement.offset(offset).limit(limit) - heroes = session.exec(statement).all() - return heroes + heros = session.exec(statement).all() + return Heros(heros=heros) + + if session is None: + engine = get_config().database.engine + with Session(engine) as session: + heros = get_heros(session, where, offset, limit) + return heros + + return get_heros(session, where, offset, limit) class HeroUpdate(SQLModel):