Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: allow adhoc columns in non-aggregate query #21729

Merged
merged 11 commits into from
Oct 10, 2022
Merged
Show file tree
Hide file tree
Changes from 7 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
91 changes: 51 additions & 40 deletions superset/connectors/sqla/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@
and_,
asc,
Boolean,
Column,
Column as SqlaColumn,
mayurnewase marked this conversation as resolved.
Show resolved Hide resolved
DateTime,
desc,
Enum,
Expand Down Expand Up @@ -119,6 +119,7 @@
from superset.superset_typing import (
AdhocColumn,
AdhocMetric,
Column,
mayurnewase marked this conversation as resolved.
Show resolved Hide resolved
Metric,
OrderBy,
QueryObjectDict,
Expand Down Expand Up @@ -232,16 +233,16 @@ class TableColumn(Model, BaseColumn, CertificationMixin):

__tablename__ = "table_columns"
__table_args__ = (UniqueConstraint("table_id", "column_name"),)
table_id = Column(Integer, ForeignKey("tables.id"))
table_id = SqlaColumn(Integer, ForeignKey("tables.id"))
table: "SqlaTable" = relationship(
"SqlaTable",
backref=backref("columns", cascade="all, delete-orphan"),
foreign_keys=[table_id],
)
is_dttm = Column(Boolean, default=False)
expression = Column(MediumText())
python_date_format = Column(String(255))
extra = Column(Text)
is_dttm = SqlaColumn(Boolean, default=False)
expression = SqlaColumn(MediumText())
python_date_format = SqlaColumn(String(255))
extra = SqlaColumn(Text)

export_fields = [
"table_id",
Expand Down Expand Up @@ -312,7 +313,7 @@ def type_generic(self) -> Optional[utils.GenericDataType]:
)
return column_spec.generic_type if column_spec else None

def get_sqla_col(self, label: Optional[str] = None) -> Column:
def get_sqla_col(self, label: Optional[str] = None) -> SqlaColumn:
label = label or self.column_name
db_engine_spec = self.db_engine_spec
column_spec = db_engine_spec.get_column_spec(self.type, db_extra=self.db_extra)
Expand Down Expand Up @@ -512,14 +513,14 @@ class SqlMetric(Model, BaseMetric, CertificationMixin):

__tablename__ = "sql_metrics"
__table_args__ = (UniqueConstraint("table_id", "metric_name"),)
table_id = Column(Integer, ForeignKey("tables.id"))
table_id = SqlaColumn(Integer, ForeignKey("tables.id"))
table = relationship(
"SqlaTable",
backref=backref("metrics", cascade="all, delete-orphan"),
foreign_keys=[table_id],
)
expression = Column(MediumText(), nullable=False)
extra = Column(Text)
expression = SqlaColumn(MediumText(), nullable=False)
extra = SqlaColumn(Text)

export_fields = [
"metric_name",
Expand All @@ -538,7 +539,7 @@ class SqlMetric(Model, BaseMetric, CertificationMixin):
def __repr__(self) -> str:
return str(self.metric_name)

def get_sqla_col(self, label: Optional[str] = None) -> Column:
def get_sqla_col(self, label: Optional[str] = None) -> SqlaColumn:
label = label or self.metric_name
tp = self.table.get_template_processor()
sqla_col: ColumnClause = literal_column(tp.process_template(self.expression))
Expand Down Expand Up @@ -644,9 +645,9 @@ def after_delete( # pylint: disable=unused-argument
sqlatable_user = Table(
"sqlatable_user",
metadata,
Column("id", Integer, primary_key=True),
Column("user_id", Integer, ForeignKey("ab_user.id")),
Column("table_id", Integer, ForeignKey("tables.id")),
SqlaColumn("id", Integer, primary_key=True),
SqlaColumn("user_id", Integer, ForeignKey("ab_user.id")),
SqlaColumn("table_id", Integer, ForeignKey("tables.id")),
)


Expand Down Expand Up @@ -694,21 +695,21 @@ class SqlaTable(Model, BaseDatasource): # pylint: disable=too-many-public-metho
# given the schema is optional.
__table_args__ = (UniqueConstraint("database_id", "schema", "table_name"),)

table_name = Column(String(250), nullable=False)
main_dttm_col = Column(String(250))
database_id = Column(Integer, ForeignKey("dbs.id"), nullable=False)
fetch_values_predicate = Column(Text)
table_name = SqlaColumn(String(250), nullable=False)
main_dttm_col = SqlaColumn(String(250))
database_id = SqlaColumn(Integer, ForeignKey("dbs.id"), nullable=False)
fetch_values_predicate = SqlaColumn(Text)
owners = relationship(owner_class, secondary=sqlatable_user, backref="tables")
database: Database = relationship(
"Database",
backref=backref("tables", cascade="all, delete-orphan"),
foreign_keys=[database_id],
)
schema = Column(String(255))
sql = Column(MediumText())
is_sqllab_view = Column(Boolean, default=False)
template_params = Column(Text)
extra = Column(Text)
schema = SqlaColumn(String(255))
sql = SqlaColumn(MediumText())
is_sqllab_view = SqlaColumn(Boolean, default=False)
template_params = SqlaColumn(Text)
extra = SqlaColumn(Text)

baselink = "tablemodelview"

Expand Down Expand Up @@ -1314,6 +1315,7 @@ def get_sqla_query( # pylint: disable=too-many-arguments,too-many-locals,too-ma
columns_by_name: Dict[str, TableColumn] = {
col.column_name: col for col in self.columns
}
print("column by names ", columns_by_name)

mayurnewase marked this conversation as resolved.
Show resolved Hide resolved
metrics_by_name: Dict[str, SqlMetric] = {m.metric_name: m for m in self.metrics}

Expand Down Expand Up @@ -1394,7 +1396,7 @@ def get_sqla_query( # pylint: disable=too-many-arguments,too-many-locals,too-ma
_("Unknown column used in orderby: %(col)s", col=orig_col)
)

select_exprs: List[Union[Column, Label]] = []
select_exprs: List[Union[SqlaColumn, Label]] = []
groupby_all_columns = {}
groupby_series_columns = {}

Expand Down Expand Up @@ -1438,15 +1440,24 @@ def get_sqla_query( # pylint: disable=too-many-arguments,too-many-locals,too-ma
select_exprs.append(outer)
elif columns:
for selected in columns:
if isinstance(selected, dict):
mayurnewase marked this conversation as resolved.
Show resolved Hide resolved
sql = selected["sqlExpression"]
column_label = selected["label"]
else:
sql = selected
column_label = selected

selected = validate_adhoc_subquery(
selected,
sql,
self.database_id,
self.schema,
)
select_exprs.append(
columns_by_name[selected].get_sqla_col()
if selected in columns_by_name
else self.make_sqla_column_compatible(literal_column(selected))
else self.make_sqla_column_compatible(
literal_column(selected), column_label
)
)
metrics_exprs = []

Expand Down Expand Up @@ -1514,7 +1525,7 @@ def get_sqla_query( # pylint: disable=too-many-arguments,too-many-locals,too-ma
val = flt.get("val")
op = flt["op"].upper()
col_obj: Optional[TableColumn] = None
sqla_col: Optional[Column] = None
sqla_col: Optional[SqlaColumn] = None
if flt_col == utils.DTTM_ALIAS and is_timeseries and dttm_col:
col_obj = dttm_col
elif is_adhoc_column(flt_col):
Expand Down Expand Up @@ -1820,7 +1831,7 @@ def _get_series_orderby(
series_limit_metric: Metric,
metrics_by_name: Dict[str, SqlMetric],
columns_by_name: Dict[str, TableColumn],
) -> Column:
) -> SqlaColumn:
if utils.is_adhoc_metric(series_limit_metric):
assert isinstance(series_limit_metric, dict)
ob = self.adhoc_metric_to_sqla(series_limit_metric, columns_by_name)
Expand Down Expand Up @@ -2568,17 +2579,17 @@ def write_shadow_dataset(
RLSFilterRoles = Table(
"rls_filter_roles",
metadata,
Column("id", Integer, primary_key=True),
Column("role_id", Integer, ForeignKey("ab_role.id"), nullable=False),
Column("rls_filter_id", Integer, ForeignKey("row_level_security_filters.id")),
SqlaColumn("id", Integer, primary_key=True),
SqlaColumn("role_id", Integer, ForeignKey("ab_role.id"), nullable=False),
SqlaColumn("rls_filter_id", Integer, ForeignKey("row_level_security_filters.id")),
)

RLSFilterTables = Table(
"rls_filter_tables",
metadata,
Column("id", Integer, primary_key=True),
Column("table_id", Integer, ForeignKey("tables.id")),
Column("rls_filter_id", Integer, ForeignKey("row_level_security_filters.id")),
SqlaColumn("id", Integer, primary_key=True),
SqlaColumn("table_id", Integer, ForeignKey("tables.id")),
SqlaColumn("rls_filter_id", Integer, ForeignKey("row_level_security_filters.id")),
)


Expand All @@ -2588,13 +2599,13 @@ class RowLevelSecurityFilter(Model, AuditMixinNullable):
"""

__tablename__ = "row_level_security_filters"
id = Column(Integer, primary_key=True)
name = Column(String(255), unique=True, nullable=False)
description = Column(Text)
filter_type = Column(
id = SqlaColumn(Integer, primary_key=True)
name = SqlaColumn(String(255), unique=True, nullable=False)
description = SqlaColumn(Text)
filter_type = SqlaColumn(
Enum(*[filter_type.value for filter_type in utils.RowLevelSecurityFilterType])
)
group_key = Column(String(255), nullable=True)
group_key = SqlaColumn(String(255), nullable=True)
roles = relationship(
security_manager.role_model,
secondary=RLSFilterRoles,
Expand All @@ -2603,4 +2614,4 @@ class RowLevelSecurityFilter(Model, AuditMixinNullable):
tables = relationship(
SqlaTable, secondary=RLSFilterTables, backref="row_level_security_filters"
)
clause = Column(Text, nullable=False)
clause = SqlaColumn(Text, nullable=False)
4 changes: 2 additions & 2 deletions superset/superset_typing.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,8 +53,8 @@ class AdhocMetric(TypedDict, total=False):

class AdhocColumn(TypedDict, total=False):
hasCustomLabel: Optional[bool]
label: Optional[str]
zhaoyongjie marked this conversation as resolved.
Show resolved Hide resolved
sqlExpression: Optional[str]
label: str
sqlExpression: str
columnType: Optional[Literal["BASE_AXIS", "SERIES"]]
timeGrain: Optional[str]

Expand Down
41 changes: 41 additions & 0 deletions tests/integration_tests/charts/data/api_tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -764,6 +764,47 @@ def test_with_virtual_table_with_colons_as_datasource(self):
assert "':xyz:qwerty'" in result["query"]
assert "':qwerty:'" in result["query"]

@pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
def test_with_table_columns_without_metrics(self):
mayurnewase marked this conversation as resolved.
Show resolved Hide resolved
request_payload = self.query_context_payload
request_payload["queries"][0]["columns"] = ["name", "gender"]
request_payload["queries"][0]["metrics"] = None
request_payload["queries"][0]["orderby"] = []

rv = self.post_assert_metric(CHART_DATA_URI, request_payload, "data")
result = rv.json["result"][0]

assert rv.status_code == 200
assert "name" in result["colnames"]
assert "gender" in result["colnames"]
assert "name" in result["query"]
assert "gender" in result["query"]
assert list(result["data"][0].keys()) == ["name", "gender"]

@pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
def test_with_adhoc_column_without_metrics(self):
request_payload = self.query_context_payload
request_payload["queries"][0]["columns"] = [
"name",
{
"label": "num divide by 10",
"sqlExpression": "num/10",
"expressionType": "SQL",
},
]
request_payload["queries"][0]["metrics"] = None
request_payload["queries"][0]["orderby"] = []

rv = self.post_assert_metric(CHART_DATA_URI, request_payload, "data")
result = rv.json["result"][0]

assert rv.status_code == 200
assert "num divide by 10" in result["colnames"]
assert "name" in result["colnames"]
assert "num divide by 10" in result["query"]
assert "name" in result["query"]
assert list(result["data"][0].keys()) == ["name", "num divide by 10"]


@pytest.mark.chart_data_flow
class TestGetChartDataApi(BaseTestChartDataApi):
Expand Down