import asyncio
import aiomysql


loop = asyncio.get_event_loop()


# async def main(all_done):
#     """Get all messages."""
#     conn = await aiomysql.connect(host='127.0.0.1', port=3306, user='test',
#                                        password='test123', db='test', loop=loop)
#
#     cur = await conn.cursor(aiomysql.DictCursor)
#     await cur.execute("SELECT receiver, message FROM messages")
#     r = await cur.fetchall()
#
#     await cur.close()
#     conn.close()
#     all_done.set_result(r)


class DBConnection(object):
    """Async context manager that makes a connection to db."""

    def __init__(self, host, port, user, password, db):
        self.conn = None
        self.cur = None
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.db = db

    async def __aenter__(self):
        """Enter async context manager."""
        self.conn = await aiomysql.connect(host=self.host, port=self.port,
            user=self.user, password=self.password, db=self.db, loop=loop)

        self.cur = await self.conn.cursor(aiomysql.DictCursor)
        return self.cur

    async def __aexit__(self, *args):
        """Exit async context manager."""
        await self.cur.close()
        self.conn.close()


async def alternative(all_done):
    """Get all messages."""

    async with DBConnection(host='127.0.0.1', port=3306, user='test',
                            password='test123', db='test') as cur:
        await cur.execute("SELECT receiver, message FROM messages")
        r = await cur.fetchall()
        all_done.set_result(r)