feat: Add Explore Database to CTAS/CVAS visualization flow

superset-frontend/src/SqlLab/components/ResultSet.jsx

@@ -194,6 +194,10 @@ export default class ResultSet extends React.PureComponent {
       this.props.search ? this.props.height - SEARCH_HEIGHT : this.props.height,
     );
     let sql;
+    let exploreDBId = query.dbId;
+    if (this.props.database && this.props.database.explore_database_id) {
+      exploreDBId = this.props.database.explore_database_id;
+    }
 
     if (this.props.showSql) {
       sql = <HighlightedSql sql={query.sql} />;
@@ -245,7 +249,7 @@ export default class ResultSet extends React.PureComponent {
               <ExploreCtasResultsButton
                 table={tmpTable}
                 schema={tmpSchema}
-                dbId={query.dbId}
+                dbId={exploreDBId}
                 database={this.props.database}
                 actions={this.props.actions}
               />

superset/models/core.py

@@ -195,6 +195,10 @@ def allows_virtual_table_explore(self) -> bool:
 
         return bool(extra.get("allows_virtual_table_explore", True))
 
+    @property
+    def explore_database_id(self) -> int:
+        return self.get_extra().get("explore_database_id", self.id)
+
     @property
     def data(self) -> Dict[str, Any]:
         return {
@@ -205,6 +209,7 @@ def data(self) -> Dict[str, Any]:
             "allows_subquery": self.allows_subquery,
             "allows_cost_estimate": self.allows_cost_estimate,
             "allows_virtual_table_explore": self.allows_virtual_table_explore,
+            "explore_database_id": self.explore_database_id,
         }
 
     @property

superset/views/database/api.py

@@ -137,6 +137,7 @@ class DatabaseRestApi(DatabaseMixin, BaseSupersetModelRestApi):
         "allows_subquery",
         "allows_cost_estimate",
         "allows_virtual_table_explore",
+        "explore_database_id",
         "backend",
         "function_names",
     ]

superset/views/database/views.py

@@ -187,7 +187,7 @@ def form_post(self, form: CsvToDatabaseForm) -> FlaskResponse:
             # E.g. if hive was used to upload a csv, presto will be a better option
             # to explore the table.
             expore_database = database
-            explore_database_id = database.get_extra().get("explore_database_id", None)
+            explore_database_id = database.explore_database_id
             if explore_database_id:
                 expore_database = (
                     db.session.query(models.Database)

tests/core_tests.py

@@ -1307,6 +1307,20 @@ def test_virtual_table_explore_visibility(self):
         database.extra = json.dumps(extra)
         self.assertEqual(database.allows_virtual_table_explore, True)
 
+    def test_explore_database_id(self):
+        database = utils.get_example_database()
+        explore_database = utils.get_example_database()
+
+        # test that explore_database_id is the regular database
+        # id if none is set in the extra
+        self.assertEqual(database.explore_database_id, database.id)
+
+        # test that explore_database_id is correct if the extra is set
+        extra = database.get_extra()
+        extra["explore_database_id"] = explore_database.id
+        database.extra = json.dumps(extra)
+        self.assertEqual(database.explore_database_id, explore_database.id)
+
 
 if __name__ == "__main__":
     unittest.main()

tests/database_api_tests.py

@@ -52,6 +52,7 @@ def test_get_items(self):
             "allows_virtual_table_explore",
             "backend",
             "database_name",
+            "explore_database_id",
             "expose_in_sqllab",
             "force_ctas_schema",
             "function_names",