61 lines
1.6 KiB
Python
61 lines
1.6 KiB
Python
import aiomysql
|
|
|
|
from core.schemas import mc_s
|
|
from core.util import logs
|
|
|
|
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
|