# Buggy code
def retry_connection(func):
    @wraps(func)
    async def wrapped(self, *args):
        try:
            return await func(self, *args)
        except (
            RuntimeError,
            pymysql.err.OperationalError,
            pymysql.err.InternalError,
            pymysql.err.InterfaceError,
        ):
            # Here we assume that a connection error has happened
            # Re-create connection and re-try the function call once only.
            await self._lock.acquire()
            logging.info("Attempting reconnect")
            try:
                self._close()
                await self.create_connection(with_db=True)
                logging.info("Reconnected")
            except Exception:
                logging.info("Failed to reconnect")
            finally:
                self._lock.release()

            return await func(self, *args)

    return wrapped


# Traceback (ConnectionResetError)
Traceback (most recent call last):
File \"/usr/lib/python3.6/asyncio/selector_events.py\", line 714, in _read_ready
data = self._sock.recv(self.max_size)
ConnectionResetError: [Errno 104] Connection reset by peer
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File \"/home/shlomy/PycharmProjects/playground/venv36/lib/python3.6/site-packages/tortoise/models.py\", line 277, in create
await instance.save(using_db=kwargs.get('using_db'))
File \"/home/shlomy/PycharmProjects/playground/venv36/lib/python3.6/site-packages/tortoise/models.py\", line 221, in save
await self._insert_instance(*args, **kwargs)
File \"/home/shlomy/PycharmProjects/playground/venv36/lib/python3.6/site-packages/tortoise/models.py\", line 210, in _insert_instance
).execute_insert(self)
File \"/home/shlomy/PycharmProjects/playground/venv36/lib/python3.6/site-packages/tortoise/backends/base/executor.py\", line 77, in execute_insert
instance.id = await self.db.execute_insert(query, values)
File \"/home/shlomy/PycharmProjects/playground/venv36/lib/python3.6/site-packages/tortoise/backends/asyncpg/client.py\", line 22, in wrapped
return await func(self, query, *args)
File \"/home/shlomy/PycharmProjects/playground/venv36/lib/python3.6/site-packages/tortoise/backends/asyncpg/client.py\", line 108, in execute_insert
stmt = await connection.prepare(query)
File \"/home/shlomy/PycharmProjects/playground/venv36/lib/python3.6/site-packages/asyncpg/connection.py\", line 403, in prepare
return await self._prepare(query, timeout=timeout, use_cache=False)
File \"/home/shlomy/PycharmProjects/playground/venv36/lib/python3.6/site-packages/asyncpg/connection.py\", line 408, in _prepare
use_cache=use_cache)
File \"/home/shlomy/PycharmProjects/playground/venv36/lib/python3.6/site-packages/asyncpg/connection.py\", line 329, in _get_statement
statement = await self._protocol.prepare(stmt_name, query, timeout)
File \"asyncpg/protocol/protocol.pyx\", line 163, in prepare
asyncpg.exceptions.ConnectionDoesNotExistError: connection was closed in the middle of operation"}

# Fixed code
def retry_connection(func):
    @wraps(func)
    async def retry_connection_(self, *args):
        try:
            return await func(self, *args)
        except (
            RuntimeError,
            pymysql.err.OperationalError,
            pymysql.err.InternalError,
            pymysql.err.InterfaceError,
        ):
            # Here we assume that a connection error has happened
            # Re-create connection and re-try the function call once only.
            if getattr(self, "_finalized", None) is False:
                raise TransactionManagementError(
                    "Connection gone away during transaction"
                )
            await self._lock.acquire()
            logging.info("Attempting reconnect")
            try:
                await self._close()
                await self.create_connection(with_db=True)
                logging.info("Reconnected")
            except Exception as e:
                logging.info("Failed to reconnect: %s", str(e))
                raise
            finally:
                self._lock.release()

            return await func(self, *args)

    return retry_connection_