Implemented alembic migrations, added cursors closing task.

app/operations.py

@@ -1,59 +1,22 @@
 from fastapi.routing import APIRouter
 from typing_extensions import Annotated
 from pydantic import Field
-from data.schemas import (
-    SelectQueryBase,
-    SelectQueryInDB,
-    SelectQuery,
-    SelectQueryIn,
-    SelectResult,
-    SelectQueryMetaData,
-    SelectQueryInResult,
-)
-from fastapi import Depends, HTTPException, status
-from sqlalchemy.ext.asyncio import AsyncSession
+from data.schemas import SelectResult, CachedCursorOut
+from fastapi import Depends
 from data.crud import (
     read_connection,
-    create_select_query,
-    read_all_select_queries,
     read_select_query,
 )
 from core.dependencies import get_db, get_current_user, get_admin_user
-from core.exceptions import QueryNotFound, ConnectionNotFound, PoolNotFound
-from utils.sql_creator import build_sql_query_text
+from core.exceptions import (
+    QueryNotFound,
+    ConnectionNotFound,
+    PoolNotFound,
+    CursorNotFound,
+)
 from dbs import mysql
 
-router = APIRouter(prefix="/select")
-
-
-@router.post("/check-query", dependencies=[Depends(get_current_user)])
-async def check_select_query(query: SelectQueryBase) -> SelectQuery:
-    sql, params = build_sql_query_text(query)
-    q = SelectQuery(**query.model_dump(), params=params, sql=sql)
-    return q
-
-
-@router.post("/query")
-async def create_select_query_endpoint(
-    query: SelectQueryBase, db=Depends(get_db), user=Depends(get_current_user)
-) -> SelectQueryInDB:
-    sql, params = build_sql_query_text(query)
-    query_in = SelectQueryIn(
-        **query.model_dump(), owner_id=user.id, params=params, sql=sql
-    )
-    return await create_select_query(db=db, query=query_in)
-
-
-@router.get("/query", dependencies=[Depends(get_current_user)])
-async def get_select_queries_endpoint(db=Depends(get_db)) -> list[SelectQueryInDB]:
-    return await read_all_select_queries(db=db)
-
-
-@router.get("/query/{query_id}", dependencies=[Depends(get_current_user)])
-async def get_select_queries_endpoint(
-    query_id: int, db=Depends(get_db)
-) -> SelectQueryInDB:
-    return await read_select_query(db=db, query_id=query_id)
+router = APIRouter()
 
 
 @router.post("/execute", dependencies=[Depends(get_current_user)])
@@ -76,17 +39,41 @@ async def execute_select(
         raise PoolNotFound
 
     raw_result, rowcount = await mysql.execute_select_query(
-        pool=pool, query=query.sql, params=query.params, fetch_num=page_size
+        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))
 
-    meta = SelectQueryMetaData(
-        cursor=None, total_number=rowcount, has_more=len(results.data) != rowcount
-    )
-
     return SelectResult(
-        meta=meta,
-        query=query,
+        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},
+    )