# mypy: ignore-errors
"""SQL DB Connector.

Manages interactions with SQL databases.
"""

from contextlib import asynccontextmanager
from typing import AsyncGenerator, Type

from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
from sqlalchemy.orm import DeclarativeMeta, declarative_base, sessionmaker

from delivery_metadata import config


class SqlConnector:
    """SQL connector."""

    base_model: Type[DeclarativeMeta]
    _db_engine: AsyncEngine
    _db_session_maker: sessionmaker

    def __init__(self, db_url):
        self.base_model = declarative_base()
        self._db_engine = create_async_engine(
            db_url,
            **config.SQLALCHEMY_ENGINE_OPTIONS,
        )
        self._db_session_maker = sessionmaker(
            bind=self._db_engine,
            class_=AsyncSession,
            expire_on_commit=False,
        )

    @asynccontextmanager
    async def db_session(
        self,
        transaction=False,
        turn_off_foreign_key_checks=False,
    ) -> AsyncGenerator[AsyncSession, None]:
        query_error = None
        cleanup_error = None
        session = self._db_session_maker()

        # do not use transaction, read only connection
        if not transaction:
            await session.connection(
                execution_options={
                    "isolation_level": "READ COMMITTED",
                    "autobegin": False,
                    "skip_autocommit_rollback": True,
                }
            )

        # give back session to query with
        try:
            if turn_off_foreign_key_checks:
                await session.execute(text("SET FOREIGN_KEY_CHECKS = 0"))
            yield session
            if transaction:
                await session.commit()
        except Exception as e:
            query_error = e
            if transaction:
                await session.rollback()
        finally:
            if turn_off_foreign_key_checks:
                await session.execute(text("SET FOREIGN_KEY_CHECKS = 1"))

        # cleanup by closing session
        try:
            await session.close()
        except Exception as e:
            cleanup_error = e

        # report errors
        if query_error:
            raise query_error
        if cleanup_error:
            raise cleanup_error

    async def close(self):
        await self._db_engine.dispose()