Implemented alembic migrations, added cursors closing task.

dbs/mysql.py

@@ -1,18 +1,59 @@
-import aiomysql, decimal, datetime
-import asyncio
+import asyncio, aiomysql, decimal, datetime, uuid, logging
+
 from typing import Literal, Any
-from data.schemas import Connection, SelectResultData
+from data.schemas import Connection, SelectResultData, SelectQuery
+from data.app_types import CachedCursor
+from core.exceptions import PoolNotFound
 
-# Database configuration
-DB_CONFIG = {
-    "host": "localhost",
-    "user": "me",  # Replace with your MySQL username
-    "password": "Passwd3.14",  # Replace with your MySQL password
-    "db": "testing",  # Replace with your database name
-    "port": 3306,  # Default MySQL port
-}
+pools: dict[str, aiomysql.Pool] = {}
+cached_cursors: dict[str, CachedCursor] = {}
+closed_cached_cursors: dict[str, CachedCursor] = {}
 
-pools: None | dict[str, aiomysql.Pool] = {}
+cached_cursors_cleaner_task = None
+
+async def close_old_cached_cursors():
+    global cached_cursors, closed_cached_cursors
+
+    for cursor_id in list(cached_cursors.keys()):
+        cursor = cached_cursors.get(cursor_id, None)
+        if cursor is None:
+            continue
+        print(cursor.close_at, datetime.datetime.now(datetime.UTC).timestamp())
+        if cursor.close_at > datetime.datetime.now(datetime.UTC).timestamp():
+            continue
+
+        try:
+            await cursor.close()
+            cached_cursors.pop(cursor_id, None)
+            closed_cached_cursors[cursor_id] = cursor
+            print(f"Closed cursor {cursor_id}")
+        except Exception as e:
+            print(f"Error closing Cursor {cursor_id} -> {e}")
+
+
+async def remove_old_closed_cached_cursors():
+    global closed_cached_cursors
+    for cursor_id in set(closed_cached_cursors.keys()):
+        closed_cursor = closed_cached_cursors.get(cursor_id, None)
+        if closed_cursor is None:
+            continue
+        if (
+            closed_cursor.close_at + closed_cursor.ttl * 5
+            > datetime.datetime.now(datetime.UTC).timestamp()
+        ):
+            continue
+
+        del closed_cached_cursors[cursor_id]
+        print(f"Removed cursor {cursor_id}")
+
+async def cached_cursors_cleaner():
+    global cached_cursors, closed_cached_cursors
+    while True:
+        print("hey")
+        await close_old_cached_cursors()
+        await remove_old_closed_cached_cursors()
+
+        await asyncio.sleep(10)
 
 
 async def pool_creator(connection: Connection, minsize=5, maxsize=10):
@@ -28,8 +69,31 @@ async def pool_creator(connection: Connection, minsize=5, maxsize=10):
         maxsize=maxsize,
     )
 
+async def create_cursor(connection_id: int, query: SelectQuery) -> CachedCursor:
+    pool = pools.get(connection_id, None)
 
-async def execute_select_query(pool: aiomysql.Pool, query: str, params: list, fetch_num:int = 100):
+    if pool is None:
+        raise PoolNotFound
+
+    connection = await pool.acquire()
+    cursor = await connection.cursor(aiomysql.SSCursor)
+    await cursor.execute(query.sql, query.params)
+
+    cached_cursor = CachedCursor(
+        id=str(uuid.uuid4()),
+        cursor=cursor,
+        connection=connection,
+        pool=pool,
+        connection_id=connection_id,
+        query=query,
+    )
+
+    return cached_cursor
+
+
+async def execute_select_query(
+    pool: aiomysql.Pool, sql_query: str, params: list, fetch_num: int = 100
+):
     """
     Executes a SELECT query on the MySQL database asynchronously and returns the results.
 
@@ -43,7 +107,7 @@ async def execute_select_query(pool: aiomysql.Pool, query: str, params: list, fe
     try:
         async with pool.acquire() as connection:
             async with connection.cursor(aiomysql.DictCursor) as cursor:
-                await cursor.execute(query, params)
+                await cursor.execute(sql_query, params)
                 result = await cursor.fetchmany(fetch_num)
 
         return result, cursor.rowcount