fix: save columns reference from sqllab save datasets flow

superset-frontend/src/SqlLab/actions/sqlLab.js

@@ -1503,11 +1503,12 @@ export function createDatasourceFailed(err) {
 export function createDatasource(vizOptions) {
   return dispatch => {
     dispatch(createDatasourceStarted());
-    const { dbId, schema, datasourceName, sql } = vizOptions;
+    const { dbId, schema, datasourceName, sql, columns } = vizOptions;
     return SupersetClient.post({
       endpoint: '/api/v1/dataset/',
       headers: { 'Content-Type': 'application/json' },
       body: JSON.stringify({
+        columns,
         database: dbId,
         schema,
         sql,

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

@@ -62,6 +62,7 @@ export type ExploreQuery = QueryResponse & {
 
 export interface ISimpleColumn {
   column_name?: string | null;
+  name?: string | null;
   type?: string | null;
   is_dttm?: boolean | null;
 }
@@ -199,6 +200,7 @@ export const SaveDatasetModal = ({
       return;
     }
     setLoading(true);
+
     const [, key] = await Promise.all([
       updateDataset(
         datasource?.dbId,
@@ -218,7 +220,9 @@ export const SaveDatasetModal = ({
         ...formDataWithDefaults,
         datasource: `${datasetToOverwrite.datasetid}__table`,
         ...(defaultVizType === 'table' && {
-          all_columns: datasource?.columns?.map(column => column.column_name),
+          all_columns: datasource?.columns?.map(
+            column => column.name || column.column_name,
+          ),
         }),
       }),
     ]);
@@ -305,7 +309,9 @@ export const SaveDatasetModal = ({
           ...formDataWithDefaults,
           datasource: `${data.id}__table`,
           ...(defaultVizType === 'table' && {
-            all_columns: selectedColumns.map(column => column.column_name),
+            all_columns: selectedColumns.map(
+              column => column.name || column.column_name,
+            ),
           }),
         }),
       )

superset/datasets/api.py

@@ -305,7 +305,7 @@ def post(self) -> Response:
 
         try:
             new_model = CreateDatasetCommand(item).run()
-            return self.response(201, id=new_model.id, result=item)
+            return self.response(201, id=new_model.id, result=item, data=new_model.data)
         except DatasetInvalidError as ex:
             return self.response_422(message=ex.normalized_messages())
         except DatasetCreateFailedError as ex:

superset/datasets/commands/create.py

@@ -22,6 +22,7 @@
 from sqlalchemy.exc import SQLAlchemyError
 
 from superset.commands.base import BaseCommand, CreateMixin
+from superset.connectors.sqla.models import SqlMetric, TableColumn
 from superset.dao.exceptions import DAOCreateFailedError
 from superset.datasets.commands.exceptions import (
     DatabaseNotFoundValidationError,
@@ -44,8 +45,23 @@ def run(self) -> Model:
         self.validate()
         try:
             # Creates SqlaTable (Dataset)
+            columns = self._properties.pop("columns", [])
             dataset = DatasetDAO.create(self._properties, commit=False)
             # Updates columns and metrics from the dataset
+            ds_columns = []
+            for column in columns:
+                column_name = column.get("column_name") or column.get("name")
+                col = TableColumn(
+                    column_name=column_name,
+                    filterable=True,
+                    groupby=True,
+                    is_dttm=column.get("is_dttm", False),
+                    type=column.get("type", False),
+                )
+                ds_columns.append(col)
+
+            dataset.columns = ds_columns
+            dataset.metrics = [SqlMetric(metric_name="count", expression="COUNT(*)")]
             dataset.fetch_metadata(commit=False)
             db.session.commit()
         except (SQLAlchemyError, DAOCreateFailedError) as ex:

superset/datasets/schemas.py

@@ -84,6 +84,7 @@ class DatasetPostSchema(Schema):
     owners = fields.List(fields.Integer())
     is_managed_externally = fields.Boolean(allow_none=True, dump_default=False)
     external_url = fields.String(allow_none=True)
+    columns = fields.List(fields.Dict())
 
 
 class DatasetPutSchema(Schema):

superset/views/core.py

@@ -2080,7 +2080,7 @@ def sqllab_viz(self) -> FlaskResponse:  # pylint: disable=no-self-use
             cols.append(col)
 
         table.columns = cols
-        table.metrics = [SqlMetric(metric_name="count", expression="count(*)")]
+        table.metrics = [SqlMetric(metric_name="count", expression="COUNT(*)")]
         db.session.commit()
 
         return json_success(