fix(sqla): labels_expected contains mutated label (#14095)

superset/connectors/sqla/models.py

@@ -883,6 +883,7 @@ class SqlaTable(  # pylint: disable=too-many-public-methods,too-many-instance-at
         if db_engine_spec.allows_alias_in_select:
             label = db_engine_spec.make_label_compatible(label_expected)
             sqla_col = sqla_col.label(label)
+        sqla_col.key = label_expected
         return sqla_col
 
     def make_orderby_compatible(
@@ -1136,7 +1137,7 @@ class SqlaTable(  # pylint: disable=too-many-public-methods,too-many-instance-at
         )
 
         # Expected output columns
-        labels_expected = [c.name for c in select_exprs]
+        labels_expected = [c.key for c in select_exprs]
 
         # Order by columns are "hidden" columns, some databases require them
         # always be present in SELECT if an aggregation function is used

tests/sqla_models_tests.py

@@ -22,6 +22,7 @@ import pytest
 
 from superset import db
 from superset.connectors.sqla.models import SqlaTable, TableColumn
+from superset.db_engine_specs.bigquery import BigQueryEngineSpec
 from superset.db_engine_specs.druid import DruidEngineSpec
 from superset.exceptions import QueryObjectValidationError
 from superset.models.core import Database
@@ -327,3 +328,36 @@ class TestDatabaseModel(SupersetTestCase):
         assert cols["mycase"].expression == ""
         assert VIRTUAL_TABLE_STRING_TYPES[backend].match(cols["mycase"].type)
         assert cols["expr"].expression == "case when 1 then 1 else 0 end"
+
+    @patch("superset.models.core.Database.db_engine_spec", BigQueryEngineSpec)
+    def test_labels_expected_on_mutated_query(self):
+        query_obj = {
+            "granularity": None,
+            "from_dttm": None,
+            "to_dttm": None,
+            "groupby": ["user"],
+            "metrics": [
+                {
+                    "expressionType": "SIMPLE",
+                    "column": {"column_name": "user"},
+                    "aggregate": "COUNT_DISTINCT",
+                    "label": "COUNT_DISTINCT(user)",
+                }
+            ],
+            "is_timeseries": False,
+            "filter": [],
+            "extras": {},
+        }
+
+        database = Database(database_name="testdb", sqlalchemy_uri="sqlite://")
+        table = SqlaTable(table_name="bq_table", database=database)
+        db.session.add(database)
+        db.session.add(table)
+        db.session.commit()
+        sqlaq = table.get_sqla_query(**query_obj)
+        assert sqlaq.labels_expected == ["user", "COUNT_DISTINCT(user)"]
+        sql = table.database.compile_sqla_query(sqlaq.sqla_query)
+        assert "COUNT_DISTINCT_user__00db1" in sql
+        db.session.delete(table)
+        db.session.delete(database)
+        db.session.commit()