aminghadersohi commented on code in PR #38416:
URL: https://github.com/apache/superset/pull/38416#discussion_r2980541558
##########
superset/mcp_service/dataset/tool/get_dataset_info.py:
##########
@@ -117,6 +118,11 @@ async def get_dataset_info(
result = tool.run_tool(request.identifier)
if isinstance(result, DatasetInfo):
+ # Compute popularity_score for single dataset retrieval
+ if result.id is not None:
+ scores = compute_dataset_popularity([result.id])
+ result.popularity_score = scores.get(result.id, 0.0)
Review Comment:
Fixed — popularity_score is now computed only when select_columns explicitly
includes 'popularity_score' or when select_columns is empty (return
everything). A new select_columns field on GetDatasetInfoRequest enables this
opt-in behavior.
##########
superset/mcp_service/dataset/tool/get_dataset_info.py:
##########
@@ -117,6 +118,17 @@ async def get_dataset_info(
result = tool.run_tool(request.identifier)
if isinstance(result, DatasetInfo):
+ # Compute popularity_score for single dataset retrieval
+ # Isolated so a failure here doesn't prevent returning dataset info
+ if result.id is not None:
+ try:
+ scores = compute_dataset_popularity([result.id])
+ result.popularity_score = scores.get(result.id, 0.0)
+ except (ValueError, TypeError, AttributeError) as e:
+ await ctx.warning(
+ "Failed to compute popularity score: %s" % (str(e),)
+ )
Review Comment:
Fixed — added SQLAlchemyError (from sqlalchemy.exc) to the except clause, so
database errors during popularity computation are caught and logged as warnings
rather than crashing the tool.
##########
superset/mcp_service/dataset/tool/list_datasets.py:
##########
@@ -124,32 +130,62 @@ def _serialize_dataset(
"""Serialize dataset (filtering via model_serializer)."""
return serialize_dataset_object(obj)
- # Create tool with standard serialization
- tool = ModelListCore(
- dao_class=DatasetDAO,
- output_schema=DatasetInfo,
- item_serializer=_serialize_dataset,
- filter_type=DatasetFilter,
- default_columns=DEFAULT_DATASET_COLUMNS,
- search_columns=["schema", "sql", "table_name", "uuid"],
- list_field_name="datasets",
- output_list_schema=DatasetList,
- all_columns=all_columns,
- sortable_columns=DATASET_SORTABLE_COLUMNS,
- logger=logger,
+ # Preserve original select_columns before any mutation for response
filtering
+ original_select_columns = (
+ list(request.select_columns) if request.select_columns else None
)
- with event_logger.log_context(action="mcp.list_datasets.query"):
- result = tool.run_tool(
- filters=request.filters,
- search=request.search,
- select_columns=request.select_columns,
- order_column=request.order_column,
- order_direction=request.order_direction,
- page=max(request.page - 1, 0),
- page_size=request.page_size,
+ # Two-pass approach when sorting by popularity_score
+ if request.order_column == "popularity_score":
+ with
event_logger.log_context(action="mcp.list_datasets.popularity_sort"):
+ result = _list_datasets_by_popularity(
+ request, DatasetDAO, _serialize_dataset, all_columns, ctx
+ )
+ else:
+ # Create tool with standard serialization
+ list_core = ModelListCore(
+ dao_class=DatasetDAO,
+ output_schema=DatasetInfo,
+ item_serializer=_serialize_dataset,
+ filter_type=DatasetFilter,
+ default_columns=DEFAULT_DATASET_COLUMNS,
+ search_columns=DATASET_SEARCH_COLUMNS,
+ list_field_name="datasets",
+ output_list_schema=DatasetList,
+ all_columns=all_columns,
+ sortable_columns=DATASET_SORTABLE_COLUMNS,
+ logger=logger,
)
+ # Strip computed fields before passing to DAO query
+ dao_columns = request.select_columns
+ if dao_columns:
+ dao_columns = [c for c in dao_columns if c !=
"popularity_score"]
+ # Ensure id is loaded when popularity_score was requested
+ # (scores are keyed by id)
+ if (
+ "popularity_score" in request.select_columns
+ and "id" not in dao_columns
+ ):
+ dao_columns = ["id"] + dao_columns
+
+ with event_logger.log_context(action="mcp.list_datasets.query"):
+ result = list_core.run_tool(
+ filters=request.filters,
+ search=request.search,
+ select_columns=dao_columns,
+ order_column=request.order_column,
+ order_direction=request.order_direction,
+ page=max(request.page - 1, 0),
+ page_size=request.page_size,
+ )
+
+ # Attach popularity scores if requested in select_columns
+ if request.select_columns and "popularity_score" in
request.select_columns:
+ if ds_ids := [d.id for d in result.datasets if d.id is not
None]:
+ scores = compute_dataset_popularity(ds_ids)
+ attach_popularity_scores(result.datasets, scores)
Review Comment:
Fixed — the response now preserves the original user-specified
select_columns by restoring columns_requested after run_tool() returns, before
any internal mutations (like injecting 'id' for score computation) affect the
metadata. Applied to list_datasets, list_charts, and list_dashboards.
##########
superset/mcp_service/dataset/tool/list_datasets.py:
##########
@@ -184,3 +228,84 @@ def _serialize_dataset(
)
)
raise
+
+
+def _list_datasets_by_popularity(
+ request: ListDatasetsRequest,
+ dao_class: Any,
+ serializer: Callable[..., dict[str, Any] | None],
+ all_columns: list[str],
+ ctx: Context,
+) -> DatasetList:
+ """Two-pass listing: sort all matching datasets by popularity score."""
+ from superset.mcp_service.common.schema_discovery import
DATASET_SORTABLE_COLUMNS
+
+ sorted_ids, scores, total_count = get_popularity_sorted_ids(
+ compute_fn=compute_dataset_popularity,
+ dao_class=dao_class,
+ filters=request.filters,
+ search=request.search,
+ search_columns=DATASET_SEARCH_COLUMNS,
+ order_direction=request.order_direction,
+ )
+
+ # Apply pagination to sorted IDs
+ page = max(request.page - 1, 0)
+ page_size = request.page_size
+ start = page * page_size
+ end = start + page_size
+
+ # Fetch full models for page IDs
+ if page_ids := sorted_ids[start:end]:
+ items = dao_class.find_by_ids(page_ids)
+ id_to_item = {item.id: item for item in items}
+ ordered_items = [id_to_item[pid] for pid in page_ids if pid in
id_to_item]
+ else:
+ ordered_items = []
Review Comment:
Fixed — added eager loading options (subqueryload for columns/metrics,
joinedload for database) to the find_by_ids call in the popularity sorting
path. Also added a query_options parameter to BaseDAO.find_by_ids() to support
this pattern.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
