From 1ee08fc2165a0c4ccbdeb12c9ccbcf56b8a806f6 Mon Sep 17 00:00:00 2001
From: John Bodley <4567245+john-bodley@users.noreply.github.com>
Date: Mon, 8 Oct 2018 10:32:40 -0700
Subject: [PATCH] [select-star] Adding optional schema to view (#6051)

---
 superset/db_engine_specs.py           | 18 ++++++++++++++++--
 superset/templates/superset/ajah.html |  1 -
 superset/views/core.py                | 17 +++++++++++------
 tests/core_tests.py                   |  5 +++++
 4 files changed, 32 insertions(+), 9 deletions(-)
 delete mode 100644 superset/templates/superset/ajah.html

diff --git a/superset/db_engine_specs.py b/superset/db_engine_specs.py
index 6c8e07061f6a..95cf6d86a9f3 100644
--- a/superset/db_engine_specs.py
+++ b/superset/db_engine_specs.py
@@ -32,7 +32,7 @@
 import pandas
 from past.builtins import basestring
 import sqlalchemy as sqla
-from sqlalchemy import select
+from sqlalchemy import Column, select
 from sqlalchemy.engine import create_engine
 from sqlalchemy.engine.url import make_url
 from sqlalchemy.sql import quoted_name, text
@@ -852,6 +852,20 @@ def _partition_query(
         """).format(**locals())
         return sql
 
+    @classmethod
+    def where_latest_partition(
+            cls, table_name, schema, database, qry, columns=None):
+        try:
+            col_name, value = cls.latest_partition(
+                table_name, schema, database, show_first=True)
+        except Exception:
+            # table is not partitioned
+            return False
+        for c in columns:
+            if c.get('name') == col_name:
+                return qry.where(Column(col_name) == value)
+        return False
+
     @classmethod
     def _latest_partition_from_df(cls, df):
         recs = df.to_records(index=False)
@@ -1180,7 +1194,7 @@ def where_latest_partition(
             cls, table_name, schema, database, qry, columns=None):
         try:
             col_name, value = cls.latest_partition(
-                table_name, schema, database)
+                table_name, schema, database, show_first=True)
         except Exception:
             # table is not partitioned
             return False
diff --git a/superset/templates/superset/ajah.html b/superset/templates/superset/ajah.html
deleted file mode 100644
index b5d122680d71..000000000000
--- a/superset/templates/superset/ajah.html
+++ /dev/null
@@ -1 +0,0 @@
-{{ content |safe }}
diff --git a/superset/views/core.py b/superset/views/core.py
index f811fd8733f4..3753ffdbe1ca 100755
--- a/superset/views/core.py
+++ b/superset/views/core.py
@@ -2332,7 +2332,7 @@ def table(self, database_id, table_name, schema):
             'columns': payload_columns,
             'selectStar': mydb.select_star(
                 table_name, schema=schema, show_cols=True, indent=True,
-                cols=columns, latest_partition=False),
+                cols=columns, latest_partition=True),
             'primaryKey': primary_key,
             'foreignKeys': foreign_keys,
             'indexes': keys,
@@ -2350,14 +2350,19 @@ def extra_table_metadata(self, database_id, table_name, schema):
         return json_success(json.dumps(payload))
 
     @has_access
-    @expose('/select_star/<database_id>/<table_name>/')
+    @expose('/select_star/<database_id>/<table_name>')
+    @expose('/select_star/<database_id>/<table_name>/<schema>')
     @log_this
-    def select_star(self, database_id, table_name):
+    def select_star(self, database_id, table_name, schema=None):
         mydb = db.session.query(
             models.Database).filter_by(id=database_id).first()
-        return self.render_template(
-            'superset/ajah.html',
-            content=mydb.select_star(table_name, show_cols=True),
+        return json_success(
+            mydb.select_star(
+                table_name,
+                schema,
+                latest_partition=True,
+                show_cols=True,
+            ),
         )
 
     @expose('/theme/')
diff --git a/tests/core_tests.py b/tests/core_tests.py
index d922953e845e..d0eeacf2296a 100644
--- a/tests/core_tests.py
+++ b/tests/core_tests.py
@@ -727,6 +727,11 @@ def test_schemas_access_for_csv_upload_endpoint(self,
                 .format(db_id=dbobj.id))
         assert data == ['this_schema_is_allowed_too']
 
+    def test_select_star(self):
+        self.login(username='admin')
+        resp = self.get_resp('/superset/select_star/1/birth_names')
+        self.assertIn('gender', resp)
+
 
 if __name__ == '__main__':
     unittest.main()