Allow specifying isolation_level with Postgres transactions.

Fixes #2790

docs/peewee/api.rst

@@ -936,6 +936,15 @@ Database
  currently connected, the attached database will be detached from the
  open connection.
 
+ .. py:method:: atomic([lock_type=None])
+
+ :param str lock_type: Locking strategy: DEFERRED, IMMEDIATE, EXCLUSIVE.
+
+ Create an atomic context-manager, optionally using the specified
+ locking strategy (if unspecified, DEFERRED is used).
+
+ .. note:: Lock type only applies to the outermost ``atomic()`` block.
+
  .. py:method:: transaction([lock_type=None])
 
  :param str lock_type: Locking strategy: DEFERRED, IMMEDIATE, EXCLUSIVE.
@@ -963,6 +972,22 @@ Database
  Set the timezone on the current connection. If no connection is open,
  then one will be opened.
 
+ .. py:method:: atomic([isolation_level=None])
+
+ :param str isolation_level: Isolation strategy: SERIALIZABLE, READ COMMITTED, REPEATABLE READ, READ UNCOMMITTED
+
+ Create an atomic context-manager, optionally using the specified
+ isolation level (if unspecified, the server default will be used).
+
+ .. note:: Isolation level only applies to the outermost ``atomic()`` block.
+
+ .. py:method:: transaction([isolation_level=None])
+
+ :param str isolation_level: Isolation strategy: SERIALIZABLE, READ COMMITTED, REPEATABLE READ, READ UNCOMMITTED
+
+ Create a transaction context-manager, optionally using the specified
+ isolation level (if unspecified, the server default will be used).
+
 
 .. py:class:: MySQLDatabase(database[, **kwargs])
 

peewee.py

@@ -3969,6 +3969,16 @@ def rows_affected(self, cursor):
  except AttributeError:
  return cursor.cursor.rowcount
 
+ def begin(self, isolation_level=None):
+ if self.is_closed():
+ self.connect()
+ if isolation_level:
+ stmt = 'BEGIN TRANSACTION ISOLATION LEVEL %s' % isolation_level
+ else:
+ stmt = 'BEGIN'
+ with __exception_wrapper__:
+ self.cursor().execute(stmt)
+
  def get_tables(self, schema=None):
  query = ('SELECT tablename FROM pg_catalog.pg_tables '
  'WHERE schemaname = %s ORDER BY tablename')

tests/transactions.py

@@ -3,6 +3,7 @@
 from .base import DatabaseTestCase
 from .base import IS_CRDB
 from .base import IS_CRDB_NESTED_TX
+from .base import IS_POSTGRESQL
 from .base import IS_SQLITE
 from .base import ModelTestCase
 from .base import db
@@ -398,3 +399,39 @@ def test_lock_type(self):
  with db2.transaction(lock_type='IMMEDIATE') as t2:
  self._save(6)
  self.assertRegister([2, 4, 5])
+
+
+@skip_unless(IS_POSTGRESQL, 'requires postgresql for txn isolation level')
+class TestTransactionIsolationLevel(BaseTransactionTestCase):
+ def test_isolation_level(self):
+ db2 = new_connection()
+ db2.connect()
+
+ with db2.atomic(isolation_level='SERIALIZABLE'):
+ with db.atomic(isolation_level='SERIALIZABLE'):
+ self._save(1)
+ self.assertDB2(db2, [])
+ self.assertDB2(db2, [])
+ self.assertDB2(db2, [1])
+
+ with db2.atomic(isolation_level='READ COMMITTED'):
+ with db.atomic():
+ self._save(2)
+ self.assertDB2(db2, [1])
+ self.assertDB2(db2, [1, 2])
+ self.assertDB2(db2, [1, 2])
+
+ # NB: Read Uncommitted is treated as Read Committed by PG, so we don't
+ # test it here.
+
+ with db2.atomic(isolation_level='REPEATABLE READ'):
+ with db.atomic(isolation_level='REPEATABLE READ'):
+ self._save(3)
+ self.assertDB2(db2, [1, 2])
+ self.assertDB2(db2, [1, 2])
+ self.assertDB2(db2, [1, 2, 3])
+
+ def assertDB2(self, db2, vals):
+ q = 'select "value" from "register" order by "value"'
+ actual = [v for v, in db2.execute_sql(q).fetchall()]
+ self.assertEqual(actual, vals)