from sqlalchemy.testing.assertions import eq_, assert_raises
from sqlalchemy.testing import fixtures
from sqlalchemy import exc, testing
from sqlalchemy.dialects.mysql import insert
from sqlalchemy import Table, Column, Boolean, Integer, String, func


class OnDuplicateTest(fixtures.TablesTest):
    __only_on__ = 'mysql',
    __backend__ = True
    run_define_tables = 'each'

    @classmethod
    def define_tables(cls, metadata):
        Table(
            'foos', metadata,
            Column('id', Integer, primary_key=True, autoincrement=True),
            Column('bar', String(10)),
            Column('baz', String(10)),
            Column('updated_once', Boolean, default=False),
        )

    def test_bad_args(self):
        assert_raises(
            ValueError,
            insert(self.tables.foos, values={}).on_duplicate_key_update
        )
        assert_raises(
            exc.ArgumentError,
            insert(self.tables.foos, values={}).on_duplicate_key_update,
            {'id': 1, 'bar': 'b'},
            id=1,
            bar='b',
        )
        assert_raises(
            exc.ArgumentError,
            insert(self.tables.foos, values={}).on_duplicate_key_update,
            {'id': 1, 'bar': 'b'},
            {'id': 2, 'bar': 'baz'},
        )

    def test_on_duplicate_key_update(self):
        foos = self.tables.foos
        with testing.db.connect() as conn:
            conn.execute(insert(foos, dict(id=1, bar='b', baz='bz')))
            stmt = insert(foos).values(
                [dict(id=1, bar='ab'), dict(id=2, bar='b')])
            stmt = stmt.on_duplicate_key_update(bar=stmt.inserted.bar)
            result = conn.execute(stmt)
            eq_(result.inserted_primary_key, [2])
            eq_(
                conn.execute(foos.select().where(foos.c.id == 1)).fetchall(),
                [(1, 'ab', 'bz', False)]
            )

    def test_on_duplicate_key_update_preserve_order(self):
        foos = self.tables.foos
        with testing.db.connect() as conn:
            conn.execute(insert(foos,
                [dict(id=1, bar='b', baz='bz'), dict(id=2, bar='b', baz='bz2')]))

            stmt = insert(foos)
            update_condition = (foos.c.updated_once == False)

            # The following statements show importance of the columns update ordering
            # as old values being referenced in UPDATE clause are getting replaced one
            # by one from left to right with their new values.
            stmt1 = stmt.on_duplicate_key_update([
                ('bar', func.if_(update_condition, func.values(foos.c.bar), foos.c.bar)),
                ('updated_once', func.if_(update_condition, True, foos.c.updated_once)),
            ])
            stmt2 = stmt.on_duplicate_key_update([
                ('updated_once', func.if_(update_condition, True, foos.c.updated_once)),
                ('bar', func.if_(update_condition, func.values(foos.c.bar), foos.c.bar)),
            ])
            # First statement should succeed updating column bar
            conn.execute(stmt1, dict(id=1, bar='ab'))
            eq_(
                conn.execute(foos.select().where(foos.c.id == 1)).fetchall(),
                [(1, 'ab', 'bz', True)],
            )
            # Second statement will do noop update of column bar
            conn.execute(stmt2, dict(id=2, bar='ab'))
            eq_(
                conn.execute(foos.select().where(foos.c.id == 2)).fetchall(),
                [(2, 'b', 'bz2', True)]
            )

    def test_last_inserted_id(self):
        foos = self.tables.foos
        with testing.db.connect() as conn:
            stmt = insert(foos).values({"bar": "b", "baz": "bz"})
            result = conn.execute(
                stmt.on_duplicate_key_update(
                    bar=stmt.inserted.bar, baz="newbz")
            )
            eq_(result.inserted_primary_key, [1])

            stmt = insert(foos).values({"id": 1, "bar": "b", "baz": "bz"})
            result = conn.execute(
                stmt.on_duplicate_key_update(
                    bar=stmt.inserted.bar, baz="newbz")
            )
            eq_(result.inserted_primary_key, [1])