import unittest

from superset import sql_parse


class SupersetTestCase(unittest.TestCase):

    def extract_tables(self, query):
        sq = sql_parse.ParsedQuery(query)
        return sq.tables

    def test_simple_select(self):
        query = 'SELECT * FROM tbname'
        self.assertEquals({'tbname'}, self.extract_tables(query))

        # underscores
        query = 'SELECT * FROM tb_name'
        self.assertEquals({'tb_name'},
                          self.extract_tables(query))

        # quotes
        query = 'SELECT * FROM "tbname"'
        self.assertEquals({'tbname'}, self.extract_tables(query))

        # unicode encoding
        query = 'SELECT * FROM "tb_name" WHERE city = "Lübeck"'
        self.assertEquals({'tb_name'}, self.extract_tables(query))

        # schema
        self.assertEquals(
            {'schemaname.tbname'},
            self.extract_tables('SELECT * FROM schemaname.tbname'))

        # quotes
        query = 'SELECT field1, field2 FROM tb_name'
        self.assertEquals({'tb_name'}, self.extract_tables(query))

        query = 'SELECT t1.f1, t2.f2 FROM t1, t2'
        self.assertEquals({'t1', 't2'}, self.extract_tables(query))

    def test_select_named_table(self):
        query = 'SELECT a.date, a.field FROM left_table a LIMIT 10'
        self.assertEquals(
            {'left_table'}, self.extract_tables(query))

    def test_reverse_select(self):
        query = 'FROM t1 SELECT field'
        self.assertEquals({'t1'}, self.extract_tables(query))

    def test_subselect(self):
        query = """
          SELECT sub.*
              FROM (
                    SELECT *
                      FROM s1.t1
                     WHERE day_of_week = 'Friday'
                   ) sub, s2.t2
          WHERE sub.resolution = 'NONE'
        """
        self.assertEquals({'s1.t1', 's2.t2'},
                          self.extract_tables(query))

        query = """
          SELECT sub.*
              FROM (
                    SELECT *
                      FROM s1.t1
                     WHERE day_of_week = 'Friday'
                   ) sub
          WHERE sub.resolution = 'NONE'
        """
        self.assertEquals({'s1.t1'}, self.extract_tables(query))

        query = """
            SELECT * FROM t1
            WHERE s11 > ANY
             (SELECT COUNT(*) /* no hint */ FROM t2
               WHERE NOT EXISTS
                (SELECT * FROM t3
                  WHERE ROW(5*t2.s1,77)=
                    (SELECT 50,11*s1 FROM t4)));
        """
        self.assertEquals({'t1', 't2', 't3', 't4'},
                          self.extract_tables(query))

    def test_select_in_expression(self):
        query = 'SELECT f1, (SELECT count(1) FROM t2) FROM t1'
        self.assertEquals({'t1', 't2'}, self.extract_tables(query))

    def test_union(self):
        query = 'SELECT * FROM t1 UNION SELECT * FROM t2'
        self.assertEquals({'t1', 't2'}, self.extract_tables(query))

        query = 'SELECT * FROM t1 UNION ALL SELECT * FROM t2'
        self.assertEquals({'t1', 't2'}, self.extract_tables(query))

        query = 'SELECT * FROM t1 INTERSECT ALL SELECT * FROM t2'
        self.assertEquals({'t1', 't2'}, self.extract_tables(query))

    def test_select_from_values(self):
        query = 'SELECT * FROM VALUES (13, 42)'
        self.assertFalse(self.extract_tables(query))

    def test_select_array(self):
        query = """
            SELECT ARRAY[1, 2, 3] AS my_array
            FROM t1 LIMIT 10
        """
        self.assertEquals({'t1'}, self.extract_tables(query))

    def test_select_if(self):
        query = """
            SELECT IF(CARDINALITY(my_array) >= 3, my_array[3], NULL)
            FROM t1 LIMIT 10
        """
        self.assertEquals({'t1'}, self.extract_tables(query))

    # SHOW TABLES ((FROM | IN) qualifiedName)? (LIKE pattern=STRING)?
    def test_show_tables(self):
        query = "SHOW TABLES FROM s1 like '%order%'"
        # TODO: figure out what should code do here
        self.assertEquals({'s1'}, self.extract_tables(query))

    # SHOW COLUMNS (FROM | IN) qualifiedName
    def test_show_columns(self):
        query = 'SHOW COLUMNS FROM t1'
        self.assertEquals({'t1'}, self.extract_tables(query))

    def test_where_subquery(self):
        query = """
          SELECT name
            FROM t1
            WHERE regionkey = (SELECT max(regionkey) FROM t2)
        """
        self.assertEquals({'t1', 't2'}, self.extract_tables(query))

        query = """
          SELECT name
            FROM t1
            WHERE regionkey IN (SELECT regionkey FROM t2)
        """
        self.assertEquals({'t1', 't2'}, self.extract_tables(query))

        query = """
          SELECT name
            FROM t1
            WHERE regionkey EXISTS (SELECT regionkey FROM t2)
        """
        self.assertEquals({'t1', 't2'}, self.extract_tables(query))

    # DESCRIBE | DESC qualifiedName
    def test_describe(self):
        self.assertEquals({'t1'}, self.extract_tables('DESCRIBE t1'))
        self.assertEquals({'t1'}, self.extract_tables('DESC t1'))

    # SHOW PARTITIONS FROM qualifiedName (WHERE booleanExpression)?
    # (ORDER BY sortItem (',' sortItem)*)? (LIMIT limit=(INTEGER_VALUE | ALL))?
    def test_show_partitions(self):
        query = """
            SHOW PARTITIONS FROM orders
            WHERE ds >= '2013-01-01' ORDER BY ds DESC;
        """
        self.assertEquals({'orders'}, self.extract_tables(query))

    def test_join(self):
        query = 'SELECT t1.*, t2.* FROM t1 JOIN t2 ON t1.a = t2.a;'
        self.assertEquals({'t1', 't2'}, self.extract_tables(query))

        # subquery + join
        query = """
            SELECT a.date, b.name FROM
                left_table a
                JOIN (
                  SELECT
                    CAST((b.year) as VARCHAR) date,
                    name
                  FROM right_table
                ) b
                ON a.date = b.date
        """
        self.assertEquals({'left_table', 'right_table'},
                          self.extract_tables(query))

        query = """
            SELECT a.date, b.name FROM
                left_table a
                LEFT INNER JOIN (
                  SELECT
                    CAST((b.year) as VARCHAR) date,
                    name
                  FROM right_table
                ) b
                ON a.date = b.date
        """
        self.assertEquals({'left_table', 'right_table'},
                          self.extract_tables(query))

        query = """
            SELECT a.date, b.name FROM
                left_table a
                RIGHT OUTER JOIN (
                  SELECT
                    CAST((b.year) as VARCHAR) date,
                    name
                  FROM right_table
                ) b
                ON a.date = b.date
        """
        self.assertEquals({'left_table', 'right_table'},
                          self.extract_tables(query))

        query = """
            SELECT a.date, b.name FROM
                left_table a
                FULL OUTER JOIN (
                  SELECT
                    CAST((b.year) as VARCHAR) date,
                    name
                  FROM right_table
                ) b
                ON a.date = b.date
        """
        self.assertEquals({'left_table', 'right_table'},
                          self.extract_tables(query))

        # TODO: add SEMI join support, SQL Parse does not handle it.
        # query = """
        #     SELECT a.date, b.name FROM
        #         left_table a
        #         LEFT SEMI JOIN (
        #           SELECT
        #             CAST((b.year) as VARCHAR) date,
        #             name
        #           FROM right_table
        #         ) b
        #         ON a.date = b.date
        # """
        # self.assertEquals({'left_table', 'right_table'},
        #                   sql_parse.extract_tables(query))

    def test_combinations(self):
        query = """
            SELECT * FROM t1
            WHERE s11 > ANY
             (SELECT * FROM t1 UNION ALL SELECT * FROM (
               SELECT t6.*, t3.* FROM t6 JOIN t3 ON t6.a = t3.a) tmp_join
               WHERE NOT EXISTS
                (SELECT * FROM t3
                  WHERE ROW(5*t3.s1,77)=
                    (SELECT 50,11*s1 FROM t4)));
        """
        self.assertEquals({'t1', 't3', 't4', 't6'},
                          self.extract_tables(query))

        query = """
        SELECT * FROM (SELECT * FROM (SELECT * FROM (SELECT * FROM EmployeeS)
            AS S1) AS S2) AS S3;
        """
        self.assertEquals({'EmployeeS'}, self.extract_tables(query))

    def test_with(self):
        query = """
            WITH
              x AS (SELECT a FROM t1),
              y AS (SELECT a AS b FROM t2),
              z AS (SELECT b AS c FROM t3)
            SELECT c FROM z;
        """
        self.assertEquals({'t1', 't2', 't3'},
                          self.extract_tables(query))

        query = """
            WITH
              x AS (SELECT a FROM t1),
              y AS (SELECT a AS b FROM x),
              z AS (SELECT b AS c FROM y)
            SELECT c FROM z;
        """
        self.assertEquals({'t1'}, self.extract_tables(query))

    def test_reusing_aliases(self):
        query = """
            with q1 as ( select key from q2 where key = '5'),
            q2 as ( select key from src where key = '5')
            select * from (select key from q1) a;
        """
        self.assertEquals({'src'}, self.extract_tables(query))

    def test_multistatement(self):
        query = 'SELECT * FROM t1; SELECT * FROM t2'
        self.assertEquals({'t1', 't2'}, self.extract_tables(query))

        query = 'SELECT * FROM t1; SELECT * FROM t2;'
        self.assertEquals({'t1', 't2'}, self.extract_tables(query))

    def test_update_not_select(self):
        sql = sql_parse.ParsedQuery('UPDATE t1 SET col1 = NULL')
        self.assertEquals(False, sql.is_select())
        self.assertEquals(False, sql.is_readonly())

    def test_explain(self):
        sql = sql_parse.ParsedQuery('EXPLAIN SELECT 1')

        self.assertEquals(True, sql.is_explain())
        self.assertEquals(False, sql.is_select())
        self.assertEquals(True, sql.is_readonly())

    def test_complex_extract_tables(self):
        query = """SELECT sum(m_examples) AS "sum__m_example"
            FROM
              (SELECT COUNT(DISTINCT id_userid) AS m_examples,
                      some_more_info
               FROM my_b_table b
               JOIN my_t_table t ON b.ds=t.ds
               JOIN my_l_table l ON b.uid=l.uid
               WHERE b.rid IN
                   (SELECT other_col
                    FROM inner_table)
                 AND l.bla IN ('x', 'y')
               GROUP BY 2
               ORDER BY 2 ASC) AS "meh"
            ORDER BY "sum__m_example" DESC
            LIMIT 10;"""
        self.assertEquals(
            {'my_l_table', 'my_b_table', 'my_t_table', 'inner_table'},
            self.extract_tables(query))

    def test_complex_extract_tables2(self):
        query = """SELECT *
            FROM table_a AS a, table_b AS b, table_c as c
            WHERE a.id = b.id and b.id = c.id"""
        self.assertEquals(
            {'table_a', 'table_b', 'table_c'},
            self.extract_tables(query))

    def test_complex_extract_tables3(self):
        query = """SELECT somecol AS somecol
            FROM
              (WITH bla AS
                 (SELECT col_a
                  FROM a
                  WHERE 1=1
                    AND column_of_choice NOT IN
                      ( SELECT interesting_col
                       FROM b ) ),
                    rb AS
                 ( SELECT yet_another_column
                  FROM
                    ( SELECT a
                     FROM c
                     GROUP BY the_other_col ) not_table
                  LEFT JOIN bla foo ON foo.prop = not_table.bad_col0
                  WHERE 1=1
                  GROUP BY not_table.bad_col1 ,
                           not_table.bad_col2 ,
                  ORDER BY not_table.bad_col_3 DESC , not_table.bad_col4 ,
                                                    not_table.bad_col5) SELECT random_col
               FROM d
               WHERE 1=1
               UNION ALL SELECT even_more_cols
               FROM e
               WHERE 1=1
               UNION ALL SELECT lets_go_deeper
               FROM f
               WHERE 1=1
            WHERE 2=2
            GROUP BY last_col
            LIMIT 50000;"""
        self.assertEquals(
            {'a', 'b', 'c', 'd', 'e', 'f'},
            self.extract_tables(query))

    def test_basic_breakdown_statements(self):
        multi_sql = """
        SELECT * FROM ab_user;
        SELECT * FROM ab_user LIMIT 1;
        """
        parsed = sql_parse.ParsedQuery(multi_sql)
        statements = parsed.get_statements()
        self.assertEquals(len(statements), 2)
        expected = [
            'SELECT * FROM ab_user',
            'SELECT * FROM ab_user LIMIT 1',
        ]
        self.assertEquals(statements, expected)

    def test_messy_breakdown_statements(self):
        multi_sql = """
        SELECT 1;\t\n\n\n  \t
        \t\nSELECT 2;
        SELECT * FROM ab_user;;;
        SELECT * FROM ab_user LIMIT 1
        """
        parsed = sql_parse.ParsedQuery(multi_sql)
        statements = parsed.get_statements()
        self.assertEquals(len(statements), 4)
        expected = [
            'SELECT 1',
            'SELECT 2',
            'SELECT * FROM ab_user',
            'SELECT * FROM ab_user LIMIT 1',
        ]
        self.assertEquals(statements, expected)