"""
Transaction semantics, savepoints, session management, and isolation level
tests.

All test cases use the Register model (a single IntegerField) from
base_models, via BaseTransactionTestCase.

Test case ordering:

* Transaction commit/rollback, nesting, savepoints
* Session (context manager) behavior
* Lock type
* Isolation level
"""
import threading

from peewee import *

from .base import DatabaseTestCase
from .base import IS_CRDB
from .base import IS_CRDB_NESTED_TX
from .base import IS_MYSQL
from .base import IS_POSTGRESQL
from .base import IS_SQLITE
from .base import ModelTestCase
from .base import db
from .base import new_connection
from .base import skip_if
from .base import skip_unless
from .base_models import Register


class BaseTransactionTestCase(ModelTestCase):
    requires = [Register]

    def assertRegister(self, vals):
        query = Register.select().order_by(Register.value)
        self.assertEqual([register.value for register in query], vals)

    def _save(self, *vals):
        Register.insert([{Register.value: val} for val in vals]).execute()


def requires_nested(fn):
    return skip_if(IS_CRDB and not IS_CRDB_NESTED_TX,
                   'nested transaction support is required')(fn)


# ===========================================================================
# Transaction commit/rollback, nesting, and savepoints
# ===========================================================================

class TestTransaction(BaseTransactionTestCase):
    def test_simple(self):
        self.assertFalse(db.in_transaction())
        with db.atomic():
            self.assertTrue(db.in_transaction())
            self._save(1)

        self.assertFalse(db.in_transaction())
        self.assertRegister([1])

        # Explicit rollback, implicit commit.
        with db.atomic() as txn:
            self._save(2)
            txn.rollback()
            self.assertTrue(db.in_transaction())
            self._save(3)

        self.assertFalse(db.in_transaction())
        self.assertRegister([1, 3])

        # Explicit rollbacks.
        with db.atomic() as txn:
            self._save(4)
            txn.rollback()
            self._save(5)
            txn.rollback()

        self.assertRegister([1, 3])

        # Explicit commit, implicit rollback.
        try:
            with db.atomic() as txn:
                self._save(6)
                txn.commit()
                self.assertTrue(db.in_transaction())  # New txn begun.
                self._save(7)
                raise ValueError('oops')
        except ValueError:
            pass

        self.assertRegister([1, 3, 6])
        self.assertFalse(db.in_transaction())

    @requires_nested
    def test_simple_nested(self):
        self.assertFalse(db.in_transaction())
        with db.atomic():
            with db.atomic():
                self._save(1)
                self.assertTrue(db.in_transaction())
            self.assertTrue(db.in_transaction())

        self.assertFalse(db.in_transaction())
        self.assertRegister([1])

        # Explicit rollback, implicit commit.
        with db.atomic() as txn:
            with db.atomic() as sp:
                self._save(2)
                sp.rollback()
                self.assertTrue(db.in_transaction())
                self._save(3)

        self.assertFalse(db.in_transaction())
        self.assertRegister([1, 3])

        # Explicit rollbacks.
        with db.atomic() as txn:
            with db.atomic() as sp:
                self._save(4)
                sp.rollback()
                self._save(5)
                sp.rollback()

        self.assertRegister([1, 3])

        # Explicit commit, implicit rollback.
        with db.atomic() as txn:
            try:
                with db.atomic() as sp:
                    self._save(6)
                    sp.commit()
                    self.assertTrue(db.in_transaction())  # New txn begun.
                    self._save(7)
                    raise ValueError('oops')
            except ValueError:
                pass

        self.assertRegister([1, 3, 6])
        self.assertFalse(db.in_transaction())

    @requires_nested
    def test_transactions(self):
        self.assertFalse(db.in_transaction())

        with db.atomic():
            self.assertTrue(db.in_transaction())
            self._save(1)

        self.assertRegister([1])

        with db.atomic() as txn:
            self._save(2)
            txn.rollback()
            self._save(3)
            with db.atomic() as sp1:
                self._save(4)
                with db.atomic() as sp2:
                    self._save(5)
                    sp2.rollback()
                with db.atomic() as sp3:
                    self._save(6)
                    with db.atomic() as sp4:
                        self._save(7)
                        with db.atomic() as sp5:
                            self._save(8)
                        self.assertRegister([1, 3, 4, 6, 7, 8])
                        sp4.rollback()

                    self.assertRegister([1, 3, 4, 6])

        self.assertRegister([1, 3, 4, 6])

    def test_commit_rollback(self):
        with db.atomic() as txn:
            self._save(1)
            txn.commit()
            self._save(2)
            txn.rollback()

        self.assertRegister([1])

        with db.atomic() as txn:
            self._save(3)
            txn.rollback()
            self._save(4)

        self.assertRegister([1, 4])

    @requires_nested
    def test_commit_rollback_nested(self):
        with db.atomic() as txn:
            self.test_commit_rollback()
            self.assertRegister([1, 4])
            txn.rollback()
        self.assertRegister([])

        with db.atomic():
            self.test_commit_rollback()
        self.assertRegister([1, 4])

    def test_nesting_transaction_obj(self):
        self.assertRegister([])

        with db.transaction() as txn:
            self._save(1)
            with db.transaction() as txn2:
                self._save(2)
                txn2.rollback()  # Actually issues a rollback.
                self.assertRegister([])
            self._save(3)
        self.assertRegister([3])

        with db.transaction() as txn:
            self._save(4)
            with db.transaction() as txn2:
                with db.transaction() as txn3:
                    self._save(5)
                    txn3.commit()  # Actually commits.
            self._save(6)
            txn2.rollback()

        self.assertRegister([3, 4, 5])

        with db.transaction() as txn:
            self._save(6)
            try:
                with db.transaction() as txn2:
                    self._save(7)
                    raise ValueError()
            except ValueError:
                pass

        self.assertRegister([3, 4, 5, 6, 7])

    @requires_nested
    def test_savepoint_commit(self):
        with db.atomic() as txn:
            self._save(1)
            txn.rollback()

            self._save(2)
            txn.commit()

            with db.atomic() as sp:
                self._save(3)
                sp.rollback()

                self._save(4)
                sp.commit()

        self.assertRegister([2, 4])

    def test_atomic_decorator(self):
        @db.atomic()
        def save(i):
            self._save(i)

        save(1)
        self.assertRegister([1])

    def text_atomic_exception(self):
        def will_fail(self):
            with db.atomic():
                self._save(1)
                self._save(None)

        self.assertRaises(IntegrityError, will_fail)
        self.assertRegister([])

        def user_error(self):
            with db.atomic():
                self._save(2)
                raise ValueError

        self.assertRaises(ValueError, user_error)
        self.assertRegister([])

    def test_manual_commit(self):
        with db.manual_commit():
            db.begin()
            self._save(1)
            db.rollback()

            db.begin()
            self._save(2)
            db.commit()

            with db.manual_commit():
                db.begin()
                self._save(3)
                db.rollback()

            db.begin()
            self._save(4)
            db.commit()

        self.assertRegister([2, 4])

    def test_manual_commit_as_decorator(self):
        @db.manual_commit()
        def do_work():
            db.begin()
            self._save(100)
            db.rollback()
            db.begin()
            self._save(200)
            db.commit()

        do_work()
        self.assertRegister([200])

    def test_mixing_manual_atomic(self):
        @db.manual_commit()
        def will_fail():
            pass

        @db.atomic()
        def also_fails():
            pass

        with db.atomic():
            self.assertRaises(ValueError, will_fail)

        with db.manual_commit():
            self.assertRaises(ValueError, also_fails)

        with db.manual_commit():
            with self.assertRaises(ValueError):
                with db.atomic(): pass
        with db.atomic():
            with self.assertRaises(ValueError):
                with db.manual_commit(): pass

    def test_closing_db_in_transaction(self):
        with db.atomic():
            self.assertRaises(OperationalError, db.close)

    @requires_nested
    def test_db_context_manager(self):
        db.close()
        self.assertTrue(db.is_closed())

        with db:
            self.assertFalse(db.is_closed())
            self._save(1)
            with db:
                self._save(2)
                try:
                    with db:
                        self._save(3)
                        raise ValueError('xxx')
                except ValueError:
                    pass
                self._save(4)

            try:
                with db:
                    self._save(5)
                    with db:
                        self._save(6)
                    raise ValueError('yyy')
            except ValueError:
                pass

            self.assertFalse(db.is_closed())

        self.assertTrue(db.is_closed())
        self.assertRegister([1, 2, 4])

    def test_transaction_concurrency(self):
        barrier = threading.Barrier(5)
        accum = []

        def run_thread():
            barrier.wait(timeout=2)
            for i in range(10):
                try:
                    with db.atomic() as tx:
                        for j in range(10):
                            try:
                                with db.atomic() as sp:
                                    sp.commit()
                                    if j % 2 == 0:
                                        raise ValueError()
                            except ValueError:
                                pass
                        if i % 1 == 0:
                            raise ValueError()
                except ValueError:
                    pass
            accum.append(True)

        threads = [threading.Thread(target=run_thread) for _ in range(4)]
        for t in threads: t.start()
        barrier.wait(timeout=2)
        for t in threads: t.join()

        self.assertEqual(accum, [True, True, True, True])


