Implement connection.transaction()

Author: 1st1

asyncpg/__init__.py

@@ -1,4 +1,5 @@
 import asyncio
+import enum
 import getpass
 import os
 
@@ -12,16 +13,27 @@
 
 
 class Connection:
+
+    __slots__ = ('_protocol', '_transport', '_loop', '_types_stmt',
+                 '_type_by_name_stmt', '_top_xact', '_uid')
+
     def __init__(self, protocol, transport, loop):
         self._protocol = protocol
         self._transport = transport
         self._loop = loop
         self._types_stmt = None
         self._type_by_name_stmt = None
+        self._top_xact = None
+        self._uid = 0
 
     def get_settings(self):
         return self._protocol.get_settings()
 
+    def transaction(self, *, isolation='read_committed', readonly=False,
+                    deferrable=False):
+
+        return Transaction(self, isolation, readonly, deferrable)
+
     async def execute_script(self, script):
         await self._protocol.query(script)
 
@@ -82,8 +94,15 @@ async def set_type_codec(self, typename, *,
     def close(self):
         self._transport.close()
 
+    def _get_unique_id(self):
+        self._uid += 1
+        return 'id{}'.format(self._uid)
+
 
 class PreparedStatement:
+
+    __slots__ = ('_connection', '_state')
+
     def __init__(self, connection, state):
         self._connection = connection
         self._state = state
@@ -99,6 +118,133 @@ async def execute(self, *args):
         return await protocol.execute(self._state, args)
 
 
+class TransactionState(enum.Enum):
+    NEW             = 0
+    STARTED         = 1
+    COMMITTED       = 2
+    ROLLEDBACK      = 3
+    FAILED          = 4
+
+
+class Transaction:
+
+    ISOLATION_LEVELS = {'read_committed', 'serializable', 'repeatable_read'}
+
+    __slots__ = ('_connection', '_isolation', '_readonly', '_deferrable',
+                 '_state', '_nested', '_id')
+
+    def __init__(self, connection, isolation, readonly, deferrable):
+        if isolation not in self.ISOLATION_LEVELS:
+            raise ValueError(
+                'isolation is expected to be either of {}, '
+                'got {!r}'.format(self.ISOLATION_LEVELS, isolation))
+
+        if isolation != 'serializable':
+            if readonly:
+                raise ValueError(
+                    '"readonly" is only supported for '
+                    'serializable transactions')
+
+            if deferrable and not readonly:
+                raise ValueError(
+                    '"deferrable" is only supported for '
+                    'serializable readonly transactions')
+
+        self._connection = connection
+        self._isolation = isolation
+        self._readonly = readonly
+        self._deferrable = deferrable
+        self._state = TransactionState.NEW
+        self._nested = False
+        self._id = None
+
+    async def __aenter__(self):
+        await self.start()
+
+    async def __aexit__(self, extype, ex, tb):
+        if extype is not None:
+            await self.rollback()
+
+    async def start(self):
+        if self._state is not TransactionState.NEW:
+            raise FatalError('cannot start transaction: inconsistent state')
+
+        con = self._connection
+
+        if con._top_xact is None:
+            con._top_xact = self
+        else:
+            # Nested transaction block
+            top_xact = con._top_xact
+            if self._isolation != top_xact._isolation:
+                raise FatalError(
+                    'nested transaction has different isolation level: '
+                    'current {!r} != outer {!r}'.format(
+                        self._isolation, top_xact._isolation))
+            self._nested = True
+
+        if self._nested:
+            self._id = con._get_unique_id()
+            query = 'SAVEPOINT {};'.format(self._id)
+        else:
+            if self._isolation == 'read_committed':
+                query = 'BEGIN;'
+            elif self._isolation == 'repeatable_read':
+                query = 'BEGIN ISOLATION LEVEL REPEATABLE READ;'
+            else:
+                query = 'BEGIN ISOLATION LEVEL SERIALIZABLE'
+                if self._readonly:
+                    query += ' READ ONLY'
+                if self._deferrable:
+                    query += ' DEFERRABLE'
+                query += ';'
+
+        try:
+            await self._connection.execute_script(query)
+        except:
+            self._state = TransactionState.FAILED
+            raise
+        else:
+            self._state = TransactionState.STARTED
+
+    async def commit(self):
+        if self._state is not TransactionState.STARTED:
+            raise FatalError('cannot commit transaction: inconsistent state')
+
+        if self._nested:
+            query = 'RELEASE SAVEPOINT {};'.format(self._id)
+        else:
+            query = 'COMMIT;'
+
+        try:
+            await self._connection.execute_script(query)
+        except:
+            self._state = TransactionState.FAILED
+            raise
+        else:
+            self._state = TransactionState.COMMITTED
+
+    async def rollback(self):
+        if self._connection._top_xact is self:
+            self._connection._top_xact = None
+
+        if self._state is not TransactionState.STARTED:
+            raise FatalError('cannot rollback transaction: inconsistent state')
+
+        if self._nested:
+            query = 'ROLLBACK TO {};'.format(self._id)
+        else:
+            query = 'ROLLBACK;'
+
+        try:
+            await self._connection.execute_script(query)
+        except:
+            self._state = TransactionState.FAILED
+            raise
+        else:
+            self._state = TransactionState.ROLLEDBACK
+
+
 async def connect(iri=None, *,
                   host=None, port=None,
                   user=None, password=None,

tests/test_transaction.py

@@ -0,0 +1,72 @@
+import asyncpg
+
+from asyncpg import _testbase as tb
+
+
+class TestTransaction(tb.ConnectedTestCase):
+
+    async def test_transaction_regular(self):
+
+        try:
+            async with self.con.transaction():
+
+                await self.con.execute_script('''
+                    CREATE TABLE mytab (a int);
+                ''')
+
+                1 / 0
+
+        except ZeroDivisionError:
+            pass
+        else:
+            self.fail('ZeroDivisionError was not raised')
+
+        with self.assertRaisesRegex(asyncpg.Error, '"mytab" does not exist'):
+            await self.con.prepare('''
+                SELECT * FROM mytab
+            ''')
+
+    async def test_transaction_nested(self):
+
+        try:
+            async with self.con.transaction():
+
+                await self.con.execute_script('''
+                    CREATE TABLE mytab (a int);
+                ''')
+
+                async with self.con.transaction():
+
+                    await self.con.execute_script('''
+                        INSERT INTO mytab (a) VALUES (1), (2);
+                    ''')
+
+                try:
+                    async with self.con.transaction():
+
+                        await self.con.execute_script('''
+                            INSERT INTO mytab (a) VALUES (3), (4);
+                        ''')
+
+                        1 / 0
+                except ZeroDivisionError:
+                    pass
+                else:
+                    self.fail('ZeroDivisionError was not raised')
+
+                res = await self.con.execute('SELECT * FROM mytab;')
+                self.assertEqual(len(res), 2)
+                self.assertEqual(res[0][0], 1)
+                self.assertEqual(res[1][0], 2)
+
+                1 / 0
+
+        except ZeroDivisionError:
+            pass
+        else:
+            self.fail('ZeroDivisionError was not raised')
+
+        with self.assertRaisesRegex(asyncpg.Error, '"mytab" does not exist'):
+            await self.con.prepare('''
+                SELECT * FROM mytab
+            ''')