fix: save columns reference from sqllab save datasets flow

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

@@ -251,8 +251,7 @@ export const CtasEnum = {
 };
 
 export type QueryColumn = {
-  name: string;
-  column_name?: string;
+  column_name: string;
   type: string | null;
   is_dttm: boolean;
 };
@@ -384,17 +383,17 @@ export const testQuery: Query = {
   type: DatasourceType.Query,
   columns: [
     {
-      name: 'Column 1',
+      column_name: 'Column 1',
       type: 'STRING',
       is_dttm: false,
     },
     {
-      name: 'Column 3',
+      column_name: 'Column 3',
       type: 'STRING',
       is_dttm: false,
     },
     {
-      name: 'Column 2',
+      column_name: 'Column 2',
       type: 'TIMESTAMP',
       is_dttm: true,
     },

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/ResultSet/index.tsx

@@ -205,7 +205,7 @@ const ResultSet = ({
         ...EXPLORE_CHART_DEFAULT,
         datasource: `${results.query_id}__query`,
         ...{
-          all_columns: results.columns.map(column => column.name),
+          all_columns: results.columns.map(column => column.column_name),
         },
       });
       const url = mountExploreUrl(null, {
@@ -491,7 +491,7 @@ const ResultSet = ({
     }
     if (data && data.length > 0) {
       const expandedColumns = results.expanded_columns
-        ? results.expanded_columns.map(col => col.name)
+        ? results.expanded_columns.map(col => col.column_name)
         : [];
       return (
         <>
@@ -500,7 +500,7 @@ const ResultSet = ({
           {sql}
           <FilterableTable
             data={data}
-            orderedColumnKeys={results.columns.map(col => col.name)}
+            orderedColumnKeys={results.columns.map(col => col.column_name)}
             height={rowsHeight}
             filterText={searchText}
             expandedColumns={expandedColumns}

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")
+                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/queries/dao.py

@@ -59,6 +59,8 @@ def update_saved_query_exec_info(query_id: int) -> None:
     def save_metadata(query: Query, payload: Dict[str, Any]) -> None:
         # pull relevant data from payload and store in extra_json
         columns = payload.get("columns", {})
+        for col in columns:
+            col["column_name"] = col.pop("name")
         db.session.add(query)
         query.set_extra_json_key("columns", columns)
 

tests/integration_tests/datasets/api_tests.py

@@ -542,6 +542,31 @@ def test_create_dataset_item(self):
         assert model.table_name == table_data["table_name"]
         assert model.database_id == table_data["database"]
 
+    def test_create_dataset_item_with_columns(self):
+        """
+        Dataset API: Test create dataset item with columns
+        """
+        if backend() == "sqlite":
+            return
+
+        main_db = get_main_database()
+        self.login(username="admin")
+        table_data = {
+            "database": main_db.id,
+            "schema": None,
+            "table_name": "ab_permission",
+            "columns": [{"column_name": "foo", "type": "STRING", "is_dttm": False}],
+        }
+        uri = "api/v1/dataset/"
+        rv = self.post_assert_metric(uri, table_data, "post")
+        assert rv.status_code == 201
+        data = json.loads(rv.data.decode("utf-8"))
+        table_id = data.get("id")
+        model = db.session.query(SqlaTable).get(table_id)
+        assert model.table_name == table_data["table_name"]
+        assert model.database_id == table_data["database"]
+        assert len(model.columns) == 1
+
         # Assert that columns were created
         columns = (
             db.session.query(TableColumn)