import pytest

from helpers.cluster import ClickHouseCluster
from helpers.test_tools import assert_eq_with_retry

cluster = ClickHouseCluster(__file__)

node1 = cluster.add_instance('node1', main_configs=['configs/remote_servers.xml'],
                             user_configs=['configs/user_good_restricted.xml'], with_zookeeper=True)
node2 = cluster.add_instance('node2', main_configs=['configs/remote_servers.xml'],
                             user_configs=['configs/user_good_restricted.xml'], with_zookeeper=True)
node3 = cluster.add_instance('node3', main_configs=['configs/remote_servers.xml'],
                             user_configs=['configs/user_good_allowed.xml'], with_zookeeper=True)
node4 = cluster.add_instance('node4', main_configs=['configs/remote_servers.xml'],
                             user_configs=['configs/user_good_allowed.xml'], with_zookeeper=True)


@pytest.fixture(scope="module")
def started_cluster():
    try:
        cluster.start()

        for node in [node1, node2]:
            node.query('''
            CREATE TABLE sometable(date Date, id UInt32, value Int32)
    ENGINE = ReplicatedMergeTree('/clickhouse/tables/0/sometable', '{replica}', date, id, 8192);
                '''.format(replica=node.name), user='awesome')

        for node in [node3, node4]:
            node.query('''
            CREATE TABLE someothertable(date Date, id UInt32, value Int32)
    ENGINE = ReplicatedMergeTree('/clickhouse/tables/0/someothertable', '{replica}', date, id, 8192);
                '''.format(replica=node.name), user='good')

        yield cluster

    finally:
        cluster.shutdown()


@pytest.mark.parametrize("table,query,expected,n1,n2", [
    pytest.param("sometable", "ALTER TABLE sometable DROP PARTITION 201706", '1', node1, node2, id="case1"),
    pytest.param("sometable", "TRUNCATE TABLE sometable", '0', node1, node2, id="case2"),
    pytest.param("sometable", "OPTIMIZE TABLE sometable", '4', node1, node2, id="case3"),
    pytest.param("someothertable", "ALTER TABLE someothertable DROP PARTITION 201706", '1', node3, node4, id="case4"),
    pytest.param("someothertable", "TRUNCATE TABLE someothertable", '0', node3, node4, id="case5"),
    pytest.param("someothertable", "OPTIMIZE TABLE someothertable", '4', node3, node4, id="case6"),
])
def test_alter_table_drop_partition(started_cluster, table, query, expected, n1, n2):
    to_insert = '''\
2017-06-16	111	0
2017-06-16	222	1
2017-06-16	333	2
2017-07-16	444	3
'''
    n1.query("INSERT INTO {} FORMAT TSV".format(table), stdin=to_insert, user='good')

    assert_eq_with_retry(n1, "SELECT COUNT(*) from {}".format(table), '4', user='good')
    assert_eq_with_retry(n2, "SELECT COUNT(*) from {}".format(table), '4', user='good')

    ### It maybe leader and everything will be ok
    n1.query(query, user='good')

    assert_eq_with_retry(n1, "SELECT COUNT(*) from {}".format(table), expected, user='good')
    assert_eq_with_retry(n2, "SELECT COUNT(*) from {}".format(table), expected, user='good')

    n1.query("INSERT INTO {} FORMAT TSV".format(table), stdin=to_insert, user='good')

    assert_eq_with_retry(n1, "SELECT COUNT(*) from {}".format(table), '4', user='good')
    assert_eq_with_retry(n2, "SELECT COUNT(*) from {}".format(table), '4', user='good')

    ### If node1 is leader than node2 will be slave
    n2.query(query, user='good')

    assert_eq_with_retry(n1, "SELECT COUNT(*) from {}".format(table), expected, user='good')
    assert_eq_with_retry(n2, "SELECT COUNT(*) from {}".format(table), expected, user='good')