Refactor and implement queueing messages

core/db/sql.py

@@ -0,0 +1,63 @@
+import aiomysql
+
+from core.util import logs
+from core.schemas import mc_s
+
+mysql_pool = None
+
+log = logs.get_logger("sql")
+
+DB_URL = "giadb"
+async def init_mysql_pool():
+    """
+    Initialize the MySQL connection pool.
+    """
+    global mysql_pool
+    mysql_pool = await aiomysql.create_pool(
+        host=DB_URL,
+        port=9306,
+        db="Manticore",
+        minsize=1,
+        maxsize=10
+    )
+
+async def close_mysql_pool():
+    """Close the MySQL connection pool properly."""
+    global mysql_pool
+    if mysql_pool:
+        mysql_pool.close()
+        await mysql_pool.wait_closed()
+
+
+async def create_index():
+    schemas = {
+        "main": mc_s.schema_main,
+        # "rule_storage": mc_s.schema_rule_storage,
+        # "meta": mc_s.schema_meta,
+        # "internal": mc_s.schema_int,
+    }
+    try:
+        async with mysql_pool.acquire() as conn:
+            async with conn.cursor() as cur:
+                for name, schema in schemas.items():
+                    schema_types = ", ".join([f"{k} {v}" for k, v in schema.items()])
+
+                    create_query = (
+                        f"create table if not exists {name}({schema_types}) engine='columnar'"
+                    )
+                    log.info(f"Schema types {create_query}")
+                    await cur.execute(create_query) # SQLi
+    except aiomysql.Error as e:
+        log.error(f"MySQL error: {e}")
+
+
+async def main():
+    await init_mysql_pool()
+    created = False
+    while not created:
+        try:
+            await create_index()
+            created = True
+        except Exception as e:
+            log.error(f"Error creating index: {e}")
+            await asyncio.sleep(1) # Block the thread, just wait for the DB