139 lines
4.1 KiB
Python
139 lines
4.1 KiB
Python
from fastapi import WebSocket
|
|
from fastapi.routing import APIRouter
|
|
from fastapi.responses import StreamingResponse
|
|
from typing_extensions import Annotated
|
|
from pydantic import Field
|
|
from data.schemas import SelectResult, CachedCursorOut
|
|
from fastapi import Depends
|
|
from data.crud import (
|
|
read_connection,
|
|
read_select_query,
|
|
)
|
|
from core.dependencies import get_db, get_current_user, get_user_from_api_key
|
|
from core.exceptions import (
|
|
QueryNotFound,
|
|
ConnectionNotFound,
|
|
PoolNotFound,
|
|
CursorNotFound,
|
|
)
|
|
from dbs import mysql
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
@router.post("/execute", dependencies=[Depends(get_current_user)])
|
|
async def execute_select(
|
|
query_id: int,
|
|
connection_id: int,
|
|
page_size: Annotated[int, Field(ge=1, le=100)] = 50,
|
|
db=Depends(get_db),
|
|
) -> SelectResult:
|
|
query = await read_select_query(db=db, query_id=query_id)
|
|
|
|
if query is None:
|
|
raise QueryNotFound
|
|
connection = await read_connection(db=db, connection_id=connection_id)
|
|
|
|
if connection is None:
|
|
raise ConnectionNotFound
|
|
pool = mysql.pools.get(connection.id, None)
|
|
if pool is None:
|
|
raise PoolNotFound
|
|
|
|
raw_result, rowcount = await mysql.execute_select_query(
|
|
pool=pool, sql_query=query.sql, params=query.params, fetch_num=page_size
|
|
)
|
|
|
|
results = mysql.dict_result_to_list(result=mysql.serializer(raw_result=raw_result))
|
|
|
|
return SelectResult(
|
|
cursor=CachedCursorOut(
|
|
id=None,
|
|
connection_id=connection_id,
|
|
query=query,
|
|
row_count=rowcount,
|
|
fetched_rows=len(results.data),
|
|
is_closed=True,
|
|
has_more=len(results.data) != rowcount,
|
|
ttl=-1,
|
|
close_at=-1,
|
|
),
|
|
results=results,
|
|
)
|
|
|
|
|
|
@router.get(path="/fetch_cursor", dependencies=[Depends(get_current_user)])
|
|
async def fetch_cursor(
|
|
cursor_id: str,
|
|
page_size: Annotated[int, Field(ge=1, le=1000)] = 50,
|
|
) -> SelectResult:
|
|
cached_cursor = mysql.cached_cursors.get(cursor_id, None)
|
|
if cached_cursor is None:
|
|
raise CursorNotFound
|
|
result = await cached_cursor.fetch_many(size=page_size)
|
|
|
|
if cached_cursor.done:
|
|
mysql.cached_cursors.pop(cursor_id, None)
|
|
|
|
return SelectResult(
|
|
cursor=cached_cursor,
|
|
results={"columns": cached_cursor.query.columns, "data": result},
|
|
)
|
|
|
|
|
|
@router.get(
|
|
"/sse-stream-cursor", dependencies=[Depends(get_current_user)], status_code=200
|
|
)
|
|
async def server_side_events_stream_cursor(
|
|
cursor_id: str,
|
|
page_size: Annotated[int, Field(ge=1, le=1000)] = 50,
|
|
):
|
|
cached_cursor = mysql.cached_cursors.get(cursor_id, None)
|
|
if cached_cursor is None:
|
|
raise CursorNotFound
|
|
|
|
async def stream():
|
|
while True:
|
|
result = await fetch_cursor(cursor_id=cursor_id, page_size=page_size)
|
|
|
|
serialized_result = result.model_dump_json()
|
|
yield f"data: {serialized_result}\n\n" # Format as Server-Sent Event (SSE)
|
|
if result.cursor.is_closed:
|
|
break
|
|
|
|
return StreamingResponse(stream(), media_type="text/event-stream")
|
|
|
|
|
|
@router.websocket(
|
|
path="/ws-stream/{cursor_id}",
|
|
)
|
|
async def websocket_stream_cursor(
|
|
websocket: WebSocket,
|
|
cursor_id: str,
|
|
page_size: Annotated[int, Field(ge=1, le=1000)] = 50,
|
|
db=Depends(get_db),
|
|
):
|
|
await websocket.accept()
|
|
api_key = websocket.headers.get("Authorization")
|
|
user = await get_user_from_api_key(db=db, api_key=api_key)
|
|
if user is None:
|
|
await websocket.close(reason="Invalid credentials", code=1008)
|
|
return
|
|
|
|
cached_cursor = mysql.cached_cursors.get(cursor_id, None)
|
|
if cached_cursor is None:
|
|
e = CursorNotFound()
|
|
await websocket.close(reason=str(e.detail), code=1011)
|
|
return
|
|
|
|
while True:
|
|
result = await fetch_cursor(cursor_id=cursor_id, page_size=page_size)
|
|
|
|
serialized_result = result.model_dump_json()
|
|
await websocket.send_text(
|
|
f"data: {serialized_result}\n\n"
|
|
) # Format as Server-Sent Event (SSE)
|
|
if result.cursor.is_closed:
|
|
break
|
|
await websocket.close(reason="Done")
|