aminghadersohi commented on code in PR #38416:
URL: https://github.com/apache/superset/pull/38416#discussion_r2931761389
##########
superset/mcp_service/dataset/tool/list_datasets.py:
##########
@@ -152,9 +190,16 @@ def _serialize_dataset(
)
# Apply field filtering via serialization context
- # Always use columns_requested (either explicit select_columns or
defaults)
- # This triggers DatasetInfo._filter_fields_by_context for each dataset
columns_to_filter = result.columns_requested
+ # Re-add popularity_score if it was originally requested
+ # (it was stripped before the DAO query since it's computed)
+ if (
+ request.select_columns
+ and "popularity_score" in request.select_columns
+ and columns_to_filter
+ and "popularity_score" not in columns_to_filter
+ ):
+ columns_to_filter = list(columns_to_filter) + ["popularity_score"]
Review Comment:
Fixed in bad796b. Response filtering now uses the original
`request.select_columns` (saved before DAO mutation) instead of
`result.columns_requested` (which reflected DAO-expanded columns with injected
`id`). When no explicit columns requested, falls back to
`result.columns_requested` with DAO defaults.
##########
superset/mcp_service/dataset/tool/list_datasets.py:
##########
@@ -176,3 +221,78 @@ 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 = []
+
+ # Serialize
+ select_columns = request.select_columns or DEFAULT_DATASET_COLUMNS
+ if "popularity_score" not in select_columns:
+ select_columns = list(select_columns) + ["popularity_score"]
+
+ ds_objs = []
+ for item in ordered_items:
+ obj = serializer(item, select_columns)
+ if obj is not None:
+ ds_objs.append(obj)
+
+ attach_popularity_scores(ds_objs, scores)
+
+ total_pages = (total_count + page_size - 1) // page_size if page_size > 0
else 0
+ pagination_info = PaginationInfo(
+ page=page,
+ page_size=page_size,
+ total_count=total_count,
+ total_pages=total_pages,
+ has_next=page < total_pages - 1,
+ has_previous=page > 0,
+ )
+
+ return DatasetList(
+ datasets=ds_objs,
+ count=len(ds_objs),
+ total_count=total_count,
+ page=page,
+ page_size=page_size,
+ total_pages=total_pages,
+ has_previous=page > 0,
+ has_next=page < total_pages - 1,
+ columns_requested=select_columns,
+ columns_loaded=select_columns,
+ columns_available=all_columns,
+ sortable_columns=DATASET_SORTABLE_COLUMNS,
+ filters_applied=request.filters if isinstance(request.filters, list)
else [],
+ pagination=pagination_info,
+ timestamp=datetime.now(timezone.utc),
+ )
Review Comment:
Fixed in bad796b. In all three `_list_*_by_popularity` functions,
`columns_requested` now reflects the original `request.select_columns` (or
defaults). The internal expansion for attach step is isolated to
`select_columns`/`columns_loaded` only.
--
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]
