ISSUE #43

* Implement circuit breaker for `aiopg` connection failures.

Author: Sergio García Prado

packages/plugins/minos-database-aiopg/minos/plugins/aiopg/clients.py

@@ -3,8 +3,12 @@
 )
 
 import logging
+from asyncio import (
+    TimeoutError,
+)
 from collections.abc import (
     AsyncIterator,
+    Iterable,
 )
 from typing import (
     Optional,
@@ -22,6 +26,7 @@
 )
 
 from minos.common import (
+    CircuitBreakerMixin,
     ConnectionException,
     DatabaseClient,
     IntegrityException,
@@ -35,7 +40,7 @@
 logger = logging.getLogger(__name__)
 
 
-class AiopgDatabaseClient(DatabaseClient):
+class AiopgDatabaseClient(DatabaseClient, CircuitBreakerMixin):
     """Aiopg Database Client class."""
 
     _connection: Optional[Connection]
@@ -48,10 +53,17 @@ def __init__(
         port: Optional[int] = None,
         user: Optional[str] = None,
         password: Optional[str] = None,
+        circuit_breaker_exceptions: Iterable[type] = tuple(),
+        connection_timeout: Optional[float] = None,
+        cursor_timeout: Optional[float] = None,
         *args,
         **kwargs,
     ):
-        super().__init__(*args, **kwargs)
+        super().__init__(
+            *args,
+            **kwargs,
+            circuit_breaker_exceptions=(OperationalError, TimeoutError, *circuit_breaker_exceptions),
+        )
 
         if host is None:
             host = "localhost"
@@ -61,13 +73,20 @@ def __init__(
             user = "postgres"
         if password is None:
             password = ""
+        if connection_timeout is None:
+            connection_timeout = 1
+        if cursor_timeout is None:
+            cursor_timeout = 60
 
         self._database = database
         self._host = host
         self._port = port
         self._user = user
         self._password = password
 
+        self._connection_timeout = connection_timeout
+        self._cursor_timeout = cursor_timeout
+
         self._connection = None
         self._cursor = None
 
@@ -80,19 +99,22 @@ async def _destroy(self) -> None:
         await self._close_connection()
 
     async def _create_connection(self):
-        try:
-            self._connection = await aiopg.connect(
-                host=self.host, port=self.port, dbname=self.database, user=self.user, password=self.password
-            )
-        except OperationalError as exc:
-            msg = f"There was an {exc!r} while trying to get a database connection."
-            logger.warning(msg)
-            raise ConnectionException(msg)
+        self._connection = await self.with_circuit_breaker(self._connect)
 
         logger.debug(f"Created {self.database!r} database connection identified by {id(self._connection)}!")
 
+    async def _connect(self) -> Connection:
+        return await aiopg.connect(
+            timeout=self._connection_timeout,
+            host=self.host,
+            port=self.port,
+            dbname=self.database,
+            user=self.user,
+            password=self.password,
+        )
+
     async def _close_connection(self):
-        if self._connection is not None and not self._connection.closed:
+        if await self.is_valid():
             await self._connection.close()
         self._connection = None
         logger.debug(f"Destroyed {self.database!r} database connection identified by {id(self._connection)}!")
@@ -114,12 +136,18 @@ async def _reset(self, **kwargs) -> None:
 
     # noinspection PyUnusedLocal
     async def _fetch_all(self) -> AsyncIterator[tuple]:
-        await self._create_cursor()
+        if self._cursor is None:
+            raise ProgrammingException("An operation must be executed before fetching any value.")
+
         try:
             async for row in self._cursor:
                 yield row
         except ProgrammingError as exc:
             raise ProgrammingException(str(exc))
+        except OperationalError as exc:
+            msg = f"There was an {exc!r} while trying to connect to the database."
+            logger.warning(msg)
+            raise ConnectionException(msg)
 
     # noinspection PyUnusedLocal
     async def _execute(self, operation: AiopgDatabaseOperation) -> None:
@@ -131,10 +159,14 @@ async def _execute(self, operation: AiopgDatabaseOperation) -> None:
             await self._cursor.execute(operation=operation.query, parameters=operation.parameters)
         except IntegrityError as exc:
             raise IntegrityException(f"The requested operation raised a integrity error: {exc!r}")
+        except OperationalError as exc:
+            msg = f"There was an {exc!r} while trying to connect to the database."
+            logger.warning(msg)
+            raise ConnectionException(msg)
 
     async def _create_cursor(self):
         if self._cursor is None:
-            self._cursor = await self._connection.cursor()
+            self._cursor = await self._connection.cursor(timeout=self._cursor_timeout)
 
     async def _destroy_cursor(self, **kwargs):
         if self._cursor is not None:

packages/plugins/minos-database-aiopg/tests/test_aiopg/test_clients.py

@@ -13,6 +13,7 @@
 from psycopg2 import (
     IntegrityError,
     OperationalError,
+    ProgrammingError,
 )
 
 from minos.common import (
@@ -30,7 +31,7 @@
 )
 
 
-# noinspection SqlNoDataSourceInspection
+# noinspection SqlNoDataSourceInspection,SqlDialectInspection
 class TestAiopgDatabaseClient(AiopgTestCase):
     def setUp(self):
         super().setUp()
@@ -79,10 +80,14 @@ async def test_connection(self):
         self.assertIsNone(client.connection)
 
     async def test_connection_raises(self):
-        with patch.object(aiopg, "connect", new_callable=PropertyMock, side_effect=OperationalError):
-            with self.assertRaises(ConnectionException):
-                async with AiopgDatabaseClient.from_config(self.config):
-                    pass
+        async with AiopgDatabaseClient.from_config(self.config) as c1:
+
+            async def _fn():
+                return c1.connection
+
+            with patch.object(aiopg, "connect", new_callable=PropertyMock, side_effect=(OperationalError, _fn())):
+                async with AiopgDatabaseClient.from_config(self.config) as c2:
+                    self.assertEqual(c1.connection, c2.connection)
 
     async def test_cursor(self):
         client = AiopgDatabaseClient.from_config(self.config)
@@ -125,17 +130,37 @@ async def test_execute_raises_integrity(self):
                 with self.assertRaises(IntegrityException):
                     await client.execute(self.operation)
 
+    async def test_execute_raises_operational(self):
+        async with AiopgDatabaseClient.from_config(self.config) as client:
+            with patch.object(Cursor, "execute", side_effect=OperationalError):
+                with self.assertRaises(ConnectionException):
+                    await client.execute(self.operation)
+
     async def test_fetch_one(self):
         async with AiopgDatabaseClient.from_config(self.config) as client:
             await client.execute(self.operation)
             observed = await client.fetch_one()
         self.assertIsInstance(observed, tuple)
 
-    async def test_fetch_one_raises(self):
+    async def test_fetch_one_raises_programming_empty(self):
         async with AiopgDatabaseClient.from_config(self.config) as client:
             with self.assertRaises(ProgrammingException):
                 await client.fetch_one()
 
+    async def test_fetch_one_raises_programming(self):
+        async with AiopgDatabaseClient.from_config(self.config) as client:
+            await client.execute(self.operation)
+            with patch.object(Cursor, "fetchone", side_effect=ProgrammingError):
+                with self.assertRaises(ProgrammingException):
+                    await client.fetch_one()
+
+    async def test_fetch_one_raises_operational(self):
+        async with AiopgDatabaseClient.from_config(self.config) as client:
+            await client.execute(self.operation)
+            with patch.object(Cursor, "fetchone", side_effect=OperationalError):
+                with self.assertRaises(ConnectionException):
+                    await client.fetch_one()
+
     async def test_fetch_all(self):
         async with AiopgDatabaseClient.from_config(self.config) as client:
             await client.execute(self.operation)