Implemented alembic migrations, added cursors closing task.

app/__init__.py

@@ -1,8 +1,10 @@
 from fastapi import APIRouter
 
 from app.connections import router as connections_router
-from app.operations import router as router
+from app.operations import router as operations_router
 from app.users import router as user_router
+from app.cursors import router as cursors_router
+from app.queries import router as queries_router
 
 api_router = APIRouter()
 api_router.include_router(router=user_router, prefix="/users", tags=["Users"])
@@ -10,5 +12,11 @@ api_router.include_router(
     router=connections_router, prefix="/connections", tags=["Connections"]
 )
 api_router.include_router(
-    router=router, prefix='/operations', tags=["Operations"]
+    router=cursors_router, prefix='/cursors', tags=['Cursors']
+)
+api_router.include_router(
+    router=queries_router, prefix='/queries', tags=['Queries']
+)
+api_router.include_router(
+    router=operations_router, prefix='/operations', tags=["Operations"]
 )

app/cursors.py

@@ -0,0 +1,78 @@
+from fastapi.routing import APIRouter
+from data.schemas import (
+    CachedCursorOut,
+)
+from fastapi import Depends, status
+from data.crud import (
+    read_connection,
+    read_select_query,
+)
+from core.dependencies import get_db, get_current_user, get_admin_user
+from core.exceptions import (
+    QueryNotFound,
+    ConnectionNotFound,
+    CursorNotFound,
+)
+from dbs import mysql
+
+router = APIRouter()
+
+
+@router.post("/", dependencies=[Depends(get_current_user)])
+async def create_cursor_endpoint(
+    query_id: int,
+    connection_id: int,
+    db=Depends(get_db),
+) -> CachedCursorOut:
+    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
+
+    cached_cursor = await mysql.create_cursor(query=query, connection_id=connection_id)
+
+    mysql.cached_cursors[cached_cursor.id] = cached_cursor
+    print(mysql.cached_cursors)
+    return cached_cursor
+
+
+@router.get("/", dependencies=[Depends(get_current_user)])
+async def get_all_cursors() -> list[CachedCursorOut]:
+    return mysql.cached_cursors.values()
+
+
+@router.get("/{cursor_id}", dependencies=[Depends(get_current_user)])
+async def get_cursors(cursor_id: str) -> CachedCursorOut:
+    try:
+        return mysql.cached_cursors[cursor_id]
+    except KeyError:
+        raise CursorNotFound
+
+
+@router.delete(
+    "/",
+    dependencies=[Depends(get_admin_user)],
+    status_code=status.HTTP_204_NO_CONTENT,
+)
+async def close_all_cursor() -> None:
+    for cached_cursor in mysql.cached_cursors.values():
+        await cached_cursor.close()
+    mysql.cached_cursors.clear()
+
+
+@router.delete(
+    "/{cursor_id}",
+    dependencies=[Depends(get_current_user)],
+    status_code=status.HTTP_204_NO_CONTENT,
+)
+async def close_cursor(cursor_id: str) -> None:
+    cached_cursor = mysql.cached_cursors.get(cursor_id, None)
+    if cached_cursor is None:
+        raise CursorNotFound
+
+    await cached_cursor.close()
+    del mysql.cached_cursors[cursor_id]

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},
+    )

app/queries.py

@@ -0,0 +1,50 @@
+
+from fastapi.routing import APIRouter
+
+from data.schemas import (
+    SelectQueryBase,
+    SelectQueryInDB,
+    SelectQuery,
+    SelectQueryIn,
+)
+from fastapi import Depends
+from data.crud import (
+    create_select_query,
+    read_all_select_queries,
+    read_select_query,
+)
+from core.dependencies import get_db, get_current_user, get_admin_user
+from utils.mysql_scripts import build_sql_query_text
+
+router = APIRouter()
+
+
+
+@router.post("/check", 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("/")
+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("/", 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_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)