You are viewing a plain text version of this content. The canonical link for it is here.
Posted to notifications@superset.apache.org by "aminghadersohi (via GitHub)" <gi...@apache.org> on 2026/04/28 15:31:57 UTC
[PR] feat(mcp): add query_dataset tool to query datasets using semantic layer [superset]
aminghadersohi opened a new pull request, #39727:
URL: https://github.com/apache/superset/pull/39727
## Summary
Add a new MCP tool `query_dataset` that allows querying a dataset directly using its semantic layer (saved metrics, calculated columns, dimensions) without requiring a saved chart.
### Changes
- **New tool**: `query_dataset` in `superset/mcp_service/dataset/tool/query_dataset.py`
- **New schemas**: `QueryDatasetFilter`, `QueryDatasetRequest`, `QueryDatasetResponse` in `dataset/schemas.py`
- **Registration**: Tool registered in `app.py` with instructions and workflow documentation
- **Tests**: 10 unit tests covering happy path, validation errors, time range, filters, UUID lookup
### Motivation
Today the MCP service has `execute_sql` (raw SQL), `get_chart_data` (requires saved chart), and `get_dataset_info` (metadata only). There's no way to query a dataset's semantic layer ad-hoc. The new tool fills this gap by accepting dataset_id + metrics/columns/filters/time_range and returning tabular data.
### Implementation
Reuses the existing `QueryContextFactory` + `ChartDataCommand` pipeline, which ensures:
- Row Level Security (RLS) enforcement
- Dataset permission checks
- Saved metric resolution
- Cache support
### Test plan
- [x] 10 unit tests pass (happy path, not found, invalid metric/column with suggestions, no metrics/columns validation, time range handling, filters, empty results, UUID lookup)
- [ ] CI passes
- [ ] Manual MCP client test against staging
--
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: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] feat(mcp): add query_dataset tool to query datasets using semantic layer [superset]
Posted by "netlify[bot] (via GitHub)" <gi...@apache.org>.
netlify[bot] commented on PR #39727:
URL: https://github.com/apache/superset/pull/39727#issuecomment-4353401596
### <span aria-hidden="true">✅</span> Deploy Preview for *superset-docs-preview* ready!
| Name | Link |
|:-:|------------------------|
|<span aria-hidden="true">🔨</span> Latest commit | 934e48a0f492ef866d8d215907d96b3fe7dcc9f6 |
|<span aria-hidden="true">🔍</span> Latest deploy log | https://app.netlify.com/projects/superset-docs-preview/deploys/69f369428ab7d90008d6d05a |
|<span aria-hidden="true">😎</span> Deploy Preview | [https://deploy-preview-39727--superset-docs-preview.netlify.app](https://deploy-preview-39727--superset-docs-preview.netlify.app) |
|<span aria-hidden="true">📱</span> Preview on mobile | <details><summary> Toggle QR Code... </summary><br /><br /><br /><br />_Use your smartphone camera to open QR code link._</details> |
|<span aria-hidden="true">🤖</span> Make changes | [Run an agent on this branch](https://app.netlify.com/projects/superset-docs-preview/agent-runs#DcrLDYAgDADQXXrWSPwguIFjFItighYBD8a4u9ze4b0QIh8hwwRQgYl4Lq4Yj_3cHJKNid3eXLeNT02YMdlcZ2ZfMtng-ZmpdKnXTuq-VWhG0kIIRZLEgPD9) |
---
<!-- [superset-docs-preview Preview](https://deploy-preview-39727--superset-docs-preview.netlify.app) -->
_To edit notification comments on pull requests, go to your [Netlify project configuration](https://app.netlify.com/projects/superset-docs-preview/configuration/notifications#deploy-notifications)._
--
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: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] feat(mcp): add query_dataset tool to query datasets using semantic layer [superset]
Posted by "codeant-ai-for-open-source[bot] (via GitHub)" <gi...@apache.org>.
codeant-ai-for-open-source[bot] commented on code in PR #39727:
URL: https://github.com/apache/superset/pull/39727#discussion_r3164411873
##########
superset/mcp_service/dataset/tool/query_dataset.py:
##########
@@ -0,0 +1,457 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+MCP tool: query_dataset
+
+Query a dataset using its semantic layer (saved metrics, calculated columns,
+dimensions) without requiring a saved chart.
+"""
+
+import difflib
+import logging
+import time
+from typing import Any
+
+from fastmcp import Context
+from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.orm import joinedload, subqueryload
+from superset_core.mcp.decorators import tool, ToolAnnotations
+
+from superset.commands.exceptions import CommandException
+from superset.exceptions import OAuth2Error, OAuth2RedirectError, SupersetException
+from superset.extensions import event_logger
+from superset.mcp_service.chart.schemas import DataColumn, PerformanceMetadata
+from superset.mcp_service.dataset.schemas import (
+ DatasetError,
+ QueryDatasetRequest,
+ QueryDatasetResponse,
+)
+from superset.mcp_service.utils import _is_uuid
+from superset.mcp_service.utils.cache_utils import get_cache_status_from_result
+from superset.mcp_service.utils.oauth2_utils import build_oauth2_redirect_message
+
+logger = logging.getLogger(__name__)
+
+
+def _resolve_dataset(identifier: int | str, eager_options: list[Any]) -> Any | None:
+ """Resolve a dataset by int ID or UUID string.
+
+ Replicates the identifier resolution logic from ModelGetInfoCore._find_object().
+ """
+ from superset.daos.dataset import DatasetDAO
+
+ opts = eager_options or None
+
+ if isinstance(identifier, int):
+ return DatasetDAO.find_by_id(identifier, query_options=opts)
+
+ # Try parsing as int
+ try:
+ id_val = int(identifier)
+ return DatasetDAO.find_by_id(id_val, query_options=opts)
+ except (ValueError, TypeError):
+ pass
+
+ # Try UUID
+ if _is_uuid(str(identifier)):
+ return DatasetDAO.find_by_id(identifier, id_column="uuid", query_options=opts)
+
+ return None
+
+
+def _validate_names(
+ requested: list[str],
+ valid: set[str],
+ kind: str,
+) -> list[str]:
+ """Return list of error messages for names not found in *valid*.
+
+ Includes close-match suggestions when available.
+ """
+ errors: list[str] = []
+ for name in requested:
+ if name not in valid:
+ suggestions = difflib.get_close_matches(name, valid, n=3, cutoff=0.6)
+ msg = f"Unknown {kind}: '{name}'"
+ if suggestions:
+ msg += f". Did you mean: {', '.join(suggestions)}?"
+ errors.append(msg)
+ return errors
+
+
+@tool(
+ tags=["data"],
+ class_permission_name="Dataset",
+ annotations=ToolAnnotations(
+ title="Query dataset",
+ readOnlyHint=True,
+ destructiveHint=False,
+ ),
+)
+async def query_dataset( # noqa: C901
+ request: QueryDatasetRequest, ctx: Context
+) -> QueryDatasetResponse | DatasetError:
Review Comment:
**Suggestion:** This tool is missing the same data-model privacy gate used by other dataset metadata tools, so users who are blocked from dataset metadata can still infer schema details (column/metric names) via validation errors and close-match suggestions. Add the privacy access decorator/check here as well to prevent metadata leakage through this endpoint. [security]
<details>
<summary><b>Severity Level:</b> Critical 🚨</summary>
```mdx
- ❌ Data-model metadata privacy bypass in query_dataset MCP tool.
- ⚠️ Restricted users can infer hidden column and metric names.
- ⚠️ LLM agents can enumerate schema despite metadata gating.
```
</details>
<details>
<summary><b>Steps of Reproduction ✅ </b></summary>
```mdx
1. Configure a Superset environment where some users have Dataset read access for chart
rendering but lack data-model metadata permissions (no "can_get_drill_info",
"can_get_or_create_dataset", or "can_write" on "Dataset"), so
`user_can_view_data_model_metadata()` in `superset/mcp_service/privacy.py:110-129` returns
False for them.
2. As such a restricted user, call the MCP `get_dataset_info` tool
(`superset/mcp_service/dataset/tool/get_dataset_info.py:49-61`) for a known dataset
identifier and observe it is decorated with `@requires_data_model_metadata_access` on line
58 and explicitly denies access when `user_can_view_data_model_metadata()` is False (lines
102-108), returning a `DatasetError` with type `DataModelMetadataRestricted`.
3. Using the same user and MCP client, call the new `query_dataset` tool defined at
`superset/mcp_service/dataset/tool/query_dataset.py:97-108` with a valid `dataset_id` plus
at least one mistyped column or metric name (e.g. `"customeer_id"`) in `request.columns`
or `request.metrics`.
4. In `query_dataset()`, the code at `query_dataset.py:181-203` builds
`valid_columns`/`valid_metrics` from `dataset.columns` and `dataset.metrics`, then
`_validate_names()` (defined at `query_dataset.py:77-94`) uses `difflib.get_close_matches`
to generate error strings like `Unknown column: 'customeer_id'. Did you mean:
customer_id?`, which are concatenated and returned via `DatasetError.create(...,
error_type="ValidationError")` at `query_dataset.py:204-210`. This leaks real
column/metric names to a user for whom `get_dataset_info` and `list_datasets` (see
`list_datasets.py:80-90` with `@requires_data_model_metadata_access` and the runtime check
at `list_datasets.py:134-139`) would have hidden data-model metadata, meaning the privacy
gate is not applied consistently to this new tool.
```
</details>
[Fix in Cursor](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt=This%20is%20a%20comment%20left%20during%20a%20code%20review.%0A%0A%2A%2APath%3A%2A%2A%20superset%2Fmcp_service%2Fdataset%2Ftool%2Fquery_dataset.py%0A%2A%2ALine%3A%2A%2A%2097%3A108%0A%2A%2AComment%3A%2A%2A%0A%09%2ASecurity%3A%20This%20tool%20is%20missing%20the%20same%20data-model%20privacy%20gate%20used%20by%20other%20dataset%20metadata%20tools%2C%20so%20users%20who%20are%20blocked%20from%20dataset%20metadata%20can%20still%20infer%20schema%20details%20%28column%2Fmetric%20names%29%20via%20validation%20errors%20and%20close-match%20suggestions.%20Add%20the%20privacy%20access%20decorator%2Fcheck%20here%20as%20well%20to%20prevent%20metadata%20leakage%20through%20this%20endpoint.%0A%0AValidate%20the%20correctness%20of%20the%20flagged%20issue.%20If%20correct%2C%20How%20can%20I%20resolve%20this%3F%20If%20you%20propose%20a%20fix%2C%20implement%20it%20and%20please%20make%20it%20concise.%0AOnce%20fix%20is%20implemented%2C%2
0also%20check%20other%20comments%20on%20the%20same%20PR%2C%20and%20ask%20user%20if%20the%20user%20wants%20to%20fix%20the%20rest%20of%20the%20comments%20as%20well.%20if%20said%20yes%2C%20then%20fetch%20all%20the%20comments%20validate%20the%20correctness%20and%20implement%20a%20minimal%20fix%0A) | [Fix in VSCode Claude](https://app.codeant.ai/fix-in-ide?tool=vscode-claude&prompt=This%20is%20a%20comment%20left%20during%20a%20code%20review.%0A%0A%2A%2APath%3A%2A%2A%20superset%2Fmcp_service%2Fdataset%2Ftool%2Fquery_dataset.py%0A%2A%2ALine%3A%2A%2A%2097%3A108%0A%2A%2AComment%3A%2A%2A%0A%09%2ASecurity%3A%20This%20tool%20is%20missing%20the%20same%20data-model%20privacy%20gate%20used%20by%20other%20dataset%20metadata%20tools%2C%20so%20users%20who%20are%20blocked%20from%20dataset%20metadata%20can%20still%20infer%20schema%20details%20%28column%2Fmetric%20names%29%20via%20validation%20errors%20and%20close-match%20suggestions.%20Add%20the%20privacy%20access%20decorator%2Fcheck%20here%20as%20well
%20to%20prevent%20metadata%20leakage%20through%20this%20endpoint.%0A%0AValidate%20the%20correctness%20of%20the%20flagged%20issue.%20If%20correct%2C%20How%20can%20I%20resolve%20this%3F%20If%20you%20propose%20a%20fix%2C%20implement%20it%20and%20please%20make%20it%20concise.%0AOnce%20fix%20is%20implemented%2C%20also%20check%20other%20comments%20on%20the%20same%20PR%2C%20and%20ask%20user%20if%20the%20user%20wants%20to%20fix%20the%20rest%20of%20the%20comments%20as%20well.%20if%20said%20yes%2C%20then%20fetch%20all%20the%20comments%20validate%20the%20correctness%20and%20implement%20a%20minimal%20fix%0A)
*(Use Cmd/Ctrl + Click for best experience)*
<details>
<summary><b>Prompt for AI Agent 🤖 </b></summary>
```mdx
This is a comment left during a code review.
**Path:** superset/mcp_service/dataset/tool/query_dataset.py
**Line:** 97:108
**Comment:**
*Security: This tool is missing the same data-model privacy gate used by other dataset metadata tools, so users who are blocked from dataset metadata can still infer schema details (column/metric names) via validation errors and close-match suggestions. Add the privacy access decorator/check here as well to prevent metadata leakage through this endpoint.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
```
</details>
<a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F39727&comment_hash=65abe2f9dbfb62593b44724037a870a3b6ab3113543baada5ce68fd2a4dd4167&reaction=like'>👍</a> | <a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F39727&comment_hash=65abe2f9dbfb62593b44724037a870a3b6ab3113543baada5ce68fd2a4dd4167&reaction=dislike'>👎</a>
##########
superset/mcp_service/dataset/tool/query_dataset.py:
##########
@@ -0,0 +1,457 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+MCP tool: query_dataset
+
+Query a dataset using its semantic layer (saved metrics, calculated columns,
+dimensions) without requiring a saved chart.
+"""
+
+import difflib
+import logging
+import time
+from typing import Any
+
+from fastmcp import Context
+from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.orm import joinedload, subqueryload
+from superset_core.mcp.decorators import tool, ToolAnnotations
+
+from superset.commands.exceptions import CommandException
+from superset.exceptions import OAuth2Error, OAuth2RedirectError, SupersetException
+from superset.extensions import event_logger
+from superset.mcp_service.chart.schemas import DataColumn, PerformanceMetadata
+from superset.mcp_service.dataset.schemas import (
+ DatasetError,
+ QueryDatasetRequest,
+ QueryDatasetResponse,
+)
+from superset.mcp_service.utils import _is_uuid
+from superset.mcp_service.utils.cache_utils import get_cache_status_from_result
+from superset.mcp_service.utils.oauth2_utils import build_oauth2_redirect_message
+
+logger = logging.getLogger(__name__)
+
+
+def _resolve_dataset(identifier: int | str, eager_options: list[Any]) -> Any | None:
+ """Resolve a dataset by int ID or UUID string.
+
+ Replicates the identifier resolution logic from ModelGetInfoCore._find_object().
+ """
+ from superset.daos.dataset import DatasetDAO
+
+ opts = eager_options or None
+
+ if isinstance(identifier, int):
+ return DatasetDAO.find_by_id(identifier, query_options=opts)
+
+ # Try parsing as int
+ try:
+ id_val = int(identifier)
+ return DatasetDAO.find_by_id(id_val, query_options=opts)
+ except (ValueError, TypeError):
+ pass
+
+ # Try UUID
+ if _is_uuid(str(identifier)):
+ return DatasetDAO.find_by_id(identifier, id_column="uuid", query_options=opts)
+
+ return None
+
+
+def _validate_names(
+ requested: list[str],
+ valid: set[str],
+ kind: str,
+) -> list[str]:
+ """Return list of error messages for names not found in *valid*.
+
+ Includes close-match suggestions when available.
+ """
+ errors: list[str] = []
+ for name in requested:
+ if name not in valid:
+ suggestions = difflib.get_close_matches(name, valid, n=3, cutoff=0.6)
+ msg = f"Unknown {kind}: '{name}'"
+ if suggestions:
+ msg += f". Did you mean: {', '.join(suggestions)}?"
+ errors.append(msg)
+ return errors
+
+
+@tool(
+ tags=["data"],
+ class_permission_name="Dataset",
+ annotations=ToolAnnotations(
+ title="Query dataset",
+ readOnlyHint=True,
+ destructiveHint=False,
+ ),
+)
+async def query_dataset( # noqa: C901
+ request: QueryDatasetRequest, ctx: Context
+) -> QueryDatasetResponse | DatasetError:
+ """Query a dataset using its semantic layer (saved metrics, dimensions, filters).
+
+ Returns tabular data without requiring a saved chart. Use this when you want
+ to compute saved metrics, group by dimensions, or apply filters directly
+ against a dataset's curated semantic layer.
+
+ Workflow:
+ 1. list_datasets -> find a dataset
+ 2. get_dataset_info -> discover available columns and metrics
+ 3. query_dataset -> query using metric names and column names
+
+ Example:
+ ```json
+ {
+ "dataset_id": 123,
+ "metrics": ["count", "avg_revenue"],
+ "columns": ["product_category"],
+ "time_range": "Last 7 days",
+ "row_limit": 100
+ }
+ ```
+ """
+ await ctx.info(
+ "Starting dataset query: dataset_id=%s, metrics=%s, columns=%s, "
+ "row_limit=%s"
+ % (
+ request.dataset_id,
+ request.metrics,
+ request.columns,
+ request.row_limit,
+ )
+ )
+
+ try:
+ from superset.commands.chart.data.get_data_command import ChartDataCommand
+ from superset.common.query_context_factory import QueryContextFactory
+ from superset.connectors.sqla.models import SqlaTable
+
+ # ------------------------------------------------------------------
+ # Step 1: Resolve dataset
+ # ------------------------------------------------------------------
+ await ctx.report_progress(1, 5, "Looking up dataset")
+ eager_options = [
+ subqueryload(SqlaTable.columns),
+ subqueryload(SqlaTable.metrics),
+ joinedload(SqlaTable.database),
+ ]
+
+ with event_logger.log_context(action="mcp.query_dataset.lookup"):
+ dataset = _resolve_dataset(request.dataset_id, eager_options)
+
+ if dataset is None:
+ await ctx.error("Dataset not found: identifier=%s" % (request.dataset_id,))
+ return DatasetError.create(
+ error=f"No dataset found with identifier: {request.dataset_id}",
+ error_type="NotFound",
+ )
+
+ dataset_name = getattr(dataset, "table_name", None) or f"Dataset {dataset.id}"
+ await ctx.info(
+ "Dataset found: id=%s, name=%s, columns=%s, metrics=%s"
+ % (
+ dataset.id,
+ dataset_name,
+ len(dataset.columns),
+ len(dataset.metrics),
+ )
+ )
+
+ # ------------------------------------------------------------------
+ # Step 2: Validate requested columns and metrics
+ # ------------------------------------------------------------------
+ await ctx.report_progress(2, 5, "Validating columns and metrics")
+ valid_columns = {c.column_name for c in dataset.columns}
+ valid_metrics = {m.metric_name for m in dataset.metrics}
+
+ validation_errors: list[str] = []
+ validation_errors.extend(
+ _validate_names(request.columns, valid_columns, "column")
+ )
+ validation_errors.extend(
+ _validate_names(request.metrics, valid_metrics, "metric")
+ )
+ # Validate filter column names against dataset columns
+ filter_cols = [f.col for f in request.filters]
+ validation_errors.extend(
+ _validate_names(filter_cols, valid_columns, "filter column")
+ )
+ # Validate order_by names against columns + metrics
+ if request.order_by:
+ valid_orderby = valid_columns | valid_metrics
+ validation_errors.extend(
+ _validate_names(request.order_by, valid_orderby, "order_by")
+ )
+
+ if validation_errors:
+ error_msg = "; ".join(validation_errors)
+ await ctx.error("Validation failed: %s" % (error_msg,))
+ return DatasetError.create(
+ error=error_msg,
+ error_type="ValidationError",
+ )
+
+ # ------------------------------------------------------------------
+ # Step 3: Build filters and time range
+ # ------------------------------------------------------------------
+ warnings: list[str] = []
+ query_filters: list[dict[str, Any]] = [
+ {"col": f.col, "op": f.op, "val": f.val} for f in request.filters
+ ]
+ granularity: str | None = None
+
+ if request.time_range:
+ temporal_col = request.time_column or getattr(
+ dataset, "main_dttm_col", None
+ )
+ if not temporal_col:
+ await ctx.error("time_range provided but no temporal column available")
+ return DatasetError.create(
+ error=(
+ "time_range was provided but no temporal column is available. "
+ "Either set time_column explicitly or ensure the dataset has "
+ "a main datetime column configured."
+ ),
+ error_type="ValidationError",
+ )
+ # Validate that the temporal column actually exists on the dataset
+ if temporal_col not in valid_columns:
+ await ctx.error("time_column '%s' not found on dataset" % temporal_col)
+ return DatasetError.create(
+ error=(
+ f"time_column '{temporal_col}' does not exist on this dataset."
+ ),
+ error_type="ValidationError",
+ )
+ # Warn if the chosen temporal column isn't marked as datetime
+ dttm_cols = {c.column_name for c in dataset.columns if c.is_dttm}
+ if temporal_col not in dttm_cols:
+ warnings.append(
+ f"Column '{temporal_col}' is not marked as a datetime "
+ f"column on this dataset. Time filtering may not work "
+ f"as expected."
+ )
+
+ query_filters.append(
+ {
+ "col": temporal_col,
+ "op": "TEMPORAL_RANGE",
+ "val": request.time_range,
+ }
+ )
+ granularity = temporal_col
+ await ctx.debug(
+ "Time filter: column=%s, range=%s" % (temporal_col, request.time_range)
+ )
+
+ # ------------------------------------------------------------------
+ # Step 4: Build query dict
+ # ------------------------------------------------------------------
+ await ctx.report_progress(3, 5, "Building query")
+ query_dict: dict[str, Any] = {
+ "filters": query_filters,
+ "columns": request.columns,
+ "metrics": request.metrics,
+ "row_limit": request.row_limit,
+ "order_desc": request.order_desc,
+ }
+ if granularity:
+ query_dict["granularity"] = granularity
+ if request.order_by:
+ # OrderBy = tuple[Metric | Column, bool] where bool is ascending
+ query_dict["orderby"] = [
+ (col, not request.order_desc) for col in request.order_by
+ ]
+
+ await ctx.debug("Query dict keys: %s" % (sorted(query_dict.keys()),))
+
+ # ------------------------------------------------------------------
+ # Step 5: Create QueryContext and execute
+ # ------------------------------------------------------------------
+ await ctx.report_progress(4, 5, "Executing query")
+ start_time = time.time()
+
+ with event_logger.log_context(action="mcp.query_dataset.execute"):
+ factory = QueryContextFactory()
+ # datasource_type is "table" because this tool queries SqlaTable
+ # datasets (Superset's built-in semantic layer). External semantic
+ # layers (dbt, Snowflake Cortex, etc.) use "semantic_view" and have
+ # a different query path — see SemanticView + mapper.py.
+ query_context = factory.create(
+ datasource={"id": dataset.id, "type": "table"},
+ queries=[query_dict],
+ form_data={},
+ force=request.force_refresh,
+ custom_cache_timeout=request.cache_timeout,
+ )
Review Comment:
**Suggestion:** The request schema includes `use_cache`, but this value is ignored during query-context creation, so setting `use_cache=false` will not actually bypass cache unless `force_refresh` is also set. Apply the same cache-control logic used elsewhere (`force` should be true when either `force_refresh` is true or `use_cache` is false). [logic error]
<details>
<summary><b>Severity Level:</b> Major ⚠️</summary>
```mdx
- ⚠️ query_dataset ignores use_cache flag for cache bypass.
- ⚠️ Callers may see stale data despite disabling cache.
- ⚠️ LLM workflows cannot reliably control query cache usage.
```
</details>
<details>
<summary><b>Steps of Reproduction ✅ </b></summary>
```mdx
1. Note that `QueryDatasetRequest` in `superset/mcp_service/dataset/schemas.py:40-99`
inherits from `QueryCacheControl` (defined in
`superset/mcp_service/common/cache_schemas.py:49-59`), which provides `use_cache: bool`,
`force_refresh: bool`, and `cache_timeout: int | None` for data queries.
2. In the `query_dataset` tool implementation
(`superset/mcp_service/dataset/tool/query_dataset.py:292-304`),
`QueryContextFactory.create` is invoked with `force=request.force_refresh` and
`custom_cache_timeout=request.cache_timeout`, but `request.use_cache` is never read or
passed through, and no call is made to `apply_cache_control_to_query_context()` in
`superset/mcp_service/utils/cache_utils.py:11-37`, which is the helper that sets
`query_context["force"] = True` when `not use_cache or force_refresh`.
3. From an MCP client, issue a `query_dataset` request with a valid `dataset_id`, some
metrics/columns, and explicitly set `use_cache=false` and `force_refresh=false` in the
JSON body; Pydantic accepts this because `QueryDatasetRequest` exposes both fields via
`QueryCacheControl`.
4. Because the tool only uses `request.force_refresh` (False) when constructing the query
context, the underlying `QueryContext` and `ChartDataCommand` will execute with
`force=False`, meaning Superset's query-result cache remains enabled and any cached result
will still be returned even though the caller requested `use_cache=false`. This can be
observed by making a first call with defaults (populating cache), mutating the underlying
table, then calling again with `use_cache=false` and `force_refresh=false` and seeing
stale data and a cache hit reported via `get_cache_status_from_result(query_result,
force_refresh=request.force_refresh)` at `query_dataset.py:387-388.
```
</details>
[Fix in Cursor](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt=This%20is%20a%20comment%20left%20during%20a%20code%20review.%0A%0A%2A%2APath%3A%2A%2A%20superset%2Fmcp_service%2Fdataset%2Ftool%2Fquery_dataset.py%0A%2A%2ALine%3A%2A%2A%20298%3A304%0A%2A%2AComment%3A%2A%2A%0A%09%2ALogic%20Error%3A%20The%20request%20schema%20includes%20%60use_cache%60%2C%20but%20this%20value%20is%20ignored%20during%20query-context%20creation%2C%20so%20setting%20%60use_cache%3Dfalse%60%20will%20not%20actually%20bypass%20cache%20unless%20%60force_refresh%60%20is%20also%20set.%20Apply%20the%20same%20cache-control%20logic%20used%20elsewhere%20%28%60force%60%20should%20be%20true%20when%20either%20%60force_refresh%60%20is%20true%20or%20%60use_cache%60%20is%20false%29.%0A%0AValidate%20the%20correctness%20of%20the%20flagged%20issue.%20If%20correct%2C%20How%20can%20I%20resolve%20this%3F%20If%20you%20propose%20a%20fix%2C%20implement%20it%20and%20please%20make%20it%20concise.%0AOnce%20fix%20is%20implemented%
2C%20also%20check%20other%20comments%20on%20the%20same%20PR%2C%20and%20ask%20user%20if%20the%20user%20wants%20to%20fix%20the%20rest%20of%20the%20comments%20as%20well.%20if%20said%20yes%2C%20then%20fetch%20all%20the%20comments%20validate%20the%20correctness%20and%20implement%20a%20minimal%20fix%0A) | [Fix in VSCode Claude](https://app.codeant.ai/fix-in-ide?tool=vscode-claude&prompt=This%20is%20a%20comment%20left%20during%20a%20code%20review.%0A%0A%2A%2APath%3A%2A%2A%20superset%2Fmcp_service%2Fdataset%2Ftool%2Fquery_dataset.py%0A%2A%2ALine%3A%2A%2A%20298%3A304%0A%2A%2AComment%3A%2A%2A%0A%09%2ALogic%20Error%3A%20The%20request%20schema%20includes%20%60use_cache%60%2C%20but%20this%20value%20is%20ignored%20during%20query-context%20creation%2C%20so%20setting%20%60use_cache%3Dfalse%60%20will%20not%20actually%20bypass%20cache%20unless%20%60force_refresh%60%20is%20also%20set.%20Apply%20the%20same%20cache-control%20logic%20used%20elsewhere%20%28%60force%60%20should%20be%20true%20when%20either%
20%60force_refresh%60%20is%20true%20or%20%60use_cache%60%20is%20false%29.%0A%0AValidate%20the%20correctness%20of%20the%20flagged%20issue.%20If%20correct%2C%20How%20can%20I%20resolve%20this%3F%20If%20you%20propose%20a%20fix%2C%20implement%20it%20and%20please%20make%20it%20concise.%0AOnce%20fix%20is%20implemented%2C%20also%20check%20other%20comments%20on%20the%20same%20PR%2C%20and%20ask%20user%20if%20the%20user%20wants%20to%20fix%20the%20rest%20of%20the%20comments%20as%20well.%20if%20said%20yes%2C%20then%20fetch%20all%20the%20comments%20validate%20the%20correctness%20and%20implement%20a%20minimal%20fix%0A)
*(Use Cmd/Ctrl + Click for best experience)*
<details>
<summary><b>Prompt for AI Agent 🤖 </b></summary>
```mdx
This is a comment left during a code review.
**Path:** superset/mcp_service/dataset/tool/query_dataset.py
**Line:** 298:304
**Comment:**
*Logic Error: The request schema includes `use_cache`, but this value is ignored during query-context creation, so setting `use_cache=false` will not actually bypass cache unless `force_refresh` is also set. Apply the same cache-control logic used elsewhere (`force` should be true when either `force_refresh` is true or `use_cache` is false).
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
```
</details>
<a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F39727&comment_hash=a9a47b3607920416bfbb4526982bf741ab6dce9631ade28645af9727f50e3a16&reaction=like'>👍</a> | <a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F39727&comment_hash=a9a47b3607920416bfbb4526982bf741ab6dce9631ade28645af9727f50e3a16&reaction=dislike'>👎</a>
--
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: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] feat(mcp): add query_dataset tool to query datasets using semantic layer [superset]
Posted by "codecov[bot] (via GitHub)" <gi...@apache.org>.
codecov[bot] commented on PR #39727:
URL: https://github.com/apache/superset/pull/39727#issuecomment-4336943870
## [Codecov](https://app.codecov.io/gh/apache/superset/pull/39727?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) Report
:x: Patch coverage is `30.50847%` with `123 lines` in your changes missing coverage. Please review.
:white_check_mark: Project coverage is 63.62%. Comparing base ([`3aa99c5`](https://app.codecov.io/gh/apache/superset/commit/3aa99c577ed3026ad4ee3148e0a11823ac638fdf?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)) to head ([`17d18fa`](https://app.codecov.io/gh/apache/superset/commit/17d18fa5acc2f119156f83459fbfb3373e850d23?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)).
:warning: Report is 1 commits behind head on master.
| [Files with missing lines](https://app.codecov.io/gh/apache/superset/pull/39727?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Patch % | Lines |
|---|---|---|
| [superset/mcp\_service/dataset/tool/query\_dataset.py](https://app.codecov.io/gh/apache/superset/pull/39727?src=pr&el=tree&filepath=superset%2Fmcp_service%2Fdataset%2Ftool%2Fquery_dataset.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-c3VwZXJzZXQvbWNwX3NlcnZpY2UvZGF0YXNldC90b29sL3F1ZXJ5X2RhdGFzZXQucHk=) | 14.89% | [120 Missing :warning: ](https://app.codecov.io/gh/apache/superset/pull/39727?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [superset/mcp\_service/dataset/schemas.py](https://app.codecov.io/gh/apache/superset/pull/39727?src=pr&el=tree&filepath=superset%2Fmcp_service%2Fdataset%2Fschemas.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-c3VwZXJzZXQvbWNwX3NlcnZpY2UvZGF0YXNldC9zY2hlbWFzLnB5) | 91.42% | [3 Missing :warning: ](https://app.codecov.io/gh/apache/superset/pull/39727?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
<details><summary>Additional details and impacted files</summary>
```diff
@@ Coverage Diff @@
## master #39727 +/- ##
==========================================
- Coverage 64.47% 63.62% -0.85%
==========================================
Files 2566 2567 +1
Lines 134027 134203 +176
Branches 31121 31141 +20
==========================================
- Hits 86410 85384 -1026
- Misses 46122 47306 +1184
- Partials 1495 1513 +18
```
| [Flag](https://app.codecov.io/gh/apache/superset/pull/39727/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Coverage Δ | |
|---|---|---|
| [hive](https://app.codecov.io/gh/apache/superset/pull/39727/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `39.70% <30.50%> (-0.03%)` | :arrow_down: |
| [mysql](https://app.codecov.io/gh/apache/superset/pull/39727/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
| [postgres](https://app.codecov.io/gh/apache/superset/pull/39727/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
| [presto](https://app.codecov.io/gh/apache/superset/pull/39727/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `41.46% <30.50%> (-0.04%)` | :arrow_down: |
| [python](https://app.codecov.io/gh/apache/superset/pull/39727/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `59.91% <30.50%> (-1.89%)` | :arrow_down: |
| [sqlite](https://app.codecov.io/gh/apache/superset/pull/39727/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `59.77% <30.50%> (-0.09%)` | :arrow_down: |
| [unit](https://app.codecov.io/gh/apache/superset/pull/39727/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#carryforward-flags-in-the-pull-request-comment) to find out more.
</details>
[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/apache/superset/pull/39727?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
<details><summary> :rocket: New features to boost your workflow: </summary>
- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
- :package: [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis): Save yourself from yourself by tracking and limiting bundle sizes in JS merges.
</details>
--
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: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] feat(mcp): add query_dataset tool to query datasets using semantic layer [superset]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #39727:
URL: https://github.com/apache/superset/pull/39727#discussion_r3169411043
##########
superset/mcp_service/dataset/tool/query_dataset.py:
##########
@@ -0,0 +1,487 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+MCP tool: query_dataset
+
+Query a dataset using its semantic layer (saved metrics, calculated columns,
+dimensions) without requiring a saved chart.
+"""
+
+import difflib
+import logging
+import time
+from typing import Any
+
+from fastmcp import Context
+from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.orm import joinedload, subqueryload
+from superset_core.mcp.decorators import tool, ToolAnnotations
+
+from superset.commands.exceptions import CommandException
+from superset.exceptions import OAuth2Error, OAuth2RedirectError, SupersetException
+from superset.extensions import event_logger
+from superset.mcp_service.chart.schemas import DataColumn, PerformanceMetadata
+from superset.mcp_service.dataset.schemas import (
+ DatasetError,
+ QueryDatasetFilter,
+ QueryDatasetRequest,
+ QueryDatasetResponse,
+)
+from superset.mcp_service.privacy import (
+ DATA_MODEL_METADATA_ERROR_TYPE,
+ requires_data_model_metadata_access,
+ user_can_view_data_model_metadata,
+)
+from superset.mcp_service.utils import _is_uuid
+from superset.mcp_service.utils.cache_utils import get_cache_status_from_result
+from superset.mcp_service.utils.oauth2_utils import build_oauth2_redirect_message
+
+logger = logging.getLogger(__name__)
+
+
+def _resolve_dataset(identifier: int | str, eager_options: list[Any]) -> Any | None:
+ """Resolve a dataset by int ID or UUID string.
+
+ Replicates the identifier resolution logic from ModelGetInfoCore._find_object().
+ """
+ from superset.daos.dataset import DatasetDAO
+
+ opts = eager_options or None
+
+ if isinstance(identifier, int):
+ return DatasetDAO.find_by_id(identifier, query_options=opts)
+
+ # Try parsing as int
+ try:
+ id_val = int(identifier)
+ return DatasetDAO.find_by_id(id_val, query_options=opts)
+ except (ValueError, TypeError):
+ pass
+
+ # Try UUID
+ if _is_uuid(str(identifier)):
+ return DatasetDAO.find_by_id(identifier, id_column="uuid", query_options=opts)
+
+ return None
+
+
+def _validate_names(
+ requested: list[str],
+ valid: set[str],
+ kind: str,
+) -> list[str]:
+ """Return list of error messages for names not found in *valid*.
+
+ Includes close-match suggestions when available.
+ """
+ errors: list[str] = []
+ for name in requested:
+ if name not in valid:
+ suggestions = difflib.get_close_matches(name, valid, n=3, cutoff=0.6)
+ msg = f"Unknown {kind}: '{name}'"
+ if suggestions:
+ msg += f". Did you mean: {', '.join(suggestions)}?"
+ errors.append(msg)
+ return errors
+
+
+@requires_data_model_metadata_access
+@tool(
+ tags=["data"],
+ class_permission_name="Dataset",
+ annotations=ToolAnnotations(
+ title="Query dataset",
+ readOnlyHint=True,
+ destructiveHint=False,
+ ),
+)
+async def query_dataset( # noqa: C901
+ request: QueryDatasetRequest, ctx: Context
+) -> QueryDatasetResponse | DatasetError:
+ """Query a dataset using its semantic layer (saved metrics, dimensions, filters).
+
+ Returns tabular data without requiring a saved chart. Use this when you want
+ to compute saved metrics, group by dimensions, or apply filters directly
+ against a dataset's curated semantic layer.
+
+ Workflow:
+ 1. list_datasets -> find a dataset
+ 2. get_dataset_info -> discover available columns and metrics
+ 3. query_dataset -> query using metric names and column names
+
+ Example:
+ ```json
+ {
+ "dataset_id": 123,
+ "metrics": ["count", "avg_revenue"],
+ "columns": ["product_category"],
+ "time_range": "Last 7 days",
+ "row_limit": 100
+ }
+ ```
+ """
+ await ctx.info(
+ "Starting dataset query: dataset_id=%s, metrics=%s, columns=%s, "
+ "row_limit=%s"
+ % (
+ request.dataset_id,
+ request.metrics,
+ request.columns,
+ request.row_limit,
+ )
+ )
+
+ try:
+ from superset.commands.chart.data.get_data_command import ChartDataCommand
+ from superset.common.query_context_factory import QueryContextFactory
+ from superset.connectors.sqla.models import SqlaTable
+
+ # ------------------------------------------------------------------
+ # Step 1: Resolve dataset
+ # ------------------------------------------------------------------
+ await ctx.report_progress(1, 5, "Looking up dataset")
+ eager_options = [
+ subqueryload(SqlaTable.columns),
+ subqueryload(SqlaTable.metrics),
+ joinedload(SqlaTable.database),
+ ]
+
+ with event_logger.log_context(action="mcp.query_dataset.lookup"):
+ dataset = _resolve_dataset(request.dataset_id, eager_options)
+
+ if dataset is None:
+ await ctx.error("Dataset not found: identifier=%s" % (request.dataset_id,))
Review Comment:
The data-model metadata permission check happens after resolving the dataset. That means metadata-restricted users can distinguish between “NotFound” vs “DataModelMetadataRestricted”, leaking dataset existence via direct calls. To align with list_datasets/get_dataset_info privacy behavior, check user_can_view_data_model_metadata() (and return the restricted error) before doing the DAO lookup / returning NotFound.
##########
tests/unit_tests/mcp_service/dataset/tool/test_query_dataset.py:
##########
@@ -0,0 +1,798 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Tests for the query_dataset MCP tool."""
+
+from __future__ import annotations
+
+import importlib
+from collections.abc import Generator
+from typing import Any
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+from fastmcp import Client, FastMCP
+
+from superset.mcp_service.app import mcp
+from superset.utils import json
+
+query_dataset_module = importlib.import_module(
+ "superset.mcp_service.dataset.tool.query_dataset"
+)
+
+
+@pytest.fixture
+def mcp_server() -> FastMCP:
+ return mcp
+
+
+@pytest.fixture(autouse=True)
+def mock_auth() -> Generator[MagicMock, None, None]:
+ """Mock authentication and metadata access for all tests."""
+ with (
+ patch("superset.mcp_service.auth.get_user_from_request") as mock_get_user,
+ patch.object(
+ query_dataset_module,
+ "user_can_view_data_model_metadata",
+ return_value=True,
+ ),
+ ):
+ mock_user = Mock()
+ mock_user.id = 1
+ mock_user.username = "admin"
+ mock_get_user.return_value = mock_user
+ yield mock_get_user
+
+
+def _make_column(name: str, is_dttm: bool = False) -> MagicMock:
+ """Build a mock SqlaTable column with the given name and datetime flag."""
+ col = MagicMock()
+ col.column_name = name
+ col.is_dttm = is_dttm
+ col.verbose_name = None
+ col.type = "VARCHAR"
+ col.groupby = True
+ col.filterable = True
+ col.description = None
+ return col
+
+
+def _make_metric(name: str, expression: str = "COUNT(*)") -> MagicMock:
+ """Build a mock SqlMetric with the given name and SQL expression."""
+ metric = MagicMock()
+ metric.metric_name = name
+ metric.verbose_name = None
+ metric.expression = expression
+ metric.description = None
+ metric.d3format = None
+ return metric
+
+
+def _make_dataset(
+ dataset_id: int = 1,
+ table_name: str = "orders",
+ columns: list[Any] | None = None,
+ metrics: list[Any] | None = None,
+ main_dttm_col: str | None = None,
+) -> MagicMock:
+ """Build a mock SqlaTable dataset with default columns and metrics."""
+ ds = MagicMock()
+ ds.id = dataset_id
+ ds.table_name = table_name
+ ds.uuid = f"test-uuid-{dataset_id}"
+ ds.main_dttm_col = main_dttm_col
+ ds.database = MagicMock()
+ ds.database.database_name = "examples"
+ ds.columns = columns or [
+ _make_column("category"),
+ _make_column("region"),
+ _make_column("order_date", is_dttm=True),
+ ]
+ ds.metrics = metrics or [
+ _make_metric("count", "COUNT(*)"),
+ _make_metric("total_revenue", "SUM(revenue)"),
+ ]
+ return ds
+
+
+def _mock_command_result(
+ data: list[dict[str, Any]] | None = None,
+ colnames: list[str] | None = None,
+) -> dict[str, Any]:
+ """Build the result dict that ChartDataCommand.run() returns."""
+ data = data or [
+ {"category": "Electronics", "count": 42},
+ {"category": "Clothing", "count": 17},
+ ]
+ colnames = colnames or ["category", "count"]
+ return {
+ "queries": [
+ {
+ "data": data,
+ "colnames": colnames,
+ "rowcount": len(data),
+ "cache_key": "abc123",
+ "is_cached": False,
+ "cached_dttm": None,
+ "cache_timeout": 300,
+ }
+ ]
+ }
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_success(mcp_server: FastMCP) -> None:
+ """Happy path: metrics + columns returns data."""
+ dataset = _make_dataset()
+ result_data = _mock_command_result()
+
+ with (
+ patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.validate",
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.run",
+ return_value=result_data,
+ ),
+ patch(
+ "superset.common.query_context_factory.QueryContextFactory.create",
+ return_value=MagicMock(),
+ ),
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metrics": ["count"],
+ "columns": ["category"],
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert data["dataset_id"] == 1
+ assert data["dataset_name"] == "orders"
+ assert data["row_count"] == 2
+ assert len(data["data"]) == 2
+ assert data["data"][0]["category"] == "Electronics"
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_not_found(mcp_server: FastMCP) -> None:
+ """Dataset ID that doesn't exist returns error."""
+ with patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=None,
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 999,
+ "metrics": ["count"],
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert data["error_type"] == "NotFound"
+ assert "999" in data["error"]
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_invalid_metric(mcp_server: FastMCP) -> None:
+ """Unknown metric name returns validation error with suggestions."""
+ dataset = _make_dataset()
+
+ with patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metrics": ["countt"], # typo
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert data["error_type"] == "ValidationError"
+ assert "countt" in data["error"]
+ # Should suggest "count" as a close match
+ assert "count" in data["error"]
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_invalid_column(mcp_server: FastMCP) -> None:
+ """Unknown column name returns validation error."""
+ dataset = _make_dataset()
+
+ with patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "columns": ["nonexistent_col"],
+ "metrics": ["count"],
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert data["error_type"] == "ValidationError"
+ assert "nonexistent_col" in data["error"]
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_no_metrics_no_columns(mcp_server: FastMCP) -> None:
+ """Providing neither metrics nor columns raises validation error."""
+ from fastmcp.exceptions import ToolError
+
+ async with Client(mcp_server) as client:
+ with pytest.raises(ToolError, match="metrics.*columns"):
+ await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metrics": [],
+ "columns": [],
+ }
+ },
+ )
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_with_time_range(mcp_server: FastMCP) -> None:
+ """time_range is converted to TEMPORAL_RANGE filter + granularity."""
+ dataset = _make_dataset(main_dttm_col="order_date")
+ result_data = _mock_command_result()
+ captured_queries: list[dict[str, Any]] = []
+
+ def capture_create(**kwargs):
+ captured_queries.extend(kwargs.get("queries", []))
+ return MagicMock()
+
+ with (
+ patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.validate",
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.run",
+ return_value=result_data,
+ ),
+ patch(
+ "superset.common.query_context_factory.QueryContextFactory.create",
+ side_effect=capture_create,
+ ),
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metrics": ["count"],
+ "time_range": "Last 7 days",
+ }
+ },
+ )
+
+ assert len(captured_queries) == 1
+ query_dict = captured_queries[0]
+ # Should have TEMPORAL_RANGE filter
+ temporal_filters = [f for f in query_dict["filters"] if f["op"] == "TEMPORAL_RANGE"]
+ assert len(temporal_filters) == 1
+ assert temporal_filters[0]["col"] == "order_date"
+ assert temporal_filters[0]["val"] == "Last 7 days"
+ # Should set granularity
+ assert query_dict["granularity"] == "order_date"
+ # applied_filters in response must include the synthesized TEMPORAL_RANGE filter
+ data = json.loads(result.content[0].text)
+ resp_filters = data["applied_filters"]
+ temporal_resp = [f for f in resp_filters if f["op"] == "TEMPORAL_RANGE"]
+ assert len(temporal_resp) == 1
+ assert temporal_resp[0]["col"] == "order_date"
+ assert temporal_resp[0]["val"] == "Last 7 days"
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_time_range_no_temporal_column(mcp_server: FastMCP) -> None:
+ """time_range without a temporal column returns error."""
+ dataset = _make_dataset(main_dttm_col=None)
+
+ with patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metrics": ["count"],
+ "time_range": "Last 7 days",
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert data["error_type"] == "ValidationError"
+ assert "temporal column" in data["error"].lower()
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_with_filters(mcp_server: FastMCP) -> None:
+ """User-provided filters are passed through to the query."""
+ dataset = _make_dataset()
+ result_data = _mock_command_result()
+ captured_queries: list[dict[str, Any]] = []
+
+ def capture_create(**kwargs):
+ captured_queries.extend(kwargs.get("queries", []))
+ return MagicMock()
+
+ with (
+ patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.validate",
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.run",
+ return_value=result_data,
+ ),
+ patch(
+ "superset.common.query_context_factory.QueryContextFactory.create",
+ side_effect=capture_create,
+ ),
+ ):
+ async with Client(mcp_server) as client:
+ await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metrics": ["count"],
+ "filters": [
+ {"col": "category", "op": "==", "val": "Electronics"}
+ ],
+ }
+ },
+ )
+
+ assert len(captured_queries) == 1
+ filters = captured_queries[0]["filters"]
+ assert len(filters) == 1
+ assert filters[0]["col"] == "category"
+ assert filters[0]["op"] == "=="
+ assert filters[0]["val"] == "Electronics"
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_empty_results(mcp_server: FastMCP) -> None:
+ """Query that returns no data gives a response with row_count=0."""
+ dataset = _make_dataset()
+ empty_result = {
+ "queries": [
+ {
+ "data": [],
+ "colnames": [],
+ "rowcount": 0,
+ "is_cached": False,
+ "cached_dttm": None,
+ "cache_timeout": 300,
+ }
+ ]
+ }
+
+ with (
+ patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.validate",
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.run",
+ return_value=empty_result,
+ ),
+ patch(
+ "superset.common.query_context_factory.QueryContextFactory.create",
+ return_value=MagicMock(),
+ ),
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metrics": ["count"],
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert data["row_count"] == 0
+ assert data["data"] == []
+ assert "no data" in data["summary"].lower()
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_by_uuid(mcp_server: FastMCP) -> None:
+ """UUID-based lookup works."""
+ dataset = _make_dataset()
+ result_data = _mock_command_result()
+
+ with (
+ patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ) as mock_resolve,
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.validate",
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.run",
+ return_value=result_data,
+ ),
+ patch(
+ "superset.common.query_context_factory.QueryContextFactory.create",
+ return_value=MagicMock(),
+ ),
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": "a1b2c3d4-5678-90ab-cdef-1234567890ab",
+ "metrics": ["count"],
+ }
+ },
+ )
+
+ # Verify the resolve function was called with the UUID
+ mock_resolve.assert_called_once()
+ call_args = mock_resolve.call_args
+ assert call_args[0][0] == "a1b2c3d4-5678-90ab-cdef-1234567890ab"
+
+ data = json.loads(result.content[0].text)
+ assert data["dataset_id"] == 1
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_permission_denied(mcp_server: FastMCP) -> None:
+ """Permission denied from ChartDataCommand.validate() returns error."""
+ from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
+ from superset.exceptions import SupersetSecurityException
+
+ dataset = _make_dataset()
+
+ with (
+ patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ),
+ patch(
+ "superset.common.query_context_factory.QueryContextFactory.create",
+ return_value=MagicMock(),
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.validate",
+ side_effect=SupersetSecurityException(
+ SupersetError(
+ message="Access denied",
+ error_type=SupersetErrorType.DATASOURCE_SECURITY_ACCESS_ERROR,
+ level=ErrorLevel.WARNING,
+ )
+ ),
+ ),
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metrics": ["count"],
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert data["error_type"] == "QueryError"
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_order_by_valid(mcp_server: FastMCP) -> None:
+ """order_by with valid column/metric names passes through."""
+ dataset = _make_dataset()
+ result_data = _mock_command_result()
+ captured_queries: list[dict[str, Any]] = []
+
+ def capture_create(**kwargs):
+ captured_queries.extend(kwargs.get("queries", []))
+ return MagicMock()
+
+ with (
+ patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.validate",
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.run",
+ return_value=result_data,
+ ),
+ patch(
+ "superset.common.query_context_factory.QueryContextFactory.create",
+ side_effect=capture_create,
+ ),
+ ):
+ async with Client(mcp_server) as client:
+ await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metrics": ["count"],
+ "columns": ["category"],
+ "order_by": ["count"],
+ "order_desc": True,
+ }
+ },
+ )
+
+ assert len(captured_queries) == 1
+ orderby = captured_queries[0].get("orderby", [])
+ assert len(orderby) == 1
+ assert orderby[0][0] == "count"
+ # order_desc=True -> ascending=False
+ assert orderby[0][1] is False
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_order_by_invalid(mcp_server: FastMCP) -> None:
+ """order_by with an unknown name returns validation error."""
+ dataset = _make_dataset()
+
+ with patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metrics": ["count"],
+ "order_by": ["nonexistent"],
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert data["error_type"] == "ValidationError"
+ assert "nonexistent" in data["error"]
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_time_column_override(mcp_server: FastMCP) -> None:
+ """Explicit time_column overrides dataset main_dttm_col."""
+ dataset = _make_dataset(main_dttm_col="order_date")
+ result_data = _mock_command_result()
+ captured_queries: list[dict[str, Any]] = []
+
+ def capture_create(**kwargs):
+ captured_queries.extend(kwargs.get("queries", []))
+ return MagicMock()
+
+ with (
+ patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.validate",
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.run",
+ return_value=result_data,
+ ),
+ patch(
+ "superset.common.query_context_factory.QueryContextFactory.create",
+ side_effect=capture_create,
+ ),
+ ):
+ async with Client(mcp_server) as client:
+ await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metrics": ["count"],
+ "time_range": "Last 30 days",
+ "time_column": "order_date",
+ }
+ },
+ )
+
+ assert len(captured_queries) == 1
+ query_dict = captured_queries[0]
+ assert query_dict["granularity"] == "order_date"
+ temporal_filters = [f for f in query_dict["filters"] if f["op"] == "TEMPORAL_RANGE"]
+ assert temporal_filters[0]["col"] == "order_date"
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_non_dttm_time_column_warns(mcp_server: FastMCP) -> None:
+ """Using a non-datetime column for time_range produces a warning."""
+ dataset = _make_dataset(main_dttm_col=None)
+ result_data = _mock_command_result()
+
+ with (
+ patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.validate",
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.run",
+ return_value=result_data,
+ ),
+ patch(
+ "superset.common.query_context_factory.QueryContextFactory.create",
+ return_value=MagicMock(),
+ ),
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metrics": ["count"],
+ "time_range": "Last 7 days",
+ "time_column": "category",
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert len(data["warnings"]) > 0
+ assert "not marked as a datetime" in data["warnings"][0]
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_invalid_filter_column(mcp_server: FastMCP) -> None:
+ """Filter on a column that doesn't exist returns validation error."""
+ dataset = _make_dataset()
+
+ with patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metrics": ["count"],
+ "filters": [
+ {
+ "col": "nonexistent",
+ "op": "==",
+ "val": "test",
+ }
+ ],
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert data["error_type"] == "ValidationError"
+ assert "nonexistent" in data["error"]
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_metadata_access_denied_no_suggestions(
+ mcp_server: FastMCP,
+) -> None:
+ """Users without data-model metadata access cannot probe column/metric names.
+
Review Comment:
Consider adding a regression test to ensure metadata-restricted users don’t get a different error for nonexistent vs existing dataset IDs (i.e., the tool should not return NotFound when metadata access is denied). This will prevent reintroducing the dataset-existence side channel when adjusting the permission check order.
--
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: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] feat(mcp): add query_dataset tool to query datasets using semantic layer [superset]
Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on PR #39727:
URL: https://github.com/apache/superset/pull/39727#issuecomment-4357051963
**Bito Automatic Review Skipped – PR Already Merged**<br/><br/>Bito scheduled an automatic review for this pull request, but the review was skipped because this PR was merged before the review could be run.<br/>No action is needed if you didn't intend to review it. To get a review, you can type `/review` in a comment and save it
--
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: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] feat(mcp): add query_dataset tool to query datasets using semantic layer [superset]
Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on code in PR #39727:
URL: https://github.com/apache/superset/pull/39727#discussion_r3164805063
##########
superset/mcp_service/dataset/tool/query_dataset.py:
##########
@@ -0,0 +1,457 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+MCP tool: query_dataset
+
+Query a dataset using its semantic layer (saved metrics, calculated columns,
+dimensions) without requiring a saved chart.
+"""
+
+import difflib
+import logging
+import time
+from typing import Any
+
+from fastmcp import Context
+from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.orm import joinedload, subqueryload
+from superset_core.mcp.decorators import tool, ToolAnnotations
+
+from superset.commands.exceptions import CommandException
+from superset.exceptions import OAuth2Error, OAuth2RedirectError, SupersetException
+from superset.extensions import event_logger
+from superset.mcp_service.chart.schemas import DataColumn, PerformanceMetadata
+from superset.mcp_service.dataset.schemas import (
+ DatasetError,
+ QueryDatasetRequest,
+ QueryDatasetResponse,
+)
+from superset.mcp_service.utils import _is_uuid
+from superset.mcp_service.utils.cache_utils import get_cache_status_from_result
+from superset.mcp_service.utils.oauth2_utils import build_oauth2_redirect_message
+
+logger = logging.getLogger(__name__)
+
+
+def _resolve_dataset(identifier: int | str, eager_options: list[Any]) -> Any | None:
+ """Resolve a dataset by int ID or UUID string.
+
+ Replicates the identifier resolution logic from ModelGetInfoCore._find_object().
+ """
+ from superset.daos.dataset import DatasetDAO
+
+ opts = eager_options or None
+
+ if isinstance(identifier, int):
+ return DatasetDAO.find_by_id(identifier, query_options=opts)
+
+ # Try parsing as int
+ try:
+ id_val = int(identifier)
+ return DatasetDAO.find_by_id(id_val, query_options=opts)
+ except (ValueError, TypeError):
+ pass
+
+ # Try UUID
+ if _is_uuid(str(identifier)):
+ return DatasetDAO.find_by_id(identifier, id_column="uuid", query_options=opts)
+
+ return None
+
+
+def _validate_names(
+ requested: list[str],
+ valid: set[str],
+ kind: str,
+) -> list[str]:
+ """Return list of error messages for names not found in *valid*.
+
+ Includes close-match suggestions when available.
+ """
+ errors: list[str] = []
+ for name in requested:
+ if name not in valid:
+ suggestions = difflib.get_close_matches(name, valid, n=3, cutoff=0.6)
+ msg = f"Unknown {kind}: '{name}'"
+ if suggestions:
+ msg += f". Did you mean: {', '.join(suggestions)}?"
+ errors.append(msg)
+ return errors
+
+
+@tool(
+ tags=["data"],
+ class_permission_name="Dataset",
+ annotations=ToolAnnotations(
+ title="Query dataset",
+ readOnlyHint=True,
+ destructiveHint=False,
+ ),
+)
+async def query_dataset( # noqa: C901
+ request: QueryDatasetRequest, ctx: Context
+) -> QueryDatasetResponse | DatasetError:
+ """Query a dataset using its semantic layer (saved metrics, dimensions, filters).
+
+ Returns tabular data without requiring a saved chart. Use this when you want
+ to compute saved metrics, group by dimensions, or apply filters directly
+ against a dataset's curated semantic layer.
+
+ Workflow:
+ 1. list_datasets -> find a dataset
+ 2. get_dataset_info -> discover available columns and metrics
+ 3. query_dataset -> query using metric names and column names
+
+ Example:
+ ```json
+ {
+ "dataset_id": 123,
+ "metrics": ["count", "avg_revenue"],
+ "columns": ["product_category"],
+ "time_range": "Last 7 days",
+ "row_limit": 100
+ }
+ ```
+ """
+ await ctx.info(
+ "Starting dataset query: dataset_id=%s, metrics=%s, columns=%s, "
+ "row_limit=%s"
+ % (
+ request.dataset_id,
+ request.metrics,
+ request.columns,
+ request.row_limit,
+ )
+ )
+
+ try:
+ from superset.commands.chart.data.get_data_command import ChartDataCommand
+ from superset.common.query_context_factory import QueryContextFactory
+ from superset.connectors.sqla.models import SqlaTable
+
+ # ------------------------------------------------------------------
+ # Step 1: Resolve dataset
+ # ------------------------------------------------------------------
+ await ctx.report_progress(1, 5, "Looking up dataset")
+ eager_options = [
+ subqueryload(SqlaTable.columns),
+ subqueryload(SqlaTable.metrics),
+ joinedload(SqlaTable.database),
+ ]
+
+ with event_logger.log_context(action="mcp.query_dataset.lookup"):
+ dataset = _resolve_dataset(request.dataset_id, eager_options)
+
+ if dataset is None:
+ await ctx.error("Dataset not found: identifier=%s" % (request.dataset_id,))
+ return DatasetError.create(
+ error=f"No dataset found with identifier: {request.dataset_id}",
+ error_type="NotFound",
+ )
+
+ dataset_name = getattr(dataset, "table_name", None) or f"Dataset {dataset.id}"
+ await ctx.info(
+ "Dataset found: id=%s, name=%s, columns=%s, metrics=%s"
+ % (
+ dataset.id,
+ dataset_name,
+ len(dataset.columns),
+ len(dataset.metrics),
+ )
+ )
+
+ # ------------------------------------------------------------------
+ # Step 2: Validate requested columns and metrics
+ # ------------------------------------------------------------------
+ await ctx.report_progress(2, 5, "Validating columns and metrics")
+ valid_columns = {c.column_name for c in dataset.columns}
+ valid_metrics = {m.metric_name for m in dataset.metrics}
+
+ validation_errors: list[str] = []
+ validation_errors.extend(
+ _validate_names(request.columns, valid_columns, "column")
+ )
+ validation_errors.extend(
+ _validate_names(request.metrics, valid_metrics, "metric")
+ )
+ # Validate filter column names against dataset columns
+ filter_cols = [f.col for f in request.filters]
+ validation_errors.extend(
+ _validate_names(filter_cols, valid_columns, "filter column")
+ )
+ # Validate order_by names against columns + metrics
+ if request.order_by:
+ valid_orderby = valid_columns | valid_metrics
+ validation_errors.extend(
+ _validate_names(request.order_by, valid_orderby, "order_by")
+ )
+
+ if validation_errors:
+ error_msg = "; ".join(validation_errors)
+ await ctx.error("Validation failed: %s" % (error_msg,))
+ return DatasetError.create(
+ error=error_msg,
+ error_type="ValidationError",
+ )
+
+ # ------------------------------------------------------------------
+ # Step 3: Build filters and time range
+ # ------------------------------------------------------------------
+ warnings: list[str] = []
+ query_filters: list[dict[str, Any]] = [
+ {"col": f.col, "op": f.op, "val": f.val} for f in request.filters
+ ]
+ granularity: str | None = None
+
+ if request.time_range:
+ temporal_col = request.time_column or getattr(
+ dataset, "main_dttm_col", None
+ )
+ if not temporal_col:
+ await ctx.error("time_range provided but no temporal column available")
+ return DatasetError.create(
+ error=(
+ "time_range was provided but no temporal column is available. "
+ "Either set time_column explicitly or ensure the dataset has "
+ "a main datetime column configured."
+ ),
+ error_type="ValidationError",
+ )
+ # Validate that the temporal column actually exists on the dataset
+ if temporal_col not in valid_columns:
+ await ctx.error("time_column '%s' not found on dataset" % temporal_col)
+ return DatasetError.create(
+ error=(
+ f"time_column '{temporal_col}' does not exist on this dataset."
+ ),
+ error_type="ValidationError",
+ )
+ # Warn if the chosen temporal column isn't marked as datetime
+ dttm_cols = {c.column_name for c in dataset.columns if c.is_dttm}
+ if temporal_col not in dttm_cols:
+ warnings.append(
+ f"Column '{temporal_col}' is not marked as a datetime "
+ f"column on this dataset. Time filtering may not work "
+ f"as expected."
+ )
+
+ query_filters.append(
+ {
+ "col": temporal_col,
+ "op": "TEMPORAL_RANGE",
+ "val": request.time_range,
+ }
+ )
+ granularity = temporal_col
+ await ctx.debug(
+ "Time filter: column=%s, range=%s" % (temporal_col, request.time_range)
+ )
+
+ # ------------------------------------------------------------------
+ # Step 4: Build query dict
+ # ------------------------------------------------------------------
+ await ctx.report_progress(3, 5, "Building query")
+ query_dict: dict[str, Any] = {
+ "filters": query_filters,
+ "columns": request.columns,
+ "metrics": request.metrics,
+ "row_limit": request.row_limit,
+ "order_desc": request.order_desc,
+ }
+ if granularity:
+ query_dict["granularity"] = granularity
+ if request.order_by:
+ # OrderBy = tuple[Metric | Column, bool] where bool is ascending
+ query_dict["orderby"] = [
+ (col, not request.order_desc) for col in request.order_by
+ ]
+
+ await ctx.debug("Query dict keys: %s" % (sorted(query_dict.keys()),))
+
+ # ------------------------------------------------------------------
+ # Step 5: Create QueryContext and execute
+ # ------------------------------------------------------------------
+ await ctx.report_progress(4, 5, "Executing query")
+ start_time = time.time()
+
+ with event_logger.log_context(action="mcp.query_dataset.execute"):
+ factory = QueryContextFactory()
+ # datasource_type is "table" because this tool queries SqlaTable
+ # datasets (Superset's built-in semantic layer). External semantic
+ # layers (dbt, Snowflake Cortex, etc.) use "semantic_view" and have
+ # a different query path — see SemanticView + mapper.py.
+ query_context = factory.create(
+ datasource={"id": dataset.id, "type": "table"},
+ queries=[query_dict],
+ form_data={},
+ force=request.force_refresh,
+ custom_cache_timeout=request.cache_timeout,
+ )
Review Comment:
Fixed in ffcd4b1d25 — changed `force=request.force_refresh` to `force=not request.use_cache or request.force_refresh` so that setting `use_cache=False` correctly bypasses the query cache.
##########
superset/mcp_service/dataset/tool/query_dataset.py:
##########
@@ -0,0 +1,457 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+MCP tool: query_dataset
+
+Query a dataset using its semantic layer (saved metrics, calculated columns,
+dimensions) without requiring a saved chart.
+"""
+
+import difflib
+import logging
+import time
+from typing import Any
+
+from fastmcp import Context
+from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.orm import joinedload, subqueryload
+from superset_core.mcp.decorators import tool, ToolAnnotations
+
+from superset.commands.exceptions import CommandException
+from superset.exceptions import OAuth2Error, OAuth2RedirectError, SupersetException
+from superset.extensions import event_logger
+from superset.mcp_service.chart.schemas import DataColumn, PerformanceMetadata
+from superset.mcp_service.dataset.schemas import (
+ DatasetError,
+ QueryDatasetRequest,
+ QueryDatasetResponse,
+)
+from superset.mcp_service.utils import _is_uuid
+from superset.mcp_service.utils.cache_utils import get_cache_status_from_result
+from superset.mcp_service.utils.oauth2_utils import build_oauth2_redirect_message
+
+logger = logging.getLogger(__name__)
+
+
+def _resolve_dataset(identifier: int | str, eager_options: list[Any]) -> Any | None:
+ """Resolve a dataset by int ID or UUID string.
+
+ Replicates the identifier resolution logic from ModelGetInfoCore._find_object().
+ """
+ from superset.daos.dataset import DatasetDAO
+
+ opts = eager_options or None
+
+ if isinstance(identifier, int):
+ return DatasetDAO.find_by_id(identifier, query_options=opts)
+
+ # Try parsing as int
+ try:
+ id_val = int(identifier)
+ return DatasetDAO.find_by_id(id_val, query_options=opts)
+ except (ValueError, TypeError):
+ pass
+
+ # Try UUID
+ if _is_uuid(str(identifier)):
+ return DatasetDAO.find_by_id(identifier, id_column="uuid", query_options=opts)
+
+ return None
+
+
+def _validate_names(
+ requested: list[str],
+ valid: set[str],
+ kind: str,
+) -> list[str]:
+ """Return list of error messages for names not found in *valid*.
+
+ Includes close-match suggestions when available.
+ """
+ errors: list[str] = []
+ for name in requested:
+ if name not in valid:
+ suggestions = difflib.get_close_matches(name, valid, n=3, cutoff=0.6)
+ msg = f"Unknown {kind}: '{name}'"
+ if suggestions:
+ msg += f". Did you mean: {', '.join(suggestions)}?"
+ errors.append(msg)
+ return errors
+
+
+@tool(
+ tags=["data"],
+ class_permission_name="Dataset",
+ annotations=ToolAnnotations(
+ title="Query dataset",
+ readOnlyHint=True,
+ destructiveHint=False,
+ ),
+)
+async def query_dataset( # noqa: C901
+ request: QueryDatasetRequest, ctx: Context
+) -> QueryDatasetResponse | DatasetError:
+ """Query a dataset using its semantic layer (saved metrics, dimensions, filters).
+
+ Returns tabular data without requiring a saved chart. Use this when you want
+ to compute saved metrics, group by dimensions, or apply filters directly
+ against a dataset's curated semantic layer.
+
+ Workflow:
+ 1. list_datasets -> find a dataset
+ 2. get_dataset_info -> discover available columns and metrics
+ 3. query_dataset -> query using metric names and column names
+
+ Example:
+ ```json
+ {
+ "dataset_id": 123,
+ "metrics": ["count", "avg_revenue"],
+ "columns": ["product_category"],
+ "time_range": "Last 7 days",
+ "row_limit": 100
+ }
+ ```
+ """
+ await ctx.info(
+ "Starting dataset query: dataset_id=%s, metrics=%s, columns=%s, "
+ "row_limit=%s"
+ % (
+ request.dataset_id,
+ request.metrics,
+ request.columns,
+ request.row_limit,
+ )
+ )
+
+ try:
+ from superset.commands.chart.data.get_data_command import ChartDataCommand
+ from superset.common.query_context_factory import QueryContextFactory
+ from superset.connectors.sqla.models import SqlaTable
+
+ # ------------------------------------------------------------------
+ # Step 1: Resolve dataset
+ # ------------------------------------------------------------------
+ await ctx.report_progress(1, 5, "Looking up dataset")
+ eager_options = [
+ subqueryload(SqlaTable.columns),
+ subqueryload(SqlaTable.metrics),
+ joinedload(SqlaTable.database),
+ ]
+
+ with event_logger.log_context(action="mcp.query_dataset.lookup"):
+ dataset = _resolve_dataset(request.dataset_id, eager_options)
+
+ if dataset is None:
+ await ctx.error("Dataset not found: identifier=%s" % (request.dataset_id,))
+ return DatasetError.create(
+ error=f"No dataset found with identifier: {request.dataset_id}",
+ error_type="NotFound",
+ )
+
+ dataset_name = getattr(dataset, "table_name", None) or f"Dataset {dataset.id}"
+ await ctx.info(
+ "Dataset found: id=%s, name=%s, columns=%s, metrics=%s"
+ % (
+ dataset.id,
+ dataset_name,
+ len(dataset.columns),
+ len(dataset.metrics),
+ )
+ )
+
+ # ------------------------------------------------------------------
+ # Step 2: Validate requested columns and metrics
+ # ------------------------------------------------------------------
+ await ctx.report_progress(2, 5, "Validating columns and metrics")
+ valid_columns = {c.column_name for c in dataset.columns}
+ valid_metrics = {m.metric_name for m in dataset.metrics}
+
+ validation_errors: list[str] = []
+ validation_errors.extend(
+ _validate_names(request.columns, valid_columns, "column")
+ )
+ validation_errors.extend(
+ _validate_names(request.metrics, valid_metrics, "metric")
+ )
+ # Validate filter column names against dataset columns
+ filter_cols = [f.col for f in request.filters]
+ validation_errors.extend(
+ _validate_names(filter_cols, valid_columns, "filter column")
+ )
+ # Validate order_by names against columns + metrics
+ if request.order_by:
+ valid_orderby = valid_columns | valid_metrics
+ validation_errors.extend(
+ _validate_names(request.order_by, valid_orderby, "order_by")
+ )
+
+ if validation_errors:
+ error_msg = "; ".join(validation_errors)
+ await ctx.error("Validation failed: %s" % (error_msg,))
+ return DatasetError.create(
+ error=error_msg,
+ error_type="ValidationError",
+ )
+
+ # ------------------------------------------------------------------
+ # Step 3: Build filters and time range
+ # ------------------------------------------------------------------
+ warnings: list[str] = []
+ query_filters: list[dict[str, Any]] = [
+ {"col": f.col, "op": f.op, "val": f.val} for f in request.filters
+ ]
+ granularity: str | None = None
+
+ if request.time_range:
+ temporal_col = request.time_column or getattr(
+ dataset, "main_dttm_col", None
+ )
+ if not temporal_col:
+ await ctx.error("time_range provided but no temporal column available")
+ return DatasetError.create(
+ error=(
+ "time_range was provided but no temporal column is available. "
+ "Either set time_column explicitly or ensure the dataset has "
+ "a main datetime column configured."
+ ),
+ error_type="ValidationError",
+ )
+ # Validate that the temporal column actually exists on the dataset
+ if temporal_col not in valid_columns:
+ await ctx.error("time_column '%s' not found on dataset" % temporal_col)
+ return DatasetError.create(
+ error=(
+ f"time_column '{temporal_col}' does not exist on this dataset."
+ ),
+ error_type="ValidationError",
+ )
+ # Warn if the chosen temporal column isn't marked as datetime
+ dttm_cols = {c.column_name for c in dataset.columns if c.is_dttm}
+ if temporal_col not in dttm_cols:
+ warnings.append(
+ f"Column '{temporal_col}' is not marked as a datetime "
+ f"column on this dataset. Time filtering may not work "
+ f"as expected."
+ )
+
+ query_filters.append(
+ {
+ "col": temporal_col,
+ "op": "TEMPORAL_RANGE",
+ "val": request.time_range,
+ }
+ )
+ granularity = temporal_col
+ await ctx.debug(
+ "Time filter: column=%s, range=%s" % (temporal_col, request.time_range)
+ )
+
+ # ------------------------------------------------------------------
+ # Step 4: Build query dict
+ # ------------------------------------------------------------------
+ await ctx.report_progress(3, 5, "Building query")
+ query_dict: dict[str, Any] = {
+ "filters": query_filters,
+ "columns": request.columns,
+ "metrics": request.metrics,
+ "row_limit": request.row_limit,
+ "order_desc": request.order_desc,
+ }
+ if granularity:
+ query_dict["granularity"] = granularity
+ if request.order_by:
+ # OrderBy = tuple[Metric | Column, bool] where bool is ascending
+ query_dict["orderby"] = [
+ (col, not request.order_desc) for col in request.order_by
+ ]
+
+ await ctx.debug("Query dict keys: %s" % (sorted(query_dict.keys()),))
+
+ # ------------------------------------------------------------------
+ # Step 5: Create QueryContext and execute
+ # ------------------------------------------------------------------
+ await ctx.report_progress(4, 5, "Executing query")
+ start_time = time.time()
+
+ with event_logger.log_context(action="mcp.query_dataset.execute"):
+ factory = QueryContextFactory()
+ # datasource_type is "table" because this tool queries SqlaTable
+ # datasets (Superset's built-in semantic layer). External semantic
+ # layers (dbt, Snowflake Cortex, etc.) use "semantic_view" and have
+ # a different query path — see SemanticView + mapper.py.
+ query_context = factory.create(
+ datasource={"id": dataset.id, "type": "table"},
+ queries=[query_dict],
+ form_data={},
+ force=request.force_refresh,
+ custom_cache_timeout=request.cache_timeout,
+ )
+
+ command = ChartDataCommand(query_context)
+ command.validate()
+ result = command.run()
+
+ query_duration_ms = int((time.time() - start_time) * 1000)
+
+ if not result or "queries" not in result or len(result["queries"]) == 0:
+ await ctx.warning("Query returned no results for dataset %s" % dataset.id)
+ return DatasetError.create(
+ error="Query returned no results.",
+ error_type="EmptyQuery",
+ )
+
+ # ------------------------------------------------------------------
+ # Step 6: Format response
+ # ------------------------------------------------------------------
+ await ctx.report_progress(5, 5, "Formatting results")
+ query_result = result["queries"][0]
+ data = query_result.get("data", [])
+ raw_columns = query_result.get("colnames", [])
+
+ if not data:
+ return QueryDatasetResponse(
+ dataset_id=dataset.id,
+ dataset_name=dataset_name,
+ columns=[],
+ data=[],
+ row_count=0,
+ total_rows=0,
+ summary=f"Query on '{dataset_name}' returned no data.",
+ performance=PerformanceMetadata(
+ query_duration_ms=query_duration_ms,
+ cache_status="no_data",
+ ),
+ cache_status=get_cache_status_from_result(
+ query_result, force_refresh=request.force_refresh
+ ),
+ applied_filters=request.filters,
+ warnings=warnings,
Review Comment:
Fixed in ffcd4b1d25 — now tracking an `effective_filters` list that starts from `request.filters` and appends the synthesized `TEMPORAL_RANGE` filter when `time_range` is provided. Both response paths (data and no-data) now use `effective_filters` instead of `request.filters`.
--
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: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] feat(mcp): add query_dataset tool to query datasets using semantic layer [superset]
Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on code in PR #39727:
URL: https://github.com/apache/superset/pull/39727#discussion_r3157414187
##########
superset/mcp_service/dataset/schemas.py:
##########
@@ -380,6 +385,146 @@ class CreateVirtualDatasetResponse(BaseModel):
)
+VALID_FILTER_OPS = Literal[
+ "==",
+ "!=",
+ ">",
+ "<",
+ ">=",
+ "<=",
+ "LIKE",
+ "NOT LIKE",
+ "ILIKE",
+ "NOT ILIKE",
+ "IN",
+ "NOT IN",
+ "IS NULL",
+ "IS NOT NULL",
+ "IS TRUE",
+ "IS FALSE",
+ "TEMPORAL_RANGE",
+]
Review Comment:
Fixed in ff3672f — updated the example in the description to use `"=="` to match `VALID_FILTER_OPS`.
##########
superset/mcp_service/dataset/schemas.py:
##########
@@ -380,6 +385,146 @@ class CreateVirtualDatasetResponse(BaseModel):
)
+VALID_FILTER_OPS = Literal[
+ "==",
+ "!=",
+ ">",
+ "<",
+ ">=",
+ "<=",
+ "LIKE",
+ "NOT LIKE",
+ "ILIKE",
+ "NOT ILIKE",
+ "IN",
+ "NOT IN",
+ "IS NULL",
+ "IS NOT NULL",
+ "IS TRUE",
+ "IS FALSE",
+ "TEMPORAL_RANGE",
+]
+
+
+class QueryDatasetFilter(BaseModel):
+ """A single filter condition for dataset queries."""
+
+ col: str = Field(..., description="Column name to filter on")
+ op: VALID_FILTER_OPS = Field(
+ ...,
+ description=(
+ 'Filter operator. Use "==" for equals, "!=" for not equals, '
+ '"IN" / "NOT IN" for membership, "IS NULL" / "IS NOT NULL", '
+ '"LIKE" for pattern matching, "TEMPORAL_RANGE" for time filters.'
+ ),
+ )
+ val: Any = Field(
+ default=None,
+ description="Filter value (omit for IS NULL/IS NOT NULL)",
+ )
+
+
+class QueryDatasetRequest(QueryCacheControl):
+ """Request schema for query_dataset tool."""
+
+ dataset_id: int | str = Field(
+ ...,
+ description="Dataset identifier — numeric ID or UUID string.",
+ )
+ metrics: List[str] = Field(
+ default_factory=list,
+ description=(
+ "Saved metric names to compute (e.g. ['count', 'avg_revenue']). "
+ "Use get_dataset_info to discover available metrics."
+ ),
+ )
+ columns: List[str] = Field(
+ default_factory=list,
+ description=(
+ "Column/dimension names for GROUP BY or SELECT "
+ "(e.g. ['category', 'region']). "
+ "Use get_dataset_info to discover available columns."
+ ),
+ )
+ filters: List[QueryDatasetFilter] = Field(
+ default_factory=list,
+ description=(
+ 'Filter conditions (e.g. [{"col": "status", "op": "=", "val": "active"}]).'
Review Comment:
Fixed in ff3672f — same fix, updated filter example from `"="` to `"=="`.
##########
superset/mcp_service/dataset/tool/query_dataset.py:
##########
@@ -0,0 +1,434 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+MCP tool: query_dataset
+
+Query a dataset using its semantic layer (saved metrics, calculated columns,
+dimensions) without requiring a saved chart.
+"""
+
+import difflib
+import logging
+import time
+from typing import Any
+
+from fastmcp import Context
+from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.orm import joinedload, subqueryload
+from superset_core.mcp.decorators import tool, ToolAnnotations
+
+from superset.commands.exceptions import CommandException
+from superset.exceptions import OAuth2Error, OAuth2RedirectError, SupersetException
+from superset.extensions import event_logger
+from superset.mcp_service.chart.schemas import DataColumn, PerformanceMetadata
+from superset.mcp_service.dataset.schemas import (
+ DatasetError,
+ QueryDatasetRequest,
+ QueryDatasetResponse,
+)
+from superset.mcp_service.utils import _is_uuid
+from superset.mcp_service.utils.cache_utils import get_cache_status_from_result
+from superset.mcp_service.utils.oauth2_utils import build_oauth2_redirect_message
+
+logger = logging.getLogger(__name__)
+
+
+def _resolve_dataset(identifier: int | str, eager_options: list[Any]) -> Any | None:
+ """Resolve a dataset by int ID or UUID string.
+
+ Replicates the identifier resolution logic from ModelGetInfoCore._find_object().
+ """
+ from superset.daos.dataset import DatasetDAO
+
+ opts = eager_options or None
+
+ if isinstance(identifier, int):
+ return DatasetDAO.find_by_id(identifier, query_options=opts)
+
+ # Try parsing as int
+ try:
+ id_val = int(identifier)
+ return DatasetDAO.find_by_id(id_val, query_options=opts)
+ except (ValueError, TypeError):
+ pass
+
+ # Try UUID
+ if _is_uuid(str(identifier)):
+ return DatasetDAO.find_by_id(identifier, id_column="uuid", query_options=opts)
+
+ return None
+
+
+def _validate_names(
+ requested: list[str],
+ valid: set[str],
+ kind: str,
+) -> list[str]:
+ """Return list of error messages for names not found in *valid*.
+
+ Includes close-match suggestions when available.
+ """
+ errors: list[str] = []
+ for name in requested:
+ if name not in valid:
+ suggestions = difflib.get_close_matches(name, valid, n=3, cutoff=0.6)
+ msg = f"Unknown {kind}: '{name}'"
+ if suggestions:
+ msg += f". Did you mean: {', '.join(suggestions)}?"
+ errors.append(msg)
+ return errors
+
+
+@tool(
+ tags=["data"],
+ class_permission_name="Dataset",
+ annotations=ToolAnnotations(
+ title="Query dataset",
+ readOnlyHint=True,
+ destructiveHint=False,
+ ),
+)
+async def query_dataset( # noqa: C901
+ request: QueryDatasetRequest, ctx: Context
+) -> QueryDatasetResponse | DatasetError:
+ """Query a dataset using its semantic layer (saved metrics, dimensions, filters).
+
+ Returns tabular data without requiring a saved chart. Use this when you want
+ to compute saved metrics, group by dimensions, or apply filters directly
+ against a dataset's curated semantic layer.
+
+ Workflow:
+ 1. list_datasets -> find a dataset
+ 2. get_dataset_info -> discover available columns and metrics
+ 3. query_dataset -> query using metric names and column names
+
+ Example:
+ ```json
+ {
+ "dataset_id": 123,
+ "metrics": ["count", "avg_revenue"],
+ "columns": ["product_category"],
+ "time_range": "Last 7 days",
+ "row_limit": 100
+ }
+ ```
+ """
+ await ctx.info(
+ "Starting dataset query: dataset_id=%s, metrics=%s, columns=%s, "
+ "row_limit=%s"
+ % (
+ request.dataset_id,
+ request.metrics,
+ request.columns,
+ request.row_limit,
+ )
+ )
+
+ try:
+ from superset.commands.chart.data.get_data_command import ChartDataCommand
+ from superset.common.query_context_factory import QueryContextFactory
+ from superset.connectors.sqla.models import SqlaTable
+
+ # ------------------------------------------------------------------
+ # Step 1: Resolve dataset
+ # ------------------------------------------------------------------
+ await ctx.report_progress(1, 5, "Looking up dataset")
+ eager_options = [
+ subqueryload(SqlaTable.columns),
+ subqueryload(SqlaTable.metrics),
+ joinedload(SqlaTable.database),
+ ]
+
+ with event_logger.log_context(action="mcp.query_dataset.lookup"):
+ dataset = _resolve_dataset(request.dataset_id, eager_options)
+
+ if dataset is None:
+ await ctx.error("Dataset not found: identifier=%s" % (request.dataset_id,))
+ return DatasetError.create(
+ error=f"No dataset found with identifier: {request.dataset_id}",
+ error_type="NotFound",
+ )
+
+ dataset_name = getattr(dataset, "table_name", None) or f"Dataset {dataset.id}"
+ await ctx.info(
+ "Dataset found: id=%s, name=%s, columns=%s, metrics=%s"
+ % (
+ dataset.id,
+ dataset_name,
+ len(dataset.columns),
+ len(dataset.metrics),
+ )
+ )
+
+ # ------------------------------------------------------------------
+ # Step 2: Validate requested columns and metrics
+ # ------------------------------------------------------------------
+ await ctx.report_progress(2, 5, "Validating columns and metrics")
+ valid_columns = {c.column_name for c in dataset.columns}
+ valid_metrics = {m.metric_name for m in dataset.metrics}
+
+ validation_errors: list[str] = []
+ validation_errors.extend(
+ _validate_names(request.columns, valid_columns, "column")
+ )
+ validation_errors.extend(
+ _validate_names(request.metrics, valid_metrics, "metric")
+ )
+ # Validate filter column names against dataset columns
+ filter_cols = [f.col for f in request.filters]
+ validation_errors.extend(
+ _validate_names(filter_cols, valid_columns, "filter column")
+ )
+ # Validate order_by names against columns + metrics
+ if request.order_by:
+ valid_orderby = valid_columns | valid_metrics
+ validation_errors.extend(
+ _validate_names(request.order_by, valid_orderby, "order_by")
+ )
+
+ if validation_errors:
+ error_msg = "; ".join(validation_errors)
+ await ctx.error("Validation failed: %s" % (error_msg,))
+ return DatasetError.create(
+ error=error_msg,
+ error_type="ValidationError",
+ )
+
+ # ------------------------------------------------------------------
+ # Step 3: Build filters and time range
+ # ------------------------------------------------------------------
+ warnings: list[str] = []
+ query_filters: list[dict[str, Any]] = [
+ {"col": f.col, "op": f.op, "val": f.val} for f in request.filters
+ ]
+ granularity: str | None = None
+
+ if request.time_range:
+ temporal_col = request.time_column or getattr(
+ dataset, "main_dttm_col", None
+ )
+ if not temporal_col:
+ await ctx.error("time_range provided but no temporal column available")
+ return DatasetError.create(
+ error=(
+ "time_range was provided but no temporal column is available. "
+ "Either set time_column explicitly or ensure the dataset has "
+ "a main datetime column configured."
+ ),
+ error_type="ValidationError",
+ )
+ # Validate that the temporal column actually exists on the dataset
+ if temporal_col not in valid_columns:
+ await ctx.error("time_column '%s' not found on dataset" % temporal_col)
+ return DatasetError.create(
+ error=(
+ f"time_column '{temporal_col}' does not exist on this dataset."
+ ),
+ error_type="ValidationError",
+ )
+ # Warn if the chosen temporal column isn't marked as datetime
+ dttm_cols = {c.column_name for c in dataset.columns if c.is_dttm}
+ if temporal_col not in dttm_cols:
+ warnings.append(
+ f"Column '{temporal_col}' is not marked as a datetime "
+ f"column on this dataset. Time filtering may not work "
+ f"as expected."
+ )
+
+ query_filters.append(
+ {
+ "col": temporal_col,
+ "op": "TEMPORAL_RANGE",
+ "val": request.time_range,
+ }
+ )
+ granularity = temporal_col
+ await ctx.debug(
+ "Time filter: column=%s, range=%s" % (temporal_col, request.time_range)
+ )
+
+ # ------------------------------------------------------------------
+ # Step 4: Build query dict
+ # ------------------------------------------------------------------
+ await ctx.report_progress(3, 5, "Building query")
+ query_dict: dict[str, Any] = {
+ "filters": query_filters,
+ "columns": request.columns,
+ "metrics": request.metrics,
+ "row_limit": request.row_limit,
+ "order_desc": request.order_desc,
+ }
+ if granularity:
+ query_dict["granularity"] = granularity
+ if request.order_by:
+ # OrderBy = tuple[Metric | Column, bool] where bool is ascending
+ query_dict["orderby"] = [
+ (col, not request.order_desc) for col in request.order_by
+ ]
+
+ await ctx.debug("Query dict keys: %s" % (sorted(query_dict.keys()),))
+
+ # ------------------------------------------------------------------
+ # Step 5: Create QueryContext and execute
+ # ------------------------------------------------------------------
+ await ctx.report_progress(4, 5, "Executing query")
+ start_time = time.time()
+
+ with event_logger.log_context(action="mcp.query_dataset.execute"):
+ factory = QueryContextFactory()
+ # datasource_type is "table" because this tool queries SqlaTable
+ # datasets (Superset's built-in semantic layer). External semantic
+ # layers (dbt, Snowflake Cortex, etc.) use "semantic_view" and have
+ # a different query path — see SemanticView + mapper.py.
+ query_context = factory.create(
+ datasource={"id": dataset.id, "type": "table"},
+ queries=[query_dict],
+ form_data={},
+ force=request.force_refresh,
+ custom_cache_timeout=request.cache_timeout,
+ )
+
+ command = ChartDataCommand(query_context)
+ command.validate()
+ result = command.run()
+
+ query_duration_ms = int((time.time() - start_time) * 1000)
+
+ if not result or "queries" not in result or len(result["queries"]) == 0:
+ await ctx.warning("Query returned no results for dataset %s" % dataset.id)
+ return DatasetError.create(
+ error="Query returned no results.",
+ error_type="EmptyQuery",
+ )
+
+ # ------------------------------------------------------------------
+ # Step 6: Format response
+ # ------------------------------------------------------------------
+ await ctx.report_progress(5, 5, "Formatting results")
+ query_result = result["queries"][0]
+ data = query_result.get("data", [])
+ raw_columns = query_result.get("colnames", [])
+
+ if not data:
+ return QueryDatasetResponse(
+ dataset_id=dataset.id,
+ dataset_name=dataset_name,
+ columns=[],
+ data=[],
+ row_count=0,
+ total_rows=0,
+ summary=f"Query on '{dataset_name}' returned no data.",
+ performance=PerformanceMetadata(
+ query_duration_ms=query_duration_ms,
+ cache_status="no_data",
+ ),
+ cache_status=get_cache_status_from_result(
+ query_result, force_refresh=request.force_refresh
+ ),
+ applied_filters=request.filters,
+ warnings=warnings,
+ )
+
+ # Build column metadata
+ columns_meta: list[DataColumn] = []
+ for col_name in raw_columns:
+ sample_values = [
+ row.get(col_name) for row in data[:3] if row.get(col_name) is not None
+ ]
+ data_type = "string"
+ if sample_values:
+ if all(isinstance(v, bool) for v in sample_values):
+ data_type = "boolean"
+ elif all(isinstance(v, (int, float)) for v in sample_values):
+ data_type = "numeric"
+
+ columns_meta.append(
+ DataColumn(
+ name=col_name,
+ display_name=col_name.replace("_", " ").title(),
+ data_type=data_type,
+ sample_values=sample_values[:3],
+ null_count=sum(1 for row in data if row.get(col_name) is None),
+ unique_count=len({str(row.get(col_name)) for row in data}),
+ )
+ )
Review Comment:
Fixed in ff3672f — stats are now computed in a single pass per column, capped at 5000 rows to avoid O(rows×cols) on large result sets.
##########
superset/mcp_service/dataset/tool/query_dataset.py:
##########
@@ -0,0 +1,434 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+MCP tool: query_dataset
+
+Query a dataset using its semantic layer (saved metrics, calculated columns,
+dimensions) without requiring a saved chart.
+"""
+
+import difflib
+import logging
+import time
+from typing import Any
+
+from fastmcp import Context
+from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.orm import joinedload, subqueryload
+from superset_core.mcp.decorators import tool, ToolAnnotations
+
+from superset.commands.exceptions import CommandException
+from superset.exceptions import OAuth2Error, OAuth2RedirectError, SupersetException
+from superset.extensions import event_logger
+from superset.mcp_service.chart.schemas import DataColumn, PerformanceMetadata
+from superset.mcp_service.dataset.schemas import (
+ DatasetError,
+ QueryDatasetRequest,
+ QueryDatasetResponse,
+)
+from superset.mcp_service.utils import _is_uuid
+from superset.mcp_service.utils.cache_utils import get_cache_status_from_result
+from superset.mcp_service.utils.oauth2_utils import build_oauth2_redirect_message
+
+logger = logging.getLogger(__name__)
+
+
+def _resolve_dataset(identifier: int | str, eager_options: list[Any]) -> Any | None:
+ """Resolve a dataset by int ID or UUID string.
+
+ Replicates the identifier resolution logic from ModelGetInfoCore._find_object().
+ """
+ from superset.daos.dataset import DatasetDAO
+
+ opts = eager_options or None
+
+ if isinstance(identifier, int):
+ return DatasetDAO.find_by_id(identifier, query_options=opts)
+
+ # Try parsing as int
+ try:
+ id_val = int(identifier)
+ return DatasetDAO.find_by_id(id_val, query_options=opts)
+ except (ValueError, TypeError):
+ pass
+
+ # Try UUID
+ if _is_uuid(str(identifier)):
+ return DatasetDAO.find_by_id(identifier, id_column="uuid", query_options=opts)
+
+ return None
+
+
+def _validate_names(
+ requested: list[str],
+ valid: set[str],
+ kind: str,
+) -> list[str]:
+ """Return list of error messages for names not found in *valid*.
+
+ Includes close-match suggestions when available.
+ """
+ errors: list[str] = []
+ for name in requested:
+ if name not in valid:
+ suggestions = difflib.get_close_matches(name, valid, n=3, cutoff=0.6)
+ msg = f"Unknown {kind}: '{name}'"
+ if suggestions:
+ msg += f". Did you mean: {', '.join(suggestions)}?"
+ errors.append(msg)
+ return errors
+
+
+@tool(
+ tags=["data"],
+ class_permission_name="Dataset",
+ annotations=ToolAnnotations(
+ title="Query dataset",
+ readOnlyHint=True,
+ destructiveHint=False,
+ ),
+)
+async def query_dataset( # noqa: C901
+ request: QueryDatasetRequest, ctx: Context
+) -> QueryDatasetResponse | DatasetError:
+ """Query a dataset using its semantic layer (saved metrics, dimensions, filters).
+
+ Returns tabular data without requiring a saved chart. Use this when you want
+ to compute saved metrics, group by dimensions, or apply filters directly
+ against a dataset's curated semantic layer.
+
+ Workflow:
+ 1. list_datasets -> find a dataset
+ 2. get_dataset_info -> discover available columns and metrics
+ 3. query_dataset -> query using metric names and column names
+
+ Example:
+ ```json
+ {
+ "dataset_id": 123,
+ "metrics": ["count", "avg_revenue"],
+ "columns": ["product_category"],
+ "time_range": "Last 7 days",
+ "row_limit": 100
+ }
+ ```
+ """
+ await ctx.info(
+ "Starting dataset query: dataset_id=%s, metrics=%s, columns=%s, "
+ "row_limit=%s"
+ % (
+ request.dataset_id,
+ request.metrics,
+ request.columns,
+ request.row_limit,
+ )
+ )
+
+ try:
+ from superset.commands.chart.data.get_data_command import ChartDataCommand
+ from superset.common.query_context_factory import QueryContextFactory
+ from superset.connectors.sqla.models import SqlaTable
+
+ # ------------------------------------------------------------------
+ # Step 1: Resolve dataset
+ # ------------------------------------------------------------------
+ await ctx.report_progress(1, 5, "Looking up dataset")
+ eager_options = [
+ subqueryload(SqlaTable.columns),
+ subqueryload(SqlaTable.metrics),
+ joinedload(SqlaTable.database),
+ ]
+
+ with event_logger.log_context(action="mcp.query_dataset.lookup"):
+ dataset = _resolve_dataset(request.dataset_id, eager_options)
+
+ if dataset is None:
+ await ctx.error("Dataset not found: identifier=%s" % (request.dataset_id,))
+ return DatasetError.create(
+ error=f"No dataset found with identifier: {request.dataset_id}",
+ error_type="NotFound",
+ )
+
+ dataset_name = getattr(dataset, "table_name", None) or f"Dataset {dataset.id}"
+ await ctx.info(
+ "Dataset found: id=%s, name=%s, columns=%s, metrics=%s"
+ % (
+ dataset.id,
+ dataset_name,
+ len(dataset.columns),
+ len(dataset.metrics),
+ )
+ )
+
+ # ------------------------------------------------------------------
+ # Step 2: Validate requested columns and metrics
+ # ------------------------------------------------------------------
+ await ctx.report_progress(2, 5, "Validating columns and metrics")
+ valid_columns = {c.column_name for c in dataset.columns}
+ valid_metrics = {m.metric_name for m in dataset.metrics}
+
+ validation_errors: list[str] = []
+ validation_errors.extend(
+ _validate_names(request.columns, valid_columns, "column")
+ )
+ validation_errors.extend(
+ _validate_names(request.metrics, valid_metrics, "metric")
+ )
+ # Validate filter column names against dataset columns
+ filter_cols = [f.col for f in request.filters]
+ validation_errors.extend(
+ _validate_names(filter_cols, valid_columns, "filter column")
+ )
+ # Validate order_by names against columns + metrics
+ if request.order_by:
+ valid_orderby = valid_columns | valid_metrics
+ validation_errors.extend(
+ _validate_names(request.order_by, valid_orderby, "order_by")
+ )
+
+ if validation_errors:
+ error_msg = "; ".join(validation_errors)
+ await ctx.error("Validation failed: %s" % (error_msg,))
+ return DatasetError.create(
+ error=error_msg,
+ error_type="ValidationError",
+ )
+
+ # ------------------------------------------------------------------
+ # Step 3: Build filters and time range
+ # ------------------------------------------------------------------
+ warnings: list[str] = []
+ query_filters: list[dict[str, Any]] = [
+ {"col": f.col, "op": f.op, "val": f.val} for f in request.filters
+ ]
+ granularity: str | None = None
+
+ if request.time_range:
+ temporal_col = request.time_column or getattr(
+ dataset, "main_dttm_col", None
+ )
+ if not temporal_col:
+ await ctx.error("time_range provided but no temporal column available")
+ return DatasetError.create(
+ error=(
+ "time_range was provided but no temporal column is available. "
+ "Either set time_column explicitly or ensure the dataset has "
+ "a main datetime column configured."
+ ),
+ error_type="ValidationError",
+ )
+ # Validate that the temporal column actually exists on the dataset
+ if temporal_col not in valid_columns:
+ await ctx.error("time_column '%s' not found on dataset" % temporal_col)
+ return DatasetError.create(
+ error=(
+ f"time_column '{temporal_col}' does not exist on this dataset."
+ ),
+ error_type="ValidationError",
+ )
+ # Warn if the chosen temporal column isn't marked as datetime
+ dttm_cols = {c.column_name for c in dataset.columns if c.is_dttm}
+ if temporal_col not in dttm_cols:
+ warnings.append(
+ f"Column '{temporal_col}' is not marked as a datetime "
+ f"column on this dataset. Time filtering may not work "
+ f"as expected."
+ )
+
+ query_filters.append(
+ {
+ "col": temporal_col,
+ "op": "TEMPORAL_RANGE",
+ "val": request.time_range,
+ }
+ )
+ granularity = temporal_col
+ await ctx.debug(
+ "Time filter: column=%s, range=%s" % (temporal_col, request.time_range)
+ )
+
+ # ------------------------------------------------------------------
+ # Step 4: Build query dict
+ # ------------------------------------------------------------------
+ await ctx.report_progress(3, 5, "Building query")
+ query_dict: dict[str, Any] = {
+ "filters": query_filters,
+ "columns": request.columns,
+ "metrics": request.metrics,
+ "row_limit": request.row_limit,
+ "order_desc": request.order_desc,
+ }
+ if granularity:
+ query_dict["granularity"] = granularity
+ if request.order_by:
+ # OrderBy = tuple[Metric | Column, bool] where bool is ascending
+ query_dict["orderby"] = [
+ (col, not request.order_desc) for col in request.order_by
+ ]
+
+ await ctx.debug("Query dict keys: %s" % (sorted(query_dict.keys()),))
+
+ # ------------------------------------------------------------------
+ # Step 5: Create QueryContext and execute
+ # ------------------------------------------------------------------
+ await ctx.report_progress(4, 5, "Executing query")
+ start_time = time.time()
+
+ with event_logger.log_context(action="mcp.query_dataset.execute"):
+ factory = QueryContextFactory()
+ # datasource_type is "table" because this tool queries SqlaTable
+ # datasets (Superset's built-in semantic layer). External semantic
+ # layers (dbt, Snowflake Cortex, etc.) use "semantic_view" and have
+ # a different query path — see SemanticView + mapper.py.
+ query_context = factory.create(
+ datasource={"id": dataset.id, "type": "table"},
+ queries=[query_dict],
+ form_data={},
+ force=request.force_refresh,
+ custom_cache_timeout=request.cache_timeout,
+ )
+
+ command = ChartDataCommand(query_context)
+ command.validate()
+ result = command.run()
+
+ query_duration_ms = int((time.time() - start_time) * 1000)
+
+ if not result or "queries" not in result or len(result["queries"]) == 0:
+ await ctx.warning("Query returned no results for dataset %s" % dataset.id)
+ return DatasetError.create(
+ error="Query returned no results.",
+ error_type="EmptyQuery",
+ )
+
+ # ------------------------------------------------------------------
+ # Step 6: Format response
+ # ------------------------------------------------------------------
+ await ctx.report_progress(5, 5, "Formatting results")
+ query_result = result["queries"][0]
+ data = query_result.get("data", [])
+ raw_columns = query_result.get("colnames", [])
+
+ if not data:
+ return QueryDatasetResponse(
+ dataset_id=dataset.id,
+ dataset_name=dataset_name,
+ columns=[],
+ data=[],
+ row_count=0,
+ total_rows=0,
+ summary=f"Query on '{dataset_name}' returned no data.",
+ performance=PerformanceMetadata(
+ query_duration_ms=query_duration_ms,
+ cache_status="no_data",
+ ),
+ cache_status=get_cache_status_from_result(
+ query_result, force_refresh=request.force_refresh
+ ),
+ applied_filters=request.filters,
+ warnings=warnings,
+ )
+
+ # Build column metadata
+ columns_meta: list[DataColumn] = []
+ for col_name in raw_columns:
+ sample_values = [
+ row.get(col_name) for row in data[:3] if row.get(col_name) is not None
+ ]
+ data_type = "string"
+ if sample_values:
+ if all(isinstance(v, bool) for v in sample_values):
+ data_type = "boolean"
+ elif all(isinstance(v, (int, float)) for v in sample_values):
+ data_type = "numeric"
+
+ columns_meta.append(
+ DataColumn(
+ name=col_name,
+ display_name=col_name.replace("_", " ").title(),
+ data_type=data_type,
+ sample_values=sample_values[:3],
+ null_count=sum(1 for row in data if row.get(col_name) is None),
+ unique_count=len({str(row.get(col_name)) for row in data}),
Review Comment:
Fixed in ff3672f — `unique_count` now excludes nulls. Single-pass loop collects null_count and unique non-null values separately.
##########
superset/mcp_service/dataset/tool/query_dataset.py:
##########
@@ -0,0 +1,434 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+MCP tool: query_dataset
+
+Query a dataset using its semantic layer (saved metrics, calculated columns,
+dimensions) without requiring a saved chart.
+"""
+
+import difflib
+import logging
+import time
+from typing import Any
+
+from fastmcp import Context
+from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.orm import joinedload, subqueryload
+from superset_core.mcp.decorators import tool, ToolAnnotations
+
+from superset.commands.exceptions import CommandException
+from superset.exceptions import OAuth2Error, OAuth2RedirectError, SupersetException
+from superset.extensions import event_logger
+from superset.mcp_service.chart.schemas import DataColumn, PerformanceMetadata
+from superset.mcp_service.dataset.schemas import (
+ DatasetError,
+ QueryDatasetRequest,
+ QueryDatasetResponse,
+)
+from superset.mcp_service.utils import _is_uuid
+from superset.mcp_service.utils.cache_utils import get_cache_status_from_result
+from superset.mcp_service.utils.oauth2_utils import build_oauth2_redirect_message
+
+logger = logging.getLogger(__name__)
+
+
+def _resolve_dataset(identifier: int | str, eager_options: list[Any]) -> Any | None:
+ """Resolve a dataset by int ID or UUID string.
+
+ Replicates the identifier resolution logic from ModelGetInfoCore._find_object().
+ """
+ from superset.daos.dataset import DatasetDAO
+
+ opts = eager_options or None
+
+ if isinstance(identifier, int):
+ return DatasetDAO.find_by_id(identifier, query_options=opts)
+
+ # Try parsing as int
+ try:
+ id_val = int(identifier)
+ return DatasetDAO.find_by_id(id_val, query_options=opts)
+ except (ValueError, TypeError):
+ pass
+
+ # Try UUID
+ if _is_uuid(str(identifier)):
+ return DatasetDAO.find_by_id(identifier, id_column="uuid", query_options=opts)
+
+ return None
+
+
+def _validate_names(
+ requested: list[str],
+ valid: set[str],
+ kind: str,
+) -> list[str]:
+ """Return list of error messages for names not found in *valid*.
+
+ Includes close-match suggestions when available.
+ """
+ errors: list[str] = []
+ for name in requested:
+ if name not in valid:
+ suggestions = difflib.get_close_matches(name, valid, n=3, cutoff=0.6)
+ msg = f"Unknown {kind}: '{name}'"
+ if suggestions:
+ msg += f". Did you mean: {', '.join(suggestions)}?"
+ errors.append(msg)
+ return errors
+
+
+@tool(
+ tags=["data"],
+ class_permission_name="Dataset",
+ annotations=ToolAnnotations(
+ title="Query dataset",
+ readOnlyHint=True,
+ destructiveHint=False,
+ ),
+)
+async def query_dataset( # noqa: C901
+ request: QueryDatasetRequest, ctx: Context
+) -> QueryDatasetResponse | DatasetError:
+ """Query a dataset using its semantic layer (saved metrics, dimensions, filters).
+
+ Returns tabular data without requiring a saved chart. Use this when you want
+ to compute saved metrics, group by dimensions, or apply filters directly
+ against a dataset's curated semantic layer.
+
+ Workflow:
+ 1. list_datasets -> find a dataset
+ 2. get_dataset_info -> discover available columns and metrics
+ 3. query_dataset -> query using metric names and column names
+
+ Example:
+ ```json
+ {
+ "dataset_id": 123,
+ "metrics": ["count", "avg_revenue"],
+ "columns": ["product_category"],
+ "time_range": "Last 7 days",
+ "row_limit": 100
+ }
+ ```
+ """
+ await ctx.info(
+ "Starting dataset query: dataset_id=%s, metrics=%s, columns=%s, "
+ "row_limit=%s"
+ % (
+ request.dataset_id,
+ request.metrics,
+ request.columns,
+ request.row_limit,
+ )
+ )
+
+ try:
+ from superset.commands.chart.data.get_data_command import ChartDataCommand
+ from superset.common.query_context_factory import QueryContextFactory
+ from superset.connectors.sqla.models import SqlaTable
+
+ # ------------------------------------------------------------------
+ # Step 1: Resolve dataset
+ # ------------------------------------------------------------------
+ await ctx.report_progress(1, 5, "Looking up dataset")
+ eager_options = [
+ subqueryload(SqlaTable.columns),
+ subqueryload(SqlaTable.metrics),
+ joinedload(SqlaTable.database),
+ ]
+
+ with event_logger.log_context(action="mcp.query_dataset.lookup"):
+ dataset = _resolve_dataset(request.dataset_id, eager_options)
+
+ if dataset is None:
+ await ctx.error("Dataset not found: identifier=%s" % (request.dataset_id,))
+ return DatasetError.create(
+ error=f"No dataset found with identifier: {request.dataset_id}",
+ error_type="NotFound",
+ )
+
+ dataset_name = getattr(dataset, "table_name", None) or f"Dataset {dataset.id}"
+ await ctx.info(
+ "Dataset found: id=%s, name=%s, columns=%s, metrics=%s"
+ % (
+ dataset.id,
+ dataset_name,
+ len(dataset.columns),
+ len(dataset.metrics),
+ )
+ )
+
+ # ------------------------------------------------------------------
+ # Step 2: Validate requested columns and metrics
+ # ------------------------------------------------------------------
+ await ctx.report_progress(2, 5, "Validating columns and metrics")
+ valid_columns = {c.column_name for c in dataset.columns}
+ valid_metrics = {m.metric_name for m in dataset.metrics}
+
+ validation_errors: list[str] = []
+ validation_errors.extend(
+ _validate_names(request.columns, valid_columns, "column")
+ )
+ validation_errors.extend(
+ _validate_names(request.metrics, valid_metrics, "metric")
+ )
+ # Validate filter column names against dataset columns
+ filter_cols = [f.col for f in request.filters]
+ validation_errors.extend(
+ _validate_names(filter_cols, valid_columns, "filter column")
+ )
+ # Validate order_by names against columns + metrics
+ if request.order_by:
+ valid_orderby = valid_columns | valid_metrics
+ validation_errors.extend(
+ _validate_names(request.order_by, valid_orderby, "order_by")
+ )
+
+ if validation_errors:
+ error_msg = "; ".join(validation_errors)
+ await ctx.error("Validation failed: %s" % (error_msg,))
+ return DatasetError.create(
+ error=error_msg,
+ error_type="ValidationError",
+ )
+
+ # ------------------------------------------------------------------
+ # Step 3: Build filters and time range
+ # ------------------------------------------------------------------
+ warnings: list[str] = []
+ query_filters: list[dict[str, Any]] = [
+ {"col": f.col, "op": f.op, "val": f.val} for f in request.filters
+ ]
+ granularity: str | None = None
+
+ if request.time_range:
+ temporal_col = request.time_column or getattr(
+ dataset, "main_dttm_col", None
+ )
+ if not temporal_col:
+ await ctx.error("time_range provided but no temporal column available")
+ return DatasetError.create(
+ error=(
+ "time_range was provided but no temporal column is available. "
+ "Either set time_column explicitly or ensure the dataset has "
+ "a main datetime column configured."
+ ),
+ error_type="ValidationError",
+ )
+ # Validate that the temporal column actually exists on the dataset
+ if temporal_col not in valid_columns:
+ await ctx.error("time_column '%s' not found on dataset" % temporal_col)
+ return DatasetError.create(
+ error=(
+ f"time_column '{temporal_col}' does not exist on this dataset."
+ ),
+ error_type="ValidationError",
+ )
+ # Warn if the chosen temporal column isn't marked as datetime
+ dttm_cols = {c.column_name for c in dataset.columns if c.is_dttm}
+ if temporal_col not in dttm_cols:
+ warnings.append(
+ f"Column '{temporal_col}' is not marked as a datetime "
+ f"column on this dataset. Time filtering may not work "
+ f"as expected."
+ )
+
+ query_filters.append(
+ {
+ "col": temporal_col,
+ "op": "TEMPORAL_RANGE",
+ "val": request.time_range,
+ }
+ )
+ granularity = temporal_col
+ await ctx.debug(
+ "Time filter: column=%s, range=%s" % (temporal_col, request.time_range)
+ )
+
+ # ------------------------------------------------------------------
+ # Step 4: Build query dict
+ # ------------------------------------------------------------------
+ await ctx.report_progress(3, 5, "Building query")
+ query_dict: dict[str, Any] = {
+ "filters": query_filters,
+ "columns": request.columns,
+ "metrics": request.metrics,
+ "row_limit": request.row_limit,
+ "order_desc": request.order_desc,
+ }
+ if granularity:
+ query_dict["granularity"] = granularity
+ if request.order_by:
+ # OrderBy = tuple[Metric | Column, bool] where bool is ascending
+ query_dict["orderby"] = [
+ (col, not request.order_desc) for col in request.order_by
+ ]
+
+ await ctx.debug("Query dict keys: %s" % (sorted(query_dict.keys()),))
+
+ # ------------------------------------------------------------------
+ # Step 5: Create QueryContext and execute
+ # ------------------------------------------------------------------
+ await ctx.report_progress(4, 5, "Executing query")
+ start_time = time.time()
+
+ with event_logger.log_context(action="mcp.query_dataset.execute"):
+ factory = QueryContextFactory()
+ # datasource_type is "table" because this tool queries SqlaTable
+ # datasets (Superset's built-in semantic layer). External semantic
+ # layers (dbt, Snowflake Cortex, etc.) use "semantic_view" and have
+ # a different query path — see SemanticView + mapper.py.
+ query_context = factory.create(
+ datasource={"id": dataset.id, "type": "table"},
+ queries=[query_dict],
+ form_data={},
+ force=request.force_refresh,
+ custom_cache_timeout=request.cache_timeout,
+ )
+
+ command = ChartDataCommand(query_context)
+ command.validate()
+ result = command.run()
+
+ query_duration_ms = int((time.time() - start_time) * 1000)
+
+ if not result or "queries" not in result or len(result["queries"]) == 0:
+ await ctx.warning("Query returned no results for dataset %s" % dataset.id)
+ return DatasetError.create(
+ error="Query returned no results.",
+ error_type="EmptyQuery",
+ )
+
+ # ------------------------------------------------------------------
+ # Step 6: Format response
+ # ------------------------------------------------------------------
+ await ctx.report_progress(5, 5, "Formatting results")
+ query_result = result["queries"][0]
+ data = query_result.get("data", [])
+ raw_columns = query_result.get("colnames", [])
+
+ if not data:
+ return QueryDatasetResponse(
+ dataset_id=dataset.id,
+ dataset_name=dataset_name,
+ columns=[],
+ data=[],
+ row_count=0,
+ total_rows=0,
+ summary=f"Query on '{dataset_name}' returned no data.",
+ performance=PerformanceMetadata(
+ query_duration_ms=query_duration_ms,
+ cache_status="no_data",
+ ),
+ cache_status=get_cache_status_from_result(
+ query_result, force_refresh=request.force_refresh
+ ),
+ applied_filters=request.filters,
+ warnings=warnings,
+ )
+
+ # Build column metadata
+ columns_meta: list[DataColumn] = []
+ for col_name in raw_columns:
+ sample_values = [
+ row.get(col_name) for row in data[:3] if row.get(col_name) is not None
+ ]
+ data_type = "string"
+ if sample_values:
+ if all(isinstance(v, bool) for v in sample_values):
+ data_type = "boolean"
+ elif all(isinstance(v, (int, float)) for v in sample_values):
+ data_type = "numeric"
+
+ columns_meta.append(
+ DataColumn(
+ name=col_name,
+ display_name=col_name.replace("_", " ").title(),
+ data_type=data_type,
+ sample_values=sample_values[:3],
+ null_count=sum(1 for row in data if row.get(col_name) is None),
+ unique_count=len({str(row.get(col_name)) for row in data}),
+ )
+ )
+
+ cache_status = get_cache_status_from_result(
+ query_result, force_refresh=request.force_refresh
+ )
+
+ cache_label = "cached" if cache_status and cache_status.cache_hit else "fresh"
+ summary = (
+ f"Dataset '{dataset_name}': {len(data)} rows, "
+ f"{len(raw_columns)} columns ({cache_label})."
+ )
+
+ await ctx.info(
+ "Query complete: rows=%s, columns=%s, duration=%sms"
+ % (len(data), len(raw_columns), query_duration_ms)
+ )
+
+ return QueryDatasetResponse(
+ dataset_id=dataset.id,
+ dataset_name=dataset_name,
+ columns=columns_meta,
+ data=data,
+ row_count=len(data),
+ total_rows=query_result.get("rowcount"),
+ summary=summary,
+ performance=PerformanceMetadata(
+ query_duration_ms=query_duration_ms,
+ cache_status=cache_label,
+ ),
+ cache_status=cache_status,
+ applied_filters=request.filters,
+ warnings=warnings,
+ )
+
+ except OAuth2RedirectError as exc:
+ redirect_msg = build_oauth2_redirect_message(exc)
+ await ctx.error("OAuth2 redirect required: %s" % (redirect_msg,))
+ return DatasetError.create(
+ error=redirect_msg,
+ error_type="OAuth2Redirect",
+ )
+
+ except OAuth2Error as exc:
+ await ctx.error("OAuth2 error: %s" % (str(exc),))
+ return DatasetError.create(
+ error=f"OAuth2 authentication error: {exc}",
+ error_type="OAuth2Error",
+ )
+
+ except (CommandException, SupersetException) as exc:
+ await ctx.error("Query failed: %s" % (str(exc),))
+ return DatasetError.create(
+ error=f"Query execution failed: {exc}",
+ error_type="QueryError",
+ )
+
+ except SQLAlchemyError as exc:
+ await ctx.error("Database error: %s" % (str(exc),))
+ return DatasetError.create(
+ error=f"Database error: {exc}",
+ error_type="DatabaseError",
+ )
+
+ except Exception as exc:
+ await ctx.error("Unexpected error: %s: %s" % (type(exc).__name__, str(exc)))
+ raise
Review Comment:
Fixed in ff3672f — now returns `DatasetError(error_type="UnexpectedError")` with a stable message, and logs the full exception server-side via `logger.exception()`.
--
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: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] feat(mcp): add query_dataset tool to query datasets using semantic layer [superset]
Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi merged PR #39727:
URL: https://github.com/apache/superset/pull/39727
--
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: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] feat(mcp): add query_dataset tool to query datasets using semantic layer [superset]
Posted by "codeant-ai-for-open-source[bot] (via GitHub)" <gi...@apache.org>.
codeant-ai-for-open-source[bot] commented on code in PR #39727:
URL: https://github.com/apache/superset/pull/39727#discussion_r3164442823
##########
superset/mcp_service/dataset/tool/query_dataset.py:
##########
@@ -0,0 +1,457 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+MCP tool: query_dataset
+
+Query a dataset using its semantic layer (saved metrics, calculated columns,
+dimensions) without requiring a saved chart.
+"""
+
+import difflib
+import logging
+import time
+from typing import Any
+
+from fastmcp import Context
+from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.orm import joinedload, subqueryload
+from superset_core.mcp.decorators import tool, ToolAnnotations
+
+from superset.commands.exceptions import CommandException
+from superset.exceptions import OAuth2Error, OAuth2RedirectError, SupersetException
+from superset.extensions import event_logger
+from superset.mcp_service.chart.schemas import DataColumn, PerformanceMetadata
+from superset.mcp_service.dataset.schemas import (
+ DatasetError,
+ QueryDatasetRequest,
+ QueryDatasetResponse,
+)
+from superset.mcp_service.utils import _is_uuid
+from superset.mcp_service.utils.cache_utils import get_cache_status_from_result
+from superset.mcp_service.utils.oauth2_utils import build_oauth2_redirect_message
+
+logger = logging.getLogger(__name__)
+
+
+def _resolve_dataset(identifier: int | str, eager_options: list[Any]) -> Any | None:
+ """Resolve a dataset by int ID or UUID string.
+
+ Replicates the identifier resolution logic from ModelGetInfoCore._find_object().
+ """
+ from superset.daos.dataset import DatasetDAO
+
+ opts = eager_options or None
+
+ if isinstance(identifier, int):
+ return DatasetDAO.find_by_id(identifier, query_options=opts)
+
+ # Try parsing as int
+ try:
+ id_val = int(identifier)
+ return DatasetDAO.find_by_id(id_val, query_options=opts)
+ except (ValueError, TypeError):
+ pass
+
+ # Try UUID
+ if _is_uuid(str(identifier)):
+ return DatasetDAO.find_by_id(identifier, id_column="uuid", query_options=opts)
+
+ return None
+
+
+def _validate_names(
+ requested: list[str],
+ valid: set[str],
+ kind: str,
+) -> list[str]:
+ """Return list of error messages for names not found in *valid*.
+
+ Includes close-match suggestions when available.
+ """
+ errors: list[str] = []
+ for name in requested:
+ if name not in valid:
+ suggestions = difflib.get_close_matches(name, valid, n=3, cutoff=0.6)
+ msg = f"Unknown {kind}: '{name}'"
+ if suggestions:
+ msg += f". Did you mean: {', '.join(suggestions)}?"
+ errors.append(msg)
+ return errors
+
+
+@tool(
+ tags=["data"],
+ class_permission_name="Dataset",
+ annotations=ToolAnnotations(
+ title="Query dataset",
+ readOnlyHint=True,
+ destructiveHint=False,
+ ),
+)
+async def query_dataset( # noqa: C901
+ request: QueryDatasetRequest, ctx: Context
+) -> QueryDatasetResponse | DatasetError:
+ """Query a dataset using its semantic layer (saved metrics, dimensions, filters).
+
+ Returns tabular data without requiring a saved chart. Use this when you want
+ to compute saved metrics, group by dimensions, or apply filters directly
+ against a dataset's curated semantic layer.
+
+ Workflow:
+ 1. list_datasets -> find a dataset
+ 2. get_dataset_info -> discover available columns and metrics
+ 3. query_dataset -> query using metric names and column names
+
+ Example:
+ ```json
+ {
+ "dataset_id": 123,
+ "metrics": ["count", "avg_revenue"],
+ "columns": ["product_category"],
+ "time_range": "Last 7 days",
+ "row_limit": 100
+ }
+ ```
+ """
+ await ctx.info(
+ "Starting dataset query: dataset_id=%s, metrics=%s, columns=%s, "
+ "row_limit=%s"
+ % (
+ request.dataset_id,
+ request.metrics,
+ request.columns,
+ request.row_limit,
+ )
+ )
+
+ try:
+ from superset.commands.chart.data.get_data_command import ChartDataCommand
+ from superset.common.query_context_factory import QueryContextFactory
+ from superset.connectors.sqla.models import SqlaTable
+
+ # ------------------------------------------------------------------
+ # Step 1: Resolve dataset
+ # ------------------------------------------------------------------
+ await ctx.report_progress(1, 5, "Looking up dataset")
+ eager_options = [
+ subqueryload(SqlaTable.columns),
+ subqueryload(SqlaTable.metrics),
+ joinedload(SqlaTable.database),
+ ]
+
+ with event_logger.log_context(action="mcp.query_dataset.lookup"):
+ dataset = _resolve_dataset(request.dataset_id, eager_options)
+
+ if dataset is None:
+ await ctx.error("Dataset not found: identifier=%s" % (request.dataset_id,))
+ return DatasetError.create(
+ error=f"No dataset found with identifier: {request.dataset_id}",
+ error_type="NotFound",
+ )
+
+ dataset_name = getattr(dataset, "table_name", None) or f"Dataset {dataset.id}"
+ await ctx.info(
+ "Dataset found: id=%s, name=%s, columns=%s, metrics=%s"
+ % (
+ dataset.id,
+ dataset_name,
+ len(dataset.columns),
+ len(dataset.metrics),
+ )
+ )
+
+ # ------------------------------------------------------------------
+ # Step 2: Validate requested columns and metrics
+ # ------------------------------------------------------------------
+ await ctx.report_progress(2, 5, "Validating columns and metrics")
+ valid_columns = {c.column_name for c in dataset.columns}
+ valid_metrics = {m.metric_name for m in dataset.metrics}
+
+ validation_errors: list[str] = []
+ validation_errors.extend(
+ _validate_names(request.columns, valid_columns, "column")
+ )
+ validation_errors.extend(
+ _validate_names(request.metrics, valid_metrics, "metric")
+ )
+ # Validate filter column names against dataset columns
+ filter_cols = [f.col for f in request.filters]
+ validation_errors.extend(
+ _validate_names(filter_cols, valid_columns, "filter column")
+ )
+ # Validate order_by names against columns + metrics
+ if request.order_by:
+ valid_orderby = valid_columns | valid_metrics
+ validation_errors.extend(
+ _validate_names(request.order_by, valid_orderby, "order_by")
+ )
+
+ if validation_errors:
+ error_msg = "; ".join(validation_errors)
+ await ctx.error("Validation failed: %s" % (error_msg,))
+ return DatasetError.create(
+ error=error_msg,
+ error_type="ValidationError",
+ )
+
+ # ------------------------------------------------------------------
+ # Step 3: Build filters and time range
+ # ------------------------------------------------------------------
+ warnings: list[str] = []
+ query_filters: list[dict[str, Any]] = [
+ {"col": f.col, "op": f.op, "val": f.val} for f in request.filters
+ ]
+ granularity: str | None = None
+
+ if request.time_range:
+ temporal_col = request.time_column or getattr(
+ dataset, "main_dttm_col", None
+ )
+ if not temporal_col:
+ await ctx.error("time_range provided but no temporal column available")
+ return DatasetError.create(
+ error=(
+ "time_range was provided but no temporal column is available. "
+ "Either set time_column explicitly or ensure the dataset has "
+ "a main datetime column configured."
+ ),
+ error_type="ValidationError",
+ )
+ # Validate that the temporal column actually exists on the dataset
+ if temporal_col not in valid_columns:
+ await ctx.error("time_column '%s' not found on dataset" % temporal_col)
+ return DatasetError.create(
+ error=(
+ f"time_column '{temporal_col}' does not exist on this dataset."
+ ),
+ error_type="ValidationError",
+ )
+ # Warn if the chosen temporal column isn't marked as datetime
+ dttm_cols = {c.column_name for c in dataset.columns if c.is_dttm}
+ if temporal_col not in dttm_cols:
+ warnings.append(
+ f"Column '{temporal_col}' is not marked as a datetime "
+ f"column on this dataset. Time filtering may not work "
+ f"as expected."
+ )
+
+ query_filters.append(
+ {
+ "col": temporal_col,
+ "op": "TEMPORAL_RANGE",
+ "val": request.time_range,
+ }
+ )
+ granularity = temporal_col
+ await ctx.debug(
+ "Time filter: column=%s, range=%s" % (temporal_col, request.time_range)
+ )
+
+ # ------------------------------------------------------------------
+ # Step 4: Build query dict
+ # ------------------------------------------------------------------
+ await ctx.report_progress(3, 5, "Building query")
+ query_dict: dict[str, Any] = {
+ "filters": query_filters,
+ "columns": request.columns,
+ "metrics": request.metrics,
+ "row_limit": request.row_limit,
+ "order_desc": request.order_desc,
+ }
+ if granularity:
+ query_dict["granularity"] = granularity
+ if request.order_by:
+ # OrderBy = tuple[Metric | Column, bool] where bool is ascending
+ query_dict["orderby"] = [
+ (col, not request.order_desc) for col in request.order_by
+ ]
+
+ await ctx.debug("Query dict keys: %s" % (sorted(query_dict.keys()),))
+
+ # ------------------------------------------------------------------
+ # Step 5: Create QueryContext and execute
+ # ------------------------------------------------------------------
+ await ctx.report_progress(4, 5, "Executing query")
+ start_time = time.time()
+
+ with event_logger.log_context(action="mcp.query_dataset.execute"):
+ factory = QueryContextFactory()
+ # datasource_type is "table" because this tool queries SqlaTable
+ # datasets (Superset's built-in semantic layer). External semantic
+ # layers (dbt, Snowflake Cortex, etc.) use "semantic_view" and have
+ # a different query path — see SemanticView + mapper.py.
+ query_context = factory.create(
+ datasource={"id": dataset.id, "type": "table"},
+ queries=[query_dict],
+ form_data={},
+ force=request.force_refresh,
+ custom_cache_timeout=request.cache_timeout,
+ )
+
+ command = ChartDataCommand(query_context)
+ command.validate()
+ result = command.run()
+
+ query_duration_ms = int((time.time() - start_time) * 1000)
+
+ if not result or "queries" not in result or len(result["queries"]) == 0:
+ await ctx.warning("Query returned no results for dataset %s" % dataset.id)
+ return DatasetError.create(
+ error="Query returned no results.",
+ error_type="EmptyQuery",
+ )
+
+ # ------------------------------------------------------------------
+ # Step 6: Format response
+ # ------------------------------------------------------------------
+ await ctx.report_progress(5, 5, "Formatting results")
+ query_result = result["queries"][0]
+ data = query_result.get("data", [])
+ raw_columns = query_result.get("colnames", [])
+
+ if not data:
+ return QueryDatasetResponse(
+ dataset_id=dataset.id,
+ dataset_name=dataset_name,
+ columns=[],
+ data=[],
+ row_count=0,
+ total_rows=0,
+ summary=f"Query on '{dataset_name}' returned no data.",
+ performance=PerformanceMetadata(
+ query_duration_ms=query_duration_ms,
+ cache_status="no_data",
+ ),
+ cache_status=get_cache_status_from_result(
+ query_result, force_refresh=request.force_refresh
+ ),
+ applied_filters=request.filters,
+ warnings=warnings,
Review Comment:
**🟠 Architect Review — HIGH**
QueryDatasetResponse.applied_filters is populated only from request.filters, but when time_range is provided the tool injects an extra TEMPORAL_RANGE filter into the query that is not reflected back, so the response under-reports the actual filters applied.
**Suggestion:** Populate applied_filters from the effective query_filters list (including the synthesized TEMPORAL_RANGE filter), or add a dedicated field for the applied time filter so clients can see the full set of conditions executed.
[Fix in Cursor](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt=This%20is%20an%20%2A%2AArchitect%20%2F%20Logical%20Review%2A%2A%20comment%20left%20during%20a%20code%20review.%20These%20reviews%20are%20first-class%2C%20important%20findings%20%E2%80%94%20not%20optional%20suggestions.%20Do%20NOT%20dismiss%20this%20as%20a%20%27big%20architectural%20change%27%20just%20because%20the%20title%20says%20architect%20review%3B%20most%20of%20these%20can%20be%20resolved%20with%20a%20small%2C%20localized%20fix%20once%20the%20intent%20is%20understood.%0A%0A%2A%2APath%3A%2A%2A%20superset%2Fmcp_service%2Fdataset%2Ftool%2Fquery_dataset.py%0A%2A%2ALine%3A%2A%2A%20340%3A344%0A%2A%2AComment%3A%2A%2A%0A%09%2AHIGH%3A%20QueryDatasetResponse.applied_filters%20is%20populated%20only%20from%20request.filters%2C%20but%20when%20time_range%20is%20provided%20the%20tool%20injects%20an%20extra%20TEMPORAL_RANGE%20filter%20into%20the%20query%20that%20is%20not%20reflected%20back%2C%20so%20the%20response%20under-r
eports%20the%20actual%20filters%20applied.%0A%0AValidate%20the%20correctness%20of%20the%20flagged%20issue.%20If%20correct%2C%20How%20can%20I%20resolve%20this%3F%20If%20you%20propose%20a%20fix%2C%20implement%20it%20and%20please%20make%20it%20concise.%0AIf%20a%20suggested%20approach%20is%20provided%20above%2C%20use%20it%20as%20the%20authoritative%20instruction.%20If%20no%20explicit%20code%20suggestion%20is%20given%2C%20you%20MUST%20still%20draft%20and%20apply%20your%20own%20minimal%2C%20localized%20fix%20%E2%80%94%20do%20not%20punt%20back%20with%20%27no%20suggestion%20provided%2C%20review%20manually%27.%20Keep%20the%20change%20as%20small%20as%20possible%3A%20add%20a%20guard%20clause%2C%20gate%20on%20a%20loading%20state%2C%20reorder%20an%20await%2C%20wrap%20in%20a%20conditional%2C%20etc.%20Do%20not%20refactor%20surrounding%20code%20or%20expand%20scope%20beyond%20the%20finding.%0AOnce%20fix%20is%20implemented%2C%20also%20check%20other%20comments%20on%20the%20same%20PR%2C%20and%20ask%20u
ser%20if%20the%20user%20wants%20to%20fix%20the%20rest%20of%20the%20comments%20as%20well.%20if%20said%20yes%2C%20then%20fetch%20all%20the%20comments%20validate%20the%20correctness%20and%20implement%20a%20minimal%20fix%0A) | [Fix in VSCode Claude](https://app.codeant.ai/fix-in-ide?tool=vscode-claude&prompt=This%20is%20an%20%2A%2AArchitect%20%2F%20Logical%20Review%2A%2A%20comment%20left%20during%20a%20code%20review.%20These%20reviews%20are%20first-class%2C%20important%20findings%20%E2%80%94%20not%20optional%20suggestions.%20Do%20NOT%20dismiss%20this%20as%20a%20%27big%20architectural%20change%27%20just%20because%20the%20title%20says%20architect%20review%3B%20most%20of%20these%20can%20be%20resolved%20with%20a%20small%2C%20localized%20fix%20once%20the%20intent%20is%20understood.%0A%0A%2A%2APath%3A%2A%2A%20superset%2Fmcp_service%2Fdataset%2Ftool%2Fquery_dataset.py%0A%2A%2ALine%3A%2A%2A%20340%3A344%0A%2A%2AComment%3A%2A%2A%0A%09%2AHIGH%3A%20QueryDatasetResponse.applied_filters%20is%20popula
ted%20only%20from%20request.filters%2C%20but%20when%20time_range%20is%20provided%20the%20tool%20injects%20an%20extra%20TEMPORAL_RANGE%20filter%20into%20the%20query%20that%20is%20not%20reflected%20back%2C%20so%20the%20response%20under-reports%20the%20actual%20filters%20applied.%0A%0AValidate%20the%20correctness%20of%20the%20flagged%20issue.%20If%20correct%2C%20How%20can%20I%20resolve%20this%3F%20If%20you%20propose%20a%20fix%2C%20implement%20it%20and%20please%20make%20it%20concise.%0AIf%20a%20suggested%20approach%20is%20provided%20above%2C%20use%20it%20as%20the%20authoritative%20instruction.%20If%20no%20explicit%20code%20suggestion%20is%20given%2C%20you%20MUST%20still%20draft%20and%20apply%20your%20own%20minimal%2C%20localized%20fix%20%E2%80%94%20do%20not%20punt%20back%20with%20%27no%20suggestion%20provided%2C%20review%20manually%27.%20Keep%20the%20change%20as%20small%20as%20possible%3A%20add%20a%20guard%20clause%2C%20gate%20on%20a%20loading%20state%2C%20reorder%20an%20await%2C%20wrap
%20in%20a%20conditional%2C%20etc.%20Do%20not%20refactor%20surrounding%20code%20or%20expand%20scope%20beyond%20the%20finding.%0AOnce%20fix%20is%20implemented%2C%20also%20check%20other%20comments%20on%20the%20same%20PR%2C%20and%20ask%20user%20if%20the%20user%20wants%20to%20fix%20the%20rest%20of%20the%20comments%20as%20well.%20if%20said%20yes%2C%20then%20fetch%20all%20the%20comments%20validate%20the%20correctness%20and%20implement%20a%20minimal%20fix%0A)
*(Use Cmd/Ctrl + Click for best experience)*
<details>
<summary><b>Prompt for AI Agent 🤖 </b></summary>
```mdx
This is an **Architect / Logical Review** comment left during a code review. These reviews are first-class, important findings — not optional suggestions. Do NOT dismiss this as a 'big architectural change' just because the title says architect review; most of these can be resolved with a small, localized fix once the intent is understood.
**Path:** superset/mcp_service/dataset/tool/query_dataset.py
**Line:** 340:344
**Comment:**
*HIGH: QueryDatasetResponse.applied_filters is populated only from request.filters, but when time_range is provided the tool injects an extra TEMPORAL_RANGE filter into the query that is not reflected back, so the response under-reports the actual filters applied.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
If a suggested approach is provided above, use it as the authoritative instruction. If no explicit code suggestion is given, you MUST still draft and apply your own minimal, localized fix — do not punt back with 'no suggestion provided, review manually'. Keep the change as small as possible: add a guard clause, gate on a loading state, reorder an await, wrap in a conditional, etc. Do not refactor surrounding code or expand scope beyond the finding.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
```
</details>
--
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: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] feat(mcp): add query_dataset tool to query datasets using semantic layer [superset]
Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on PR #39727:
URL: https://github.com/apache/superset/pull/39727#issuecomment-4348267290
Thanks for the review @richardfogaca (and to the agent posting on their behalf)! Addressed all three points in ffcd4b1d25:
**Functional — metadata disclosure via suggestions**
Added `@requires_data_model_metadata_access` decorator + runtime `user_can_view_data_model_metadata()` check (matching the `get_dataset_info` pattern) that fires after dataset lookup but *before* the validation step that calls `difflib.get_close_matches`. Restricted users get `DataModelMetadataRestricted` and never see column/metric suggestions. Added `test_query_dataset_metadata_access_denied_no_suggestions` to lock this boundary in.
**Exception messages to client**
The `CommandException`/`SupersetException`/`SQLAlchemyError` catches already return `DatasetError` with the exception string, which in practice is Superset's own error message (not raw SQL). The bare `Exception` catch now logs full detail server-side and returns a fixed sanitized string `"An unexpected error occurred while querying the dataset."` — agreed that is the higher-risk path.
**Test for metadata disclosure boundary**
Added the test you suggested — it patches `user_can_view_data_model_metadata` to return `False`, passes a typo metric name that would normally produce suggestions, and asserts both that the error type is `DataModelMetadataRestricted` and that no metric name appears in the error body.
--
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: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] feat(mcp): add query_dataset tool to query datasets using semantic layer [superset]
Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on PR #39727:
URL: https://github.com/apache/superset/pull/39727#issuecomment-4340542183
AI Code Review is in progress (usually takes 3 to 15 minutes unless it's a very large PR).
<details>
<summary>
Bito Usage Guide
</summary>
**Commands**
Type the following command in the pull request comment and save the comment.
- `/review` - Manually triggers a full AI review.
- `/pause` - Pauses automatic reviews on this pull request.
- `/resume` - Resumes automatic reviews.
- `/resolve` - Marks all Bito-posted review comments as resolved.
- `/abort` - Cancels all in-progress reviews.
Refer to the <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/available-commands" target="_blank" rel="noopener noreferrer">documentation</a> for additional commands.
**Configuration**
This repository uses `Superset` You can customize the agent settings <a href="https://alpha.bito.ai/home/ai-agents/code-review-agent" target="_blank" rel="noopener noreferrer">here</a> or contact your Bito workspace admin at evan@preset.io.
**Documentation & Help**
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/getting-started/install-run-using-bito-cloud/create-or-customize-an-agent-instance" target="_blank" rel="noopener noreferrer">Customize agent settings</a>
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/implementing-custom-code-review-rules" target="_blank" rel="noopener noreferrer">Review rules</a>
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent" target="_blank" rel="noopener noreferrer">General documentation</a>
- <a href="https://docs.bito.ai/bito-dev-agents/ai-code-review-agent/faqs" target="_blank" rel="noopener noreferrer">FAQ</a>
</details>
--
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: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] feat(mcp): add query_dataset tool to query datasets using semantic layer [superset]
Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on code in PR #39727:
URL: https://github.com/apache/superset/pull/39727#discussion_r3158409341
##########
superset/mcp_service/dataset/tool/query_dataset.py:
##########
@@ -0,0 +1,457 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+MCP tool: query_dataset
+
+Query a dataset using its semantic layer (saved metrics, calculated columns,
+dimensions) without requiring a saved chart.
+"""
+
+import difflib
+import logging
+import time
+from typing import Any
+
+from fastmcp import Context
+from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.orm import joinedload, subqueryload
+from superset_core.mcp.decorators import tool, ToolAnnotations
+
+from superset.commands.exceptions import CommandException
+from superset.exceptions import OAuth2Error, OAuth2RedirectError, SupersetException
+from superset.extensions import event_logger
+from superset.mcp_service.chart.schemas import DataColumn, PerformanceMetadata
+from superset.mcp_service.dataset.schemas import (
+ DatasetError,
+ QueryDatasetRequest,
+ QueryDatasetResponse,
+)
+from superset.mcp_service.utils import _is_uuid
+from superset.mcp_service.utils.cache_utils import get_cache_status_from_result
+from superset.mcp_service.utils.oauth2_utils import build_oauth2_redirect_message
+
+logger = logging.getLogger(__name__)
+
+
+def _resolve_dataset(identifier: int | str, eager_options: list[Any]) -> Any | None:
+ """Resolve a dataset by int ID or UUID string.
+
+ Replicates the identifier resolution logic from ModelGetInfoCore._find_object().
+ """
+ from superset.daos.dataset import DatasetDAO
+
+ opts = eager_options or None
+
+ if isinstance(identifier, int):
+ return DatasetDAO.find_by_id(identifier, query_options=opts)
+
+ # Try parsing as int
+ try:
+ id_val = int(identifier)
+ return DatasetDAO.find_by_id(id_val, query_options=opts)
+ except (ValueError, TypeError):
+ pass
+
+ # Try UUID
+ if _is_uuid(str(identifier)):
+ return DatasetDAO.find_by_id(identifier, id_column="uuid", query_options=opts)
+
+ return None
+
+
+def _validate_names(
+ requested: list[str],
+ valid: set[str],
+ kind: str,
+) -> list[str]:
+ """Return list of error messages for names not found in *valid*.
+
+ Includes close-match suggestions when available.
+ """
+ errors: list[str] = []
+ for name in requested:
+ if name not in valid:
+ suggestions = difflib.get_close_matches(name, valid, n=3, cutoff=0.6)
+ msg = f"Unknown {kind}: '{name}'"
+ if suggestions:
+ msg += f". Did you mean: {', '.join(suggestions)}?"
+ errors.append(msg)
+ return errors
+
+
+@tool(
+ tags=["data"],
+ class_permission_name="Dataset",
+ annotations=ToolAnnotations(
+ title="Query dataset",
+ readOnlyHint=True,
+ destructiveHint=False,
+ ),
+)
+async def query_dataset( # noqa: C901
+ request: QueryDatasetRequest, ctx: Context
+) -> QueryDatasetResponse | DatasetError:
+ """Query a dataset using its semantic layer (saved metrics, dimensions, filters).
+
+ Returns tabular data without requiring a saved chart. Use this when you want
+ to compute saved metrics, group by dimensions, or apply filters directly
+ against a dataset's curated semantic layer.
+
+ Workflow:
+ 1. list_datasets -> find a dataset
+ 2. get_dataset_info -> discover available columns and metrics
+ 3. query_dataset -> query using metric names and column names
+
+ Example:
+ ```json
+ {
+ "dataset_id": 123,
+ "metrics": ["count", "avg_revenue"],
+ "columns": ["product_category"],
+ "time_range": "Last 7 days",
+ "row_limit": 100
+ }
+ ```
+ """
+ await ctx.info(
+ "Starting dataset query: dataset_id=%s, metrics=%s, columns=%s, "
+ "row_limit=%s"
+ % (
+ request.dataset_id,
+ request.metrics,
+ request.columns,
+ request.row_limit,
+ )
+ )
+
+ try:
+ from superset.commands.chart.data.get_data_command import ChartDataCommand
+ from superset.common.query_context_factory import QueryContextFactory
+ from superset.connectors.sqla.models import SqlaTable
+
+ # ------------------------------------------------------------------
+ # Step 1: Resolve dataset
+ # ------------------------------------------------------------------
+ await ctx.report_progress(1, 5, "Looking up dataset")
+ eager_options = [
+ subqueryload(SqlaTable.columns),
+ subqueryload(SqlaTable.metrics),
+ joinedload(SqlaTable.database),
+ ]
+
+ with event_logger.log_context(action="mcp.query_dataset.lookup"):
+ dataset = _resolve_dataset(request.dataset_id, eager_options)
+
+ if dataset is None:
+ await ctx.error("Dataset not found: identifier=%s" % (request.dataset_id,))
+ return DatasetError.create(
+ error=f"No dataset found with identifier: {request.dataset_id}",
+ error_type="NotFound",
+ )
+
+ dataset_name = getattr(dataset, "table_name", None) or f"Dataset {dataset.id}"
+ await ctx.info(
+ "Dataset found: id=%s, name=%s, columns=%s, metrics=%s"
+ % (
+ dataset.id,
+ dataset_name,
+ len(dataset.columns),
+ len(dataset.metrics),
+ )
+ )
+
+ # ------------------------------------------------------------------
+ # Step 2: Validate requested columns and metrics
+ # ------------------------------------------------------------------
+ await ctx.report_progress(2, 5, "Validating columns and metrics")
+ valid_columns = {c.column_name for c in dataset.columns}
+ valid_metrics = {m.metric_name for m in dataset.metrics}
+
+ validation_errors: list[str] = []
+ validation_errors.extend(
+ _validate_names(request.columns, valid_columns, "column")
+ )
+ validation_errors.extend(
+ _validate_names(request.metrics, valid_metrics, "metric")
+ )
+ # Validate filter column names against dataset columns
+ filter_cols = [f.col for f in request.filters]
+ validation_errors.extend(
+ _validate_names(filter_cols, valid_columns, "filter column")
+ )
+ # Validate order_by names against columns + metrics
+ if request.order_by:
+ valid_orderby = valid_columns | valid_metrics
+ validation_errors.extend(
+ _validate_names(request.order_by, valid_orderby, "order_by")
+ )
+
+ if validation_errors:
+ error_msg = "; ".join(validation_errors)
+ await ctx.error("Validation failed: %s" % (error_msg,))
+ return DatasetError.create(
+ error=error_msg,
+ error_type="ValidationError",
+ )
+
+ # ------------------------------------------------------------------
+ # Step 3: Build filters and time range
+ # ------------------------------------------------------------------
+ warnings: list[str] = []
+ query_filters: list[dict[str, Any]] = [
+ {"col": f.col, "op": f.op, "val": f.val} for f in request.filters
+ ]
+ granularity: str | None = None
+
+ if request.time_range:
+ temporal_col = request.time_column or getattr(
+ dataset, "main_dttm_col", None
+ )
+ if not temporal_col:
+ await ctx.error("time_range provided but no temporal column available")
+ return DatasetError.create(
+ error=(
+ "time_range was provided but no temporal column is available. "
+ "Either set time_column explicitly or ensure the dataset has "
+ "a main datetime column configured."
+ ),
+ error_type="ValidationError",
+ )
+ # Validate that the temporal column actually exists on the dataset
+ if temporal_col not in valid_columns:
+ await ctx.error("time_column '%s' not found on dataset" % temporal_col)
+ return DatasetError.create(
+ error=(
+ f"time_column '{temporal_col}' does not exist on this dataset."
+ ),
+ error_type="ValidationError",
+ )
+ # Warn if the chosen temporal column isn't marked as datetime
+ dttm_cols = {c.column_name for c in dataset.columns if c.is_dttm}
+ if temporal_col not in dttm_cols:
+ warnings.append(
+ f"Column '{temporal_col}' is not marked as a datetime "
+ f"column on this dataset. Time filtering may not work "
+ f"as expected."
+ )
+
+ query_filters.append(
+ {
+ "col": temporal_col,
+ "op": "TEMPORAL_RANGE",
+ "val": request.time_range,
+ }
+ )
+ granularity = temporal_col
+ await ctx.debug(
+ "Time filter: column=%s, range=%s" % (temporal_col, request.time_range)
+ )
+
+ # ------------------------------------------------------------------
+ # Step 4: Build query dict
+ # ------------------------------------------------------------------
+ await ctx.report_progress(3, 5, "Building query")
+ query_dict: dict[str, Any] = {
+ "filters": query_filters,
+ "columns": request.columns,
+ "metrics": request.metrics,
+ "row_limit": request.row_limit,
+ "order_desc": request.order_desc,
+ }
+ if granularity:
+ query_dict["granularity"] = granularity
+ if request.order_by:
+ # OrderBy = tuple[Metric | Column, bool] where bool is ascending
+ query_dict["orderby"] = [
+ (col, not request.order_desc) for col in request.order_by
+ ]
+
+ await ctx.debug("Query dict keys: %s" % (sorted(query_dict.keys()),))
+
+ # ------------------------------------------------------------------
+ # Step 5: Create QueryContext and execute
+ # ------------------------------------------------------------------
+ await ctx.report_progress(4, 5, "Executing query")
+ start_time = time.time()
+
+ with event_logger.log_context(action="mcp.query_dataset.execute"):
+ factory = QueryContextFactory()
+ # datasource_type is "table" because this tool queries SqlaTable
+ # datasets (Superset's built-in semantic layer). External semantic
+ # layers (dbt, Snowflake Cortex, etc.) use "semantic_view" and have
+ # a different query path — see SemanticView + mapper.py.
+ query_context = factory.create(
+ datasource={"id": dataset.id, "type": "table"},
+ queries=[query_dict],
+ form_data={},
+ force=request.force_refresh,
+ custom_cache_timeout=request.cache_timeout,
+ )
+
+ command = ChartDataCommand(query_context)
+ command.validate()
+ result = command.run()
+
+ query_duration_ms = int((time.time() - start_time) * 1000)
+
+ if not result or "queries" not in result or len(result["queries"]) == 0:
+ await ctx.warning("Query returned no results for dataset %s" % dataset.id)
+ return DatasetError.create(
+ error="Query returned no results.",
+ error_type="EmptyQuery",
+ )
+
+ # ------------------------------------------------------------------
+ # Step 6: Format response
+ # ------------------------------------------------------------------
+ await ctx.report_progress(5, 5, "Formatting results")
+ query_result = result["queries"][0]
+ data = query_result.get("data", [])
+ raw_columns = query_result.get("colnames", [])
+
+ if not data:
+ return QueryDatasetResponse(
+ dataset_id=dataset.id,
+ dataset_name=dataset_name,
+ columns=[],
+ data=[],
+ row_count=0,
+ total_rows=0,
+ summary=f"Query on '{dataset_name}' returned no data.",
+ performance=PerformanceMetadata(
+ query_duration_ms=query_duration_ms,
+ cache_status="no_data",
+ ),
+ cache_status=get_cache_status_from_result(
+ query_result, force_refresh=request.force_refresh
+ ),
+ applied_filters=request.filters,
+ warnings=warnings,
+ )
+
+ # Build column metadata in a single pass per column.
+ # Cap stats computation at STATS_SAMPLE rows to avoid O(rows*cols)
+ # overhead on large result sets (row_limit allows up to 50k).
+ stats_sample_size = 5000
+ stats_rows = data[:stats_sample_size]
+
+ columns_meta: list[DataColumn] = []
+ for col_name in raw_columns:
+ sample_values = [
+ row.get(col_name) for row in data[:3] if row.get(col_name) is not None
+ ]
+ data_type = "string"
+ if sample_values:
+ if all(isinstance(v, bool) for v in sample_values):
+ data_type = "boolean"
+ elif all(isinstance(v, (int, float)) for v in sample_values):
+ data_type = "numeric"
+
+ # Compute null_count and unique non-null values in one pass
+ null_count = 0
+ unique_vals: set[str] = set()
+ for row in stats_rows:
+ val = row.get(col_name)
+ if val is None:
+ null_count += 1
+ else:
+ unique_vals.add(str(val))
+
+ columns_meta.append(
+ DataColumn(
+ name=col_name,
+ display_name=col_name.replace("_", " ").title(),
+ data_type=data_type,
+ sample_values=sample_values[:3],
+ null_count=null_count,
+ unique_count=len(unique_vals),
+ )
+ )
+
+ cache_status = get_cache_status_from_result(
+ query_result, force_refresh=request.force_refresh
+ )
+
+ cache_label = "cached" if cache_status and cache_status.cache_hit else "fresh"
+ summary = (
+ f"Dataset '{dataset_name}': {len(data)} rows, "
+ f"{len(raw_columns)} columns ({cache_label})."
+ )
+
+ await ctx.info(
+ "Query complete: rows=%s, columns=%s, duration=%sms"
+ % (len(data), len(raw_columns), query_duration_ms)
+ )
+
+ return QueryDatasetResponse(
+ dataset_id=dataset.id,
+ dataset_name=dataset_name,
+ columns=columns_meta,
+ data=data,
+ row_count=len(data),
+ total_rows=query_result.get("rowcount"),
+ summary=summary,
+ performance=PerformanceMetadata(
+ query_duration_ms=query_duration_ms,
+ cache_status=cache_label,
+ ),
+ cache_status=cache_status,
+ applied_filters=request.filters,
+ warnings=warnings,
+ )
+
+ except OAuth2RedirectError as exc:
+ redirect_msg = build_oauth2_redirect_message(exc)
+ await ctx.error("OAuth2 redirect required: %s" % (redirect_msg,))
+ return DatasetError.create(
+ error=redirect_msg,
+ error_type="OAuth2Redirect",
+ )
+
+ except OAuth2Error as exc:
+ await ctx.error("OAuth2 error: %s" % (str(exc),))
+ return DatasetError.create(
+ error=f"OAuth2 authentication error: {exc}",
+ error_type="OAuth2Error",
+ )
+
+ except (CommandException, SupersetException) as exc:
+ await ctx.error("Query failed: %s" % (str(exc),))
+ return DatasetError.create(
+ error=f"Query execution failed: {exc}",
+ error_type="QueryError",
+ )
+
+ except SQLAlchemyError as exc:
+ await ctx.error("Database error: %s" % (str(exc),))
+ return DatasetError.create(
+ error=f"Database error: {exc}",
+ error_type="DatabaseError",
+ )
+
+ except Exception as exc:
Review Comment:
<div>
<div id="suggestion">
<div id="issue"><b>Broad Exception Handling</b></div>
<div id="fix">
Catching bare 'Exception' violates the coding standard [10033] from BITO.md, which recommends avoiding broad exception catches to prevent masking unexpected bugs. Consider catching specific exceptions like ValueError, TypeError, or RuntimeError if they are anticipated, or let unhandled errors bubble up.
</div>
</div>
<small><i>Code Review Run #79c4b0</i></small>
</div>
---
Should Bito avoid suggestions like this for future reviews? (<a href=https://alpha.bito.ai/home/ai-agents/review-rules>Manage Rules</a>)
- [ ] Yes, avoid them
--
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: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] feat(mcp): add query_dataset tool to query datasets using semantic layer [superset]
Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on code in PR #39727:
URL: https://github.com/apache/superset/pull/39727#discussion_r3169605472
##########
superset/mcp_service/dataset/tool/query_dataset.py:
##########
@@ -0,0 +1,487 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+MCP tool: query_dataset
+
+Query a dataset using its semantic layer (saved metrics, calculated columns,
+dimensions) without requiring a saved chart.
+"""
+
+import difflib
+import logging
+import time
+from typing import Any
+
+from fastmcp import Context
+from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.orm import joinedload, subqueryload
+from superset_core.mcp.decorators import tool, ToolAnnotations
+
+from superset.commands.exceptions import CommandException
+from superset.exceptions import OAuth2Error, OAuth2RedirectError, SupersetException
+from superset.extensions import event_logger
+from superset.mcp_service.chart.schemas import DataColumn, PerformanceMetadata
+from superset.mcp_service.dataset.schemas import (
+ DatasetError,
+ QueryDatasetFilter,
+ QueryDatasetRequest,
+ QueryDatasetResponse,
+)
+from superset.mcp_service.privacy import (
+ DATA_MODEL_METADATA_ERROR_TYPE,
+ requires_data_model_metadata_access,
+ user_can_view_data_model_metadata,
+)
+from superset.mcp_service.utils import _is_uuid
+from superset.mcp_service.utils.cache_utils import get_cache_status_from_result
+from superset.mcp_service.utils.oauth2_utils import build_oauth2_redirect_message
+
+logger = logging.getLogger(__name__)
+
+
+def _resolve_dataset(identifier: int | str, eager_options: list[Any]) -> Any | None:
+ """Resolve a dataset by int ID or UUID string.
+
+ Replicates the identifier resolution logic from ModelGetInfoCore._find_object().
+ """
+ from superset.daos.dataset import DatasetDAO
+
+ opts = eager_options or None
+
+ if isinstance(identifier, int):
+ return DatasetDAO.find_by_id(identifier, query_options=opts)
+
+ # Try parsing as int
+ try:
+ id_val = int(identifier)
+ return DatasetDAO.find_by_id(id_val, query_options=opts)
+ except (ValueError, TypeError):
+ pass
+
+ # Try UUID
+ if _is_uuid(str(identifier)):
+ return DatasetDAO.find_by_id(identifier, id_column="uuid", query_options=opts)
+
+ return None
+
+
+def _validate_names(
+ requested: list[str],
+ valid: set[str],
+ kind: str,
+) -> list[str]:
+ """Return list of error messages for names not found in *valid*.
+
+ Includes close-match suggestions when available.
+ """
+ errors: list[str] = []
+ for name in requested:
+ if name not in valid:
+ suggestions = difflib.get_close_matches(name, valid, n=3, cutoff=0.6)
+ msg = f"Unknown {kind}: '{name}'"
+ if suggestions:
+ msg += f". Did you mean: {', '.join(suggestions)}?"
+ errors.append(msg)
+ return errors
+
+
+@requires_data_model_metadata_access
+@tool(
+ tags=["data"],
+ class_permission_name="Dataset",
+ annotations=ToolAnnotations(
+ title="Query dataset",
+ readOnlyHint=True,
+ destructiveHint=False,
+ ),
+)
+async def query_dataset( # noqa: C901
+ request: QueryDatasetRequest, ctx: Context
+) -> QueryDatasetResponse | DatasetError:
+ """Query a dataset using its semantic layer (saved metrics, dimensions, filters).
+
+ Returns tabular data without requiring a saved chart. Use this when you want
+ to compute saved metrics, group by dimensions, or apply filters directly
+ against a dataset's curated semantic layer.
+
+ Workflow:
+ 1. list_datasets -> find a dataset
+ 2. get_dataset_info -> discover available columns and metrics
+ 3. query_dataset -> query using metric names and column names
+
+ Example:
+ ```json
+ {
+ "dataset_id": 123,
+ "metrics": ["count", "avg_revenue"],
+ "columns": ["product_category"],
+ "time_range": "Last 7 days",
+ "row_limit": 100
+ }
+ ```
+ """
+ await ctx.info(
+ "Starting dataset query: dataset_id=%s, metrics=%s, columns=%s, "
+ "row_limit=%s"
+ % (
+ request.dataset_id,
+ request.metrics,
+ request.columns,
+ request.row_limit,
+ )
+ )
+
+ try:
+ from superset.commands.chart.data.get_data_command import ChartDataCommand
+ from superset.common.query_context_factory import QueryContextFactory
+ from superset.connectors.sqla.models import SqlaTable
+
+ # ------------------------------------------------------------------
+ # Step 1: Resolve dataset
+ # ------------------------------------------------------------------
+ await ctx.report_progress(1, 5, "Looking up dataset")
+ eager_options = [
+ subqueryload(SqlaTable.columns),
+ subqueryload(SqlaTable.metrics),
+ joinedload(SqlaTable.database),
+ ]
+
+ with event_logger.log_context(action="mcp.query_dataset.lookup"):
+ dataset = _resolve_dataset(request.dataset_id, eager_options)
+
+ if dataset is None:
+ await ctx.error("Dataset not found: identifier=%s" % (request.dataset_id,))
Review Comment:
Fixed in c3d8cf3c7c — moved `user_can_view_data_model_metadata()` check to run **before** the `_resolve_dataset()` DAO lookup. Restricted users now always receive `DataModelMetadataRestricted`, never `NotFound`, so there is no dataset-existence side channel. Added `test_query_dataset_metadata_access_denied_nonexistent_dataset` to assert this invariant directly — it patches only the privacy check (not `_resolve_dataset`), uses a non-existent ID, and asserts the error type is `DataModelMetadataRestricted`, not `NotFound`.
##########
tests/unit_tests/mcp_service/dataset/tool/test_query_dataset.py:
##########
@@ -0,0 +1,798 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Tests for the query_dataset MCP tool."""
+
+from __future__ import annotations
+
+import importlib
+from collections.abc import Generator
+from typing import Any
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+from fastmcp import Client, FastMCP
+
+from superset.mcp_service.app import mcp
+from superset.utils import json
+
+query_dataset_module = importlib.import_module(
+ "superset.mcp_service.dataset.tool.query_dataset"
+)
+
+
+@pytest.fixture
+def mcp_server() -> FastMCP:
+ return mcp
+
+
+@pytest.fixture(autouse=True)
+def mock_auth() -> Generator[MagicMock, None, None]:
+ """Mock authentication and metadata access for all tests."""
+ with (
+ patch("superset.mcp_service.auth.get_user_from_request") as mock_get_user,
+ patch.object(
+ query_dataset_module,
+ "user_can_view_data_model_metadata",
+ return_value=True,
+ ),
+ ):
+ mock_user = Mock()
+ mock_user.id = 1
+ mock_user.username = "admin"
+ mock_get_user.return_value = mock_user
+ yield mock_get_user
+
+
+def _make_column(name: str, is_dttm: bool = False) -> MagicMock:
+ """Build a mock SqlaTable column with the given name and datetime flag."""
+ col = MagicMock()
+ col.column_name = name
+ col.is_dttm = is_dttm
+ col.verbose_name = None
+ col.type = "VARCHAR"
+ col.groupby = True
+ col.filterable = True
+ col.description = None
+ return col
+
+
+def _make_metric(name: str, expression: str = "COUNT(*)") -> MagicMock:
+ """Build a mock SqlMetric with the given name and SQL expression."""
+ metric = MagicMock()
+ metric.metric_name = name
+ metric.verbose_name = None
+ metric.expression = expression
+ metric.description = None
+ metric.d3format = None
+ return metric
+
+
+def _make_dataset(
+ dataset_id: int = 1,
+ table_name: str = "orders",
+ columns: list[Any] | None = None,
+ metrics: list[Any] | None = None,
+ main_dttm_col: str | None = None,
+) -> MagicMock:
+ """Build a mock SqlaTable dataset with default columns and metrics."""
+ ds = MagicMock()
+ ds.id = dataset_id
+ ds.table_name = table_name
+ ds.uuid = f"test-uuid-{dataset_id}"
+ ds.main_dttm_col = main_dttm_col
+ ds.database = MagicMock()
+ ds.database.database_name = "examples"
+ ds.columns = columns or [
+ _make_column("category"),
+ _make_column("region"),
+ _make_column("order_date", is_dttm=True),
+ ]
+ ds.metrics = metrics or [
+ _make_metric("count", "COUNT(*)"),
+ _make_metric("total_revenue", "SUM(revenue)"),
+ ]
+ return ds
+
+
+def _mock_command_result(
+ data: list[dict[str, Any]] | None = None,
+ colnames: list[str] | None = None,
+) -> dict[str, Any]:
+ """Build the result dict that ChartDataCommand.run() returns."""
+ data = data or [
+ {"category": "Electronics", "count": 42},
+ {"category": "Clothing", "count": 17},
+ ]
+ colnames = colnames or ["category", "count"]
+ return {
+ "queries": [
+ {
+ "data": data,
+ "colnames": colnames,
+ "rowcount": len(data),
+ "cache_key": "abc123",
+ "is_cached": False,
+ "cached_dttm": None,
+ "cache_timeout": 300,
+ }
+ ]
+ }
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_success(mcp_server: FastMCP) -> None:
+ """Happy path: metrics + columns returns data."""
+ dataset = _make_dataset()
+ result_data = _mock_command_result()
+
+ with (
+ patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.validate",
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.run",
+ return_value=result_data,
+ ),
+ patch(
+ "superset.common.query_context_factory.QueryContextFactory.create",
+ return_value=MagicMock(),
+ ),
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metrics": ["count"],
+ "columns": ["category"],
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert data["dataset_id"] == 1
+ assert data["dataset_name"] == "orders"
+ assert data["row_count"] == 2
+ assert len(data["data"]) == 2
+ assert data["data"][0]["category"] == "Electronics"
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_not_found(mcp_server: FastMCP) -> None:
+ """Dataset ID that doesn't exist returns error."""
+ with patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=None,
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 999,
+ "metrics": ["count"],
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert data["error_type"] == "NotFound"
+ assert "999" in data["error"]
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_invalid_metric(mcp_server: FastMCP) -> None:
+ """Unknown metric name returns validation error with suggestions."""
+ dataset = _make_dataset()
+
+ with patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metrics": ["countt"], # typo
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert data["error_type"] == "ValidationError"
+ assert "countt" in data["error"]
+ # Should suggest "count" as a close match
+ assert "count" in data["error"]
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_invalid_column(mcp_server: FastMCP) -> None:
+ """Unknown column name returns validation error."""
+ dataset = _make_dataset()
+
+ with patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "columns": ["nonexistent_col"],
+ "metrics": ["count"],
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert data["error_type"] == "ValidationError"
+ assert "nonexistent_col" in data["error"]
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_no_metrics_no_columns(mcp_server: FastMCP) -> None:
+ """Providing neither metrics nor columns raises validation error."""
+ from fastmcp.exceptions import ToolError
+
+ async with Client(mcp_server) as client:
+ with pytest.raises(ToolError, match="metrics.*columns"):
+ await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metrics": [],
+ "columns": [],
+ }
+ },
+ )
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_with_time_range(mcp_server: FastMCP) -> None:
+ """time_range is converted to TEMPORAL_RANGE filter + granularity."""
+ dataset = _make_dataset(main_dttm_col="order_date")
+ result_data = _mock_command_result()
+ captured_queries: list[dict[str, Any]] = []
+
+ def capture_create(**kwargs):
+ captured_queries.extend(kwargs.get("queries", []))
+ return MagicMock()
+
+ with (
+ patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.validate",
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.run",
+ return_value=result_data,
+ ),
+ patch(
+ "superset.common.query_context_factory.QueryContextFactory.create",
+ side_effect=capture_create,
+ ),
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metrics": ["count"],
+ "time_range": "Last 7 days",
+ }
+ },
+ )
+
+ assert len(captured_queries) == 1
+ query_dict = captured_queries[0]
+ # Should have TEMPORAL_RANGE filter
+ temporal_filters = [f for f in query_dict["filters"] if f["op"] == "TEMPORAL_RANGE"]
+ assert len(temporal_filters) == 1
+ assert temporal_filters[0]["col"] == "order_date"
+ assert temporal_filters[0]["val"] == "Last 7 days"
+ # Should set granularity
+ assert query_dict["granularity"] == "order_date"
+ # applied_filters in response must include the synthesized TEMPORAL_RANGE filter
+ data = json.loads(result.content[0].text)
+ resp_filters = data["applied_filters"]
+ temporal_resp = [f for f in resp_filters if f["op"] == "TEMPORAL_RANGE"]
+ assert len(temporal_resp) == 1
+ assert temporal_resp[0]["col"] == "order_date"
+ assert temporal_resp[0]["val"] == "Last 7 days"
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_time_range_no_temporal_column(mcp_server: FastMCP) -> None:
+ """time_range without a temporal column returns error."""
+ dataset = _make_dataset(main_dttm_col=None)
+
+ with patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metrics": ["count"],
+ "time_range": "Last 7 days",
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert data["error_type"] == "ValidationError"
+ assert "temporal column" in data["error"].lower()
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_with_filters(mcp_server: FastMCP) -> None:
+ """User-provided filters are passed through to the query."""
+ dataset = _make_dataset()
+ result_data = _mock_command_result()
+ captured_queries: list[dict[str, Any]] = []
+
+ def capture_create(**kwargs):
+ captured_queries.extend(kwargs.get("queries", []))
+ return MagicMock()
+
+ with (
+ patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.validate",
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.run",
+ return_value=result_data,
+ ),
+ patch(
+ "superset.common.query_context_factory.QueryContextFactory.create",
+ side_effect=capture_create,
+ ),
+ ):
+ async with Client(mcp_server) as client:
+ await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metrics": ["count"],
+ "filters": [
+ {"col": "category", "op": "==", "val": "Electronics"}
+ ],
+ }
+ },
+ )
+
+ assert len(captured_queries) == 1
+ filters = captured_queries[0]["filters"]
+ assert len(filters) == 1
+ assert filters[0]["col"] == "category"
+ assert filters[0]["op"] == "=="
+ assert filters[0]["val"] == "Electronics"
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_empty_results(mcp_server: FastMCP) -> None:
+ """Query that returns no data gives a response with row_count=0."""
+ dataset = _make_dataset()
+ empty_result = {
+ "queries": [
+ {
+ "data": [],
+ "colnames": [],
+ "rowcount": 0,
+ "is_cached": False,
+ "cached_dttm": None,
+ "cache_timeout": 300,
+ }
+ ]
+ }
+
+ with (
+ patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.validate",
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.run",
+ return_value=empty_result,
+ ),
+ patch(
+ "superset.common.query_context_factory.QueryContextFactory.create",
+ return_value=MagicMock(),
+ ),
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metrics": ["count"],
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert data["row_count"] == 0
+ assert data["data"] == []
+ assert "no data" in data["summary"].lower()
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_by_uuid(mcp_server: FastMCP) -> None:
+ """UUID-based lookup works."""
+ dataset = _make_dataset()
+ result_data = _mock_command_result()
+
+ with (
+ patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ) as mock_resolve,
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.validate",
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.run",
+ return_value=result_data,
+ ),
+ patch(
+ "superset.common.query_context_factory.QueryContextFactory.create",
+ return_value=MagicMock(),
+ ),
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": "a1b2c3d4-5678-90ab-cdef-1234567890ab",
+ "metrics": ["count"],
+ }
+ },
+ )
+
+ # Verify the resolve function was called with the UUID
+ mock_resolve.assert_called_once()
+ call_args = mock_resolve.call_args
+ assert call_args[0][0] == "a1b2c3d4-5678-90ab-cdef-1234567890ab"
+
+ data = json.loads(result.content[0].text)
+ assert data["dataset_id"] == 1
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_permission_denied(mcp_server: FastMCP) -> None:
+ """Permission denied from ChartDataCommand.validate() returns error."""
+ from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
+ from superset.exceptions import SupersetSecurityException
+
+ dataset = _make_dataset()
+
+ with (
+ patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ),
+ patch(
+ "superset.common.query_context_factory.QueryContextFactory.create",
+ return_value=MagicMock(),
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.validate",
+ side_effect=SupersetSecurityException(
+ SupersetError(
+ message="Access denied",
+ error_type=SupersetErrorType.DATASOURCE_SECURITY_ACCESS_ERROR,
+ level=ErrorLevel.WARNING,
+ )
+ ),
+ ),
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metrics": ["count"],
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert data["error_type"] == "QueryError"
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_order_by_valid(mcp_server: FastMCP) -> None:
+ """order_by with valid column/metric names passes through."""
+ dataset = _make_dataset()
+ result_data = _mock_command_result()
+ captured_queries: list[dict[str, Any]] = []
+
+ def capture_create(**kwargs):
+ captured_queries.extend(kwargs.get("queries", []))
+ return MagicMock()
+
+ with (
+ patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.validate",
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.run",
+ return_value=result_data,
+ ),
+ patch(
+ "superset.common.query_context_factory.QueryContextFactory.create",
+ side_effect=capture_create,
+ ),
+ ):
+ async with Client(mcp_server) as client:
+ await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metrics": ["count"],
+ "columns": ["category"],
+ "order_by": ["count"],
+ "order_desc": True,
+ }
+ },
+ )
+
+ assert len(captured_queries) == 1
+ orderby = captured_queries[0].get("orderby", [])
+ assert len(orderby) == 1
+ assert orderby[0][0] == "count"
+ # order_desc=True -> ascending=False
+ assert orderby[0][1] is False
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_order_by_invalid(mcp_server: FastMCP) -> None:
+ """order_by with an unknown name returns validation error."""
+ dataset = _make_dataset()
+
+ with patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metrics": ["count"],
+ "order_by": ["nonexistent"],
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert data["error_type"] == "ValidationError"
+ assert "nonexistent" in data["error"]
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_time_column_override(mcp_server: FastMCP) -> None:
+ """Explicit time_column overrides dataset main_dttm_col."""
+ dataset = _make_dataset(main_dttm_col="order_date")
+ result_data = _mock_command_result()
+ captured_queries: list[dict[str, Any]] = []
+
+ def capture_create(**kwargs):
+ captured_queries.extend(kwargs.get("queries", []))
+ return MagicMock()
+
+ with (
+ patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.validate",
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.run",
+ return_value=result_data,
+ ),
+ patch(
+ "superset.common.query_context_factory.QueryContextFactory.create",
+ side_effect=capture_create,
+ ),
+ ):
+ async with Client(mcp_server) as client:
+ await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metrics": ["count"],
+ "time_range": "Last 30 days",
+ "time_column": "order_date",
+ }
+ },
+ )
+
+ assert len(captured_queries) == 1
+ query_dict = captured_queries[0]
+ assert query_dict["granularity"] == "order_date"
+ temporal_filters = [f for f in query_dict["filters"] if f["op"] == "TEMPORAL_RANGE"]
+ assert temporal_filters[0]["col"] == "order_date"
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_non_dttm_time_column_warns(mcp_server: FastMCP) -> None:
+ """Using a non-datetime column for time_range produces a warning."""
+ dataset = _make_dataset(main_dttm_col=None)
+ result_data = _mock_command_result()
+
+ with (
+ patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.validate",
+ ),
+ patch(
+ "superset.commands.chart.data.get_data_command.ChartDataCommand.run",
+ return_value=result_data,
+ ),
+ patch(
+ "superset.common.query_context_factory.QueryContextFactory.create",
+ return_value=MagicMock(),
+ ),
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metrics": ["count"],
+ "time_range": "Last 7 days",
+ "time_column": "category",
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert len(data["warnings"]) > 0
+ assert "not marked as a datetime" in data["warnings"][0]
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_invalid_filter_column(mcp_server: FastMCP) -> None:
+ """Filter on a column that doesn't exist returns validation error."""
+ dataset = _make_dataset()
+
+ with patch.object(
+ query_dataset_module,
+ "_resolve_dataset",
+ return_value=dataset,
+ ):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "query_dataset",
+ {
+ "request": {
+ "dataset_id": 1,
+ "metrics": ["count"],
+ "filters": [
+ {
+ "col": "nonexistent",
+ "op": "==",
+ "val": "test",
+ }
+ ],
+ }
+ },
+ )
+
+ data = json.loads(result.content[0].text)
+ assert data["error_type"] == "ValidationError"
+ assert "nonexistent" in data["error"]
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_metadata_access_denied_no_suggestions(
+ mcp_server: FastMCP,
+) -> None:
+ """Users without data-model metadata access cannot probe column/metric names.
+
Review Comment:
Added in c3d8cf3c7c — `test_query_dataset_metadata_access_denied_nonexistent_dataset` covers exactly this case. It does not patch `_resolve_dataset`, passes a non-existent dataset ID (999999), and asserts `error_type == "DataModelMetadataRestricted"` while explicitly asserting it is not `"NotFound"`. This prevents regressing the existence side channel.
--
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: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] feat(mcp): add query_dataset tool to query datasets using semantic layer [superset]
Posted by "codeant-ai-for-open-source[bot] (via GitHub)" <gi...@apache.org>.
codeant-ai-for-open-source[bot] commented on code in PR #39727:
URL: https://github.com/apache/superset/pull/39727#discussion_r3164396977
##########
tests/unit_tests/mcp_service/dataset/tool/test_query_dataset.py:
##########
@@ -0,0 +1,737 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Tests for the query_dataset MCP tool."""
+
+from __future__ import annotations
+
+import importlib
+from typing import Any
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+from fastmcp import Client
+
+from superset.mcp_service.app import mcp
+from superset.utils import json
+
+query_dataset_module = importlib.import_module(
+ "superset.mcp_service.dataset.tool.query_dataset"
+)
+
+
+@pytest.fixture
+def mcp_server():
+ return mcp
+
+
+@pytest.fixture(autouse=True)
+def mock_auth():
Review Comment:
**Suggestion:** Add a return type annotation for this fixture (including the yielded type) to satisfy full typing requirements. [custom_rule]
**Severity Level:** Minor ⚠️
<details>
<summary><b>Why it matters? 🤔 </b></summary>
The autouse fixture is a generator-based fixture and has no return/yield type annotation. This is a genuine typing violation under the stated rule.
</details>
[Fix in Cursor](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt=This%20is%20a%20comment%20left%20during%20a%20code%20review.%0A%0A%2A%2APath%3A%2A%2A%20tests%2Funit_tests%2Fmcp_service%2Fdataset%2Ftool%2Ftest_query_dataset.py%0A%2A%2ALine%3A%2A%2A%2042%3A43%0A%2A%2AComment%3A%2A%2A%0A%09%2ACustom%20Rule%3A%20Add%20a%20return%20type%20annotation%20for%20this%20fixture%20%28including%20the%20yielded%20type%29%20to%20satisfy%20full%20typing%20requirements.%0A%0AValidate%20the%20correctness%20of%20the%20flagged%20issue.%20If%20correct%2C%20How%20can%20I%20resolve%20this%3F%20If%20you%20propose%20a%20fix%2C%20implement%20it%20and%20please%20make%20it%20concise.%0AOnce%20fix%20is%20implemented%2C%20also%20check%20other%20comments%20on%20the%20same%20PR%2C%20and%20ask%20user%20if%20the%20user%20wants%20to%20fix%20the%20rest%20of%20the%20comments%20as%20well.%20if%20said%20yes%2C%20then%20fetch%20all%20the%20comments%20validate%20the%20correctness%20and%20implement%20a%20minimal%20fi
x%0A) | [Fix in VSCode Claude](https://app.codeant.ai/fix-in-ide?tool=vscode-claude&prompt=This%20is%20a%20comment%20left%20during%20a%20code%20review.%0A%0A%2A%2APath%3A%2A%2A%20tests%2Funit_tests%2Fmcp_service%2Fdataset%2Ftool%2Ftest_query_dataset.py%0A%2A%2ALine%3A%2A%2A%2042%3A43%0A%2A%2AComment%3A%2A%2A%0A%09%2ACustom%20Rule%3A%20Add%20a%20return%20type%20annotation%20for%20this%20fixture%20%28including%20the%20yielded%20type%29%20to%20satisfy%20full%20typing%20requirements.%0A%0AValidate%20the%20correctness%20of%20the%20flagged%20issue.%20If%20correct%2C%20How%20can%20I%20resolve%20this%3F%20If%20you%20propose%20a%20fix%2C%20implement%20it%20and%20please%20make%20it%20concise.%0AOnce%20fix%20is%20implemented%2C%20also%20check%20other%20comments%20on%20the%20same%20PR%2C%20and%20ask%20user%20if%20the%20user%20wants%20to%20fix%20the%20rest%20of%20the%20comments%20as%20well.%20if%20said%20yes%2C%20then%20fetch%20all%20the%20comments%20validate%20the%20correctness%20and%20implemen
t%20a%20minimal%20fix%0A)
*(Use Cmd/Ctrl + Click for best experience)*
<details>
<summary><b>Prompt for AI Agent 🤖 </b></summary>
```mdx
This is a comment left during a code review.
**Path:** tests/unit_tests/mcp_service/dataset/tool/test_query_dataset.py
**Line:** 42:43
**Comment:**
*Custom Rule: Add a return type annotation for this fixture (including the yielded type) to satisfy full typing requirements.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
```
</details>
<a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F39727&comment_hash=68f5b00bb49cd862524269fe1f15ad5aac3bacce38bda2aae66cade6162d444c&reaction=like'>👍</a> | <a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F39727&comment_hash=68f5b00bb49cd862524269fe1f15ad5aac3bacce38bda2aae66cade6162d444c&reaction=dislike'>👎</a>
##########
tests/unit_tests/mcp_service/dataset/tool/test_query_dataset.py:
##########
@@ -0,0 +1,737 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Tests for the query_dataset MCP tool."""
+
+from __future__ import annotations
+
+import importlib
+from typing import Any
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+from fastmcp import Client
+
+from superset.mcp_service.app import mcp
+from superset.utils import json
+
+query_dataset_module = importlib.import_module(
+ "superset.mcp_service.dataset.tool.query_dataset"
+)
+
+
+@pytest.fixture
+def mcp_server():
+ return mcp
+
+
+@pytest.fixture(autouse=True)
+def mock_auth():
+ """Mock authentication for all tests."""
+ with patch("superset.mcp_service.auth.get_user_from_request") as mock_get_user:
+ mock_user = Mock()
+ mock_user.id = 1
+ mock_user.username = "admin"
+ mock_get_user.return_value = mock_user
+ yield mock_get_user
+
+
+def _make_column(name: str, is_dttm: bool = False) -> MagicMock:
+ col = MagicMock()
+ col.column_name = name
+ col.is_dttm = is_dttm
+ col.verbose_name = None
+ col.type = "VARCHAR"
+ col.groupby = True
+ col.filterable = True
+ col.description = None
+ return col
+
+
+def _make_metric(name: str, expression: str = "COUNT(*)") -> MagicMock:
+ metric = MagicMock()
+ metric.metric_name = name
+ metric.verbose_name = None
+ metric.expression = expression
+ metric.description = None
+ metric.d3format = None
+ return metric
+
+
+def _make_dataset(
+ dataset_id: int = 1,
+ table_name: str = "orders",
+ columns: list[Any] | None = None,
+ metrics: list[Any] | None = None,
+ main_dttm_col: str | None = None,
+) -> MagicMock:
Review Comment:
**Suggestion:** Add a docstring for this helper constructor function to document what test object it builds. [custom_rule]
**Severity Level:** Minor ⚠️
<details>
<summary><b>Why it matters? 🤔 </b></summary>
The helper constructor lacks a docstring in the final file state. This is a real omission consistent with the suggestion.
</details>
[Fix in Cursor](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt=This%20is%20a%20comment%20left%20during%20a%20code%20review.%0A%0A%2A%2APath%3A%2A%2A%20tests%2Funit_tests%2Fmcp_service%2Fdataset%2Ftool%2Ftest_query_dataset.py%0A%2A%2ALine%3A%2A%2A%2075%3A81%0A%2A%2AComment%3A%2A%2A%0A%09%2ACustom%20Rule%3A%20Add%20a%20docstring%20for%20this%20helper%20constructor%20function%20to%20document%20what%20test%20object%20it%20builds.%0A%0AValidate%20the%20correctness%20of%20the%20flagged%20issue.%20If%20correct%2C%20How%20can%20I%20resolve%20this%3F%20If%20you%20propose%20a%20fix%2C%20implement%20it%20and%20please%20make%20it%20concise.%0AOnce%20fix%20is%20implemented%2C%20also%20check%20other%20comments%20on%20the%20same%20PR%2C%20and%20ask%20user%20if%20the%20user%20wants%20to%20fix%20the%20rest%20of%20the%20comments%20as%20well.%20if%20said%20yes%2C%20then%20fetch%20all%20the%20comments%20validate%20the%20correctness%20and%20implement%20a%20minimal%20fix%0A) | [Fix in VSCode Clau
de](https://app.codeant.ai/fix-in-ide?tool=vscode-claude&prompt=This%20is%20a%20comment%20left%20during%20a%20code%20review.%0A%0A%2A%2APath%3A%2A%2A%20tests%2Funit_tests%2Fmcp_service%2Fdataset%2Ftool%2Ftest_query_dataset.py%0A%2A%2ALine%3A%2A%2A%2075%3A81%0A%2A%2AComment%3A%2A%2A%0A%09%2ACustom%20Rule%3A%20Add%20a%20docstring%20for%20this%20helper%20constructor%20function%20to%20document%20what%20test%20object%20it%20builds.%0A%0AValidate%20the%20correctness%20of%20the%20flagged%20issue.%20If%20correct%2C%20How%20can%20I%20resolve%20this%3F%20If%20you%20propose%20a%20fix%2C%20implement%20it%20and%20please%20make%20it%20concise.%0AOnce%20fix%20is%20implemented%2C%20also%20check%20other%20comments%20on%20the%20same%20PR%2C%20and%20ask%20user%20if%20the%20user%20wants%20to%20fix%20the%20rest%20of%20the%20comments%20as%20well.%20if%20said%20yes%2C%20then%20fetch%20all%20the%20comments%20validate%20the%20correctness%20and%20implement%20a%20minimal%20fix%0A)
*(Use Cmd/Ctrl + Click for best experience)*
<details>
<summary><b>Prompt for AI Agent 🤖 </b></summary>
```mdx
This is a comment left during a code review.
**Path:** tests/unit_tests/mcp_service/dataset/tool/test_query_dataset.py
**Line:** 75:81
**Comment:**
*Custom Rule: Add a docstring for this helper constructor function to document what test object it builds.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
```
</details>
<a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F39727&comment_hash=59322f51d5d45cdabc73e235516dd18d4ff718f41ab0a25598ebdf22ed6b66e1&reaction=like'>👍</a> | <a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F39727&comment_hash=59322f51d5d45cdabc73e235516dd18d4ff718f41ab0a25598ebdf22ed6b66e1&reaction=dislike'>👎</a>
##########
tests/unit_tests/mcp_service/dataset/tool/test_query_dataset.py:
##########
@@ -0,0 +1,737 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Tests for the query_dataset MCP tool."""
+
+from __future__ import annotations
+
+import importlib
+from typing import Any
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+from fastmcp import Client
+
+from superset.mcp_service.app import mcp
+from superset.utils import json
+
+query_dataset_module = importlib.import_module(
+ "superset.mcp_service.dataset.tool.query_dataset"
+)
+
+
+@pytest.fixture
+def mcp_server():
+ return mcp
+
+
+@pytest.fixture(autouse=True)
+def mock_auth():
+ """Mock authentication for all tests."""
+ with patch("superset.mcp_service.auth.get_user_from_request") as mock_get_user:
+ mock_user = Mock()
+ mock_user.id = 1
+ mock_user.username = "admin"
+ mock_get_user.return_value = mock_user
+ yield mock_get_user
+
+
+def _make_column(name: str, is_dttm: bool = False) -> MagicMock:
Review Comment:
**Suggestion:** Add a docstring describing this helper function's purpose and return value. [custom_rule]
**Severity Level:** Minor ⚠️
<details>
<summary><b>Why it matters? 🤔 </b></summary>
This helper function is introduced without a docstring. The suggestion accurately identifies a missing docstring for the helper.
</details>
[Fix in Cursor](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt=This%20is%20a%20comment%20left%20during%20a%20code%20review.%0A%0A%2A%2APath%3A%2A%2A%20tests%2Funit_tests%2Fmcp_service%2Fdataset%2Ftool%2Ftest_query_dataset.py%0A%2A%2ALine%3A%2A%2A%2053%3A53%0A%2A%2AComment%3A%2A%2A%0A%09%2ACustom%20Rule%3A%20Add%20a%20docstring%20describing%20this%20helper%20function%27s%20purpose%20and%20return%20value.%0A%0AValidate%20the%20correctness%20of%20the%20flagged%20issue.%20If%20correct%2C%20How%20can%20I%20resolve%20this%3F%20If%20you%20propose%20a%20fix%2C%20implement%20it%20and%20please%20make%20it%20concise.%0AOnce%20fix%20is%20implemented%2C%20also%20check%20other%20comments%20on%20the%20same%20PR%2C%20and%20ask%20user%20if%20the%20user%20wants%20to%20fix%20the%20rest%20of%20the%20comments%20as%20well.%20if%20said%20yes%2C%20then%20fetch%20all%20the%20comments%20validate%20the%20correctness%20and%20implement%20a%20minimal%20fix%0A) | [Fix in VSCode Claude](https://app.codeant
.ai/fix-in-ide?tool=vscode-claude&prompt=This%20is%20a%20comment%20left%20during%20a%20code%20review.%0A%0A%2A%2APath%3A%2A%2A%20tests%2Funit_tests%2Fmcp_service%2Fdataset%2Ftool%2Ftest_query_dataset.py%0A%2A%2ALine%3A%2A%2A%2053%3A53%0A%2A%2AComment%3A%2A%2A%0A%09%2ACustom%20Rule%3A%20Add%20a%20docstring%20describing%20this%20helper%20function%27s%20purpose%20and%20return%20value.%0A%0AValidate%20the%20correctness%20of%20the%20flagged%20issue.%20If%20correct%2C%20How%20can%20I%20resolve%20this%3F%20If%20you%20propose%20a%20fix%2C%20implement%20it%20and%20please%20make%20it%20concise.%0AOnce%20fix%20is%20implemented%2C%20also%20check%20other%20comments%20on%20the%20same%20PR%2C%20and%20ask%20user%20if%20the%20user%20wants%20to%20fix%20the%20rest%20of%20the%20comments%20as%20well.%20if%20said%20yes%2C%20then%20fetch%20all%20the%20comments%20validate%20the%20correctness%20and%20implement%20a%20minimal%20fix%0A)
*(Use Cmd/Ctrl + Click for best experience)*
<details>
<summary><b>Prompt for AI Agent 🤖 </b></summary>
```mdx
This is a comment left during a code review.
**Path:** tests/unit_tests/mcp_service/dataset/tool/test_query_dataset.py
**Line:** 53:53
**Comment:**
*Custom Rule: Add a docstring describing this helper function's purpose and return value.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
```
</details>
<a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F39727&comment_hash=28b7d658a0d5708c90a6746488522e647459d538fce3d047aa627597021c85fa&reaction=like'>👍</a> | <a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F39727&comment_hash=28b7d658a0d5708c90a6746488522e647459d538fce3d047aa627597021c85fa&reaction=dislike'>👎</a>
##########
tests/unit_tests/mcp_service/dataset/tool/test_query_dataset.py:
##########
@@ -0,0 +1,737 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Tests for the query_dataset MCP tool."""
+
+from __future__ import annotations
+
+import importlib
+from typing import Any
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+from fastmcp import Client
+
+from superset.mcp_service.app import mcp
+from superset.utils import json
+
+query_dataset_module = importlib.import_module(
+ "superset.mcp_service.dataset.tool.query_dataset"
+)
+
+
+@pytest.fixture
+def mcp_server():
+ return mcp
Review Comment:
**Suggestion:** Add an explicit return type annotation to this fixture so the new function is fully typed. [custom_rule]
**Severity Level:** Minor ⚠️
<details>
<summary><b>Why it matters? 🤔 </b></summary>
The fixture is defined without any return type annotation, which matches the suggestion's typing rule. This is a real violation of the custom rule requiring explicit return typing.
</details>
[Fix in Cursor](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt=This%20is%20a%20comment%20left%20during%20a%20code%20review.%0A%0A%2A%2APath%3A%2A%2A%20tests%2Funit_tests%2Fmcp_service%2Fdataset%2Ftool%2Ftest_query_dataset.py%0A%2A%2ALine%3A%2A%2A%2037%3A39%0A%2A%2AComment%3A%2A%2A%0A%09%2ACustom%20Rule%3A%20Add%20an%20explicit%20return%20type%20annotation%20to%20this%20fixture%20so%20the%20new%20function%20is%20fully%20typed.%0A%0AValidate%20the%20correctness%20of%20the%20flagged%20issue.%20If%20correct%2C%20How%20can%20I%20resolve%20this%3F%20If%20you%20propose%20a%20fix%2C%20implement%20it%20and%20please%20make%20it%20concise.%0AOnce%20fix%20is%20implemented%2C%20also%20check%20other%20comments%20on%20the%20same%20PR%2C%20and%20ask%20user%20if%20the%20user%20wants%20to%20fix%20the%20rest%20of%20the%20comments%20as%20well.%20if%20said%20yes%2C%20then%20fetch%20all%20the%20comments%20validate%20the%20correctness%20and%20implement%20a%20minimal%20fix%0A) | [Fix in VSCode Clau
de](https://app.codeant.ai/fix-in-ide?tool=vscode-claude&prompt=This%20is%20a%20comment%20left%20during%20a%20code%20review.%0A%0A%2A%2APath%3A%2A%2A%20tests%2Funit_tests%2Fmcp_service%2Fdataset%2Ftool%2Ftest_query_dataset.py%0A%2A%2ALine%3A%2A%2A%2037%3A39%0A%2A%2AComment%3A%2A%2A%0A%09%2ACustom%20Rule%3A%20Add%20an%20explicit%20return%20type%20annotation%20to%20this%20fixture%20so%20the%20new%20function%20is%20fully%20typed.%0A%0AValidate%20the%20correctness%20of%20the%20flagged%20issue.%20If%20correct%2C%20How%20can%20I%20resolve%20this%3F%20If%20you%20propose%20a%20fix%2C%20implement%20it%20and%20please%20make%20it%20concise.%0AOnce%20fix%20is%20implemented%2C%20also%20check%20other%20comments%20on%20the%20same%20PR%2C%20and%20ask%20user%20if%20the%20user%20wants%20to%20fix%20the%20rest%20of%20the%20comments%20as%20well.%20if%20said%20yes%2C%20then%20fetch%20all%20the%20comments%20validate%20the%20correctness%20and%20implement%20a%20minimal%20fix%0A)
*(Use Cmd/Ctrl + Click for best experience)*
<details>
<summary><b>Prompt for AI Agent 🤖 </b></summary>
```mdx
This is a comment left during a code review.
**Path:** tests/unit_tests/mcp_service/dataset/tool/test_query_dataset.py
**Line:** 37:39
**Comment:**
*Custom Rule: Add an explicit return type annotation to this fixture so the new function is fully typed.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
```
</details>
<a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F39727&comment_hash=3b3c02498b4d1b339a275dc1fa212f554c460a7b8083be5db553a16427dfb0d2&reaction=like'>👍</a> | <a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F39727&comment_hash=3b3c02498b4d1b339a275dc1fa212f554c460a7b8083be5db553a16427dfb0d2&reaction=dislike'>👎</a>
##########
tests/unit_tests/mcp_service/dataset/tool/test_query_dataset.py:
##########
@@ -0,0 +1,737 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Tests for the query_dataset MCP tool."""
+
+from __future__ import annotations
+
+import importlib
+from typing import Any
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+from fastmcp import Client
+
+from superset.mcp_service.app import mcp
+from superset.utils import json
+
+query_dataset_module = importlib.import_module(
+ "superset.mcp_service.dataset.tool.query_dataset"
+)
+
+
+@pytest.fixture
+def mcp_server():
+ return mcp
+
+
+@pytest.fixture(autouse=True)
+def mock_auth():
+ """Mock authentication for all tests."""
+ with patch("superset.mcp_service.auth.get_user_from_request") as mock_get_user:
+ mock_user = Mock()
+ mock_user.id = 1
+ mock_user.username = "admin"
+ mock_get_user.return_value = mock_user
+ yield mock_get_user
+
+
+def _make_column(name: str, is_dttm: bool = False) -> MagicMock:
+ col = MagicMock()
+ col.column_name = name
+ col.is_dttm = is_dttm
+ col.verbose_name = None
+ col.type = "VARCHAR"
+ col.groupby = True
+ col.filterable = True
+ col.description = None
+ return col
+
+
+def _make_metric(name: str, expression: str = "COUNT(*)") -> MagicMock:
+ metric = MagicMock()
+ metric.metric_name = name
+ metric.verbose_name = None
+ metric.expression = expression
+ metric.description = None
+ metric.d3format = None
+ return metric
+
+
+def _make_dataset(
+ dataset_id: int = 1,
+ table_name: str = "orders",
+ columns: list[Any] | None = None,
+ metrics: list[Any] | None = None,
+ main_dttm_col: str | None = None,
+) -> MagicMock:
+ ds = MagicMock()
+ ds.id = dataset_id
+ ds.table_name = table_name
+ ds.uuid = f"test-uuid-{dataset_id}"
+ ds.main_dttm_col = main_dttm_col
+ ds.database = MagicMock()
+ ds.database.database_name = "examples"
+ ds.columns = columns or [
+ _make_column("category"),
+ _make_column("region"),
+ _make_column("order_date", is_dttm=True),
+ ]
+ ds.metrics = metrics or [
+ _make_metric("count", "COUNT(*)"),
+ _make_metric("total_revenue", "SUM(revenue)"),
+ ]
+ return ds
+
+
+def _mock_command_result(
+ data: list[dict[str, Any]] | None = None,
+ colnames: list[str] | None = None,
+) -> dict[str, Any]:
+ """Build the result dict that ChartDataCommand.run() returns."""
+ data = data or [
+ {"category": "Electronics", "count": 42},
+ {"category": "Clothing", "count": 17},
+ ]
+ colnames = colnames or ["category", "count"]
+ return {
+ "queries": [
+ {
+ "data": data,
+ "colnames": colnames,
+ "rowcount": len(data),
+ "cache_key": "abc123",
+ "is_cached": False,
+ "cached_dttm": None,
+ "cache_timeout": 300,
+ }
+ ]
+ }
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_success(mcp_server) -> None:
Review Comment:
**Suggestion:** Add a type annotation for the `mcp_server` parameter in this new test function signature. [custom_rule]
**Severity Level:** Minor ⚠️
<details>
<summary><b>Why it matters? 🤔 </b></summary>
The test function accepts `mcp_server` without a type annotation, which is exactly the kind of missing parameter typing described by the suggestion.
</details>
[Fix in Cursor](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt=This%20is%20a%20comment%20left%20during%20a%20code%20review.%0A%0A%2A%2APath%3A%2A%2A%20tests%2Funit_tests%2Fmcp_service%2Fdataset%2Ftool%2Ftest_query_dataset.py%0A%2A%2ALine%3A%2A%2A%20126%3A127%0A%2A%2AComment%3A%2A%2A%0A%09%2ACustom%20Rule%3A%20Add%20a%20type%20annotation%20for%20the%20%60mcp_server%60%20parameter%20in%20this%20new%20test%20function%20signature.%0A%0AValidate%20the%20correctness%20of%20the%20flagged%20issue.%20If%20correct%2C%20How%20can%20I%20resolve%20this%3F%20If%20you%20propose%20a%20fix%2C%20implement%20it%20and%20please%20make%20it%20concise.%0AOnce%20fix%20is%20implemented%2C%20also%20check%20other%20comments%20on%20the%20same%20PR%2C%20and%20ask%20user%20if%20the%20user%20wants%20to%20fix%20the%20rest%20of%20the%20comments%20as%20well.%20if%20said%20yes%2C%20then%20fetch%20all%20the%20comments%20validate%20the%20correctness%20and%20implement%20a%20minimal%20fix%0A) | [Fix in VSCode Cla
ude](https://app.codeant.ai/fix-in-ide?tool=vscode-claude&prompt=This%20is%20a%20comment%20left%20during%20a%20code%20review.%0A%0A%2A%2APath%3A%2A%2A%20tests%2Funit_tests%2Fmcp_service%2Fdataset%2Ftool%2Ftest_query_dataset.py%0A%2A%2ALine%3A%2A%2A%20126%3A127%0A%2A%2AComment%3A%2A%2A%0A%09%2ACustom%20Rule%3A%20Add%20a%20type%20annotation%20for%20the%20%60mcp_server%60%20parameter%20in%20this%20new%20test%20function%20signature.%0A%0AValidate%20the%20correctness%20of%20the%20flagged%20issue.%20If%20correct%2C%20How%20can%20I%20resolve%20this%3F%20If%20you%20propose%20a%20fix%2C%20implement%20it%20and%20please%20make%20it%20concise.%0AOnce%20fix%20is%20implemented%2C%20also%20check%20other%20comments%20on%20the%20same%20PR%2C%20and%20ask%20user%20if%20the%20user%20wants%20to%20fix%20the%20rest%20of%20the%20comments%20as%20well.%20if%20said%20yes%2C%20then%20fetch%20all%20the%20comments%20validate%20the%20correctness%20and%20implement%20a%20minimal%20fix%0A)
*(Use Cmd/Ctrl + Click for best experience)*
<details>
<summary><b>Prompt for AI Agent 🤖 </b></summary>
```mdx
This is a comment left during a code review.
**Path:** tests/unit_tests/mcp_service/dataset/tool/test_query_dataset.py
**Line:** 126:127
**Comment:**
*Custom Rule: Add a type annotation for the `mcp_server` parameter in this new test function signature.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
```
</details>
<a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F39727&comment_hash=cb6122249530d6bae3de7dd014c8fafe71ea8daedbec2e4417b147710a0527f0&reaction=like'>👍</a> | <a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F39727&comment_hash=cb6122249530d6bae3de7dd014c8fafe71ea8daedbec2e4417b147710a0527f0&reaction=dislike'>👎</a>
##########
superset/mcp_service/dataset/schemas.py:
##########
@@ -383,6 +386,146 @@ class CreateVirtualDatasetResponse(BaseModel):
)
+VALID_FILTER_OPS = Literal[
+ "==",
+ "!=",
+ ">",
+ "<",
+ ">=",
+ "<=",
+ "LIKE",
+ "NOT LIKE",
+ "ILIKE",
+ "NOT ILIKE",
+ "IN",
+ "NOT IN",
+ "IS NULL",
+ "IS NOT NULL",
+ "IS TRUE",
+ "IS FALSE",
+ "TEMPORAL_RANGE",
+]
+
+
+class QueryDatasetFilter(BaseModel):
+ """A single filter condition for dataset queries."""
+
+ col: str = Field(..., description="Column name to filter on")
+ op: VALID_FILTER_OPS = Field(
+ ...,
+ description=(
+ 'Filter operator. Use "==" for equals, "!=" for not equals, '
+ '"IN" / "NOT IN" for membership, "IS NULL" / "IS NOT NULL", '
+ '"LIKE" for pattern matching, "TEMPORAL_RANGE" for time filters.'
+ ),
+ )
+ val: Any = Field(
+ default=None,
+ description="Filter value (omit for IS NULL/IS NOT NULL)",
+ )
+
+
+class QueryDatasetRequest(QueryCacheControl):
+ """Request schema for query_dataset tool."""
+
+ dataset_id: int | str = Field(
+ ...,
+ description="Dataset identifier — numeric ID or UUID string.",
+ )
+ metrics: List[str] = Field(
+ default_factory=list,
+ description=(
+ "Saved metric names to compute (e.g. ['count', 'avg_revenue']). "
+ "Use get_dataset_info to discover available metrics."
+ ),
+ )
+ columns: List[str] = Field(
+ default_factory=list,
+ description=(
+ "Column/dimension names for GROUP BY or SELECT "
+ "(e.g. ['category', 'region']). "
+ "Use get_dataset_info to discover available columns."
+ ),
+ )
+ filters: List[QueryDatasetFilter] = Field(
+ default_factory=list,
+ description=(
+ 'Filter conditions (e.g. [{"col": "status", "op": "==", "val": "active"}]).'
+ ),
+ )
+ time_range: str | None = Field(
+ default=None,
+ description=(
+ "Time range filter (e.g. 'Last 7 days', 'Last month', "
+ "'2024-01-01 : 2024-12-31'). Requires a temporal column "
+ "on the dataset."
+ ),
+ )
+ time_column: str | None = Field(
+ default=None,
+ description=(
+ "Temporal column to apply time_range to. "
+ "Defaults to the dataset's main datetime column."
+ ),
+ )
+ order_by: List[str] | None = Field(
+ default=None,
+ description="Column or metric names to sort results by.",
+ )
+ order_desc: bool = Field(
+ default=True,
+ description="Sort descending (True) or ascending (False).",
+ )
+ row_limit: int = Field(
+ default=1000,
+ ge=1,
+ le=50000,
+ description="Maximum number of rows to return (default 1000, max 50000).",
+ )
+
+ @model_validator(mode="after")
+ def validate_metrics_or_columns(self) -> "QueryDatasetRequest":
+ """At least one of metrics or columns must be provided."""
+ if not self.metrics and not self.columns:
+ raise ValueError(
+ "At least one of 'metrics' or 'columns' must be provided. "
+ "Use get_dataset_info to discover available metrics and columns."
+ )
+ return self
+
+
+class QueryDatasetResponse(BaseModel):
+ """Response schema for query_dataset tool."""
+
+ model_config = ConfigDict(ser_json_timedelta="iso8601")
+
+ dataset_id: int = Field(..., description="Dataset ID")
Review Comment:
**Suggestion:** Replace the response `dataset_id` integer field with a UUID-based identifier field so the public response does not expose internal numeric primary keys. [custom_rule]
**Severity Level:** Minor ⚠️
<details>
<summary><b>Why it matters? 🤔 </b></summary>
The response schema uses an integer `dataset_id`, so it exposes an internal numeric identifier exactly as the suggestion describes. This is a real match to the stated rule violation.
</details>
[Fix in Cursor](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt=This%20is%20a%20comment%20left%20during%20a%20code%20review.%0A%0A%2A%2APath%3A%2A%2A%20superset%2Fmcp_service%2Fdataset%2Fschemas.py%0A%2A%2ALine%3A%2A%2A%20502%3A502%0A%2A%2AComment%3A%2A%2A%0A%09%2ACustom%20Rule%3A%20Replace%20the%20response%20%60dataset_id%60%20integer%20field%20with%20a%20UUID-based%20identifier%20field%20so%20the%20public%20response%20does%20not%20expose%20internal%20numeric%20primary%20keys.%0A%0AValidate%20the%20correctness%20of%20the%20flagged%20issue.%20If%20correct%2C%20How%20can%20I%20resolve%20this%3F%20If%20you%20propose%20a%20fix%2C%20implement%20it%20and%20please%20make%20it%20concise.%0AOnce%20fix%20is%20implemented%2C%20also%20check%20other%20comments%20on%20the%20same%20PR%2C%20and%20ask%20user%20if%20the%20user%20wants%20to%20fix%20the%20rest%20of%20the%20comments%20as%20well.%20if%20said%20yes%2C%20then%20fetch%20all%20the%20comments%20validate%20the%20correctness%20and%20imp
lement%20a%20minimal%20fix%0A) | [Fix in VSCode Claude](https://app.codeant.ai/fix-in-ide?tool=vscode-claude&prompt=This%20is%20a%20comment%20left%20during%20a%20code%20review.%0A%0A%2A%2APath%3A%2A%2A%20superset%2Fmcp_service%2Fdataset%2Fschemas.py%0A%2A%2ALine%3A%2A%2A%20502%3A502%0A%2A%2AComment%3A%2A%2A%0A%09%2ACustom%20Rule%3A%20Replace%20the%20response%20%60dataset_id%60%20integer%20field%20with%20a%20UUID-based%20identifier%20field%20so%20the%20public%20response%20does%20not%20expose%20internal%20numeric%20primary%20keys.%0A%0AValidate%20the%20correctness%20of%20the%20flagged%20issue.%20If%20correct%2C%20How%20can%20I%20resolve%20this%3F%20If%20you%20propose%20a%20fix%2C%20implement%20it%20and%20please%20make%20it%20concise.%0AOnce%20fix%20is%20implemented%2C%20also%20check%20other%20comments%20on%20the%20same%20PR%2C%20and%20ask%20user%20if%20the%20user%20wants%20to%20fix%20the%20rest%20of%20the%20comments%20as%20well.%20if%20said%20yes%2C%20then%20fetch%20all%20the%20commen
ts%20validate%20the%20correctness%20and%20implement%20a%20minimal%20fix%0A)
*(Use Cmd/Ctrl + Click for best experience)*
<details>
<summary><b>Prompt for AI Agent 🤖 </b></summary>
```mdx
This is a comment left during a code review.
**Path:** superset/mcp_service/dataset/schemas.py
**Line:** 502:502
**Comment:**
*Custom Rule: Replace the response `dataset_id` integer field with a UUID-based identifier field so the public response does not expose internal numeric primary keys.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
```
</details>
<a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F39727&comment_hash=18770e99617f0e4748784692bd9a7891c421974ce98dd3fcb2de1adb44791929&reaction=like'>👍</a> | <a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F39727&comment_hash=18770e99617f0e4748784692bd9a7891c421974ce98dd3fcb2de1adb44791929&reaction=dislike'>👎</a>
##########
superset/mcp_service/dataset/schemas.py:
##########
@@ -383,6 +386,146 @@ class CreateVirtualDatasetResponse(BaseModel):
)
+VALID_FILTER_OPS = Literal[
+ "==",
+ "!=",
+ ">",
+ "<",
+ ">=",
+ "<=",
+ "LIKE",
+ "NOT LIKE",
+ "ILIKE",
+ "NOT ILIKE",
+ "IN",
+ "NOT IN",
+ "IS NULL",
+ "IS NOT NULL",
+ "IS TRUE",
+ "IS FALSE",
+ "TEMPORAL_RANGE",
+]
+
+
+class QueryDatasetFilter(BaseModel):
+ """A single filter condition for dataset queries."""
+
+ col: str = Field(..., description="Column name to filter on")
+ op: VALID_FILTER_OPS = Field(
+ ...,
+ description=(
+ 'Filter operator. Use "==" for equals, "!=" for not equals, '
+ '"IN" / "NOT IN" for membership, "IS NULL" / "IS NOT NULL", '
+ '"LIKE" for pattern matching, "TEMPORAL_RANGE" for time filters.'
+ ),
+ )
+ val: Any = Field(
+ default=None,
+ description="Filter value (omit for IS NULL/IS NOT NULL)",
+ )
+
+
+class QueryDatasetRequest(QueryCacheControl):
+ """Request schema for query_dataset tool."""
+
+ dataset_id: int | str = Field(
+ ...,
+ description="Dataset identifier — numeric ID or UUID string.",
Review Comment:
**Suggestion:** Restrict externally supplied dataset identifiers to UUIDs only and remove support for numeric IDs from this request schema to avoid exposing internal integer identifiers. [custom_rule]
**Severity Level:** Minor ⚠️
<details>
<summary><b>Why it matters? 🤔 </b></summary>
The request schema explicitly accepts numeric dataset IDs as well as UUID strings, which matches the issue described by the suggestion. If the intended rule is to avoid exposing internal integer identifiers, this field violates it directly.
</details>
[Fix in Cursor](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt=This%20is%20a%20comment%20left%20during%20a%20code%20review.%0A%0A%2A%2APath%3A%2A%2A%20superset%2Fmcp_service%2Fdataset%2Fschemas.py%0A%2A%2ALine%3A%2A%2A%20431%3A433%0A%2A%2AComment%3A%2A%2A%0A%09%2ACustom%20Rule%3A%20Restrict%20externally%20supplied%20dataset%20identifiers%20to%20UUIDs%20only%20and%20remove%20support%20for%20numeric%20IDs%20from%20this%20request%20schema%20to%20avoid%20exposing%20internal%20integer%20identifiers.%0A%0AValidate%20the%20correctness%20of%20the%20flagged%20issue.%20If%20correct%2C%20How%20can%20I%20resolve%20this%3F%20If%20you%20propose%20a%20fix%2C%20implement%20it%20and%20please%20make%20it%20concise.%0AOnce%20fix%20is%20implemented%2C%20also%20check%20other%20comments%20on%20the%20same%20PR%2C%20and%20ask%20user%20if%20the%20user%20wants%20to%20fix%20the%20rest%20of%20the%20comments%20as%20well.%20if%20said%20yes%2C%20then%20fetch%20all%20the%20comments%20validate%20the%20corre
ctness%20and%20implement%20a%20minimal%20fix%0A) | [Fix in VSCode Claude](https://app.codeant.ai/fix-in-ide?tool=vscode-claude&prompt=This%20is%20a%20comment%20left%20during%20a%20code%20review.%0A%0A%2A%2APath%3A%2A%2A%20superset%2Fmcp_service%2Fdataset%2Fschemas.py%0A%2A%2ALine%3A%2A%2A%20431%3A433%0A%2A%2AComment%3A%2A%2A%0A%09%2ACustom%20Rule%3A%20Restrict%20externally%20supplied%20dataset%20identifiers%20to%20UUIDs%20only%20and%20remove%20support%20for%20numeric%20IDs%20from%20this%20request%20schema%20to%20avoid%20exposing%20internal%20integer%20identifiers.%0A%0AValidate%20the%20correctness%20of%20the%20flagged%20issue.%20If%20correct%2C%20How%20can%20I%20resolve%20this%3F%20If%20you%20propose%20a%20fix%2C%20implement%20it%20and%20please%20make%20it%20concise.%0AOnce%20fix%20is%20implemented%2C%20also%20check%20other%20comments%20on%20the%20same%20PR%2C%20and%20ask%20user%20if%20the%20user%20wants%20to%20fix%20the%20rest%20of%20the%20comments%20as%20well.%20if%20said%20yes%2C
%20then%20fetch%20all%20the%20comments%20validate%20the%20correctness%20and%20implement%20a%20minimal%20fix%0A)
*(Use Cmd/Ctrl + Click for best experience)*
<details>
<summary><b>Prompt for AI Agent 🤖 </b></summary>
```mdx
This is a comment left during a code review.
**Path:** superset/mcp_service/dataset/schemas.py
**Line:** 431:433
**Comment:**
*Custom Rule: Restrict externally supplied dataset identifiers to UUIDs only and remove support for numeric IDs from this request schema to avoid exposing internal integer identifiers.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
```
</details>
<a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F39727&comment_hash=2045719712d74cfab413ea69fac9ca5246ab8af0814569c0f1b8e14e4cd6cc66&reaction=like'>👍</a> | <a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F39727&comment_hash=2045719712d74cfab413ea69fac9ca5246ab8af0814569c0f1b8e14e4cd6cc66&reaction=dislike'>👎</a>
--
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: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] feat(mcp): add query_dataset tool to query datasets using semantic layer [superset]
Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on code in PR #39727:
URL: https://github.com/apache/superset/pull/39727#discussion_r3164803000
##########
superset/mcp_service/dataset/schemas.py:
##########
@@ -383,6 +386,146 @@ class CreateVirtualDatasetResponse(BaseModel):
)
+VALID_FILTER_OPS = Literal[
+ "==",
+ "!=",
+ ">",
+ "<",
+ ">=",
+ "<=",
+ "LIKE",
+ "NOT LIKE",
+ "ILIKE",
+ "NOT ILIKE",
+ "IN",
+ "NOT IN",
+ "IS NULL",
+ "IS NOT NULL",
+ "IS TRUE",
+ "IS FALSE",
+ "TEMPORAL_RANGE",
+]
+
+
+class QueryDatasetFilter(BaseModel):
+ """A single filter condition for dataset queries."""
+
+ col: str = Field(..., description="Column name to filter on")
+ op: VALID_FILTER_OPS = Field(
+ ...,
+ description=(
+ 'Filter operator. Use "==" for equals, "!=" for not equals, '
+ '"IN" / "NOT IN" for membership, "IS NULL" / "IS NOT NULL", '
+ '"LIKE" for pattern matching, "TEMPORAL_RANGE" for time filters.'
+ ),
+ )
+ val: Any = Field(
+ default=None,
+ description="Filter value (omit for IS NULL/IS NOT NULL)",
+ )
+
+
+class QueryDatasetRequest(QueryCacheControl):
+ """Request schema for query_dataset tool."""
+
+ dataset_id: int | str = Field(
+ ...,
+ description="Dataset identifier — numeric ID or UUID string.",
+ )
+ metrics: List[str] = Field(
+ default_factory=list,
+ description=(
+ "Saved metric names to compute (e.g. ['count', 'avg_revenue']). "
+ "Use get_dataset_info to discover available metrics."
+ ),
+ )
+ columns: List[str] = Field(
+ default_factory=list,
+ description=(
+ "Column/dimension names for GROUP BY or SELECT "
+ "(e.g. ['category', 'region']). "
+ "Use get_dataset_info to discover available columns."
+ ),
+ )
+ filters: List[QueryDatasetFilter] = Field(
+ default_factory=list,
+ description=(
+ 'Filter conditions (e.g. [{"col": "status", "op": "==", "val": "active"}]).'
+ ),
+ )
+ time_range: str | None = Field(
+ default=None,
+ description=(
+ "Time range filter (e.g. 'Last 7 days', 'Last month', "
+ "'2024-01-01 : 2024-12-31'). Requires a temporal column "
+ "on the dataset."
+ ),
+ )
+ time_column: str | None = Field(
+ default=None,
+ description=(
+ "Temporal column to apply time_range to. "
+ "Defaults to the dataset's main datetime column."
+ ),
+ )
+ order_by: List[str] | None = Field(
+ default=None,
+ description="Column or metric names to sort results by.",
+ )
+ order_desc: bool = Field(
+ default=True,
+ description="Sort descending (True) or ascending (False).",
+ )
+ row_limit: int = Field(
+ default=1000,
+ ge=1,
+ le=50000,
+ description="Maximum number of rows to return (default 1000, max 50000).",
+ )
+
+ @model_validator(mode="after")
+ def validate_metrics_or_columns(self) -> "QueryDatasetRequest":
+ """At least one of metrics or columns must be provided."""
+ if not self.metrics and not self.columns:
+ raise ValueError(
+ "At least one of 'metrics' or 'columns' must be provided. "
+ "Use get_dataset_info to discover available metrics and columns."
+ )
+ return self
+
+
+class QueryDatasetResponse(BaseModel):
+ """Response schema for query_dataset tool."""
+
+ model_config = ConfigDict(ser_json_timedelta="iso8601")
+
+ dataset_id: int = Field(..., description="Dataset ID")
Review Comment:
Intentionally keeping `dataset_id: int` in the response — this matches the pattern used across all other MCP response schemas (e.g. `DatasetInfo`, `ChartInfo`) and is consistent with how Superset exposes identifiers in its REST API responses.
##########
tests/unit_tests/mcp_service/dataset/tool/test_query_dataset.py:
##########
@@ -0,0 +1,737 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Tests for the query_dataset MCP tool."""
+
+from __future__ import annotations
+
+import importlib
+from typing import Any
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+from fastmcp import Client
+
+from superset.mcp_service.app import mcp
+from superset.utils import json
+
+query_dataset_module = importlib.import_module(
+ "superset.mcp_service.dataset.tool.query_dataset"
+)
+
+
+@pytest.fixture
+def mcp_server():
+ return mcp
Review Comment:
Fixed in ffcd4b1d25 — `mcp_server()` now annotated as `-> FastMCP`.
--
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: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] feat(mcp): add query_dataset tool to query datasets using semantic layer [superset]
Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on code in PR #39727:
URL: https://github.com/apache/superset/pull/39727#discussion_r3164802362
##########
superset/mcp_service/dataset/tool/query_dataset.py:
##########
@@ -0,0 +1,457 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+MCP tool: query_dataset
+
+Query a dataset using its semantic layer (saved metrics, calculated columns,
+dimensions) without requiring a saved chart.
+"""
+
+import difflib
+import logging
+import time
+from typing import Any
+
+from fastmcp import Context
+from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.orm import joinedload, subqueryload
+from superset_core.mcp.decorators import tool, ToolAnnotations
+
+from superset.commands.exceptions import CommandException
+from superset.exceptions import OAuth2Error, OAuth2RedirectError, SupersetException
+from superset.extensions import event_logger
+from superset.mcp_service.chart.schemas import DataColumn, PerformanceMetadata
+from superset.mcp_service.dataset.schemas import (
+ DatasetError,
+ QueryDatasetRequest,
+ QueryDatasetResponse,
+)
+from superset.mcp_service.utils import _is_uuid
+from superset.mcp_service.utils.cache_utils import get_cache_status_from_result
+from superset.mcp_service.utils.oauth2_utils import build_oauth2_redirect_message
+
+logger = logging.getLogger(__name__)
+
+
+def _resolve_dataset(identifier: int | str, eager_options: list[Any]) -> Any | None:
+ """Resolve a dataset by int ID or UUID string.
+
+ Replicates the identifier resolution logic from ModelGetInfoCore._find_object().
+ """
+ from superset.daos.dataset import DatasetDAO
+
+ opts = eager_options or None
+
+ if isinstance(identifier, int):
+ return DatasetDAO.find_by_id(identifier, query_options=opts)
+
+ # Try parsing as int
+ try:
+ id_val = int(identifier)
+ return DatasetDAO.find_by_id(id_val, query_options=opts)
+ except (ValueError, TypeError):
+ pass
+
+ # Try UUID
+ if _is_uuid(str(identifier)):
+ return DatasetDAO.find_by_id(identifier, id_column="uuid", query_options=opts)
+
+ return None
+
+
+def _validate_names(
+ requested: list[str],
+ valid: set[str],
+ kind: str,
+) -> list[str]:
+ """Return list of error messages for names not found in *valid*.
+
+ Includes close-match suggestions when available.
+ """
+ errors: list[str] = []
+ for name in requested:
+ if name not in valid:
+ suggestions = difflib.get_close_matches(name, valid, n=3, cutoff=0.6)
+ msg = f"Unknown {kind}: '{name}'"
+ if suggestions:
+ msg += f". Did you mean: {', '.join(suggestions)}?"
+ errors.append(msg)
+ return errors
+
+
+@tool(
+ tags=["data"],
+ class_permission_name="Dataset",
+ annotations=ToolAnnotations(
+ title="Query dataset",
+ readOnlyHint=True,
+ destructiveHint=False,
+ ),
+)
+async def query_dataset( # noqa: C901
+ request: QueryDatasetRequest, ctx: Context
+) -> QueryDatasetResponse | DatasetError:
+ """Query a dataset using its semantic layer (saved metrics, dimensions, filters).
+
+ Returns tabular data without requiring a saved chart. Use this when you want
+ to compute saved metrics, group by dimensions, or apply filters directly
+ against a dataset's curated semantic layer.
+
+ Workflow:
+ 1. list_datasets -> find a dataset
+ 2. get_dataset_info -> discover available columns and metrics
+ 3. query_dataset -> query using metric names and column names
+
+ Example:
+ ```json
+ {
+ "dataset_id": 123,
+ "metrics": ["count", "avg_revenue"],
+ "columns": ["product_category"],
+ "time_range": "Last 7 days",
+ "row_limit": 100
+ }
+ ```
+ """
+ await ctx.info(
+ "Starting dataset query: dataset_id=%s, metrics=%s, columns=%s, "
+ "row_limit=%s"
+ % (
+ request.dataset_id,
+ request.metrics,
+ request.columns,
+ request.row_limit,
+ )
+ )
+
+ try:
+ from superset.commands.chart.data.get_data_command import ChartDataCommand
+ from superset.common.query_context_factory import QueryContextFactory
+ from superset.connectors.sqla.models import SqlaTable
+
+ # ------------------------------------------------------------------
+ # Step 1: Resolve dataset
+ # ------------------------------------------------------------------
+ await ctx.report_progress(1, 5, "Looking up dataset")
+ eager_options = [
+ subqueryload(SqlaTable.columns),
+ subqueryload(SqlaTable.metrics),
+ joinedload(SqlaTable.database),
+ ]
+
+ with event_logger.log_context(action="mcp.query_dataset.lookup"):
+ dataset = _resolve_dataset(request.dataset_id, eager_options)
+
+ if dataset is None:
+ await ctx.error("Dataset not found: identifier=%s" % (request.dataset_id,))
+ return DatasetError.create(
+ error=f"No dataset found with identifier: {request.dataset_id}",
+ error_type="NotFound",
+ )
+
+ dataset_name = getattr(dataset, "table_name", None) or f"Dataset {dataset.id}"
+ await ctx.info(
+ "Dataset found: id=%s, name=%s, columns=%s, metrics=%s"
+ % (
+ dataset.id,
+ dataset_name,
+ len(dataset.columns),
+ len(dataset.metrics),
+ )
+ )
+
+ # ------------------------------------------------------------------
+ # Step 2: Validate requested columns and metrics
+ # ------------------------------------------------------------------
+ await ctx.report_progress(2, 5, "Validating columns and metrics")
+ valid_columns = {c.column_name for c in dataset.columns}
+ valid_metrics = {m.metric_name for m in dataset.metrics}
+
+ validation_errors: list[str] = []
+ validation_errors.extend(
+ _validate_names(request.columns, valid_columns, "column")
+ )
+ validation_errors.extend(
+ _validate_names(request.metrics, valid_metrics, "metric")
+ )
+ # Validate filter column names against dataset columns
+ filter_cols = [f.col for f in request.filters]
+ validation_errors.extend(
+ _validate_names(filter_cols, valid_columns, "filter column")
+ )
+ # Validate order_by names against columns + metrics
+ if request.order_by:
+ valid_orderby = valid_columns | valid_metrics
+ validation_errors.extend(
+ _validate_names(request.order_by, valid_orderby, "order_by")
+ )
+
+ if validation_errors:
+ error_msg = "; ".join(validation_errors)
+ await ctx.error("Validation failed: %s" % (error_msg,))
+ return DatasetError.create(
+ error=error_msg,
+ error_type="ValidationError",
+ )
+
+ # ------------------------------------------------------------------
+ # Step 3: Build filters and time range
+ # ------------------------------------------------------------------
+ warnings: list[str] = []
+ query_filters: list[dict[str, Any]] = [
+ {"col": f.col, "op": f.op, "val": f.val} for f in request.filters
+ ]
+ granularity: str | None = None
+
+ if request.time_range:
+ temporal_col = request.time_column or getattr(
+ dataset, "main_dttm_col", None
+ )
+ if not temporal_col:
+ await ctx.error("time_range provided but no temporal column available")
+ return DatasetError.create(
+ error=(
+ "time_range was provided but no temporal column is available. "
+ "Either set time_column explicitly or ensure the dataset has "
+ "a main datetime column configured."
+ ),
+ error_type="ValidationError",
+ )
+ # Validate that the temporal column actually exists on the dataset
+ if temporal_col not in valid_columns:
+ await ctx.error("time_column '%s' not found on dataset" % temporal_col)
+ return DatasetError.create(
+ error=(
+ f"time_column '{temporal_col}' does not exist on this dataset."
+ ),
+ error_type="ValidationError",
+ )
+ # Warn if the chosen temporal column isn't marked as datetime
+ dttm_cols = {c.column_name for c in dataset.columns if c.is_dttm}
+ if temporal_col not in dttm_cols:
+ warnings.append(
+ f"Column '{temporal_col}' is not marked as a datetime "
+ f"column on this dataset. Time filtering may not work "
+ f"as expected."
+ )
+
+ query_filters.append(
+ {
+ "col": temporal_col,
+ "op": "TEMPORAL_RANGE",
+ "val": request.time_range,
+ }
+ )
+ granularity = temporal_col
+ await ctx.debug(
+ "Time filter: column=%s, range=%s" % (temporal_col, request.time_range)
+ )
+
+ # ------------------------------------------------------------------
+ # Step 4: Build query dict
+ # ------------------------------------------------------------------
+ await ctx.report_progress(3, 5, "Building query")
+ query_dict: dict[str, Any] = {
+ "filters": query_filters,
+ "columns": request.columns,
+ "metrics": request.metrics,
+ "row_limit": request.row_limit,
+ "order_desc": request.order_desc,
+ }
+ if granularity:
+ query_dict["granularity"] = granularity
+ if request.order_by:
+ # OrderBy = tuple[Metric | Column, bool] where bool is ascending
+ query_dict["orderby"] = [
+ (col, not request.order_desc) for col in request.order_by
+ ]
+
+ await ctx.debug("Query dict keys: %s" % (sorted(query_dict.keys()),))
+
+ # ------------------------------------------------------------------
+ # Step 5: Create QueryContext and execute
+ # ------------------------------------------------------------------
+ await ctx.report_progress(4, 5, "Executing query")
+ start_time = time.time()
+
+ with event_logger.log_context(action="mcp.query_dataset.execute"):
+ factory = QueryContextFactory()
+ # datasource_type is "table" because this tool queries SqlaTable
+ # datasets (Superset's built-in semantic layer). External semantic
+ # layers (dbt, Snowflake Cortex, etc.) use "semantic_view" and have
+ # a different query path — see SemanticView + mapper.py.
+ query_context = factory.create(
+ datasource={"id": dataset.id, "type": "table"},
+ queries=[query_dict],
+ form_data={},
+ force=request.force_refresh,
+ custom_cache_timeout=request.cache_timeout,
+ )
+
+ command = ChartDataCommand(query_context)
+ command.validate()
+ result = command.run()
+
+ query_duration_ms = int((time.time() - start_time) * 1000)
+
+ if not result or "queries" not in result or len(result["queries"]) == 0:
+ await ctx.warning("Query returned no results for dataset %s" % dataset.id)
+ return DatasetError.create(
+ error="Query returned no results.",
+ error_type="EmptyQuery",
+ )
+
+ # ------------------------------------------------------------------
+ # Step 6: Format response
+ # ------------------------------------------------------------------
+ await ctx.report_progress(5, 5, "Formatting results")
+ query_result = result["queries"][0]
+ data = query_result.get("data", [])
+ raw_columns = query_result.get("colnames", [])
+
+ if not data:
+ return QueryDatasetResponse(
+ dataset_id=dataset.id,
+ dataset_name=dataset_name,
+ columns=[],
+ data=[],
+ row_count=0,
+ total_rows=0,
+ summary=f"Query on '{dataset_name}' returned no data.",
+ performance=PerformanceMetadata(
+ query_duration_ms=query_duration_ms,
+ cache_status="no_data",
+ ),
+ cache_status=get_cache_status_from_result(
+ query_result, force_refresh=request.force_refresh
+ ),
+ applied_filters=request.filters,
+ warnings=warnings,
+ )
+
+ # Build column metadata in a single pass per column.
+ # Cap stats computation at STATS_SAMPLE rows to avoid O(rows*cols)
+ # overhead on large result sets (row_limit allows up to 50k).
+ stats_sample_size = 5000
+ stats_rows = data[:stats_sample_size]
+
+ columns_meta: list[DataColumn] = []
+ for col_name in raw_columns:
+ sample_values = [
+ row.get(col_name) for row in data[:3] if row.get(col_name) is not None
+ ]
+ data_type = "string"
+ if sample_values:
+ if all(isinstance(v, bool) for v in sample_values):
+ data_type = "boolean"
+ elif all(isinstance(v, (int, float)) for v in sample_values):
+ data_type = "numeric"
+
+ # Compute null_count and unique non-null values in one pass
+ null_count = 0
+ unique_vals: set[str] = set()
+ for row in stats_rows:
+ val = row.get(col_name)
+ if val is None:
+ null_count += 1
+ else:
+ unique_vals.add(str(val))
+
+ columns_meta.append(
+ DataColumn(
+ name=col_name,
+ display_name=col_name.replace("_", " ").title(),
+ data_type=data_type,
+ sample_values=sample_values[:3],
+ null_count=null_count,
+ unique_count=len(unique_vals),
+ )
+ )
+
+ cache_status = get_cache_status_from_result(
+ query_result, force_refresh=request.force_refresh
+ )
+
+ cache_label = "cached" if cache_status and cache_status.cache_hit else "fresh"
+ summary = (
+ f"Dataset '{dataset_name}': {len(data)} rows, "
+ f"{len(raw_columns)} columns ({cache_label})."
+ )
+
+ await ctx.info(
+ "Query complete: rows=%s, columns=%s, duration=%sms"
+ % (len(data), len(raw_columns), query_duration_ms)
+ )
+
+ return QueryDatasetResponse(
+ dataset_id=dataset.id,
+ dataset_name=dataset_name,
+ columns=columns_meta,
+ data=data,
+ row_count=len(data),
+ total_rows=query_result.get("rowcount"),
+ summary=summary,
+ performance=PerformanceMetadata(
+ query_duration_ms=query_duration_ms,
+ cache_status=cache_label,
+ ),
+ cache_status=cache_status,
+ applied_filters=request.filters,
+ warnings=warnings,
+ )
+
+ except OAuth2RedirectError as exc:
+ redirect_msg = build_oauth2_redirect_message(exc)
+ await ctx.error("OAuth2 redirect required: %s" % (redirect_msg,))
+ return DatasetError.create(
+ error=redirect_msg,
+ error_type="OAuth2Redirect",
+ )
+
+ except OAuth2Error as exc:
+ await ctx.error("OAuth2 error: %s" % (str(exc),))
+ return DatasetError.create(
+ error=f"OAuth2 authentication error: {exc}",
+ error_type="OAuth2Error",
+ )
+
+ except (CommandException, SupersetException) as exc:
+ await ctx.error("Query failed: %s" % (str(exc),))
+ return DatasetError.create(
+ error=f"Query execution failed: {exc}",
+ error_type="QueryError",
+ )
+
+ except SQLAlchemyError as exc:
+ await ctx.error("Database error: %s" % (str(exc),))
+ return DatasetError.create(
+ error=f"Database error: {exc}",
+ error_type="DatabaseError",
+ )
+
+ except Exception as exc:
Review Comment:
Fixed in ffcd4b1d25 — the bare `Exception` catch now logs the full exception server-side via `logger.exception()` and returns a structured `DatasetError(error_type="UnexpectedError")` with a sanitized client message, rather than re-raising.
##########
superset/mcp_service/dataset/schemas.py:
##########
@@ -383,6 +386,146 @@ class CreateVirtualDatasetResponse(BaseModel):
)
+VALID_FILTER_OPS = Literal[
+ "==",
+ "!=",
+ ">",
+ "<",
+ ">=",
+ "<=",
+ "LIKE",
+ "NOT LIKE",
+ "ILIKE",
+ "NOT ILIKE",
+ "IN",
+ "NOT IN",
+ "IS NULL",
+ "IS NOT NULL",
+ "IS TRUE",
+ "IS FALSE",
+ "TEMPORAL_RANGE",
+]
+
+
+class QueryDatasetFilter(BaseModel):
+ """A single filter condition for dataset queries."""
+
+ col: str = Field(..., description="Column name to filter on")
+ op: VALID_FILTER_OPS = Field(
+ ...,
+ description=(
+ 'Filter operator. Use "==" for equals, "!=" for not equals, '
+ '"IN" / "NOT IN" for membership, "IS NULL" / "IS NOT NULL", '
+ '"LIKE" for pattern matching, "TEMPORAL_RANGE" for time filters.'
+ ),
+ )
+ val: Any = Field(
+ default=None,
+ description="Filter value (omit for IS NULL/IS NOT NULL)",
+ )
+
+
+class QueryDatasetRequest(QueryCacheControl):
+ """Request schema for query_dataset tool."""
+
+ dataset_id: int | str = Field(
+ ...,
+ description="Dataset identifier — numeric ID or UUID string.",
Review Comment:
Intentionally keeping `int | str` — all other dataset tools (`get_dataset_info`, `list_datasets`) accept both integer IDs and UUIDs as a consistent pattern across the Superset MCP service. Restricting to UUID-only here would break that consistency and require callers to first resolve an integer ID before querying.
--
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: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] feat(mcp): add query_dataset tool to query datasets using semantic layer [superset]
Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on code in PR #39727:
URL: https://github.com/apache/superset/pull/39727#discussion_r3164804466
##########
tests/unit_tests/mcp_service/dataset/tool/test_query_dataset.py:
##########
@@ -0,0 +1,737 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Tests for the query_dataset MCP tool."""
+
+from __future__ import annotations
+
+import importlib
+from typing import Any
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+from fastmcp import Client
+
+from superset.mcp_service.app import mcp
+from superset.utils import json
+
+query_dataset_module = importlib.import_module(
+ "superset.mcp_service.dataset.tool.query_dataset"
+)
+
+
+@pytest.fixture
+def mcp_server():
+ return mcp
+
+
+@pytest.fixture(autouse=True)
+def mock_auth():
+ """Mock authentication for all tests."""
+ with patch("superset.mcp_service.auth.get_user_from_request") as mock_get_user:
+ mock_user = Mock()
+ mock_user.id = 1
+ mock_user.username = "admin"
+ mock_get_user.return_value = mock_user
+ yield mock_get_user
+
+
+def _make_column(name: str, is_dttm: bool = False) -> MagicMock:
+ col = MagicMock()
+ col.column_name = name
+ col.is_dttm = is_dttm
+ col.verbose_name = None
+ col.type = "VARCHAR"
+ col.groupby = True
+ col.filterable = True
+ col.description = None
+ return col
+
+
+def _make_metric(name: str, expression: str = "COUNT(*)") -> MagicMock:
+ metric = MagicMock()
+ metric.metric_name = name
+ metric.verbose_name = None
+ metric.expression = expression
+ metric.description = None
+ metric.d3format = None
+ return metric
+
+
+def _make_dataset(
+ dataset_id: int = 1,
+ table_name: str = "orders",
+ columns: list[Any] | None = None,
+ metrics: list[Any] | None = None,
+ main_dttm_col: str | None = None,
+) -> MagicMock:
+ ds = MagicMock()
+ ds.id = dataset_id
+ ds.table_name = table_name
+ ds.uuid = f"test-uuid-{dataset_id}"
+ ds.main_dttm_col = main_dttm_col
+ ds.database = MagicMock()
+ ds.database.database_name = "examples"
+ ds.columns = columns or [
+ _make_column("category"),
+ _make_column("region"),
+ _make_column("order_date", is_dttm=True),
+ ]
+ ds.metrics = metrics or [
+ _make_metric("count", "COUNT(*)"),
+ _make_metric("total_revenue", "SUM(revenue)"),
+ ]
+ return ds
+
+
+def _mock_command_result(
+ data: list[dict[str, Any]] | None = None,
+ colnames: list[str] | None = None,
+) -> dict[str, Any]:
+ """Build the result dict that ChartDataCommand.run() returns."""
+ data = data or [
+ {"category": "Electronics", "count": 42},
+ {"category": "Clothing", "count": 17},
+ ]
+ colnames = colnames or ["category", "count"]
+ return {
+ "queries": [
+ {
+ "data": data,
+ "colnames": colnames,
+ "rowcount": len(data),
+ "cache_key": "abc123",
+ "is_cached": False,
+ "cached_dttm": None,
+ "cache_timeout": 300,
+ }
+ ]
+ }
+
+
+@pytest.mark.asyncio
+async def test_query_dataset_success(mcp_server) -> None:
Review Comment:
Fixed in ffcd4b1d25 — all test function signatures now annotate `mcp_server: FastMCP`.
##########
superset/mcp_service/dataset/tool/query_dataset.py:
##########
@@ -0,0 +1,457 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+MCP tool: query_dataset
+
+Query a dataset using its semantic layer (saved metrics, calculated columns,
+dimensions) without requiring a saved chart.
+"""
+
+import difflib
+import logging
+import time
+from typing import Any
+
+from fastmcp import Context
+from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.orm import joinedload, subqueryload
+from superset_core.mcp.decorators import tool, ToolAnnotations
+
+from superset.commands.exceptions import CommandException
+from superset.exceptions import OAuth2Error, OAuth2RedirectError, SupersetException
+from superset.extensions import event_logger
+from superset.mcp_service.chart.schemas import DataColumn, PerformanceMetadata
+from superset.mcp_service.dataset.schemas import (
+ DatasetError,
+ QueryDatasetRequest,
+ QueryDatasetResponse,
+)
+from superset.mcp_service.utils import _is_uuid
+from superset.mcp_service.utils.cache_utils import get_cache_status_from_result
+from superset.mcp_service.utils.oauth2_utils import build_oauth2_redirect_message
+
+logger = logging.getLogger(__name__)
+
+
+def _resolve_dataset(identifier: int | str, eager_options: list[Any]) -> Any | None:
+ """Resolve a dataset by int ID or UUID string.
+
+ Replicates the identifier resolution logic from ModelGetInfoCore._find_object().
+ """
+ from superset.daos.dataset import DatasetDAO
+
+ opts = eager_options or None
+
+ if isinstance(identifier, int):
+ return DatasetDAO.find_by_id(identifier, query_options=opts)
+
+ # Try parsing as int
+ try:
+ id_val = int(identifier)
+ return DatasetDAO.find_by_id(id_val, query_options=opts)
+ except (ValueError, TypeError):
+ pass
+
+ # Try UUID
+ if _is_uuid(str(identifier)):
+ return DatasetDAO.find_by_id(identifier, id_column="uuid", query_options=opts)
+
+ return None
+
+
+def _validate_names(
+ requested: list[str],
+ valid: set[str],
+ kind: str,
+) -> list[str]:
+ """Return list of error messages for names not found in *valid*.
+
+ Includes close-match suggestions when available.
+ """
+ errors: list[str] = []
+ for name in requested:
+ if name not in valid:
+ suggestions = difflib.get_close_matches(name, valid, n=3, cutoff=0.6)
+ msg = f"Unknown {kind}: '{name}'"
+ if suggestions:
+ msg += f". Did you mean: {', '.join(suggestions)}?"
+ errors.append(msg)
+ return errors
+
+
+@tool(
+ tags=["data"],
+ class_permission_name="Dataset",
+ annotations=ToolAnnotations(
+ title="Query dataset",
+ readOnlyHint=True,
+ destructiveHint=False,
+ ),
+)
+async def query_dataset( # noqa: C901
+ request: QueryDatasetRequest, ctx: Context
+) -> QueryDatasetResponse | DatasetError:
Review Comment:
Fixed in ffcd4b1d25 — added `@requires_data_model_metadata_access` decorator and a runtime `user_can_view_data_model_metadata()` check immediately after dataset lookup but before the validation step that returns close-match suggestions. This matches the pattern in `get_dataset_info`. Also added `test_query_dataset_metadata_access_denied_no_suggestions` to verify restricted users receive `DataModelMetadataRestricted` before any schema suggestions are returned.
--
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: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] feat(mcp): add query_dataset tool to query datasets using semantic layer [superset]
Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on PR #39727:
URL: https://github.com/apache/superset/pull/39727#issuecomment-4354082779
Thanks for the second pass @richardfogaca. Both caught — fixed in 3242c4c4cd:
**ctx.info() metadata disclosure**
Moved the `user_can_view_data_model_metadata()` check to before the `ctx.info()` that logged dataset name and column/metric counts. Restricted callers now get the `DataModelMetadataRestricted` error before any dataset-specific info is sent to the client.
**applied_filters under-report in success path**
The non-empty success response was still returning `request.filters` — the empty-data branch was correctly fixed earlier but this path was missed. Changed to `effective_filters` so both response paths now include the synthesized `TEMPORAL_RANGE` entry when `time_range` is provided. Also added an assertion to `test_query_dataset_with_time_range` that verifies `applied_filters` in the response includes the `TEMPORAL_RANGE` filter.
--
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: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] feat(mcp): add query_dataset tool to query datasets using semantic layer [superset]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #39727:
URL: https://github.com/apache/superset/pull/39727#discussion_r3155961421
##########
superset/mcp_service/dataset/schemas.py:
##########
@@ -380,6 +385,146 @@ class CreateVirtualDatasetResponse(BaseModel):
)
+VALID_FILTER_OPS = Literal[
+ "==",
+ "!=",
+ ">",
+ "<",
+ ">=",
+ "<=",
+ "LIKE",
+ "NOT LIKE",
+ "ILIKE",
+ "NOT ILIKE",
+ "IN",
+ "NOT IN",
+ "IS NULL",
+ "IS NOT NULL",
+ "IS TRUE",
+ "IS FALSE",
+ "TEMPORAL_RANGE",
+]
Review Comment:
The schema examples/documentation advertise `op: "="`, but `VALID_FILTER_OPS` does not allow `"="` (only `"=="`). This will cause client confusion and immediate validation failures for a request that follows the example. Update the example to use `"=="`, or explicitly support `"="` (and normalize it to `"=="`) to match the documented contract.
##########
superset/mcp_service/dataset/tool/query_dataset.py:
##########
@@ -0,0 +1,434 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+MCP tool: query_dataset
+
+Query a dataset using its semantic layer (saved metrics, calculated columns,
+dimensions) without requiring a saved chart.
+"""
+
+import difflib
+import logging
+import time
+from typing import Any
+
+from fastmcp import Context
+from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.orm import joinedload, subqueryload
+from superset_core.mcp.decorators import tool, ToolAnnotations
+
+from superset.commands.exceptions import CommandException
+from superset.exceptions import OAuth2Error, OAuth2RedirectError, SupersetException
+from superset.extensions import event_logger
+from superset.mcp_service.chart.schemas import DataColumn, PerformanceMetadata
+from superset.mcp_service.dataset.schemas import (
+ DatasetError,
+ QueryDatasetRequest,
+ QueryDatasetResponse,
+)
+from superset.mcp_service.utils import _is_uuid
+from superset.mcp_service.utils.cache_utils import get_cache_status_from_result
+from superset.mcp_service.utils.oauth2_utils import build_oauth2_redirect_message
+
+logger = logging.getLogger(__name__)
+
+
+def _resolve_dataset(identifier: int | str, eager_options: list[Any]) -> Any | None:
+ """Resolve a dataset by int ID or UUID string.
+
+ Replicates the identifier resolution logic from ModelGetInfoCore._find_object().
+ """
+ from superset.daos.dataset import DatasetDAO
+
+ opts = eager_options or None
+
+ if isinstance(identifier, int):
+ return DatasetDAO.find_by_id(identifier, query_options=opts)
+
+ # Try parsing as int
+ try:
+ id_val = int(identifier)
+ return DatasetDAO.find_by_id(id_val, query_options=opts)
+ except (ValueError, TypeError):
+ pass
+
+ # Try UUID
+ if _is_uuid(str(identifier)):
+ return DatasetDAO.find_by_id(identifier, id_column="uuid", query_options=opts)
+
+ return None
+
+
+def _validate_names(
+ requested: list[str],
+ valid: set[str],
+ kind: str,
+) -> list[str]:
+ """Return list of error messages for names not found in *valid*.
+
+ Includes close-match suggestions when available.
+ """
+ errors: list[str] = []
+ for name in requested:
+ if name not in valid:
+ suggestions = difflib.get_close_matches(name, valid, n=3, cutoff=0.6)
+ msg = f"Unknown {kind}: '{name}'"
+ if suggestions:
+ msg += f". Did you mean: {', '.join(suggestions)}?"
+ errors.append(msg)
+ return errors
+
+
+@tool(
+ tags=["data"],
+ class_permission_name="Dataset",
+ annotations=ToolAnnotations(
+ title="Query dataset",
+ readOnlyHint=True,
+ destructiveHint=False,
+ ),
+)
+async def query_dataset( # noqa: C901
+ request: QueryDatasetRequest, ctx: Context
+) -> QueryDatasetResponse | DatasetError:
+ """Query a dataset using its semantic layer (saved metrics, dimensions, filters).
+
+ Returns tabular data without requiring a saved chart. Use this when you want
+ to compute saved metrics, group by dimensions, or apply filters directly
+ against a dataset's curated semantic layer.
+
+ Workflow:
+ 1. list_datasets -> find a dataset
+ 2. get_dataset_info -> discover available columns and metrics
+ 3. query_dataset -> query using metric names and column names
+
+ Example:
+ ```json
+ {
+ "dataset_id": 123,
+ "metrics": ["count", "avg_revenue"],
+ "columns": ["product_category"],
+ "time_range": "Last 7 days",
+ "row_limit": 100
+ }
+ ```
+ """
+ await ctx.info(
+ "Starting dataset query: dataset_id=%s, metrics=%s, columns=%s, "
+ "row_limit=%s"
+ % (
+ request.dataset_id,
+ request.metrics,
+ request.columns,
+ request.row_limit,
+ )
+ )
+
+ try:
+ from superset.commands.chart.data.get_data_command import ChartDataCommand
+ from superset.common.query_context_factory import QueryContextFactory
+ from superset.connectors.sqla.models import SqlaTable
+
+ # ------------------------------------------------------------------
+ # Step 1: Resolve dataset
+ # ------------------------------------------------------------------
+ await ctx.report_progress(1, 5, "Looking up dataset")
+ eager_options = [
+ subqueryload(SqlaTable.columns),
+ subqueryload(SqlaTable.metrics),
+ joinedload(SqlaTable.database),
+ ]
+
+ with event_logger.log_context(action="mcp.query_dataset.lookup"):
+ dataset = _resolve_dataset(request.dataset_id, eager_options)
+
+ if dataset is None:
+ await ctx.error("Dataset not found: identifier=%s" % (request.dataset_id,))
+ return DatasetError.create(
+ error=f"No dataset found with identifier: {request.dataset_id}",
+ error_type="NotFound",
+ )
+
+ dataset_name = getattr(dataset, "table_name", None) or f"Dataset {dataset.id}"
+ await ctx.info(
+ "Dataset found: id=%s, name=%s, columns=%s, metrics=%s"
+ % (
+ dataset.id,
+ dataset_name,
+ len(dataset.columns),
+ len(dataset.metrics),
+ )
+ )
+
+ # ------------------------------------------------------------------
+ # Step 2: Validate requested columns and metrics
+ # ------------------------------------------------------------------
+ await ctx.report_progress(2, 5, "Validating columns and metrics")
+ valid_columns = {c.column_name for c in dataset.columns}
+ valid_metrics = {m.metric_name for m in dataset.metrics}
+
+ validation_errors: list[str] = []
+ validation_errors.extend(
+ _validate_names(request.columns, valid_columns, "column")
+ )
+ validation_errors.extend(
+ _validate_names(request.metrics, valid_metrics, "metric")
+ )
+ # Validate filter column names against dataset columns
+ filter_cols = [f.col for f in request.filters]
+ validation_errors.extend(
+ _validate_names(filter_cols, valid_columns, "filter column")
+ )
+ # Validate order_by names against columns + metrics
+ if request.order_by:
+ valid_orderby = valid_columns | valid_metrics
+ validation_errors.extend(
+ _validate_names(request.order_by, valid_orderby, "order_by")
+ )
+
+ if validation_errors:
+ error_msg = "; ".join(validation_errors)
+ await ctx.error("Validation failed: %s" % (error_msg,))
+ return DatasetError.create(
+ error=error_msg,
+ error_type="ValidationError",
+ )
+
+ # ------------------------------------------------------------------
+ # Step 3: Build filters and time range
+ # ------------------------------------------------------------------
+ warnings: list[str] = []
+ query_filters: list[dict[str, Any]] = [
+ {"col": f.col, "op": f.op, "val": f.val} for f in request.filters
+ ]
+ granularity: str | None = None
+
+ if request.time_range:
+ temporal_col = request.time_column or getattr(
+ dataset, "main_dttm_col", None
+ )
+ if not temporal_col:
+ await ctx.error("time_range provided but no temporal column available")
+ return DatasetError.create(
+ error=(
+ "time_range was provided but no temporal column is available. "
+ "Either set time_column explicitly or ensure the dataset has "
+ "a main datetime column configured."
+ ),
+ error_type="ValidationError",
+ )
+ # Validate that the temporal column actually exists on the dataset
+ if temporal_col not in valid_columns:
+ await ctx.error("time_column '%s' not found on dataset" % temporal_col)
+ return DatasetError.create(
+ error=(
+ f"time_column '{temporal_col}' does not exist on this dataset."
+ ),
+ error_type="ValidationError",
+ )
+ # Warn if the chosen temporal column isn't marked as datetime
+ dttm_cols = {c.column_name for c in dataset.columns if c.is_dttm}
+ if temporal_col not in dttm_cols:
+ warnings.append(
+ f"Column '{temporal_col}' is not marked as a datetime "
+ f"column on this dataset. Time filtering may not work "
+ f"as expected."
+ )
+
+ query_filters.append(
+ {
+ "col": temporal_col,
+ "op": "TEMPORAL_RANGE",
+ "val": request.time_range,
+ }
+ )
+ granularity = temporal_col
+ await ctx.debug(
+ "Time filter: column=%s, range=%s" % (temporal_col, request.time_range)
+ )
+
+ # ------------------------------------------------------------------
+ # Step 4: Build query dict
+ # ------------------------------------------------------------------
+ await ctx.report_progress(3, 5, "Building query")
+ query_dict: dict[str, Any] = {
+ "filters": query_filters,
+ "columns": request.columns,
+ "metrics": request.metrics,
+ "row_limit": request.row_limit,
+ "order_desc": request.order_desc,
+ }
+ if granularity:
+ query_dict["granularity"] = granularity
+ if request.order_by:
+ # OrderBy = tuple[Metric | Column, bool] where bool is ascending
+ query_dict["orderby"] = [
+ (col, not request.order_desc) for col in request.order_by
+ ]
+
+ await ctx.debug("Query dict keys: %s" % (sorted(query_dict.keys()),))
+
+ # ------------------------------------------------------------------
+ # Step 5: Create QueryContext and execute
+ # ------------------------------------------------------------------
+ await ctx.report_progress(4, 5, "Executing query")
+ start_time = time.time()
+
+ with event_logger.log_context(action="mcp.query_dataset.execute"):
+ factory = QueryContextFactory()
+ # datasource_type is "table" because this tool queries SqlaTable
+ # datasets (Superset's built-in semantic layer). External semantic
+ # layers (dbt, Snowflake Cortex, etc.) use "semantic_view" and have
+ # a different query path — see SemanticView + mapper.py.
+ query_context = factory.create(
+ datasource={"id": dataset.id, "type": "table"},
+ queries=[query_dict],
+ form_data={},
+ force=request.force_refresh,
+ custom_cache_timeout=request.cache_timeout,
+ )
+
+ command = ChartDataCommand(query_context)
+ command.validate()
+ result = command.run()
+
+ query_duration_ms = int((time.time() - start_time) * 1000)
+
+ if not result or "queries" not in result or len(result["queries"]) == 0:
+ await ctx.warning("Query returned no results for dataset %s" % dataset.id)
+ return DatasetError.create(
+ error="Query returned no results.",
+ error_type="EmptyQuery",
+ )
+
+ # ------------------------------------------------------------------
+ # Step 6: Format response
+ # ------------------------------------------------------------------
+ await ctx.report_progress(5, 5, "Formatting results")
+ query_result = result["queries"][0]
+ data = query_result.get("data", [])
+ raw_columns = query_result.get("colnames", [])
+
+ if not data:
+ return QueryDatasetResponse(
+ dataset_id=dataset.id,
+ dataset_name=dataset_name,
+ columns=[],
+ data=[],
+ row_count=0,
+ total_rows=0,
+ summary=f"Query on '{dataset_name}' returned no data.",
+ performance=PerformanceMetadata(
+ query_duration_ms=query_duration_ms,
+ cache_status="no_data",
+ ),
+ cache_status=get_cache_status_from_result(
+ query_result, force_refresh=request.force_refresh
+ ),
+ applied_filters=request.filters,
+ warnings=warnings,
+ )
+
+ # Build column metadata
+ columns_meta: list[DataColumn] = []
+ for col_name in raw_columns:
+ sample_values = [
+ row.get(col_name) for row in data[:3] if row.get(col_name) is not None
+ ]
+ data_type = "string"
+ if sample_values:
+ if all(isinstance(v, bool) for v in sample_values):
+ data_type = "boolean"
+ elif all(isinstance(v, (int, float)) for v in sample_values):
+ data_type = "numeric"
+
+ columns_meta.append(
+ DataColumn(
+ name=col_name,
+ display_name=col_name.replace("_", " ").title(),
+ data_type=data_type,
+ sample_values=sample_values[:3],
+ null_count=sum(1 for row in data if row.get(col_name) is None),
+ unique_count=len({str(row.get(col_name)) for row in data}),
+ )
+ )
Review Comment:
`null_count` and `unique_count` do a full scan of `data` for each returned column, and `unique_count` materializes a full set, making response formatting O(rows × cols) with potentially large memory overhead (row_limit allows up to 50k). Consider computing these stats only on a bounded sample (e.g., first N rows), making them optional behind a request flag, or computing them in a single pass over rows to avoid repeated scans and large per-column sets.
##########
superset/mcp_service/dataset/schemas.py:
##########
@@ -380,6 +385,146 @@ class CreateVirtualDatasetResponse(BaseModel):
)
+VALID_FILTER_OPS = Literal[
+ "==",
+ "!=",
+ ">",
+ "<",
+ ">=",
+ "<=",
+ "LIKE",
+ "NOT LIKE",
+ "ILIKE",
+ "NOT ILIKE",
+ "IN",
+ "NOT IN",
+ "IS NULL",
+ "IS NOT NULL",
+ "IS TRUE",
+ "IS FALSE",
+ "TEMPORAL_RANGE",
+]
+
+
+class QueryDatasetFilter(BaseModel):
+ """A single filter condition for dataset queries."""
+
+ col: str = Field(..., description="Column name to filter on")
+ op: VALID_FILTER_OPS = Field(
+ ...,
+ description=(
+ 'Filter operator. Use "==" for equals, "!=" for not equals, '
+ '"IN" / "NOT IN" for membership, "IS NULL" / "IS NOT NULL", '
+ '"LIKE" for pattern matching, "TEMPORAL_RANGE" for time filters.'
+ ),
+ )
+ val: Any = Field(
+ default=None,
+ description="Filter value (omit for IS NULL/IS NOT NULL)",
+ )
+
+
+class QueryDatasetRequest(QueryCacheControl):
+ """Request schema for query_dataset tool."""
+
+ dataset_id: int | str = Field(
+ ...,
+ description="Dataset identifier — numeric ID or UUID string.",
+ )
+ metrics: List[str] = Field(
+ default_factory=list,
+ description=(
+ "Saved metric names to compute (e.g. ['count', 'avg_revenue']). "
+ "Use get_dataset_info to discover available metrics."
+ ),
+ )
+ columns: List[str] = Field(
+ default_factory=list,
+ description=(
+ "Column/dimension names for GROUP BY or SELECT "
+ "(e.g. ['category', 'region']). "
+ "Use get_dataset_info to discover available columns."
+ ),
+ )
+ filters: List[QueryDatasetFilter] = Field(
+ default_factory=list,
+ description=(
+ 'Filter conditions (e.g. [{"col": "status", "op": "=", "val": "active"}]).'
Review Comment:
The schema examples/documentation advertise `op: "="`, but `VALID_FILTER_OPS` does not allow `"="` (only `"=="`). This will cause client confusion and immediate validation failures for a request that follows the example. Update the example to use `"=="`, or explicitly support `"="` (and normalize it to `"=="`) to match the documented contract.
```suggestion
'Filter conditions (e.g. [{"col": "status", "op": "==", "val": "active"}]).'
```
##########
superset/mcp_service/dataset/tool/query_dataset.py:
##########
@@ -0,0 +1,434 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+MCP tool: query_dataset
+
+Query a dataset using its semantic layer (saved metrics, calculated columns,
+dimensions) without requiring a saved chart.
+"""
+
+import difflib
+import logging
+import time
+from typing import Any
+
+from fastmcp import Context
+from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.orm import joinedload, subqueryload
+from superset_core.mcp.decorators import tool, ToolAnnotations
+
+from superset.commands.exceptions import CommandException
+from superset.exceptions import OAuth2Error, OAuth2RedirectError, SupersetException
+from superset.extensions import event_logger
+from superset.mcp_service.chart.schemas import DataColumn, PerformanceMetadata
+from superset.mcp_service.dataset.schemas import (
+ DatasetError,
+ QueryDatasetRequest,
+ QueryDatasetResponse,
+)
+from superset.mcp_service.utils import _is_uuid
+from superset.mcp_service.utils.cache_utils import get_cache_status_from_result
+from superset.mcp_service.utils.oauth2_utils import build_oauth2_redirect_message
+
+logger = logging.getLogger(__name__)
+
+
+def _resolve_dataset(identifier: int | str, eager_options: list[Any]) -> Any | None:
+ """Resolve a dataset by int ID or UUID string.
+
+ Replicates the identifier resolution logic from ModelGetInfoCore._find_object().
+ """
+ from superset.daos.dataset import DatasetDAO
+
+ opts = eager_options or None
+
+ if isinstance(identifier, int):
+ return DatasetDAO.find_by_id(identifier, query_options=opts)
+
+ # Try parsing as int
+ try:
+ id_val = int(identifier)
+ return DatasetDAO.find_by_id(id_val, query_options=opts)
+ except (ValueError, TypeError):
+ pass
+
+ # Try UUID
+ if _is_uuid(str(identifier)):
+ return DatasetDAO.find_by_id(identifier, id_column="uuid", query_options=opts)
+
+ return None
+
+
+def _validate_names(
+ requested: list[str],
+ valid: set[str],
+ kind: str,
+) -> list[str]:
+ """Return list of error messages for names not found in *valid*.
+
+ Includes close-match suggestions when available.
+ """
+ errors: list[str] = []
+ for name in requested:
+ if name not in valid:
+ suggestions = difflib.get_close_matches(name, valid, n=3, cutoff=0.6)
+ msg = f"Unknown {kind}: '{name}'"
+ if suggestions:
+ msg += f". Did you mean: {', '.join(suggestions)}?"
+ errors.append(msg)
+ return errors
+
+
+@tool(
+ tags=["data"],
+ class_permission_name="Dataset",
+ annotations=ToolAnnotations(
+ title="Query dataset",
+ readOnlyHint=True,
+ destructiveHint=False,
+ ),
+)
+async def query_dataset( # noqa: C901
+ request: QueryDatasetRequest, ctx: Context
+) -> QueryDatasetResponse | DatasetError:
+ """Query a dataset using its semantic layer (saved metrics, dimensions, filters).
+
+ Returns tabular data without requiring a saved chart. Use this when you want
+ to compute saved metrics, group by dimensions, or apply filters directly
+ against a dataset's curated semantic layer.
+
+ Workflow:
+ 1. list_datasets -> find a dataset
+ 2. get_dataset_info -> discover available columns and metrics
+ 3. query_dataset -> query using metric names and column names
+
+ Example:
+ ```json
+ {
+ "dataset_id": 123,
+ "metrics": ["count", "avg_revenue"],
+ "columns": ["product_category"],
+ "time_range": "Last 7 days",
+ "row_limit": 100
+ }
+ ```
+ """
+ await ctx.info(
+ "Starting dataset query: dataset_id=%s, metrics=%s, columns=%s, "
+ "row_limit=%s"
+ % (
+ request.dataset_id,
+ request.metrics,
+ request.columns,
+ request.row_limit,
+ )
+ )
+
+ try:
+ from superset.commands.chart.data.get_data_command import ChartDataCommand
+ from superset.common.query_context_factory import QueryContextFactory
+ from superset.connectors.sqla.models import SqlaTable
+
+ # ------------------------------------------------------------------
+ # Step 1: Resolve dataset
+ # ------------------------------------------------------------------
+ await ctx.report_progress(1, 5, "Looking up dataset")
+ eager_options = [
+ subqueryload(SqlaTable.columns),
+ subqueryload(SqlaTable.metrics),
+ joinedload(SqlaTable.database),
+ ]
+
+ with event_logger.log_context(action="mcp.query_dataset.lookup"):
+ dataset = _resolve_dataset(request.dataset_id, eager_options)
+
+ if dataset is None:
+ await ctx.error("Dataset not found: identifier=%s" % (request.dataset_id,))
+ return DatasetError.create(
+ error=f"No dataset found with identifier: {request.dataset_id}",
+ error_type="NotFound",
+ )
+
+ dataset_name = getattr(dataset, "table_name", None) or f"Dataset {dataset.id}"
+ await ctx.info(
+ "Dataset found: id=%s, name=%s, columns=%s, metrics=%s"
+ % (
+ dataset.id,
+ dataset_name,
+ len(dataset.columns),
+ len(dataset.metrics),
+ )
+ )
+
+ # ------------------------------------------------------------------
+ # Step 2: Validate requested columns and metrics
+ # ------------------------------------------------------------------
+ await ctx.report_progress(2, 5, "Validating columns and metrics")
+ valid_columns = {c.column_name for c in dataset.columns}
+ valid_metrics = {m.metric_name for m in dataset.metrics}
+
+ validation_errors: list[str] = []
+ validation_errors.extend(
+ _validate_names(request.columns, valid_columns, "column")
+ )
+ validation_errors.extend(
+ _validate_names(request.metrics, valid_metrics, "metric")
+ )
+ # Validate filter column names against dataset columns
+ filter_cols = [f.col for f in request.filters]
+ validation_errors.extend(
+ _validate_names(filter_cols, valid_columns, "filter column")
+ )
+ # Validate order_by names against columns + metrics
+ if request.order_by:
+ valid_orderby = valid_columns | valid_metrics
+ validation_errors.extend(
+ _validate_names(request.order_by, valid_orderby, "order_by")
+ )
+
+ if validation_errors:
+ error_msg = "; ".join(validation_errors)
+ await ctx.error("Validation failed: %s" % (error_msg,))
+ return DatasetError.create(
+ error=error_msg,
+ error_type="ValidationError",
+ )
+
+ # ------------------------------------------------------------------
+ # Step 3: Build filters and time range
+ # ------------------------------------------------------------------
+ warnings: list[str] = []
+ query_filters: list[dict[str, Any]] = [
+ {"col": f.col, "op": f.op, "val": f.val} for f in request.filters
+ ]
+ granularity: str | None = None
+
+ if request.time_range:
+ temporal_col = request.time_column or getattr(
+ dataset, "main_dttm_col", None
+ )
+ if not temporal_col:
+ await ctx.error("time_range provided but no temporal column available")
+ return DatasetError.create(
+ error=(
+ "time_range was provided but no temporal column is available. "
+ "Either set time_column explicitly or ensure the dataset has "
+ "a main datetime column configured."
+ ),
+ error_type="ValidationError",
+ )
+ # Validate that the temporal column actually exists on the dataset
+ if temporal_col not in valid_columns:
+ await ctx.error("time_column '%s' not found on dataset" % temporal_col)
+ return DatasetError.create(
+ error=(
+ f"time_column '{temporal_col}' does not exist on this dataset."
+ ),
+ error_type="ValidationError",
+ )
+ # Warn if the chosen temporal column isn't marked as datetime
+ dttm_cols = {c.column_name for c in dataset.columns if c.is_dttm}
+ if temporal_col not in dttm_cols:
+ warnings.append(
+ f"Column '{temporal_col}' is not marked as a datetime "
+ f"column on this dataset. Time filtering may not work "
+ f"as expected."
+ )
+
+ query_filters.append(
+ {
+ "col": temporal_col,
+ "op": "TEMPORAL_RANGE",
+ "val": request.time_range,
+ }
+ )
+ granularity = temporal_col
+ await ctx.debug(
+ "Time filter: column=%s, range=%s" % (temporal_col, request.time_range)
+ )
+
+ # ------------------------------------------------------------------
+ # Step 4: Build query dict
+ # ------------------------------------------------------------------
+ await ctx.report_progress(3, 5, "Building query")
+ query_dict: dict[str, Any] = {
+ "filters": query_filters,
+ "columns": request.columns,
+ "metrics": request.metrics,
+ "row_limit": request.row_limit,
+ "order_desc": request.order_desc,
+ }
+ if granularity:
+ query_dict["granularity"] = granularity
+ if request.order_by:
+ # OrderBy = tuple[Metric | Column, bool] where bool is ascending
+ query_dict["orderby"] = [
+ (col, not request.order_desc) for col in request.order_by
+ ]
+
+ await ctx.debug("Query dict keys: %s" % (sorted(query_dict.keys()),))
+
+ # ------------------------------------------------------------------
+ # Step 5: Create QueryContext and execute
+ # ------------------------------------------------------------------
+ await ctx.report_progress(4, 5, "Executing query")
+ start_time = time.time()
+
+ with event_logger.log_context(action="mcp.query_dataset.execute"):
+ factory = QueryContextFactory()
+ # datasource_type is "table" because this tool queries SqlaTable
+ # datasets (Superset's built-in semantic layer). External semantic
+ # layers (dbt, Snowflake Cortex, etc.) use "semantic_view" and have
+ # a different query path — see SemanticView + mapper.py.
+ query_context = factory.create(
+ datasource={"id": dataset.id, "type": "table"},
+ queries=[query_dict],
+ form_data={},
+ force=request.force_refresh,
+ custom_cache_timeout=request.cache_timeout,
+ )
+
+ command = ChartDataCommand(query_context)
+ command.validate()
+ result = command.run()
+
+ query_duration_ms = int((time.time() - start_time) * 1000)
+
+ if not result or "queries" not in result or len(result["queries"]) == 0:
+ await ctx.warning("Query returned no results for dataset %s" % dataset.id)
+ return DatasetError.create(
+ error="Query returned no results.",
+ error_type="EmptyQuery",
+ )
+
+ # ------------------------------------------------------------------
+ # Step 6: Format response
+ # ------------------------------------------------------------------
+ await ctx.report_progress(5, 5, "Formatting results")
+ query_result = result["queries"][0]
+ data = query_result.get("data", [])
+ raw_columns = query_result.get("colnames", [])
+
+ if not data:
+ return QueryDatasetResponse(
+ dataset_id=dataset.id,
+ dataset_name=dataset_name,
+ columns=[],
+ data=[],
+ row_count=0,
+ total_rows=0,
+ summary=f"Query on '{dataset_name}' returned no data.",
+ performance=PerformanceMetadata(
+ query_duration_ms=query_duration_ms,
+ cache_status="no_data",
+ ),
+ cache_status=get_cache_status_from_result(
+ query_result, force_refresh=request.force_refresh
+ ),
+ applied_filters=request.filters,
+ warnings=warnings,
+ )
+
+ # Build column metadata
+ columns_meta: list[DataColumn] = []
+ for col_name in raw_columns:
+ sample_values = [
+ row.get(col_name) for row in data[:3] if row.get(col_name) is not None
+ ]
+ data_type = "string"
+ if sample_values:
+ if all(isinstance(v, bool) for v in sample_values):
+ data_type = "boolean"
+ elif all(isinstance(v, (int, float)) for v in sample_values):
+ data_type = "numeric"
+
+ columns_meta.append(
+ DataColumn(
+ name=col_name,
+ display_name=col_name.replace("_", " ").title(),
+ data_type=data_type,
+ sample_values=sample_values[:3],
+ null_count=sum(1 for row in data if row.get(col_name) is None),
+ unique_count=len({str(row.get(col_name)) for row in data}),
Review Comment:
`unique_count` currently includes nulls (because `row.get(col_name)` may be `None`, which becomes `"None"`), which will overcount distinct values compared to typical “distinct non-null” semantics. Exclude `None` (and potentially missing keys) from the unique set so `unique_count` aligns with `null_count` and expected behavior.
```suggestion
unique_count=len(
{
str(value)
for row in data
if (value := row.get(col_name)) is not None
}
),
```
##########
superset/mcp_service/dataset/tool/query_dataset.py:
##########
@@ -0,0 +1,434 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+MCP tool: query_dataset
+
+Query a dataset using its semantic layer (saved metrics, calculated columns,
+dimensions) without requiring a saved chart.
+"""
+
+import difflib
+import logging
+import time
+from typing import Any
+
+from fastmcp import Context
+from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.orm import joinedload, subqueryload
+from superset_core.mcp.decorators import tool, ToolAnnotations
+
+from superset.commands.exceptions import CommandException
+from superset.exceptions import OAuth2Error, OAuth2RedirectError, SupersetException
+from superset.extensions import event_logger
+from superset.mcp_service.chart.schemas import DataColumn, PerformanceMetadata
+from superset.mcp_service.dataset.schemas import (
+ DatasetError,
+ QueryDatasetRequest,
+ QueryDatasetResponse,
+)
+from superset.mcp_service.utils import _is_uuid
+from superset.mcp_service.utils.cache_utils import get_cache_status_from_result
+from superset.mcp_service.utils.oauth2_utils import build_oauth2_redirect_message
+
+logger = logging.getLogger(__name__)
+
+
+def _resolve_dataset(identifier: int | str, eager_options: list[Any]) -> Any | None:
+ """Resolve a dataset by int ID or UUID string.
+
+ Replicates the identifier resolution logic from ModelGetInfoCore._find_object().
+ """
+ from superset.daos.dataset import DatasetDAO
+
+ opts = eager_options or None
+
+ if isinstance(identifier, int):
+ return DatasetDAO.find_by_id(identifier, query_options=opts)
+
+ # Try parsing as int
+ try:
+ id_val = int(identifier)
+ return DatasetDAO.find_by_id(id_val, query_options=opts)
+ except (ValueError, TypeError):
+ pass
+
+ # Try UUID
+ if _is_uuid(str(identifier)):
+ return DatasetDAO.find_by_id(identifier, id_column="uuid", query_options=opts)
+
+ return None
+
+
+def _validate_names(
+ requested: list[str],
+ valid: set[str],
+ kind: str,
+) -> list[str]:
+ """Return list of error messages for names not found in *valid*.
+
+ Includes close-match suggestions when available.
+ """
+ errors: list[str] = []
+ for name in requested:
+ if name not in valid:
+ suggestions = difflib.get_close_matches(name, valid, n=3, cutoff=0.6)
+ msg = f"Unknown {kind}: '{name}'"
+ if suggestions:
+ msg += f". Did you mean: {', '.join(suggestions)}?"
+ errors.append(msg)
+ return errors
+
+
+@tool(
+ tags=["data"],
+ class_permission_name="Dataset",
+ annotations=ToolAnnotations(
+ title="Query dataset",
+ readOnlyHint=True,
+ destructiveHint=False,
+ ),
+)
+async def query_dataset( # noqa: C901
+ request: QueryDatasetRequest, ctx: Context
+) -> QueryDatasetResponse | DatasetError:
+ """Query a dataset using its semantic layer (saved metrics, dimensions, filters).
+
+ Returns tabular data without requiring a saved chart. Use this when you want
+ to compute saved metrics, group by dimensions, or apply filters directly
+ against a dataset's curated semantic layer.
+
+ Workflow:
+ 1. list_datasets -> find a dataset
+ 2. get_dataset_info -> discover available columns and metrics
+ 3. query_dataset -> query using metric names and column names
+
+ Example:
+ ```json
+ {
+ "dataset_id": 123,
+ "metrics": ["count", "avg_revenue"],
+ "columns": ["product_category"],
+ "time_range": "Last 7 days",
+ "row_limit": 100
+ }
+ ```
+ """
+ await ctx.info(
+ "Starting dataset query: dataset_id=%s, metrics=%s, columns=%s, "
+ "row_limit=%s"
+ % (
+ request.dataset_id,
+ request.metrics,
+ request.columns,
+ request.row_limit,
+ )
+ )
+
+ try:
+ from superset.commands.chart.data.get_data_command import ChartDataCommand
+ from superset.common.query_context_factory import QueryContextFactory
+ from superset.connectors.sqla.models import SqlaTable
+
+ # ------------------------------------------------------------------
+ # Step 1: Resolve dataset
+ # ------------------------------------------------------------------
+ await ctx.report_progress(1, 5, "Looking up dataset")
+ eager_options = [
+ subqueryload(SqlaTable.columns),
+ subqueryload(SqlaTable.metrics),
+ joinedload(SqlaTable.database),
+ ]
+
+ with event_logger.log_context(action="mcp.query_dataset.lookup"):
+ dataset = _resolve_dataset(request.dataset_id, eager_options)
+
+ if dataset is None:
+ await ctx.error("Dataset not found: identifier=%s" % (request.dataset_id,))
+ return DatasetError.create(
+ error=f"No dataset found with identifier: {request.dataset_id}",
+ error_type="NotFound",
+ )
+
+ dataset_name = getattr(dataset, "table_name", None) or f"Dataset {dataset.id}"
+ await ctx.info(
+ "Dataset found: id=%s, name=%s, columns=%s, metrics=%s"
+ % (
+ dataset.id,
+ dataset_name,
+ len(dataset.columns),
+ len(dataset.metrics),
+ )
+ )
+
+ # ------------------------------------------------------------------
+ # Step 2: Validate requested columns and metrics
+ # ------------------------------------------------------------------
+ await ctx.report_progress(2, 5, "Validating columns and metrics")
+ valid_columns = {c.column_name for c in dataset.columns}
+ valid_metrics = {m.metric_name for m in dataset.metrics}
+
+ validation_errors: list[str] = []
+ validation_errors.extend(
+ _validate_names(request.columns, valid_columns, "column")
+ )
+ validation_errors.extend(
+ _validate_names(request.metrics, valid_metrics, "metric")
+ )
+ # Validate filter column names against dataset columns
+ filter_cols = [f.col for f in request.filters]
+ validation_errors.extend(
+ _validate_names(filter_cols, valid_columns, "filter column")
+ )
+ # Validate order_by names against columns + metrics
+ if request.order_by:
+ valid_orderby = valid_columns | valid_metrics
+ validation_errors.extend(
+ _validate_names(request.order_by, valid_orderby, "order_by")
+ )
+
+ if validation_errors:
+ error_msg = "; ".join(validation_errors)
+ await ctx.error("Validation failed: %s" % (error_msg,))
+ return DatasetError.create(
+ error=error_msg,
+ error_type="ValidationError",
+ )
+
+ # ------------------------------------------------------------------
+ # Step 3: Build filters and time range
+ # ------------------------------------------------------------------
+ warnings: list[str] = []
+ query_filters: list[dict[str, Any]] = [
+ {"col": f.col, "op": f.op, "val": f.val} for f in request.filters
+ ]
+ granularity: str | None = None
+
+ if request.time_range:
+ temporal_col = request.time_column or getattr(
+ dataset, "main_dttm_col", None
+ )
+ if not temporal_col:
+ await ctx.error("time_range provided but no temporal column available")
+ return DatasetError.create(
+ error=(
+ "time_range was provided but no temporal column is available. "
+ "Either set time_column explicitly or ensure the dataset has "
+ "a main datetime column configured."
+ ),
+ error_type="ValidationError",
+ )
+ # Validate that the temporal column actually exists on the dataset
+ if temporal_col not in valid_columns:
+ await ctx.error("time_column '%s' not found on dataset" % temporal_col)
+ return DatasetError.create(
+ error=(
+ f"time_column '{temporal_col}' does not exist on this dataset."
+ ),
+ error_type="ValidationError",
+ )
+ # Warn if the chosen temporal column isn't marked as datetime
+ dttm_cols = {c.column_name for c in dataset.columns if c.is_dttm}
+ if temporal_col not in dttm_cols:
+ warnings.append(
+ f"Column '{temporal_col}' is not marked as a datetime "
+ f"column on this dataset. Time filtering may not work "
+ f"as expected."
+ )
+
+ query_filters.append(
+ {
+ "col": temporal_col,
+ "op": "TEMPORAL_RANGE",
+ "val": request.time_range,
+ }
+ )
+ granularity = temporal_col
+ await ctx.debug(
+ "Time filter: column=%s, range=%s" % (temporal_col, request.time_range)
+ )
+
+ # ------------------------------------------------------------------
+ # Step 4: Build query dict
+ # ------------------------------------------------------------------
+ await ctx.report_progress(3, 5, "Building query")
+ query_dict: dict[str, Any] = {
+ "filters": query_filters,
+ "columns": request.columns,
+ "metrics": request.metrics,
+ "row_limit": request.row_limit,
+ "order_desc": request.order_desc,
+ }
+ if granularity:
+ query_dict["granularity"] = granularity
+ if request.order_by:
+ # OrderBy = tuple[Metric | Column, bool] where bool is ascending
+ query_dict["orderby"] = [
+ (col, not request.order_desc) for col in request.order_by
+ ]
+
+ await ctx.debug("Query dict keys: %s" % (sorted(query_dict.keys()),))
+
+ # ------------------------------------------------------------------
+ # Step 5: Create QueryContext and execute
+ # ------------------------------------------------------------------
+ await ctx.report_progress(4, 5, "Executing query")
+ start_time = time.time()
+
+ with event_logger.log_context(action="mcp.query_dataset.execute"):
+ factory = QueryContextFactory()
+ # datasource_type is "table" because this tool queries SqlaTable
+ # datasets (Superset's built-in semantic layer). External semantic
+ # layers (dbt, Snowflake Cortex, etc.) use "semantic_view" and have
+ # a different query path — see SemanticView + mapper.py.
+ query_context = factory.create(
+ datasource={"id": dataset.id, "type": "table"},
+ queries=[query_dict],
+ form_data={},
+ force=request.force_refresh,
+ custom_cache_timeout=request.cache_timeout,
+ )
+
+ command = ChartDataCommand(query_context)
+ command.validate()
+ result = command.run()
+
+ query_duration_ms = int((time.time() - start_time) * 1000)
+
+ if not result or "queries" not in result or len(result["queries"]) == 0:
+ await ctx.warning("Query returned no results for dataset %s" % dataset.id)
+ return DatasetError.create(
+ error="Query returned no results.",
+ error_type="EmptyQuery",
+ )
+
+ # ------------------------------------------------------------------
+ # Step 6: Format response
+ # ------------------------------------------------------------------
+ await ctx.report_progress(5, 5, "Formatting results")
+ query_result = result["queries"][0]
+ data = query_result.get("data", [])
+ raw_columns = query_result.get("colnames", [])
+
+ if not data:
+ return QueryDatasetResponse(
+ dataset_id=dataset.id,
+ dataset_name=dataset_name,
+ columns=[],
+ data=[],
+ row_count=0,
+ total_rows=0,
+ summary=f"Query on '{dataset_name}' returned no data.",
+ performance=PerformanceMetadata(
+ query_duration_ms=query_duration_ms,
+ cache_status="no_data",
+ ),
+ cache_status=get_cache_status_from_result(
+ query_result, force_refresh=request.force_refresh
+ ),
+ applied_filters=request.filters,
+ warnings=warnings,
+ )
+
+ # Build column metadata
+ columns_meta: list[DataColumn] = []
+ for col_name in raw_columns:
+ sample_values = [
+ row.get(col_name) for row in data[:3] if row.get(col_name) is not None
+ ]
+ data_type = "string"
+ if sample_values:
+ if all(isinstance(v, bool) for v in sample_values):
+ data_type = "boolean"
+ elif all(isinstance(v, (int, float)) for v in sample_values):
+ data_type = "numeric"
+
+ columns_meta.append(
+ DataColumn(
+ name=col_name,
+ display_name=col_name.replace("_", " ").title(),
+ data_type=data_type,
+ sample_values=sample_values[:3],
+ null_count=sum(1 for row in data if row.get(col_name) is None),
+ unique_count=len({str(row.get(col_name)) for row in data}),
+ )
+ )
+
+ cache_status = get_cache_status_from_result(
+ query_result, force_refresh=request.force_refresh
+ )
+
+ cache_label = "cached" if cache_status and cache_status.cache_hit else "fresh"
+ summary = (
+ f"Dataset '{dataset_name}': {len(data)} rows, "
+ f"{len(raw_columns)} columns ({cache_label})."
+ )
+
+ await ctx.info(
+ "Query complete: rows=%s, columns=%s, duration=%sms"
+ % (len(data), len(raw_columns), query_duration_ms)
+ )
+
+ return QueryDatasetResponse(
+ dataset_id=dataset.id,
+ dataset_name=dataset_name,
+ columns=columns_meta,
+ data=data,
+ row_count=len(data),
+ total_rows=query_result.get("rowcount"),
+ summary=summary,
+ performance=PerformanceMetadata(
+ query_duration_ms=query_duration_ms,
+ cache_status=cache_label,
+ ),
+ cache_status=cache_status,
+ applied_filters=request.filters,
+ warnings=warnings,
+ )
+
+ except OAuth2RedirectError as exc:
+ redirect_msg = build_oauth2_redirect_message(exc)
+ await ctx.error("OAuth2 redirect required: %s" % (redirect_msg,))
+ return DatasetError.create(
+ error=redirect_msg,
+ error_type="OAuth2Redirect",
+ )
+
+ except OAuth2Error as exc:
+ await ctx.error("OAuth2 error: %s" % (str(exc),))
+ return DatasetError.create(
+ error=f"OAuth2 authentication error: {exc}",
+ error_type="OAuth2Error",
+ )
+
+ except (CommandException, SupersetException) as exc:
+ await ctx.error("Query failed: %s" % (str(exc),))
+ return DatasetError.create(
+ error=f"Query execution failed: {exc}",
+ error_type="QueryError",
+ )
+
+ except SQLAlchemyError as exc:
+ await ctx.error("Database error: %s" % (str(exc),))
+ return DatasetError.create(
+ error=f"Database error: {exc}",
+ error_type="DatabaseError",
+ )
+
+ except Exception as exc:
+ await ctx.error("Unexpected error: %s: %s" % (type(exc).__name__, str(exc)))
+ raise
Review Comment:
Re-raising the generic exception makes the tool fail with an unstructured error response, which can reduce MCP reliability and consistency (especially for clients expecting `DatasetError`). Consider returning a `DatasetError` with a stable `error_type` (e.g., `"UnexpectedError"`) while logging the exception details server-side.
```suggestion
logger.exception(
"Unexpected error while querying dataset: %s: %s",
type(exc).__name__,
str(exc),
)
await ctx.error("Unexpected error: %s: %s" % (type(exc).__name__, str(exc)))
return DatasetError.create(
error="An unexpected error occurred while querying the dataset.",
error_type="UnexpectedError",
)
```
--
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: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] feat(mcp): add query_dataset tool to query datasets using semantic layer [superset]
Posted by "aminghadersohi (via GitHub)" <gi...@apache.org>.
aminghadersohi commented on code in PR #39727:
URL: https://github.com/apache/superset/pull/39727#discussion_r3164803573
##########
tests/unit_tests/mcp_service/dataset/tool/test_query_dataset.py:
##########
@@ -0,0 +1,737 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Tests for the query_dataset MCP tool."""
+
+from __future__ import annotations
+
+import importlib
+from typing import Any
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+from fastmcp import Client
+
+from superset.mcp_service.app import mcp
+from superset.utils import json
+
+query_dataset_module = importlib.import_module(
+ "superset.mcp_service.dataset.tool.query_dataset"
+)
+
+
+@pytest.fixture
+def mcp_server():
+ return mcp
+
+
+@pytest.fixture(autouse=True)
+def mock_auth():
Review Comment:
Fixed in ffcd4b1d25 — `mock_auth()` now annotated as `-> Generator[MagicMock, None, None]`.
##########
tests/unit_tests/mcp_service/dataset/tool/test_query_dataset.py:
##########
@@ -0,0 +1,737 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Tests for the query_dataset MCP tool."""
+
+from __future__ import annotations
+
+import importlib
+from typing import Any
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+from fastmcp import Client
+
+from superset.mcp_service.app import mcp
+from superset.utils import json
+
+query_dataset_module = importlib.import_module(
+ "superset.mcp_service.dataset.tool.query_dataset"
+)
+
+
+@pytest.fixture
+def mcp_server():
+ return mcp
+
+
+@pytest.fixture(autouse=True)
+def mock_auth():
+ """Mock authentication for all tests."""
+ with patch("superset.mcp_service.auth.get_user_from_request") as mock_get_user:
+ mock_user = Mock()
+ mock_user.id = 1
+ mock_user.username = "admin"
+ mock_get_user.return_value = mock_user
+ yield mock_get_user
+
+
+def _make_column(name: str, is_dttm: bool = False) -> MagicMock:
Review Comment:
Fixed in ffcd4b1d25 — added docstring to `_make_column`.
##########
tests/unit_tests/mcp_service/dataset/tool/test_query_dataset.py:
##########
@@ -0,0 +1,737 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Tests for the query_dataset MCP tool."""
+
+from __future__ import annotations
+
+import importlib
+from typing import Any
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+from fastmcp import Client
+
+from superset.mcp_service.app import mcp
+from superset.utils import json
+
+query_dataset_module = importlib.import_module(
+ "superset.mcp_service.dataset.tool.query_dataset"
+)
+
+
+@pytest.fixture
+def mcp_server():
+ return mcp
+
+
+@pytest.fixture(autouse=True)
+def mock_auth():
+ """Mock authentication for all tests."""
+ with patch("superset.mcp_service.auth.get_user_from_request") as mock_get_user:
+ mock_user = Mock()
+ mock_user.id = 1
+ mock_user.username = "admin"
+ mock_get_user.return_value = mock_user
+ yield mock_get_user
+
+
+def _make_column(name: str, is_dttm: bool = False) -> MagicMock:
+ col = MagicMock()
+ col.column_name = name
+ col.is_dttm = is_dttm
+ col.verbose_name = None
+ col.type = "VARCHAR"
+ col.groupby = True
+ col.filterable = True
+ col.description = None
+ return col
+
+
+def _make_metric(name: str, expression: str = "COUNT(*)") -> MagicMock:
+ metric = MagicMock()
+ metric.metric_name = name
+ metric.verbose_name = None
+ metric.expression = expression
+ metric.description = None
+ metric.d3format = None
+ return metric
+
+
+def _make_dataset(
+ dataset_id: int = 1,
+ table_name: str = "orders",
+ columns: list[Any] | None = None,
+ metrics: list[Any] | None = None,
+ main_dttm_col: str | None = None,
+) -> MagicMock:
Review Comment:
Fixed in ffcd4b1d25 — added docstrings to `_make_column`, `_make_metric`, and `_make_dataset`.
--
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: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] feat(mcp): add query_dataset tool to query datasets using semantic layer [superset]
Posted by "codeant-ai-for-open-source[bot] (via GitHub)" <gi...@apache.org>.
codeant-ai-for-open-source[bot] commented on code in PR #39727:
URL: https://github.com/apache/superset/pull/39727#discussion_r3164804086
##########
superset/mcp_service/dataset/schemas.py:
##########
@@ -383,6 +386,146 @@ class CreateVirtualDatasetResponse(BaseModel):
)
+VALID_FILTER_OPS = Literal[
+ "==",
+ "!=",
+ ">",
+ "<",
+ ">=",
+ "<=",
+ "LIKE",
+ "NOT LIKE",
+ "ILIKE",
+ "NOT ILIKE",
+ "IN",
+ "NOT IN",
+ "IS NULL",
+ "IS NOT NULL",
+ "IS TRUE",
+ "IS FALSE",
+ "TEMPORAL_RANGE",
+]
+
+
+class QueryDatasetFilter(BaseModel):
+ """A single filter condition for dataset queries."""
+
+ col: str = Field(..., description="Column name to filter on")
+ op: VALID_FILTER_OPS = Field(
+ ...,
+ description=(
+ 'Filter operator. Use "==" for equals, "!=" for not equals, '
+ '"IN" / "NOT IN" for membership, "IS NULL" / "IS NOT NULL", '
+ '"LIKE" for pattern matching, "TEMPORAL_RANGE" for time filters.'
+ ),
+ )
+ val: Any = Field(
+ default=None,
+ description="Filter value (omit for IS NULL/IS NOT NULL)",
+ )
+
+
+class QueryDatasetRequest(QueryCacheControl):
+ """Request schema for query_dataset tool."""
+
+ dataset_id: int | str = Field(
+ ...,
+ description="Dataset identifier — numeric ID or UUID string.",
Review Comment:
✅ **Customized review instruction saved!**
**Instruction:**
> Do not flag `int | str` dataset identifiers in MCP dataset request schemas when they intentionally match the identifier handling used by other dataset tools in the same service.
**Applied to:**
- `superset/mcp_service/dataset/**`
---
💡 *To manage or update this instruction, visit: [CodeAnt AI Settings](https://app.codeant.ai/settings/learnings)*
##########
superset/mcp_service/dataset/schemas.py:
##########
@@ -383,6 +386,146 @@ class CreateVirtualDatasetResponse(BaseModel):
)
+VALID_FILTER_OPS = Literal[
+ "==",
+ "!=",
+ ">",
+ "<",
+ ">=",
+ "<=",
+ "LIKE",
+ "NOT LIKE",
+ "ILIKE",
+ "NOT ILIKE",
+ "IN",
+ "NOT IN",
+ "IS NULL",
+ "IS NOT NULL",
+ "IS TRUE",
+ "IS FALSE",
+ "TEMPORAL_RANGE",
+]
+
+
+class QueryDatasetFilter(BaseModel):
+ """A single filter condition for dataset queries."""
+
+ col: str = Field(..., description="Column name to filter on")
+ op: VALID_FILTER_OPS = Field(
+ ...,
+ description=(
+ 'Filter operator. Use "==" for equals, "!=" for not equals, '
+ '"IN" / "NOT IN" for membership, "IS NULL" / "IS NOT NULL", '
+ '"LIKE" for pattern matching, "TEMPORAL_RANGE" for time filters.'
+ ),
+ )
+ val: Any = Field(
+ default=None,
+ description="Filter value (omit for IS NULL/IS NOT NULL)",
+ )
+
+
+class QueryDatasetRequest(QueryCacheControl):
+ """Request schema for query_dataset tool."""
+
+ dataset_id: int | str = Field(
+ ...,
+ description="Dataset identifier — numeric ID or UUID string.",
+ )
+ metrics: List[str] = Field(
+ default_factory=list,
+ description=(
+ "Saved metric names to compute (e.g. ['count', 'avg_revenue']). "
+ "Use get_dataset_info to discover available metrics."
+ ),
+ )
+ columns: List[str] = Field(
+ default_factory=list,
+ description=(
+ "Column/dimension names for GROUP BY or SELECT "
+ "(e.g. ['category', 'region']). "
+ "Use get_dataset_info to discover available columns."
+ ),
+ )
+ filters: List[QueryDatasetFilter] = Field(
+ default_factory=list,
+ description=(
+ 'Filter conditions (e.g. [{"col": "status", "op": "==", "val": "active"}]).'
+ ),
+ )
+ time_range: str | None = Field(
+ default=None,
+ description=(
+ "Time range filter (e.g. 'Last 7 days', 'Last month', "
+ "'2024-01-01 : 2024-12-31'). Requires a temporal column "
+ "on the dataset."
+ ),
+ )
+ time_column: str | None = Field(
+ default=None,
+ description=(
+ "Temporal column to apply time_range to. "
+ "Defaults to the dataset's main datetime column."
+ ),
+ )
+ order_by: List[str] | None = Field(
+ default=None,
+ description="Column or metric names to sort results by.",
+ )
+ order_desc: bool = Field(
+ default=True,
+ description="Sort descending (True) or ascending (False).",
+ )
+ row_limit: int = Field(
+ default=1000,
+ ge=1,
+ le=50000,
+ description="Maximum number of rows to return (default 1000, max 50000).",
+ )
+
+ @model_validator(mode="after")
+ def validate_metrics_or_columns(self) -> "QueryDatasetRequest":
+ """At least one of metrics or columns must be provided."""
+ if not self.metrics and not self.columns:
+ raise ValueError(
+ "At least one of 'metrics' or 'columns' must be provided. "
+ "Use get_dataset_info to discover available metrics and columns."
+ )
+ return self
+
+
+class QueryDatasetResponse(BaseModel):
+ """Response schema for query_dataset tool."""
+
+ model_config = ConfigDict(ser_json_timedelta="iso8601")
+
+ dataset_id: int = Field(..., description="Dataset ID")
Review Comment:
✅ **Customized review instruction saved!**
**Instruction:**
> Keep `dataset_id` as an integer in MCP response schemas when they intentionally follow the existing Superset API identifier pattern; do not flag this as an issue solely because it exposes a numeric primary key.
**Applied to:**
- `superset/mcp_service/dataset/schemas.py`
---
💡 *To manage or update this instruction, visit: [CodeAnt AI Settings](https://app.codeant.ai/settings/learnings)*
--
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: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] feat(mcp): add query_dataset tool to query datasets using semantic layer [superset]
Posted by "richardfogaca (via GitHub)" <gi...@apache.org>.
richardfogaca commented on PR #39727:
URL: https://github.com/apache/superset/pull/39727#issuecomment-4354263089
_Posting on Richard's behalf - this is his PR reviewer agent._\n\nI did one more pass against HEAD `3242c4c4cd990f55041b12d7808422c5bfe873f3`, specifically looking for bugs/regressions, side effects, security/privacy risk, overall design fit, DRY/SRP/KISS, and project conventions. The previous privacy/filter concerns look addressed, the tool stays aligned with the existing `QueryContextFactory` + `ChartDataCommand` path, and the focused `query_dataset` tests pass locally (`17 passed`). No new findings from me - good job tightening this up.
--
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: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] feat(mcp): add query_dataset tool to query datasets using semantic layer [superset]
Posted by "bito-code-review[bot] (via GitHub)" <gi...@apache.org>.
bito-code-review[bot] commented on PR #39727:
URL: https://github.com/apache/superset/pull/39727#issuecomment-4336749425
AI Code Review is in progress (usually takes 3 to 15 minutes unless it's a very large PR).
--
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: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org