From 50ea4f1e76022b84eb753b27e6981e20a2b4e82f Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sun, 22 Mar 2020 01:46:20 +1300
Subject: [PATCH] Fixed async connection lock to be awaitable

---
 psycopg3/connection.py | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/psycopg3/connection.py b/psycopg3/connection.py
index 12fbfcbe7..fd14c3052 100644
--- a/psycopg3/connection.py
+++ b/psycopg3/connection.py
@@ -5,7 +5,8 @@ psycopg3 connection objects
 # Copyright (C) 2020 The Psycopg Team
 
 import logging
-from threading import Lock
+import asyncio
+import threading
 
 from . import pq
 from . import exceptions as exc
@@ -25,7 +26,6 @@ class BaseConnection:
 
     def __init__(self, pgconn):
         self.pgconn = pgconn
-        self.lock = Lock()
 
     @classmethod
     def _connect_gen(cls, conninfo):
@@ -111,6 +111,10 @@ class Connection(BaseConnection):
     This class implements a DBAPI-compliant interface.
     """
 
+    def __init__(self, pgconn):
+        super().__init__(pgconn)
+        self.lock = threading.Lock()
+
     @classmethod
     def connect(cls, conninfo, connection_factory=None, **kwargs):
         if connection_factory is not None:
@@ -153,6 +157,10 @@ class AsyncConnection(BaseConnection):
     methods implemented as coroutines.
     """
 
+    def __init__(self, pgconn):
+        super().__init__(pgconn)
+        self.lock = asyncio.Lock()
+
     @classmethod
     async def connect(cls, conninfo, **kwargs):
         conninfo = make_conninfo(conninfo, **kwargs)
@@ -167,7 +175,7 @@ class AsyncConnection(BaseConnection):
         await self._exec_commit_rollback(b"rollback")
 
     async def _exec_commit_rollback(self, command):
-        with self.lock:
+        with await self.lock:
             status = self.pgconn.transaction_status
             if status == pq.TransactionStatus.IDLE:
                 return
-- 
2.47.3