meshview/startdb.py

import asyncio
import json
import datetime
import logging
from sqlalchemy import delete
from meshview import mqtt_reader
from meshview import mqtt_database
from meshview import mqtt_store
from meshview import models
from meshview.config import CONFIG

# -------------------------
# Logging for cleanup
# -------------------------
cleanup_logger = logging.getLogger("dbcleanup")
cleanup_logger.setLevel(logging.INFO)
file_handler = logging.FileHandler("dbcleanup.log")
file_handler.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s [%(levelname)s] %(message)s')
file_handler.setFormatter(formatter)
cleanup_logger.addHandler(file_handler)

# -------------------------
# Helper functions
# -------------------------
def get_bool(config, section, key, default=False):
    return str(config.get(section, {}).get(key, default)).lower() in ("1", "true", "yes", "on")

def get_int(config, section, key, default=0):
    try:
        return int(config.get(section, {}).get(key, default))
    except ValueError:
        return default

# -------------------------
# Shared DB lock
# -------------------------
db_lock = asyncio.Lock()

# -------------------------
# Database cleanup using ORM
# -------------------------
async def daily_cleanup_at(
    hour: int = 2,
    minute: int = 0,
    days_to_keep: int = 14,
    vacuum_db: bool = False
):
    while True:
        now = datetime.datetime.now()
        next_run = now.replace(hour=hour, minute=minute, second=0, microsecond=0)
        if next_run <= now:
            next_run += datetime.timedelta(days=1)
        delay = (next_run - now).total_seconds()
        cleanup_logger.info(f"Next cleanup scheduled at {next_run}")
        await asyncio.sleep(delay)

        cutoff = datetime.datetime.now() - datetime.timedelta(days=days_to_keep)
        cleanup_logger.info(f"Running cleanup for records older than {cutoff}...")

        try:
            async with db_lock:   # <--- Pause ingestion during cleanup
                cleanup_logger.info("Ingestion paused for cleanup.")

                async with mqtt_database.async_session() as session:
                    # Packet
                    result = await session.execute(
                        delete(models.Packet).where(models.Packet.import_time < cutoff)
                    )
                    cleanup_logger.info(f"Deleted {result.rowcount} rows from Packet")

                    # PacketSeen
                    result = await session.execute(
                        delete(models.PacketSeen).where(models.PacketSeen.import_time < cutoff)
                    )
                    cleanup_logger.info(f"Deleted {result.rowcount} rows from PacketSeen")

                    # Traceroute
                    result = await session.execute(
                        delete(models.Traceroute).where(models.Traceroute.import_time < cutoff)
                    )
                    cleanup_logger.info(f"Deleted {result.rowcount} rows from Traceroute")

                    # Node
                    result = await session.execute(
                        delete(models.Node).where(models.Node.last_update < cutoff)
                    )
                    cleanup_logger.info(f"Deleted {result.rowcount} rows from Node")

                    await session.commit()

                if vacuum_db:
                    cleanup_logger.info("Running VACUUM...")
                    async with mqtt_database.engine.begin() as conn:
                        await conn.exec_driver_sql("VACUUM;")
                    cleanup_logger.info("VACUUM completed.")

                cleanup_logger.info("Cleanup completed successfully.")
                cleanup_logger.info("Ingestion resumed after cleanup.")

        except Exception as e:
            cleanup_logger.error(f"Error during cleanup: {e}")

# -------------------------
# MQTT loading
# -------------------------
async def load_database_from_mqtt(
    mqtt_server: str,
    mqtt_port: int,
    topics: list,
    mqtt_user: str | None = None,
    mqtt_passwd: str | None = None
):
    async for topic, env in mqtt_reader.get_topic_envelopes(
        mqtt_server, mqtt_port, topics, mqtt_user, mqtt_passwd
    ):
        async with db_lock:   # <--- Block here if cleanup is running
            await mqtt_store.process_envelope(topic, env)

# -------------------------
# Main function
# -------------------------
async def main():
    # Initialize database
    mqtt_database.init_database(CONFIG["database"]["connection_string"])
    await mqtt_database.create_tables()

    mqtt_user = CONFIG["mqtt"].get("username") or None
    mqtt_passwd = CONFIG["mqtt"].get("password") or None
    mqtt_topics = json.loads(CONFIG["mqtt"]["topics"])

    cleanup_enabled = get_bool(CONFIG, "cleanup", "enabled", False)
    cleanup_days = get_int(CONFIG, "cleanup", "days_to_keep", 14)
    vacuum_db = get_bool(CONFIG, "cleanup", "vacuum", False)
    cleanup_hour = get_int(CONFIG, "cleanup", "hour", 2)
    cleanup_minute = get_int(CONFIG, "cleanup", "minute", 0)

    async with asyncio.TaskGroup() as tg:
        tg.create_task(
            load_database_from_mqtt(
                CONFIG["mqtt"]["server"],
                int(CONFIG["mqtt"]["port"]),
                mqtt_topics,
                mqtt_user,
                mqtt_passwd,
            )
        )

        if cleanup_enabled:
            tg.create_task(
                daily_cleanup_at(cleanup_hour, cleanup_minute, cleanup_days, vacuum_db)
            )
        else:
            cleanup_logger.info("Daily cleanup is disabled by configuration.")

# -------------------------
# Entry point
# -------------------------
if __name__ == '__main__':
    asyncio.run(main())