from contextlib import contextmanager
from sqlalchemy.ext.compiler import compiles, deregister
from sqlalchemy.sql import Insert

from src.db.base import session


def check_health() -> bool:
    """Check DB connection."""
    result = session.execute("SELECT 1").first()
    return result and result[0] == 1


def insert_ignore(insert, compiler, **kw):
    """Custom insert compiler to ignore errors on insert."""
    return compiler.visit_insert(insert.prefix_with("IGNORE"), **kw)


@contextmanager
def set_compiler(compiler, class_, *specs):
    """Context manager to set custom compiler for some statement."""

    compiles(class_, *specs)(compiler)
    try:
        yield
    finally:
        deregister(class_)


def set_ignore_on_insert():
    """Context manger to ignore errors on update inside, is useful for bulk operations."""
    return set_compiler(insert_ignore, Insert, "mysql")