bito-code-review[bot] commented on code in PR #40303: URL: https://github.com/apache/superset/pull/40303#discussion_r3312889348
########## superset/mcp_service/annotation_layer/tool/list_layer_annotations.py: ########## @@ -0,0 +1,149 @@ +# 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. + +"""List annotations within a layer FastMCP tool.""" + +import logging +from datetime import datetime, timezone + +from fastmcp import Context +from superset_core.mcp.decorators import tool, ToolAnnotations + +from superset.daos.base import ColumnOperator, ColumnOperatorEnum +from superset.extensions import event_logger +from superset.mcp_service.annotation_layer.schemas import ( + AnnotationFilter, + AnnotationInfo, + AnnotationLayerError, + AnnotationList, + DEFAULT_ANNOTATION_COLUMNS, + ListLayerAnnotationsRequest, + serialize_annotation, +) +from superset.mcp_service.mcp_core import ModelListCore + +logger = logging.getLogger(__name__) + +_ALL_ANNOTATION_COLUMNS = [ + "id", + "short_descr", + "long_descr", + "start_dttm", + "end_dttm", + "json_metadata", + "layer_id", +] +_SORTABLE_ANNOTATION_COLUMNS = ["id", "short_descr", "start_dttm", "end_dttm"] + + +@tool( + tags=["core"], + class_permission_name="Annotation", + annotations=ToolAnnotations( + title="List annotations in a layer", + readOnlyHint=True, + destructiveHint=False, + ), +) +async def list_layer_annotations( + request: ListLayerAnnotationsRequest, + ctx: Context, +) -> AnnotationList | AnnotationLayerError: + """List annotations within a specific annotation layer. + + The layer_id parameter is required and scopes all results to that layer. + + Sortable columns for order_column: id, short_descr, start_dttm, end_dttm + + Example: + ```json + {"layer_id": 1, "page": 1, "page_size": 25} + ``` + """ + await ctx.info( + "Listing annotations: layer_id=%s, page=%s, page_size=%s, search=%s" + % (request.layer_id, request.page, request.page_size, request.search) + ) + + try: + from superset.daos.annotation_layer import AnnotationDAO, AnnotationLayerDAO + + # Verify the layer exists before listing + layer = AnnotationLayerDAO.find_by_id(request.layer_id) + if layer is None: + await ctx.warning("Annotation layer not found: id=%s" % (request.layer_id,)) + return AnnotationLayerError.create( + error=f"Annotation layer with id '{request.layer_id}' not found", + error_type="not_found", + ) + + # Prepend the layer_id filter so results are scoped to this layer + layer_filter = ColumnOperator( + col="layer_id", opr=ColumnOperatorEnum.eq, value=request.layer_id + ) + combined_filters: list[ColumnOperator] = [layer_filter] + list(request.filters) + + def _serialize(obj: object, cols: list[str] | None) -> AnnotationInfo | None: + return serialize_annotation(obj) + + list_tool = ModelListCore( + dao_class=AnnotationDAO, + output_schema=AnnotationInfo, + item_serializer=_serialize, + filter_type=AnnotationFilter, + default_columns=DEFAULT_ANNOTATION_COLUMNS, + search_columns=["short_descr", "long_descr"], + list_field_name="annotations", + output_list_schema=AnnotationList, + all_columns=_ALL_ANNOTATION_COLUMNS, + sortable_columns=_SORTABLE_ANNOTATION_COLUMNS, + logger=logger, + ) + + with event_logger.log_context(action="mcp.list_layer_annotations.query"): + result = list_tool.run_tool( + filters=combined_filters, + search=request.search, + select_columns=request.select_columns, + order_column=request.order_column, + order_direction=request.order_direction, + page=max(request.page - 1, 0), + page_size=request.page_size, + ) + + result.layer_id = request.layer_id + + await ctx.info( + "Annotations listed: layer_id=%s, count=%s, total_count=%s" + % ( + request.layer_id, + len(result.annotations) if hasattr(result, "annotations") else 0, + getattr(result, "total_count", None), + ) + ) + return result + + except Exception as e: + await ctx.error( + "Annotation listing failed: layer_id=%s, error=%s, error_type=%s" + % (request.layer_id, str(e), type(e).__name__) + ) + return AnnotationLayerError( + error=f"Failed to list annotations: {str(e)}", + error_type="InternalError", + timestamp=datetime.now(timezone.utc), + ) Review Comment: <!-- Bito Reply --> The suggestion to change the exception handler from returning an `AnnotationLayerError` to re-raising the exception is valid and appropriate. It ensures consistent error handling with other list tools (`list_annotation_layers`, `list_charts`), which re-raise after logging. This change aligns with the middleware (`GlobalErrorHandlerMiddleware`) expectations and prevents returning a 200 OK status with error content, which could mislead callers. **superset/mcp_service/annotation_layer/tool/list_layer_annotations.py** ``` except Exception as e: await ctx.error( "Annotation listing failed: layer_id=%s, error=%s, error_type=%s" % (request.layer_id, str(e), type(e).__name__) ) raise ``` ########## tests/unit_tests/mcp_service/annotation_layer/tool/test_annotation_layer_tools.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. + +import logging +from unittest.mock import MagicMock, patch + +import pytest +from fastmcp import Client +from pydantic import ValidationError + +from superset.mcp_service.annotation_layer.schemas import ( + AnnotationFilter, + AnnotationLayerFilter, + ListAnnotationLayersRequest, + ListLayerAnnotationsRequest, +) +from superset.mcp_service.app import mcp +from superset.utils import json + +logging.basicConfig(level=logging.DEBUG) +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def make_layer( + layer_id: int = 1, name: str = "My Layer", descr: str = "desc" +) -> MagicMock: + obj = MagicMock() + obj.id = layer_id + obj.name = name + obj.descr = descr + obj.changed_on = None + obj.created_on = None + return obj + + +def make_annotation( + annotation_id: int = 10, + layer_id: int = 1, + short_descr: str = "Deploy", + long_descr: str = "Deployment annotation", +) -> MagicMock: + obj = MagicMock() + obj.id = annotation_id + obj.layer_id = layer_id + obj.short_descr = short_descr + obj.long_descr = long_descr + obj.start_dttm = None + obj.end_dttm = None + obj.json_metadata = None + return obj + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + [email protected] +def mcp_server(): + return mcp + + [email protected](autouse=True) +def mock_auth(): + from unittest.mock import Mock + + 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 + + +# --------------------------------------------------------------------------- +# Schema validation tests +# --------------------------------------------------------------------------- + + +class TestAnnotationLayerFilterSchema: + def test_valid_name_filter(self): + f = AnnotationLayerFilter(col="name", opr="eq", value="My Layer") + assert f.col == "name" + + def test_invalid_column_rejected(self): + with pytest.raises(ValidationError): + AnnotationLayerFilter(col="descr", opr="eq", value="x") + + def test_search_and_filters_mutual_exclusion(self): + with pytest.raises(ValidationError): + ListAnnotationLayersRequest( + search="foo", + filters=[{"col": "name", "opr": "eq", "value": "bar"}], + ) + + +class TestAnnotationFilterSchema: + def test_valid_short_descr_filter(self): + f = AnnotationFilter(col="short_descr", opr="eq", value="Deploy") + assert f.col == "short_descr" + + def test_invalid_column_rejected(self): + with pytest.raises(ValidationError): + AnnotationFilter(col="layer_id", opr="eq", value=1) + + def test_search_and_filters_mutual_exclusion(self): + with pytest.raises(ValidationError): + ListLayerAnnotationsRequest( + layer_id=1, + search="foo", + filters=[{"col": "short_descr", "opr": "eq", "value": "bar"}], + ) + + +# --------------------------------------------------------------------------- +# list_annotation_layers tests +# --------------------------------------------------------------------------- + + +@patch("superset.daos.annotation_layer.AnnotationLayerDAO.list") [email protected] +async def test_list_annotation_layers_basic(mock_list, mcp_server): + """Basic listing returns structured response with annotation layers.""" + layer = make_layer() + mock_list.return_value = ([layer], 1) + + async with Client(mcp_server) as client: + result = await client.call_tool( + "list_annotation_layers", + {"request": {"page": 1, "page_size": 10}}, + ) + + data = json.loads(result.content[0].text) + assert data["annotation_layers"] is not None + assert len(data["annotation_layers"]) == 1 + assert data["annotation_layers"][0]["id"] == 1 + assert data["annotation_layers"][0]["name"] == "My Layer" + + +@patch("superset.daos.annotation_layer.AnnotationLayerDAO.list") [email protected] +async def test_list_annotation_layers_empty(mock_list, mcp_server): + """Empty result set returns zero count.""" + mock_list.return_value = ([], 0) + + async with Client(mcp_server) as client: + result = await client.call_tool("list_annotation_layers", {}) + + data = json.loads(result.content[0].text) + assert data["annotation_layers"] == [] + assert data["total_count"] == 0 + + +@patch("superset.daos.annotation_layer.AnnotationLayerDAO.list") [email protected] +async def test_list_annotation_layers_search(mock_list, mcp_server): + """Search parameter is passed through to DAO.""" + layer = make_layer(name="Release Events") + mock_list.return_value = ([layer], 1) + + async with Client(mcp_server) as client: + result = await client.call_tool( + "list_annotation_layers", + {"request": {"search": "release"}}, + ) + + data = json.loads(result.content[0].text) + assert data["annotation_layers"][0]["name"] == "Release Events" + call_kwargs = mock_list.call_args.kwargs + assert call_kwargs["search"] == "release" + + +@patch("superset.daos.annotation_layer.AnnotationLayerDAO.list") [email protected] +async def test_list_annotation_layers_pagination(mock_list, mcp_server): + """Pagination metadata is correctly computed.""" + mock_list.return_value = ([], 50) + + async with Client(mcp_server) as client: + result = await client.call_tool( + "list_annotation_layers", + {"request": {"page": 2, "page_size": 25}}, + ) + + data = json.loads(result.content[0].text) + assert data["page"] == 2 + assert data["page_size"] == 25 + assert data["total_count"] == 50 + assert data["total_pages"] == 2 + # Page 2 of 2, so no next page + assert data["has_next"] is False + assert data["has_previous"] is True + + +# --------------------------------------------------------------------------- +# get_annotation_layer_info tests +# --------------------------------------------------------------------------- + + +@patch("superset.daos.annotation_layer.AnnotationLayerDAO.find_by_id") [email protected] +async def test_get_annotation_layer_info_found(mock_find, mcp_server): + """Returns annotation layer data when found.""" + mock_find.return_value = make_layer(layer_id=5, name="Prod Events") + + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_annotation_layer_info", + {"request": {"id": 5}}, + ) + + data = json.loads(result.content[0].text) + assert data["id"] == 5 + assert data["name"] == "Prod Events" + mock_find.assert_called_once_with(5, query_options=None) + + +@patch("superset.daos.annotation_layer.AnnotationLayerDAO.find_by_id") [email protected] +async def test_get_annotation_layer_info_not_found(mock_find, mcp_server): + """Returns error response when layer is not found.""" + mock_find.return_value = None + + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_annotation_layer_info", + {"request": {"id": 999}}, + ) + + data = json.loads(result.content[0].text) + assert data["error_type"] == "not_found" + assert "999" in data["error"] + + +# --------------------------------------------------------------------------- +# list_layer_annotations tests +# --------------------------------------------------------------------------- + + +@patch("superset.daos.annotation_layer.AnnotationLayerDAO.find_by_id") +@patch("superset.daos.annotation_layer.AnnotationDAO.list") [email protected] +async def test_list_layer_annotations_basic(mock_list, mock_layer_find, mcp_server): + """Annotations are listed and scoped to the specified layer.""" + mock_layer_find.return_value = make_layer(layer_id=1) + ann = make_annotation(annotation_id=10, layer_id=1) + mock_list.return_value = ([ann], 1) + + async with Client(mcp_server) as client: + result = await client.call_tool( + "list_layer_annotations", + {"request": {"layer_id": 1, "page": 1, "page_size": 10}}, + ) + + data = json.loads(result.content[0].text) + assert data["layer_id"] == 1 + assert len(data["annotations"]) == 1 + assert data["annotations"][0]["id"] == 10 + assert data["annotations"][0]["layer_id"] == 1 + + +@patch("superset.daos.annotation_layer.AnnotationLayerDAO.find_by_id") +@patch("superset.daos.annotation_layer.AnnotationDAO.list") [email protected] +async def test_list_layer_annotations_layer_id_filter_prepended( + mock_list, mock_layer_find, mcp_server +): + """The layer_id filter is always prepended to DAO column_operators.""" + mock_layer_find.return_value = make_layer(layer_id=3) + mock_list.return_value = ([], 0) + + async with Client(mcp_server) as client: + await client.call_tool( + "list_layer_annotations", + {"request": {"layer_id": 3}}, + ) + + call_kwargs = mock_list.call_args.kwargs + filters = call_kwargs.get("column_operators", []) + # First filter must be the layer_id eq filter + assert filters, "Expected at least one filter (layer_id)" + first = filters[0] + col = first.get("col") if isinstance(first, dict) else getattr(first, "col", None) + val = ( + first.get("value") if isinstance(first, dict) else getattr(first, "value", None) + ) + assert col == "layer_id" + assert val == 3 + + +@patch("superset.daos.annotation_layer.AnnotationLayerDAO.find_by_id") [email protected] +async def test_list_layer_annotations_layer_not_found(mock_layer_find, mcp_server): + """Returns error when the layer does not exist.""" + mock_layer_find.return_value = None + + async with Client(mcp_server) as client: + result = await client.call_tool( + "list_layer_annotations", + {"request": {"layer_id": 42}}, + ) + + data = json.loads(result.content[0].text) + assert data["error_type"] == "not_found" + assert "42" in data["error"] + + +@patch("superset.daos.annotation_layer.AnnotationLayerDAO.find_by_id") +@patch("superset.daos.annotation_layer.AnnotationDAO.list") [email protected] +async def test_list_layer_annotations_only_returns_own_layer( + mock_list, mock_layer_find, mcp_server +): + """Results are filtered to the requested layer only — wrong layer_id is rejected.""" + mock_layer_find.return_value = make_layer(layer_id=1) + # Simulate DAO returning annotations — the layer_id filter is applied at DB level + ann_wrong = make_annotation(annotation_id=99, layer_id=2) + mock_list.return_value = ([ann_wrong], 1) + + async with Client(mcp_server) as client: + result = await client.call_tool( + "list_layer_annotations", + {"request": {"layer_id": 1}}, + ) + + data = json.loads(result.content[0].text) + # layer_id in response header must still be 1 (the requested layer) + assert data["layer_id"] == 1 Review Comment: <!-- Bito Reply --> The suggestion in the review comment is valid and should be applied. It improves the test by adding an explicit assertion to verify that annotations with a different `layer_id` are filtered out, which aligns with the business logic rather than just checking the response header. The change ensures the test is more robust and meaningful. **tests/unit_tests/mcp_service/annotation_layer/tool/test_annotation_layer_tools.py** ``` # Verify annotation with wrong layer_id is NOT in results assert len(data["annotations"]) == 0 or all(a["layer_id"] == 1 for a in data["annotations"]) ``` ########## superset/mcp_service/annotation_layer/tool/get_layer_annotation_info.py: ########## @@ -0,0 +1,130 @@ +# 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. + +"""Get a single annotation within a layer FastMCP tool.""" + +import logging +from datetime import datetime, timezone + +from fastmcp import Context +from superset_core.mcp.decorators import tool, ToolAnnotations + +from superset.extensions import event_logger +from superset.mcp_service.annotation_layer.schemas import ( + AnnotationInfo, + AnnotationLayerError, + GetLayerAnnotationInfoRequest, + serialize_annotation, +) + +logger = logging.getLogger(__name__) + + +@tool( + tags=["discovery"], + class_permission_name="Annotation", + annotations=ToolAnnotations( + title="Get annotation info", + readOnlyHint=True, + destructiveHint=False, + ), +) +async def get_layer_annotation_info( + request: GetLayerAnnotationInfoRequest, + ctx: Context, +) -> AnnotationInfo | AnnotationLayerError: + """Get detailed information about a specific annotation within a layer. + + Both layer_id and annotation_id are required. Returns an error if the + annotation does not belong to the specified layer. + + Example: + ```json + {"layer_id": 1, "annotation_id": 42} + ``` + """ + await ctx.info( + "Retrieving annotation: layer_id=%s, annotation_id=%s" + % (request.layer_id, request.annotation_id) + ) + + try: + from superset.daos.annotation_layer import AnnotationDAO, AnnotationLayerDAO + + # Verify the layer exists + with event_logger.log_context( + action="mcp.get_layer_annotation_info.layer_lookup" + ): + layer = AnnotationLayerDAO.find_by_id(request.layer_id) + + if layer is None: + await ctx.warning("Annotation layer not found: id=%s" % (request.layer_id,)) + return AnnotationLayerError.create( + error=f"Annotation layer with id '{request.layer_id}' not found", + error_type="not_found", + ) + + # Fetch the annotation + with event_logger.log_context( + action="mcp.get_layer_annotation_info.annotation_lookup" + ): + annotation = AnnotationDAO.find_by_id(request.annotation_id) + + if annotation is None: + await ctx.warning( + "Annotation not found: annotation_id=%s" % (request.annotation_id,) + ) + return AnnotationLayerError.create( + error=f"Annotation with id '{request.annotation_id}' not found", + error_type="not_found", + ) + + # Verify the annotation belongs to the requested layer + if getattr(annotation, "layer_id", None) != request.layer_id: + await ctx.warning( + "Annotation %s does not belong to layer %s" + % (request.annotation_id, request.layer_id) + ) + return AnnotationLayerError.create( + error=( + f"Annotation '{request.annotation_id}' does not belong to " + f"layer '{request.layer_id}'" + ), + error_type="not_found", + ) + + result = serialize_annotation(annotation) + await ctx.info( + "Annotation retrieved: id=%s, short_descr=%s" + % (result.id if result else None, result.short_descr if result else None) + ) + return result or AnnotationLayerError.create( + error="Failed to serialize annotation", + error_type="SerializationError", + ) + + except Exception as e: Review Comment: <!-- Bito Reply --> The suggestion to avoid catching broad `Exception` types is valid in this context. The current code catches `(ValueError, KeyError, AttributeError, TypeError)`, which is more specific than the broader `Exception` catch the reviewer is concerned about. This aligns with best practices for error handling by explicitly catching only the expected exception types. The suggestion would improve clarity and debugging by ensuring only relevant exceptions are caught, avoiding silent failures from unexpected errors. **superset/mcp_service/annotation_layer/tool/get_layer_annotation_info.py** ``` except (ValueError, KeyError, AttributeError, TypeError) as e: ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
