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


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_null(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(updated_once=None)
            result = conn.execute(stmt)
            eq_(result.inserted_primary_key, [2])
            eq_(
                conn.execute(foos.select().where(foos.c.id == 1)).fetchall(),
                [(1, "b", "bz", None)],
            )

    def test_on_duplicate_key_update_expression(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=func.concat(stmt.inserted.bar, "_foo")
            )
            result = conn.execute(stmt)
            eq_(result.inserted_primary_key, [2])
            eq_(
                conn.execute(foos.select().where(foos.c.id == 1)).fetchall(),
                [(1, "ab_foo", "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])