feat!: pass datasource_type and datasource_id to form_data (#19981)

* pass datasource_type and datasource_id to form_data

* add datasource_type to delete command

* add datasource_type to delete command

* fix old keys implementation

* add more tests

superset-frontend/src/constants.ts

@@ -67,10 +67,18 @@ export const URL_PARAMS = {
     name: 'slice_id',
     type: 'string',
   },
+  datasourceId: {
+    name: 'datasource_id',
+    type: 'string',
+  },
   datasetId: {
     name: 'dataset_id',
     type: 'string',
   },
+  datasourceType: {
+    name: 'datasource_type',
+    type: 'string',
+  },
   dashboardId: {
     name: 'dashboard_id',
     type: 'string',
@@ -88,6 +96,8 @@ export const URL_PARAMS = {
 export const RESERVED_CHART_URL_PARAMS: string[] = [
   URL_PARAMS.formDataKey.name,
   URL_PARAMS.sliceId.name,
+  URL_PARAMS.datasourceId.name,
+  URL_PARAMS.datasourceType.name,
   URL_PARAMS.datasetId.name,
 ];
 export const RESERVED_DASHBOARD_URL_PARAMS: string[] = [

superset-frontend/src/dashboard/components/gridComponents/Chart.jsx

@@ -272,6 +272,7 @@ export default class Chart extends React.Component {
         : undefined;
       const key = await postFormData(
         this.props.datasource.id,
+        this.props.datasource.type,
         this.props.formData,
         this.props.slice.slice_id,
         nextTabId,

superset-frontend/src/explore/components/ExploreViewContainer/ExploreViewContainer.test.tsx

@@ -92,7 +92,7 @@ test('generates a new form_data param when none is available', async () => {
   expect(replaceState).toHaveBeenCalledWith(
     expect.anything(),
     undefined,
-    expect.stringMatching('dataset_id'),
+    expect.stringMatching('datasource_id'),
   );
   replaceState.mockRestore();
 });
@@ -109,7 +109,7 @@ test('generates a different form_data param when one is provided and is mounting
   expect(replaceState).toHaveBeenCalledWith(
     expect.anything(),
     undefined,
-    expect.stringMatching('dataset_id'),
+    expect.stringMatching('datasource_id'),
   );
   replaceState.mockRestore();
 });

superset-frontend/src/explore/components/ExploreViewContainer/index.jsx

@@ -152,14 +152,24 @@ const ExplorePanelContainer = styled.div`
 `;
 
 const updateHistory = debounce(
-  async (formData, datasetId, isReplace, standalone, force, title, tabId) => {
+  async (
+    formData,
+    datasourceId,
+    datasourceType,
+    isReplace,
+    standalone,
+    force,
+    title,
+    tabId,
+  ) => {
     const payload = { ...formData };
     const chartId = formData.slice_id;
     const additionalParam = {};
     if (chartId) {
       additionalParam[URL_PARAMS.sliceId.name] = chartId;
     } else {
-      additionalParam[URL_PARAMS.datasetId.name] = datasetId;
+      additionalParam[URL_PARAMS.datasourceId.name] = datasourceId;
+      additionalParam[URL_PARAMS.datasourceType.name] = datasourceType;
     }
 
     const urlParams = payload?.url_params || {};
@@ -173,11 +183,24 @@ const updateHistory = debounce(
       let key;
       let stateModifier;
       if (isReplace) {
-        key = await postFormData(datasetId, formData, chartId, tabId);
+        key = await postFormData(
+          datasourceId,
+          datasourceType,
+          formData,
+          chartId,
+          tabId,
+        );
         stateModifier = 'replaceState';
       } else {
         key = getUrlParam(URL_PARAMS.formDataKey);
-        await putFormData(datasetId, key, formData, chartId, tabId);
+        await putFormData(
+          datasourceId,
+          datasourceType,
+          key,
+          formData,
+          chartId,
+          tabId,
+        );
         stateModifier = 'pushState';
       }
       const url = mountExploreUrl(
@@ -229,11 +252,12 @@ function ExploreViewContainer(props) {
             dashboardId: props.dashboardId,
           }
         : props.form_data;
-      const datasetId = props.datasource.id;
+      const { id: datasourceId, type: datasourceType } = props.datasource;
 
       updateHistory(
         formData,
-        datasetId,
+        datasourceId,
+        datasourceType,
         isReplace,
         props.standalone,
         props.force,
@@ -245,6 +269,7 @@ function ExploreViewContainer(props) {
       props.dashboardId,
       props.form_data,
       props.datasource.id,
+      props.datasource.type,
       props.standalone,
       props.force,
       tabId,

superset-frontend/src/explore/components/controls/DatasourceControl/index.jsx

@@ -189,9 +189,9 @@ class DatasourceControl extends React.PureComponent {
     const isMissingDatasource = datasource.id == null;
     let isMissingParams = false;
     if (isMissingDatasource) {
-      const datasetId = getUrlParam(URL_PARAMS.datasetId);
+      const datasourceId = getUrlParam(URL_PARAMS.datasourceId);
       const sliceId = getUrlParam(URL_PARAMS.sliceId);
-      if (!datasetId && !sliceId) {
+      if (!datasourceId && !sliceId) {
         isMissingParams = true;
       }
     }

superset-frontend/src/explore/exploreUtils/formData.ts

@@ -20,7 +20,8 @@ import { omit } from 'lodash';
 import { SupersetClient, JsonObject } from '@superset-ui/core';
 
 type Payload = {
-  dataset_id: number;
+  datasource_id: number;
+  datasource_type: string;
   form_data: string;
   chart_id?: number;
 };
@@ -42,12 +43,14 @@ const assembleEndpoint = (key?: string, tabId?: string) => {
 };
 
 const assemblePayload = (
-  datasetId: number,
+  datasourceId: number,
+  datasourceType: string,
   formData: JsonObject,
   chartId?: number,
 ) => {
   const payload: Payload = {
-    dataset_id: datasetId,
+    datasource_id: datasourceId,
+    datasource_type: datasourceType,
     form_data: JSON.stringify(sanitizeFormData(formData)),
   };
   if (chartId) {
@@ -57,24 +60,36 @@ const assemblePayload = (
 };
 
 export const postFormData = (
-  datasetId: number,
+  datasourceId: number,
+  datasourceType: string,
   formData: JsonObject,
   chartId?: number,
   tabId?: string,
 ): Promise<string> =>
   SupersetClient.post({
     endpoint: assembleEndpoint(undefined, tabId),
-    jsonPayload: assemblePayload(datasetId, formData, chartId),
+    jsonPayload: assemblePayload(
+      datasourceId,
+      datasourceType,
+      formData,
+      chartId,
+    ),
   }).then(r => r.json.key);
 
 export const putFormData = (
-  datasetId: number,
+  datasourceId: number,
+  datasourceType: string,
   key: string,
   formData: JsonObject,
   chartId?: number,
   tabId?: string,
 ): Promise<string> =>
   SupersetClient.put({
     endpoint: assembleEndpoint(key, tabId),
-    jsonPayload: assemblePayload(datasetId, formData, chartId),
+    jsonPayload: assemblePayload(
+      datasourceId,
+      datasourceType,
+      formData,
+      chartId,
+    ),
   }).then(r => r.json.message);

superset/cachekeys/schemas.py

@@ -22,6 +22,7 @@
     datasource_type_description,
     datasource_uid_description,
 )
+from superset.utils.core import DatasourceType
 
 
 class Datasource(Schema):
@@ -36,7 +37,7 @@ class Datasource(Schema):
     )
     datasource_type = fields.String(
         description=datasource_type_description,
-        validate=validate.OneOf(choices=("druid", "table", "view")),
+        validate=validate.OneOf(choices=[ds.value for ds in DatasourceType]),
         required=True,
     )
 

superset/charts/schemas.py

@@ -31,6 +31,7 @@
 from superset.utils import pandas_postprocessing, schema as utils
 from superset.utils.core import (
     AnnotationType,
+    DatasourceType,
     FilterOperator,
     PostProcessingBoxplotWhiskerType,
     PostProcessingContributionOrientation,
@@ -198,7 +199,7 @@ class ChartPostSchema(Schema):
     datasource_id = fields.Integer(description=datasource_id_description, required=True)
     datasource_type = fields.String(
         description=datasource_type_description,
-        validate=validate.OneOf(choices=("druid", "table", "view")),
+        validate=validate.OneOf(choices=[ds.value for ds in DatasourceType]),
         required=True,
     )
     datasource_name = fields.String(
@@ -244,7 +245,7 @@ class ChartPutSchema(Schema):
     )
     datasource_type = fields.String(
         description=datasource_type_description,
-        validate=validate.OneOf(choices=("druid", "table", "view")),
+        validate=validate.OneOf(choices=[ds.value for ds in DatasourceType]),
         allow_none=True,
     )
     dashboards = fields.List(fields.Integer(description=dashboards_description))
@@ -983,7 +984,7 @@ class ChartDataDatasourceSchema(Schema):
     )
     type = fields.String(
         description="Datasource type",
-        validate=validate.OneOf(choices=("druid", "table")),
+        validate=validate.OneOf(choices=[ds.value for ds in DatasourceType]),
     )
 
 

superset/commands/exceptions.py

@@ -115,8 +115,24 @@ def __init__(self) -> None:
         super().__init__([_("Some roles do not exist")], field_name="roles")
 
 
+class DatasourceTypeInvalidError(ValidationError):
+    status = 422
+
+    def __init__(self) -> None:
+        super().__init__(
+            [_("Datasource type is invalid")], field_name="datasource_type"
+        )
+
+
 class DatasourceNotFoundValidationError(ValidationError):
     status = 404
 
     def __init__(self) -> None:
-        super().__init__([_("Dataset does not exist")], field_name="datasource_id")
+        super().__init__([_("Datasource does not exist")], field_name="datasource_id")
+
+
+class QueryNotFoundValidationError(ValidationError):
+    status = 404
+
+    def __init__(self) -> None:
+        super().__init__([_("Query does not exist")], field_name="datasource_id")

superset/dao/datasource/dao.py

@@ -39,11 +39,11 @@
 class DatasourceDAO(BaseDAO):
 
     sources: Dict[DatasourceType, Type[Datasource]] = {
-        DatasourceType.SQLATABLE: SqlaTable,
+        DatasourceType.TABLE: SqlaTable,
         DatasourceType.QUERY: Query,
         DatasourceType.SAVEDQUERY: SavedQuery,
         DatasourceType.DATASET: Dataset,
-        DatasourceType.TABLE: Table,
+        DatasourceType.SLTABLE: Table,
     }
 
     @classmethod
@@ -66,7 +66,7 @@ def get_datasource(
 
     @classmethod
     def get_all_sqlatables_datasources(cls, session: Session) -> List[Datasource]:
-        source_class = DatasourceDAO.sources[DatasourceType.SQLATABLE]
+        source_class = DatasourceDAO.sources[DatasourceType.TABLE]
         qry = session.query(source_class)
         qry = source_class.default_query(qry)
         return qry.all()

superset/databases/dao.py

@@ -24,6 +24,7 @@
 from superset.models.dashboard import Dashboard
 from superset.models.slice import Slice
 from superset.models.sql_lab import TabState
+from superset.utils.core import DatasourceType
 
 logger = logging.getLogger(__name__)
 
@@ -75,7 +76,8 @@ def get_related_objects(cls, database_id: int) -> Dict[str, Any]:
         charts = (
             db.session.query(Slice)
             .filter(
-                Slice.datasource_id.in_(dataset_ids), Slice.datasource_type == "table"
+                Slice.datasource_id.in_(dataset_ids),
+                Slice.datasource_type == DatasourceType.TABLE,
             )
             .all()
         )

superset/datasets/dao.py

@@ -26,6 +26,7 @@
 from superset.models.core import Database
 from superset.models.dashboard import Dashboard
 from superset.models.slice import Slice
+from superset.utils.core import DatasourceType
 from superset.views.base import DatasourceFilter
 
 logger = logging.getLogger(__name__)
@@ -56,7 +57,8 @@ def get_related_objects(database_id: int) -> Dict[str, Any]:
         charts = (
             db.session.query(Slice)
             .filter(
-                Slice.datasource_id == database_id, Slice.datasource_type == "table"
+                Slice.datasource_id == database_id,
+                Slice.datasource_type == DatasourceType.TABLE,
             )
             .all()
         )

superset/examples/birth_names.py

@@ -29,6 +29,7 @@
 from superset.models.core import Database
 from superset.models.dashboard import Dashboard
 from superset.models.slice import Slice
+from superset.utils.core import DatasourceType
 
 from ..utils.database import get_example_database
 from .helpers import (
@@ -205,13 +206,16 @@ def create_slices(tbl: SqlaTable, admin_owner: bool) -> Tuple[List[Slice], List[
     if admin_owner:
         slice_props = dict(
             datasource_id=tbl.id,
-            datasource_type="table",
+            datasource_type=DatasourceType.TABLE,
             owners=[admin],
             created_by=admin,
         )
     else:
         slice_props = dict(
-            datasource_id=tbl.id, datasource_type="table", owners=[], created_by=admin
+            datasource_id=tbl.id,
+            datasource_type=DatasourceType.TABLE,
+            owners=[],
+            created_by=admin,
         )
 
     print("Creating some slices")

superset/examples/country_map.py

@@ -24,6 +24,7 @@
 from superset import db
 from superset.connectors.sqla.models import SqlMetric
 from superset.models.slice import Slice
+from superset.utils.core import DatasourceType
 
 from .helpers import (
     get_example_data,
@@ -112,7 +113,7 @@ def load_country_map_data(only_metadata: bool = False, force: bool = False) -> N
     slc = Slice(
         slice_name="Birth in France by department in 2016",
         viz_type="country_map",
-        datasource_type="table",
+        datasource_type=DatasourceType.TABLE,
         datasource_id=tbl.id,
         params=get_slice_json(slice_data),
     )