import logging
import os
from typing import Optional

from sqlalchemy import create_engine
from sqlalchemy.engine import Engine
from sqlalchemy.orm import (
    sessionmaker,
    Session,
)

from fidesops.core.config import config

logger = logging.getLogger(__name__)


def get_db_engine(database_uri: Optional[str] = None) -> Engine:
    """
    Return a database engine. If the TESTING environment var is set the
    database engine returned will be connected to the test DB.
    """

    if database_uri is None:
        # Don't override any database_uri explicity passed in
        if os.getenv("TESTING"):
            database_uri = config.database.SQLALCHEMY_TEST_DATABASE_URI
        else:
            database_uri = config.database.SQLALCHEMY_DATABASE_URI

    return create_engine(database_uri, pool_pre_ping=True)


def get_db_session(engine: Optional[Engine] = None) -> sessionmaker:
    """Return a database SessionLocal"""
    return sessionmaker(
        autocommit=False,
        autoflush=False,
        bind=engine or get_db_engine(),
        class_=ExtendedSession,
    )


class ExtendedSession(Session):
    """This class wraps the SQLAlchemy Session to provide some error handling on commits."""

    def commit(self) -> None:
        """Provide the option to automatically rollback failed transactions"""
        try:
            return super().commit()
        except Exception as exc:
            logger.error(exc)
            # Rollback the current transaction after each failed commit
            self.rollback()
            raise
