ASE/load_orchestrator.py

#!.venv/bin/python

# Import necessary libraries
import logging
import importlib
import asyncio
import os
import aiomysql
import contextvars

# Import custom modules for configuration and database connection
from utils.config import loader_load_data as setting
from utils.database import CSV_RECEIVED

# Crea una context variable per identificare il worker
worker_context = contextvars.ContextVar('worker_id', default='00')

# Formatter personalizzato che include il worker_id
class WorkerFormatter(logging.Formatter):
    def format(self, record):
        record.worker_id = worker_context.get()
        return super().format(record)

# Initialize the logger for this module
logger = logging.getLogger()

# Delay tra un processamento CSV e il successivo (in secondi)
CSV_PROCESSING_DELAY = 0.2
# Tempo di attesa se non ci sono record da elaborare
NO_RECORD_SLEEP = 60

async def get_next_csv_atomic(pool, table_name):
    """Preleva atomicamente il prossimo CSV da elaborare"""
    async with pool.acquire() as conn:
        # IMPORTANTE: Disabilita autocommit per questa transazione
        await conn.begin()

        try:
            async with conn.cursor() as cur:
                # Usa SELECT FOR UPDATE per lock atomico
                await cur.execute(f"""
                    SELECT id, unit_type, tool_type, unit_name, tool_name
                    FROM {table_name}
                    WHERE locked = 0 AND status = %s
                    ORDER BY id
                    LIMIT 1
                    FOR UPDATE SKIP LOCKED
                """, (CSV_RECEIVED,))

                result = await cur.fetchone()
                if result:
                    await cur.execute(f"""
                        UPDATE {table_name}
                        SET locked = 1
                        WHERE id = %s
                    """, (result[0],))

                # Commit esplicito per rilasciare il lock
                await conn.commit()
                return result

        except Exception as e:
            # Rollback in caso di errore
            await conn.rollback()
            raise e

async def worker(worker_id: int, cfg: object, pool) -> None:
    # Imposta il context per questo worker
    worker_context.set(f"W{worker_id}")

    debug_mode = (logging.getLogger().getEffectiveLevel() == logging.DEBUG)
    logger.info("Avviato")

    while True:
        try:
            logger.info("Inizio elaborazione")

            record = await get_next_csv_atomic(pool, cfg.dbrectable)

            if record:
                success = await load_csv(record, cfg, pool)
                if not success:
                    logger.error("Errore durante l'elaborazione")
                await asyncio.sleep(CSV_PROCESSING_DELAY)
            else:
                logger.debug("Nessun record disponibile")
                await asyncio.sleep(NO_RECORD_SLEEP)

        except Exception as e:
            logger.error(f"Errore durante l'esecuzione: {e}", exc_info=debug_mode)
            await asyncio.sleep(1)

async def load_csv(record: tuple, cfg: object, pool) -> bool:
    debug_mode = (logging.getLogger().getEffectiveLevel() == logging.DEBUG)
    logger.debug("Inizio ricerca nuovo CSV da elaborare")

    id, unit_type, tool_type, unit_name, tool_name = [x.lower().replace(" ", "_") if isinstance(x, str) else x for x in record]
    logger.info(f'Trovato CSV da elaborare: ID={id}, Tipo={unit_type}_{tool_type}, Nome={unit_name}_{tool_name}')

    # Costruisce il nome del modulo da caricare dinamicamente
    module_names = [f'utils.parsers.by_name.{unit_name}_{tool_name}',
                    f'utils.parsers.by_name.{unit_name}_{tool_type}',
                    f'utils.parsers.by_name.{unit_name}_all',
                    f'utils.parsers.by_type.{unit_type}_{tool_type}']
    modulo = None
    for module_name in module_names:
        try:
            logger.debug(f"Caricamento dinamico del modulo: {module_name}")
            modulo = importlib.import_module(module_name)
            logger.info(f"Funzione 'main_loader' caricata dal modulo {module_name}")
            break
        except (ImportError, AttributeError) as e:
            logger.debug(f"Modulo {module_name} non presente o non valido. {e}", exc_info=debug_mode)

    if not modulo:
        logger.error(f"Nessun modulo trovato {module_names}")
        return False

    # Ottiene la funzione 'main_loader' dal modulo
    funzione = getattr(modulo, "main_loader")

    # Esegui la funzione
    logger.info(f"Elaborazione con modulo {modulo} per ID={id}")
    await funzione(cfg, id, pool)
    logger.info(f"Elaborazione completata per ID={id}")
    return True

async def main():
    """Main function: avvia i worker e gestisce il ciclo principale."""
    logger.info("Avvio del sistema...")

    cfg = setting.Config()
    logger.info("Configurazione caricata correttamente")

    try:
        # Configura il logging globale
        log_level = os.getenv("LOG_LEVEL", "INFO").upper()
        debug_mode = (logging.getLogger().getEffectiveLevel() == logging.DEBUG)

        # Configura il logging con il formatter personalizzato
        handler = logging.FileHandler(cfg.logfilename)
        formatter = WorkerFormatter(
            "%(asctime)s - PID: %(process)d.Worker-%(worker_id)s.%(name)s.%(funcName)s.%(levelname)s: %(message)s"
        )
        handler.setFormatter(formatter)

        # Rimuovi eventuali handler esistenti e aggiungi il nostro
        logger.handlers.clear()
        logger.addHandler(handler)
        logger.setLevel(getattr(logging, log_level))

        logger.info("Logging configurato correttamente")

        # Numero massimo di worker concorrenti
        logger.info(f"Avvio di {cfg.max_threads} worker concorrenti")

        pool = await aiomysql.create_pool(
            host=cfg.dbhost,
            user=cfg.dbuser,
            password=cfg.dbpass,
            db=cfg.dbname,
            minsize=4,
            maxsize=cfg.max_threads*4,
            pool_recycle=3600
        )

        # Avvia i worker
        workers = [
            asyncio.create_task(worker(i, cfg, pool))
            for i in range(cfg.max_threads)
        ]

        logger.info("Sistema avviato correttamente. In attesa di nuovi task...")

        try:
            await asyncio.gather(*workers, return_exceptions=debug_mode)
        finally:
            pool.close()
            await pool.wait_closed()

    except KeyboardInterrupt:
        logger.info("Info: Shutdown richiesto... chiusura in corso")

    except Exception as e:
        logger.error(f"Errore principale: {e}", exc_info=debug_mode)

if __name__ == "__main__":
    asyncio.run(main())