[+] improve websocket api

This commit is contained in:
Siarhei Siniak 2025-07-14 11:41:02 +03:00
parent 612d807bc4
commit 7d6ce1eaee

@ -1,4 +1,5 @@
import fastapi import fastapi
import pydantic
import functools import functools
import logging import logging
import copy import copy
@ -8,7 +9,7 @@ import sys
from .settings import Settings as APISettings from .settings import Settings as APISettings
from typing import (Any, Optional,) from typing import (Any, Optional, Literal, Annotated,)
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -28,11 +29,41 @@ class WebsocketAPI:
self.connections.add(client) self.connections.add(client)
class Subscribe(pydantic.BaseModel):
action: Literal['subscribe']
class message_t(pydantic.BaseModel):
asset_id: Annotated[
int,
pydantic.Field(alias='assetId')
]
message: message_t
class Assets(pydantic.BaseModel):
action: Literal['assets']
class message_t(pydantic.BaseModel):
pass
message: Annotated[
message_t,
pydantic.Field(
default_factory=message_t,
)
]
Message = pydantic.RootModel[
Assets | Subscribe
]
async def on_message( async def on_message(
self, self,
client: fastapi.WebSocket, client: fastapi.WebSocket,
msg: Any msg_raw: str
) -> None: ) -> None:
msg = self.Message.model_validate_json(
msg_raw
).root
raise NotImplementedError raise NotImplementedError
async def disconnect(self, client: fastapi.WebSocket) -> None: async def disconnect(self, client: fastapi.WebSocket) -> None:
@ -55,6 +86,7 @@ async def websocket_tickers(
# websocket_api.disconnect(websocket) # websocket_api.disconnect(websocket)
except: except:
logger.exception('') logger.exception('')
raise
finally: finally:
await websocket_api.disconnect(websocket) await websocket_api.disconnect(websocket)
@ -68,7 +100,7 @@ def create_app() -> fastapi.FastAPI:
)( )(
functools.partial( functools.partial(
websocket_tickers, websocket_tickers,
websocket_api=websocket_api, websocket_api=fastapi.Depends(lambda : websocket_api),
) )
) )