Cancel current DB operation on asyncio timeout

CHANGES.txt

@@ -8,6 +8,8 @@ CHANGES
 
 * Support pools with unlimited size #59
 
+* Cancel current DB operation on asyncio timeout
+
 
 0.7.0 (2015-04-22)
 ^^^^^^^^^^^^^^^^^^

aiopg/connection.py

@@ -286,7 +286,10 @@ def cancel(self, timeout=None):
         self._conn.cancel()
         if timeout is None:
             timeout = self._timeout
-        yield from self._poll(waiter, timeout)
+        try:
+            yield from self._poll(waiter, timeout)
+        except psycopg2.extensions.QueryCanceledError:
+            pass
 
     @asyncio.coroutine
     def reset(self):

aiopg/cursor.py

@@ -108,8 +108,12 @@ def execute(self, operation, parameters=None, *, timeout=None):
         except:
             self._conn._waiter = None
             raise
-        else:
+        try:
             yield from self._conn._poll(waiter, timeout)
+        except asyncio.TimeoutError:
+            yield from self._conn.cancel()
+            self._impl.close()
+            raise
 
     @asyncio.coroutine
     def executemany(self, operation, seq_of_parameters):

tests/test_connection.py

@@ -654,3 +654,18 @@ def go():
             self.assertEqual('hello', val.payload)
 
         self.loop.run_until_complete(go())
+
+    def test_close_cursor_on_timeout_error(self):
+        @asyncio.coroutine
+        def go():
+            conn = yield from self.connect()
+            cur = yield from conn.cursor(timeout=0.01)
+            with self.assertRaises(asyncio.TimeoutError):
+                yield from cur.execute("SELECT pg_sleep(10)")
+
+            self.assertTrue(cur.closed)
+            self.assertFalse(conn.closed)
+
+            yield from conn.close()
+
+        self.loop.run_until_complete(go())