# ===========================================================================
# Session (context manager) behavior
# ===========================================================================

@requires_nested
class TestSession(BaseTransactionTestCase):
    def test_session(self):
        self.assertTrue(db.session_start())
        self.assertTrue(db.session_start())
        self.assertEqual(db.transaction_depth(), 2)

        self._save(1)
        self.assertTrue(db.session_commit())
        self.assertEqual(db.transaction_depth(), 1)

        self._save(2)  # Now we're in autocommit mode.
        self.assertTrue(db.session_rollback())
        self.assertEqual(db.transaction_depth(), 0)

        self.assertTrue(db.session_start())
        self._save(3)
        self.assertTrue(db.session_rollback())
        self.assertRegister([1])

    def test_session_with_closed_db(self):
        db.close()
        self.assertTrue(db.session_start())
        self.assertFalse(db.is_closed())
        self.assertRaises(OperationalError, db.close)
        self._save(1)
        self.assertTrue(db.session_rollback())
        self.assertRegister([])

    def test_session_inside_context_manager(self):
        with db.atomic():
            self.assertTrue(db.session_start())
            self._save(1)
            self.assertTrue(db.session_commit())
            self._save(2)
            self.assertTrue(db.session_rollback())
            db.session_start()
            self._save(3)

        self.assertRegister([1, 3])

    def test_commit_rollback_mix(self):
        db.session_start()

        with db.atomic() as txn:  # Will be a savepoint.
            self._save(1)
            with db.atomic() as t2:
                self._save(2)
                with db.atomic() as t3:
                    self._save(3)
                t2.rollback()

            txn.commit()
            self._save(4)
            txn.rollback()

        self.assertTrue(db.session_commit())
        self.assertRegister([1])

    def test_session_rollback(self):
        db.session_start()

        self._save(1)
        with db.atomic() as txn:
            self._save(2)
            with db.atomic() as t2:
                self._save(3)

        self.assertRegister([1, 2, 3])
        self.assertTrue(db.session_rollback())
        self.assertRegister([])

        db.session_start()
        self._save(1)

        with db.transaction() as txn:
            self._save(2)
            with db.transaction() as t2:
                self._save(3)
                t2.rollback()  # Rolls back everything, starts new txn.

        db.session_commit()
        self.assertRegister([])

    def test_session_commit(self):
        db.session_start()

        self._save(1)
        with db.transaction() as txn:
            self._save(2)
            with db.transaction() as t2:
                self._save(3)
                t2.commit()  # Saves everything, starts new txn.
            txn.rollback()

        self.assertTrue(db.session_rollback())
        self.assertRegister([1, 2, 3])


