feat: Visualize SqlLab.Query model data in Explore 📈

superset-frontend/packages/superset-ui-core/src/query/DatasourceKey.ts

@@ -29,6 +29,7 @@ export default class DatasourceKey {
     this.id = parseInt(idStr, 10);
     this.type =
       typeStr === 'table' ? DatasourceType.Table : DatasourceType.Druid;
+    this.type = typeStr === 'query' ? DatasourceType.Query : this.type;
   }
 
   public toString() {

superset-frontend/packages/superset-ui-core/src/query/types/Column.ts

@@ -38,7 +38,7 @@ export type PhysicalColumn = string;
  * Column information defined in datasource.
  */
 export interface Column {
-  id: number;
+  id?: number;
   type?: string;
   type_generic?: GenericDataType;
   column_name: string;

superset-frontend/packages/superset-ui-core/src/query/types/Datasource.ts

@@ -47,11 +47,17 @@ export interface Datasource {
   };
 }
 
-export const DEFAULT_METRICS = [
+export const DEFAULT_METRICS: Metric[] = [
   {
     metric_name: 'COUNT(*)',
     expression: 'COUNT(*)',
   },
 ];
 
+export const isValidDatasourceType = (datasource: DatasourceType) =>
+  datasource === DatasourceType.Dataset ||
+  datasource === DatasourceType.SlTable ||
+  datasource === DatasourceType.SavedQuery ||
+  datasource === DatasourceType.Query;
+
 export default {};

superset-frontend/packages/superset-ui-core/src/query/types/Query.ts

@@ -336,6 +336,7 @@ export type QueryResults = {
     expanded_columns: QueryColumn[];
     selected_columns: QueryColumn[];
     query: { limit: number };
+    query_id?: number;
   };
 };
 

superset-frontend/src/SqlLab/components/ResultSet/index.tsx

@@ -23,7 +23,11 @@ import Button from 'src/components/Button';
 import shortid from 'shortid';
 import { styled, t, QueryResponse } from '@superset-ui/core';
 import ErrorMessageWithStackTrace from 'src/components/ErrorMessage/ErrorMessageWithStackTrace';
-import { SaveDatasetModal } from 'src/SqlLab/components/SaveDatasetModal';
+import {
+  ISaveableDataset,
+  ISimpleColumn,
+  SaveDatasetModal,
+} from 'src/SqlLab/components/SaveDatasetModal';
 import { UserWithPermissionsAndRoles } from 'src/types/bootstrapTypes';
 import ProgressBar from 'src/components/ProgressBar';
 import Loading from 'src/components/Loading';
@@ -220,6 +224,15 @@ export default class ResultSet extends React.PureComponent<
       const { showSaveDatasetModal } = this.state;
       const { query } = this.props;
 
+      const dataset: ISaveableDataset = {
+        columns: query.columns as ISimpleColumn[],
+        name: query?.tab || 'Untitled',
+        dbId: 1,
+        sql: query.sql,
+        templateParams: query.templateParams,
+        schema: query.schema,
+      };
+
       return (
         <ResultSetControls>
           <SaveDatasetModal
@@ -230,14 +243,24 @@ export default class ResultSet extends React.PureComponent<
             modalDescription={t(
               'Save this query as a virtual dataset to continue exploring',
             )}
-            datasource={query}
+            datasource={dataset}
           />
           <ResultSetButtons>
             {this.props.visualize &&
               this.props.database?.allows_virtual_table_explore && (
                 <ExploreResultsButton
                   database={this.props.database}
-                  onClick={() => this.setState({ showSaveDatasetModal: true })}
+                  onClick={() => {
+                    // There is currently redux / state issue where sometimes a query will have serverId
+                    // and other times it will not.  We need this attribute consistently for this to work
+                    const qid = this.props?.query?.results?.query_id;
+                    if (qid) {
+                      // This will open expolore using the query as datasource
+                      window.location.href = `/superset/explore/query/${qid}`;
+                    } else {
+                      this.setState({ showSaveDatasetModal: true });
+                    }
+                  }}
                 />
               )}
             {this.props.csv && (

superset-frontend/src/SqlLab/components/SaveDatasetModal/SaveDatasetModal.test.tsx

@@ -17,16 +17,42 @@
  * under the License.
  */
 import React from 'react';
-import { QueryResponse, testQuery } from '@superset-ui/core';
-import { SaveDatasetModal } from 'src/SqlLab/components/SaveDatasetModal';
+import {
+  ISaveableDataset,
+  SaveDatasetModal,
+} from 'src/SqlLab/components/SaveDatasetModal';
 import { render, screen } from 'spec/helpers/testing-library';
+import { DatasourceType } from '@superset-ui/core';
+
+const testQuery: ISaveableDataset = {
+  name: 'unimportant',
+  dbId: 1,
+  sql: 'SELECT *',
+  columns: [
+    {
+      name: 'Column 1',
+      type: DatasourceType.Query,
+      is_dttm: false,
+    },
+    {
+      name: 'Column 3',
+      type: DatasourceType.Query,
+      is_dttm: false,
+    },
+    {
+      name: 'Column 2',
+      type: DatasourceType.Query,
+      is_dttm: true,
+    },
+  ],
+};
 
 const mockedProps = {
   visible: true,
   onHide: () => {},
   buttonTextOnSave: 'Save',
   buttonTextOnOverwrite: 'Overwrite',
-  datasource: testQuery as QueryResponse,
+  datasource: testQuery,
 };
 
 describe('SaveDatasetModal RTL', () => {
@@ -36,6 +62,7 @@ describe('SaveDatasetModal RTL', () => {
     const saveRadioBtn = screen.getByRole('radio', {
       name: /save as new unimportant/i,
     });
+
     const fieldLabel = screen.getByText(/save as new/i);
     const inputField = screen.getByRole('textbox');
     const inputFieldText = screen.getByDisplayValue(/unimportant/i);

superset-frontend/src/SqlLab/components/SaveDatasetModal/index.tsx

@@ -45,17 +45,41 @@ import {
   DatasetOptionAutocomplete,
   SqlLabExploreRootState,
   getInitialState,
-  ExploreDatasource,
 } from 'src/SqlLab/types';
+
 import { exploreChart } from 'src/explore/exploreUtils';
+import { isString } from 'lodash';
+
+interface QueryDatabase {
+  id?: number;
+}
+
+export type ExploreQuery = QueryResponse & {
+  database?: QueryDatabase | null | undefined;
+};
+
+export interface ISimpleColumn {
+  name?: string | null;
+  type?: string | null;
+  is_dttm?: boolean | null;
+}
+
+export interface ISaveableDataset {
+  columns: ISimpleColumn[];
+  name: string;
+  dbId: number;
+  sql: string;
+  templateParams?: string | object | null;
+  schema?: string | null;
+}
 
 interface SaveDatasetModalProps {
   visible: boolean;
   onHide: () => void;
   buttonTextOnSave: string;
   buttonTextOnOverwrite: string;
   modalDescription?: string;
-  datasource: ExploreDatasource;
+  datasource: ISaveableDataset;
 }
 
 const Styles = styled.div`
@@ -106,6 +130,8 @@ const updateDataset = async (
   return data.json.result;
 };
 
+const UNTITLED = t('Untitled Dataset');
+
 // eslint-disable-next-line no-empty-pattern
 export const SaveDatasetModal: FunctionComponent<SaveDatasetModalProps> = ({
   visible,
@@ -115,9 +141,8 @@ export const SaveDatasetModal: FunctionComponent<SaveDatasetModalProps> = ({
   modalDescription,
   datasource,
 }) => {
-  const query = datasource as QueryResponse;
   const getDefaultDatasetName = () =>
-    `${query.tab} ${moment().format('MM/DD/YYYY HH:mm:ss')}`;
+    `${datasource?.name || UNTITLED} ${moment().format('MM/DD/YYYY HH:mm:ss')}`;
   const [datasetName, setDatasetName] = useState(getDefaultDatasetName());
   const [newOrOverwrite, setNewOrOverwrite] = useState(
     DatasetRadioState.SAVE_NEW,
@@ -138,10 +163,10 @@ export const SaveDatasetModal: FunctionComponent<SaveDatasetModalProps> = ({
 
   const handleOverwriteDataset = async () => {
     await updateDataset(
-      query.dbId,
+      datasource.dbId,
       datasetToOverwrite.datasetId,
-      query.sql,
-      query.results.selected_columns.map(
+      datasource.sql,
+      datasource?.columns?.map(
         (d: { name: string; type: string; is_dttm: boolean }) => ({
           column_name: d.name,
           type: d.type,
@@ -159,7 +184,8 @@ export const SaveDatasetModal: FunctionComponent<SaveDatasetModalProps> = ({
     exploreChart({
       ...EXPLORE_CHART_DEFAULT,
       datasource: `${datasetToOverwrite.datasetId}__table`,
-      selected_columns: query.results.selected_columns,
+      all_columns: datasource?.columns?.map?.((d: ISimpleColumn) => d.name),
+      selected_columns: datasource?.columns,
     });
   };
 
@@ -209,28 +235,29 @@ export const SaveDatasetModal: FunctionComponent<SaveDatasetModalProps> = ({
       return;
     }
 
-    const selectedColumns = query.results.selected_columns || [];
+    const selectedColumns = datasource?.columns ?? [];
 
     // The filters param is only used to test jinja templates.
     // Remove the special filters entry from the templateParams
     // before saving the dataset.
-    if (query.templateParams) {
-      const p = JSON.parse(query.templateParams);
+    let templateParams;
+    if (isString(datasource.templateParams)) {
+      const p = JSON.parse(datasource.templateParams);
       /* eslint-disable-next-line no-underscore-dangle */
       if (p._filters) {
         /* eslint-disable-next-line no-underscore-dangle */
         delete p._filters;
         // eslint-disable-next-line no-param-reassign
-        query.templateParams = JSON.stringify(p);
+        templateParams = JSON.stringify(p);
       }
     }
 
     dispatch(
       createDatasource({
-        schema: query.schema,
-        sql: query.sql,
-        dbId: query.dbId,
-        templateParams: query.templateParams,
+        schema: datasource.schema,
+        sql: datasource.sql,
+        dbId: datasource.dbId,
+        templateParams,
         datasourceName: datasetName,
         columns: selectedColumns,
       }),

superset-frontend/src/SqlLab/components/ScheduleQueryButton/index.tsx

@@ -16,13 +16,13 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import React, { FunctionComponent, useState } from 'react';
+import React, { FunctionComponent, useState, useRef } from 'react';
 import SchemaForm, { FormProps, FormValidation } from 'react-jsonschema-form';
 import { Row, Col } from 'src/components';
 import { Input, TextArea } from 'src/components/Input';
 import { t, styled } from '@superset-ui/core';
 import * as chrono from 'chrono-node';
-import ModalTrigger from 'src/components/ModalTrigger';
+import ModalTrigger, { ModalTriggerRef } from 'src/components/ModalTrigger';
 import { Form, FormItem } from 'src/components/Form';
 import Button from 'src/components/Button';
 
@@ -143,7 +143,7 @@ const ScheduleQueryButton: FunctionComponent<ScheduleQueryButtonProps> = ({
   const [description, setDescription] = useState('');
   const [label, setLabel] = useState(defaultLabel);
   const [showSchedule, setShowSchedule] = useState(false);
-  let saveModal: ModalTrigger | null;
+  const saveModal: ModalTriggerRef | null = useRef() as ModalTriggerRef;
 
   const onScheduleSubmit = ({
     formData,
@@ -159,7 +159,7 @@ const ScheduleQueryButton: FunctionComponent<ScheduleQueryButtonProps> = ({
       extra_json: JSON.stringify({ schedule_info: formData }),
     };
     onSchedule(query);
-    saveModal?.close();
+    saveModal?.current?.close();
   };
 
   const renderModalBody = () => (
@@ -225,9 +225,7 @@ const ScheduleQueryButton: FunctionComponent<ScheduleQueryButtonProps> = ({
   return (
     <span className="ScheduleQueryButton">
       <ModalTrigger
-        ref={ref => {
-          saveModal = ref;
-        }}
+        ref={saveModal}
         modalTitle={t('Schedule query')}
         modalBody={renderModalBody()}
         triggerNode={

superset-frontend/src/SqlLab/components/TableElement/index.tsx

@@ -152,11 +152,7 @@ const TableElement = ({ table, actions, ...props }: TableElementProps) => {
     if (table?.indexes?.length) {
       keyLink = (
         <ModalTrigger
-          modalTitle={
-            <div>
-              {t('Keys for table')} <strong>{table.name}</strong>
-            </div>
-          }
+          modalTitle={`${t('Keys for table')} ${table.name}`}
           modalBody={table.indexes.map((ix, i) => (
             <pre key={i}>{JSON.stringify(ix, null, '  ')}</pre>
           ))}

superset-frontend/src/SqlLab/types.ts

@@ -19,11 +19,8 @@
 import { SupersetError } from 'src/components/ErrorMessage/types';
 import { UserWithPermissionsAndRoles } from 'src/types/bootstrapTypes';
 import { ToastType } from 'src/components/MessageToasts/types';
-import { Dataset } from '@superset-ui/chart-controls';
-import { Query, QueryResponse } from '@superset-ui/core';
 import { ExploreRootState } from 'src/explore/types';
-
-export type ExploreDatasource = Dataset | QueryResponse;
+import { Query, QueryResponse } from '@superset-ui/core';
 
 // Object as Dictionary (associative array) with Query id as the key and type Query as the value
 export type QueryDictionary = {

superset-frontend/src/components/Modal/Modal.tsx

@@ -247,7 +247,13 @@ const CustomModal = ({
   const draggableRef = useRef<HTMLDivElement>(null);
   const [bounds, setBounds] = useState<DraggableBounds>();
   const [dragDisabled, setDragDisabled] = useState<boolean>(true);
-  const modalFooter = isNil(footer)
+  let FooterComponent;
+  if (React.isValidElement(footer)) {
+    // If a footer component is provided inject a closeModal function
+    // so the footer can provide a "close" button if desired
+    FooterComponent = React.cloneElement(footer, { closeModal: onHide });
+  }
+  const modalFooter = isNil(FooterComponent)
     ? [
         <Button key="back" onClick={onHide} cta data-test="modal-cancel-button">
           {t('Cancel')}
@@ -264,7 +270,7 @@ const CustomModal = ({
           {primaryButtonName}
         </Button>,
       ]
-    : footer;
+    : FooterComponent;
 
   const modalWidth = width || (responsive ? '100vw' : '600px');
   const shouldShowMask = !(resizable || draggable);

superset-frontend/src/components/ModalTrigger/ModalTrigger.stories.tsx

@@ -20,11 +20,11 @@ import React from 'react';
 import ModalTrigger from '.';
 
 interface IModalTriggerProps {
-  triggerNode: React.ReactNode;
+  triggerNode: JSX.Element;
   dialogClassName?: string;
-  modalTitle?: React.ReactNode;
-  modalBody?: React.ReactNode;
-  modalFooter?: React.ReactNode;
+  modalTitle?: string;
+  modalBody?: JSX.Element;
+  modalFooter?: JSX.Element;
   beforeOpen?: () => void;
   onExit?: () => void;
   isButton?: boolean;