improved setup.py, new enum for transaction isolation levels

funbringer · funbringer · commit 4f00dbb09b63 · 2017-11-17T15:07:00.000+03:00
diff --git a/setup.py b/setup.py
@@ -1,4 +1,16 @@
-from distutils.core import setup
+import sys
+
+try:
+    from setuptools import setup
+except ImportError:
+    from distutils.core import setup
+
+# Basic dependencies
+install_requires = ["pg8000", "six", "port-for"]
+
+# Add compatibility enum class
+if sys.version_info < (3, 4):
+    install_requires.append("enum34")
 
 setup(
     name='testgres',
@@ -11,4 +23,4 @@
     url='https://github.com/postgrespro/testgres',
     keywords=['testing', 'postgresql'],
     classifiers=[],
-    install_requires=["pg8000", "six", "port-for", "enum34"])
+    install_requires=install_requires)
diff --git a/testgres/testgres.py b/testgres/testgres.py
@@ -171,19 +171,12 @@ def stopped(self):
         return self.stop_event.isSet()
 
 
-def log_watch(node_name, pg_logname):
+class IsolationLevel(Enum):
     """
-    Starts thread for node that redirects
-    postgresql logs to python logging system
+    Transaction isolation level for NodeConnection
     """
 
-    reader = TestgresLogger(node_name, open(pg_logname, 'r'))
-    reader.start()
-
-    global util_threads
-    util_threads.append(reader)
-
-    return reader
+    ReadUncommitted, ReadCommitted, RepeatableRead, Serializable = range(4)
 
 
 class NodeConnection(object):
@@ -218,7 +211,7 @@ def __enter__(self):
     def __exit__(self, type, value, traceback):
         self.close()
 
-    def begin(self, isolation_level=0):
+    def begin(self, isolation_level=IsolationLevel.ReadCommitted):
         # yapf: disable
         levels = [
             'read uncommitted',
@@ -227,40 +220,51 @@ def begin(self, isolation_level=0):
             'serializable'
         ]
 
-        # Check if level is int [0..3]
-        if (isinstance(isolation_level, int) and
-                isolation_level in range(0, 4)):
+        # Check if level is an IsolationLevel
+        if (isinstance(isolation_level, IsolationLevel)):
 
-            # Replace index with isolation level type
-            isolation_level = levels[isolation_level]
+            # Get index of isolation level
+            level_idx = isolation_level.value
+            assert(level_idx in range(4))
 
-        # Or it might be a string
-        elif (isinstance(isolation_level, six.text_type) and
-              isolation_level.lower() in levels):
+            # Replace isolation level with its name
+            isolation_level = levels[level_idx]
 
-            # Nothing to do here
-            pass
-
-        # Something is wrong, emit exception
         else:
-            raise QueryException(
-                'Invalid isolation level "{}"'.format(isolation_level))
+            # Get name of isolation level
+            level_str = str(isolation_level).lower()
+
+            # Validate level string
+            if level_str not in levels:
+                error = 'Invalid isolation level "{}"'
+                raise QueryException(error.format(level_str))
+
+            # Replace isolation level with its name
+            isolation_level = level_str
 
-        self.cursor.execute(
-            'SET TRANSACTION ISOLATION LEVEL {}'.format(isolation_level))
+        # Set isolation level
+        cmd = 'SET TRANSACTION ISOLATION LEVEL {}'
+        self.cursor.execute(cmd.format(isolation_level))
+
+        return self
 
     def commit(self):
         self.connection.commit()
 
+        return self
+
     def rollback(self):
         self.connection.rollback()
 
+        return self
+
     def execute(self, query, *args):
         self.cursor.execute(query, args)
 
         try:
             res = self.cursor.fetchall()
 
+            # pg8000 might return tuples
             if isinstance(res, tuple):
                 res = [tuple(t) for t in res]
 
@@ -1311,3 +1315,18 @@ def configure_testgres(**options):
 
     for key, option in options.items():
         setattr(TestgresConfig, key, option)
+
+
+def log_watch(node_name, pg_logname):
+    """
+    Start thread for node that redirects
+    PostgreSQL logs to python logging system.
+    """
+
+    reader = TestgresLogger(node_name, open(pg_logname, 'r'))
+    reader.start()
+
+    global util_threads
+    util_threads.append(reader)
+
+    return reader
diff --git a/tests/test_simple.py b/tests/test_simple.py
@@ -23,6 +23,7 @@
 from testgres import get_new_node, get_pg_config, configure_testgres
 from testgres import bound_ports
 from testgres import NodeStatus
+from testgres import IsolationLevel
 
 
 class SimpleTest(unittest.TestCase):
@@ -442,6 +443,30 @@ def test_configure(self):
         # return to the base state
         configure_testgres(cache_initdb=True, cache_pg_config=True)
 
+    def test_isolation_levels(self):
+        with get_new_node('node').init().start() as node:
+            with node.connect() as con:
+                # string levels
+                con.begin('Read Uncommitted').commit()
+                con.begin('Read Committed').commit()
+                con.begin('Repeatable Read').commit()
+                con.begin('Serializable').commit()
+
+                # enum levels
+                con.begin(IsolationLevel.ReadUncommitted).commit()
+                con.begin(IsolationLevel.ReadCommitted).commit()
+                con.begin(IsolationLevel.RepeatableRead).commit()
+                con.begin(IsolationLevel.Serializable).commit()
+
+                # check wrong level
+                got_exception = False
+                try:
+                    con.begin('Garbage').commit()
+                except QueryException:
+                    got_exception = True
+
+                self.assertTrue(got_exception)
+
 
 if __name__ == '__main__':
     unittest.main()