# ===========================================================================
# Lock type and isolation level
# ===========================================================================

@skip_unless(IS_SQLITE, 'requires sqlite for transaction lock type')
class TestTransactionLockType(BaseTransactionTestCase):
    def test_lock_type(self):
        db2 = new_connection(timeout=0.0001)
        db2.connect()

        with self.database.atomic(lock_type='EXCLUSIVE') as txn:
            with self.assertRaises(OperationalError):
                with db2.atomic(lock_type='IMMEDIATE') as t2:
                    self._save(1)
            self._save(2)
        self.assertRegister([2])

        with self.database.atomic('IMMEDIATE') as txn:
            with self.assertRaises(OperationalError):
                with db2.atomic('EXCLUSIVE') as t2:
                    self._save(3)
            self._save(4)
        self.assertRegister([2, 4])

        with self.database.transaction(lock_type='DEFERRED') as txn:
            self._save(5)  # Deferred -> Exclusive after our write.
            with self.assertRaises(OperationalError):
                with db2.transaction(lock_type='IMMEDIATE') as t2:
                    self._save(6)
        self.assertRegister([2, 4, 5])


class TestTransactionIsolationLevel(BaseTransactionTestCase):
    @skip_unless(IS_POSTGRESQL, 'requires postgresql')
    def test_isolation_level_pg(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])

    @skip_unless(IS_MYSQL, 'requires mysql')
    def test_isolation_level_mysql(self):
        db2 = new_connection()
        db2.connect()

        with db2.atomic():
            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])

        with db2.atomic(isolation_level='READ UNCOMMITTED'):
            with db.atomic():
                self._save(3)
                self.assertDB2(db2, [1, 2, 3])
            self.assertDB2(db2, [1, 2, 3])
        self.assertDB2(db2, [1, 2, 3])

        with db2.atomic(isolation_level='REPEATABLE READ'):
            with db.atomic(isolation_level='REPEATABLE READ'):
                self._save(4)
                self.assertDB2(db2, [1, 2, 3])
            self.assertDB2(db2, [1, 2, 3])
        self.assertDB2(db2, [1, 2, 3, 4])

    def assertDB2(self, db2, vals):
        with Register.bind_ctx(db2):
            q = Register.select().order_by(Register.value)
            self.assertEqual([r.value for r in q], vals)


class TestConnectionContextDecorator(BaseTransactionTestCase):
    def test_connection_context_as_decorator(self):
        db2 = new_connection()

        @db2.connection_context()
        def do_work():
            self.assertFalse(db2.is_closed())
            return db2.execute_sql('SELECT 1').fetchone()

        self.assertTrue(db2.is_closed())
        result = do_work()
        self.assertEqual(result, (1,))
        self.assertTrue(db2.is_closed())
