Fixing tests

superset/sql_lab.py

@@ -81,8 +81,7 @@ def handle_error(msg):
             query.tmp_table_name = 'tmp_{}_table_{}'.format(
                 query.user_id,
                 start_dttm.strftime('%Y_%m_%d_%H_%M_%S'))
-        executed_sql = superset_query.as_create_table(
-            executed_sql, query.tmp_table_name)
+        executed_sql = superset_query.as_create_table(query.tmp_table_name)
         query.select_as_cta_used = True
     elif (
             query.limit and superset_query.is_select() and

superset/sql_parse.py

@@ -75,24 +75,25 @@ def __process_identifier(self, identifier):
                 self._alias_names.add(identifier.tokens[0].value)
         self.__extract_from_token(identifier)
 
-    def as_create_table(self, table_name, override=False):
+    def as_create_table(self, table_name, overwrite=False):
         """Reformats the query into the create table as query.
 
         Works only for the single select SQL statements, in all other cases
         the sql query is not modified.
         :param superset_query: string, sql query that will be executed
-        :param table_name: string, will contain the results of the query execution
-        :param override, boolean, table table_name will be dropped if true
+        :param table_name: string, will contain the results of the
+            query execution
+        :param overwrite, boolean, table table_name will be dropped if true
         :return: string, create table as query
         """
-        # TODO(bkyryliuk): enforce that all the columns have names. Presto requires it
-        #                  for the CTA operation.
+        # TODO(bkyryliuk): enforce that all the columns have names.
+        # Presto requires it for the CTA operation.
         # TODO(bkyryliuk): drop table if allowed, check the namespace and
         #                  the permissions.
         # TODO raise if multi-statement
         exec_sql = ''
         sql = self.stripped()
-        if override:
+        if overwrite:
             exec_sql = 'DROP TABLE IF EXISTS {table_name};\n'
         exec_sql += "CREATE TABLE {table_name} AS \n{sql}"
         return exec_sql.format(**locals())

superset/utils.py

@@ -401,6 +401,3 @@ def ping_connection(dbapi_connection, connection_record, connection_proxy):
         except:
             raise exc.DisconnectionError()
         cursor.close()
-
-def strip_sql(sql):
-    return sql

tests/celery_tests.py

@@ -12,8 +12,9 @@
 
 import pandas as pd
 
-from superset import app, appbuilder, cli, db, models, sql_lab, dataframe
+from superset import app, appbuilder, cli, db, models, dataframe
 from superset.security import sync_role_definitions
+from superset.sql_parse import SupersetQuery
 
 from .base_tests import SupersetTestCase
 
@@ -35,38 +36,33 @@ class UtilityFunctionTests(SupersetTestCase):
 
     # TODO(bkyryliuk): support more cases in CTA function.
     def test_create_table_as(self):
-        select_query = "SELECT * FROM outer_space;"
-        updated_select_query = sql_lab.create_table_as(
-            select_query, "tmp")
+        q = SupersetQuery("SELECT * FROM outer_space;")
+
         self.assertEqual(
-            "CREATE TABLE tmp AS \nSELECT * FROM outer_space;",
-            updated_select_query)
+            "CREATE TABLE tmp AS \nSELECT * FROM outer_space",
+            q.as_create_table("tmp"))
 
-        updated_select_query_with_drop = sql_lab.create_table_as(
-            select_query, "tmp", override=True)
         self.assertEqual(
             "DROP TABLE IF EXISTS tmp;\n"
-            "CREATE TABLE tmp AS \nSELECT * FROM outer_space;",
-            updated_select_query_with_drop)
+            "CREATE TABLE tmp AS \nSELECT * FROM outer_space",
+            q.as_create_table("tmp", overwrite=True))
 
-        select_query_no_semicolon = "SELECT * FROM outer_space"
-        updated_select_query_no_semicolon = sql_lab.create_table_as(
-            select_query_no_semicolon, "tmp")
+        # now without a semicolon
+        q = SupersetQuery("SELECT * FROM outer_space")
         self.assertEqual(
             "CREATE TABLE tmp AS \nSELECT * FROM outer_space",
-            updated_select_query_no_semicolon)
+            q.as_create_table("tmp"))
 
+        # now a multi-line query
         multi_line_query = (
             "SELECT * FROM planets WHERE\n"
-            "Luke_Father = 'Darth Vader';")
-        updated_multi_line_query = sql_lab.create_table_as(
-            multi_line_query, "tmp")
-        expected_updated_multi_line_query = (
-            "CREATE TABLE tmp AS \nSELECT * FROM planets WHERE\n"
-            "Luke_Father = 'Darth Vader';")
+            "Luke_Father = 'Darth Vader'")
+        q = SupersetQuery(multi_line_query)
         self.assertEqual(
-            expected_updated_multi_line_query,
-            updated_multi_line_query)
+            "CREATE TABLE tmp AS \nSELECT * FROM planets WHERE\n"
+            "Luke_Father = 'Darth Vader'",
+            q.as_create_table("tmp")
+        )
 
 
 class CeleryTestCase(SupersetTestCase):