From e06a0cd89bc84f3a7e75ee6f03df8c9c3a2badeb Mon Sep 17 00:00:00 2001
From: vera-liu <vera.liu@airbnb.com>
Date: Thu, 15 Dec 2016 13:19:54 -0800
Subject: [PATCH] Add force_ctas_schema to query model when enabled (#1825)

* Add force_ctas_schema to query model when enabled

* Add schema to temp_table_name

* Remove extra arg in create_table_as
---
 .../javascripts/SqlLab/components/QueryTable.jsx       |  5 ++++-
 superset/sql_lab.py                                    |  6 ++----
 superset/views.py                                      | 10 +++++++++-
 3 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/superset/assets/javascripts/SqlLab/components/QueryTable.jsx b/superset/assets/javascripts/SqlLab/components/QueryTable.jsx
index 70ad7e5dc4c..04c6b79ac33 100644
--- a/superset/assets/javascripts/SqlLab/components/QueryTable.jsx
+++ b/superset/assets/javascripts/SqlLab/components/QueryTable.jsx
@@ -129,7 +129,10 @@ class QueryTable extends React.PureComponent {
           />
         );
       } else {
-        q.output = [q.schema, q.tempTable].filter((v) => (v)).join('.');
+        // if query was run using ctas and force_ctas_schema was set
+        // tempTable will have the schema
+        const schemaUsed = q.ctas && q.tempTable.includes('.') ? '' : q.schema;
+        q.output = [schemaUsed, q.tempTable].filter((v) => (v)).join('.');
       }
       q.progress = (
         <ProgressBar
diff --git a/superset/sql_lab.py b/superset/sql_lab.py
index 6dcda0bcac9..1bb1d0b7340 100644
--- a/superset/sql_lab.py
+++ b/superset/sql_lab.py
@@ -40,7 +40,7 @@ def dedup(l, suffix='__'):
     return new_l
 
 
-def create_table_as(sql, table_name, schema=None, override=False):
+def create_table_as(sql, table_name, override=False):
     """Reformats the query into the create table as query.
 
     Works only for the single select SQL statements, in all other cases
@@ -55,8 +55,6 @@ def create_table_as(sql, table_name, schema=None, override=False):
     # TODO(bkyryliuk): drop table if allowed, check the namespace and
     #                  the permissions.
     # TODO raise if multi-statement
-    if schema:
-        table_name = schema + '.' + table_name
     exec_sql = ''
     if override:
         exec_sql = 'DROP TABLE IF EXISTS {table_name};\n'
@@ -105,7 +103,7 @@ def get_sql_results(self, query_id, return_results=True, store_results=False):
                 query.user_id,
                 start_dttm.strftime('%Y_%m_%d_%H_%M_%S'))
         executed_sql = create_table_as(
-            executed_sql, query.tmp_table_name, database.force_ctas_schema)
+            executed_sql, query.tmp_table_name)
         query.select_as_cta_used = True
     elif (
             query.limit and superset_query.is_select() and
diff --git a/superset/views.py b/superset/views.py
index 27be04527c4..85fe9a11d28 100755
--- a/superset/views.py
+++ b/superset/views.py
@@ -2392,6 +2392,14 @@ class Superset(BaseSupersetView):
                 get_datasource_access_error_msg('{}'.format(rejected_tables)))
         session.commit()
 
+        select_as_cta = request.form.get('select_as_cta') == 'true'
+        tmp_table_name = request.form.get('tmp_table_name')
+        if select_as_cta and mydb.force_ctas_schema:
+            tmp_table_name = '{}.{}'.format(
+                mydb.force_ctas_schema,
+                tmp_table_name
+            )
+
         query = models.Query(
             database_id=int(database_id),
             limit=int(app.config.get('SQL_MAX_ROW', None)),
@@ -2402,7 +2410,7 @@ class Superset(BaseSupersetView):
             tab_name=request.form.get('tab'),
             status=QueryStatus.PENDING if async else QueryStatus.RUNNING,
             sql_editor_id=request.form.get('sql_editor_id'),
-            tmp_table_name=request.form.get('tmp_table_name'),
+            tmp_table_name=tmp_table_name,
             user_id=int(g.user.get_id()),
             client_id=request.form.get('client_id'),
         )