Extract and test query execution

Author: MasterMind2k

flask_phpbb3/backends/base.py

@@ -25,11 +25,11 @@ def __init__(
     ):
         # type: (...) -> None
         self._functions = {}  # type: typing.Dict[str, str]
+        self._connection = None
         self._cache = cache
         self._config = config
 
         self._prepare_statements()
-        self._setup_connection()
 
     def _setup_connection(self):
         # type: () -> None

flask_phpbb3/backends/psycopg2.py

@@ -29,7 +29,7 @@ def _setup_connection(self):
     def _db(self):
         # type: () -> psycopg2.extensions.connection
         if not self._connection:
-            raise NotImplementedError
+            self._setup_connection()
 
         return self._connection
 
@@ -130,45 +130,58 @@ def _sql_query(
                     # Woops :S
                     pass
 
-        # FIXME Driver specific code!
-        c = self._db.cursor()
-
         if operation == 'fetch':
-            # Add skip and limit
-            query += ' OFFSET {:d}'.format(skip)
-            if limit:
-                query += 'LIMIT {:d}'.format(limit)
+            query = self._paginate_query(query, skip, limit)
+
+        output = self._execute_operation(operation, query, kwargs)
+
+        if cache_key:
+            try:
+                self._cache.set(cache_key,
+                                json.dumps(output),
+                                cache_ttl)
+            except ValueError:
+                # Woops :S
+                pass
+
+        return output
+
+    def _paginate_query(self, query, skip, limit):
+        # type: (str, int, typing.Optional[int]) -> str
+        output = query + ' OFFSET {:d}'.format(skip)
+        if limit:
+            output += ' LIMIT {:d}'.format(limit)
+        return output
 
-        c.execute(
+    def _execute_operation(
+        self,
+        operation,  # type: str
+        query,  # type: str
+        params  # type: typing.Dict[str, typing.Union[str, int]]
+    ):
+        # type: (...) -> typing.Any
+        cursor = self._db.cursor()
+
+        cursor.execute(
             query.format(TABLE_PREFIX=self._config['TABLE_PREFIX']),
-            kwargs
+            params
         )
 
         if operation == 'get':
-            output = c.fetchone()
+            output = cursor.fetchone()
             if output is not None:
                 output = dict(output)
         elif operation == 'has':
-            output = bool(c.fetchone())
+            output = bool(cursor.fetchone())
         elif operation == 'fetch':
-            # FIXME a more performant option
-            output = [dict(i) for i in c]
+            output = [dict(i) for i in cursor]
         elif operation == 'set':
             # It is an update
-            output = c.statusmessage
+            output = cursor.statusmessage
             self._db.commit()
         else:
-            raise ValueError('Unknown operation {}'.format(operation))
-        c.close()
-
-        if cache_key:
-            try:
-                self._cache.set(cache_key,
-                                json.dumps(output),
-                                cache_ttl)
-            except ValueError:
-                # Woops :S
-                pass
+            output = None
+        cursor.close()
 
         return output
 

tests/unit/backends/test_psycopg2.py

@@ -0,0 +1,108 @@
+from __future__ import absolute_import
+
+import unittest
+
+import flask_phpbb3.backends.psycopg2
+
+import mock
+
+import werkzeug.contrib.cache
+
+
+@mock.patch('flask_phpbb3.backends.psycopg2.Psycopg2Backend._db')
+class TestExecuteOperation(unittest.TestCase):
+    def setUp(self):
+        # type: () -> None
+        self.connection = flask_phpbb3.backends.psycopg2.Psycopg2Backend(
+            werkzeug.contrib.cache.SimpleCache(),
+            {
+                'TABLE_PREFIX': '',
+            }
+        )
+
+    def test_get_query(self, mocked_db):
+        # type: (mock.Mock) -> None
+        parameters = mock.Mock()
+        cursor = mock.Mock()
+        cursor.fetchone.return_value = {'key': 'value'}
+        mocked_db.cursor.return_value = cursor
+
+        actual_value = self.connection._execute_operation(
+            'get',
+            'select * from somewhere',
+            parameters
+        )
+
+        self.assertEqual(actual_value, {'key': 'value'})
+        cursor.execute.assert_called_once_with(
+            'select * from somewhere',
+            parameters
+        )
+
+    def test_has_query_negative(self, mocked_db):
+        # type: (mock.Mock) -> None
+        cursor = mock.Mock()
+        cursor.fetchone.return_value = 0
+        mocked_db.cursor.return_value = cursor
+
+        actual_value = self.connection._execute_operation(
+            'has',
+            'select * from somewhere',
+            {}
+        )
+
+        self.assertFalse(actual_value)
+
+    def test_has_query_positive(self, mocked_db):
+        # type: (mock.Mock) -> None
+        cursor = mock.Mock()
+        cursor.fetchone.return_value = 1
+        mocked_db.cursor.return_value = cursor
+
+        actual_value = self.connection._execute_operation(
+            'has',
+            'select * from somewhere',
+            {}
+        )
+
+        self.assertTrue(actual_value)
+
+    def test_fetch_query(self, mocked_db):
+        # type: (mock.Mock) -> None
+        expected_value = [{'key': 1}, {'key': 2}]
+        cursor = mock.Mock()
+        cursor.__iter__ = mock.Mock(return_value=iter(expected_value))
+        mocked_db.cursor.return_value = cursor
+
+        actual_value = self.connection._execute_operation(
+            'fetch',
+            'select * from somewhere',
+            {}
+        )
+
+        self.assertEqual(actual_value, expected_value)
+
+    def test_set_query(self, mocked_db):
+        # type: (mock.Mock) -> None
+        cursor = mock.Mock()
+        mocked_db.cursor.return_value = cursor
+
+        actual_value = self.connection._execute_operation(
+            'set',
+            'select * from somewhere',
+            {}
+        )
+
+        self.assertEqual(actual_value, cursor.statusmessage)
+        mocked_db.commit.assert_called_once_with()
+
+    def test_unknown_query(self, mocked_db):
+        # type: (mock.Mock) -> None
+        actual_value = self.connection._execute_operation(
+            'unknown_op',
+            'select * from somewhere',
+            {}
+        )
+
+        # Internal function, this should never happen
+        self.assertIsNone(actual_value)