import logging
from sqlite3 import Connection
from types import TracebackType
from typing import (
    Callable,
    Generic,
    Optional,
    Type,
    TypeVar,
)


logger = logging.getLogger(__name__)

T = TypeVar('T')


class ConnectionContextManager(Generic[T]):
    def __init__(
        self,
        connection_allocator: Callable[[], Connection],
        object_creator: Callable[[Connection], T],
    ) -> None:
        self.connection_allocator = connection_allocator
        self.object_creator = object_creator

        self.conn = None  # type: Optional[Connection]

    def __enter__(self) -> T:
        logger.debug("Opening connection")
        self.conn = self.connection_allocator()
        return self.object_creator(self.conn)

    def __exit__(
        self,
        exc_type: Optional[Type[Exception]],
        exc_val: Optional[Exception],
        exc_tb: Optional[TracebackType],
    ) -> None:

        if exc_val is None:
            logger.debug("Commit")
            self.conn.commit()
        else:
            logger.debug("Rollback")
            self.conn.rollback()

        self.conn.close()
