Skip to content

Commit 09d88d7

Browse files
fix(billing-cost-management-mcp-server): storage-lens validation errors carry error_type; state Trino dialect and string CSV columns (#4950)
The two local validation failures in storage_lens_run_query (a query that is not a single read-only SELECT, and a missing manifest location) returned format_response('error', {}, message) with no error_type, so the response had error_type=None and the retry guidance fell back to do_not_retry. They now pass error_type='validation_error', operation='storage_lens_run_query' and service='Storage Lens', so callers get retry_strategy=modify_request like the other tools' argument errors. The description said "Use standard SQL syntax for Athena queries". That line now says the dialect is Athena engine version 3 (Trino SQL) and that every column of a CSV export is a string, so numbers need a cast, e.g. SUM(CAST(metric_value AS BIGINT)). The module docstring pointed at resources/storage_lens_metrics_reference.md, which the package does not ship; it now points at the tool description. No parameter, type or default changes.
1 parent 4fa757d commit 09d88d7

2 files changed

Lines changed: 78 additions & 3 deletions

File tree

‎src/billing-cost-management-mcp-server/awslabs/billing_cost_management_mcp_server/tools/storage_lens_tools.py‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
"""AWS S3 Storage Lens tools for the AWS Billing and Cost Management MCP server.
1616
1717
This module provides functionality to create and query Athena tables for S3 Storage Lens data.
18-
See the resources/storage_lens_metrics_reference.md file for detailed metrics information and sample queries.
18+
The storage-lens tool description lists the export columns and sample queries.
1919
"""
2020

2121
import asyncio
@@ -797,7 +797,7 @@ async def query_storage_lens(
797797
798798
IMPORTANT USAGE GUIDELINES:
799799
- Before using this tool, provide a 1-3 sentence explanation starting with "EXPLANATION:"
800-
- Use standard SQL syntax for Athena queries
800+
- Dialect: Athena engine version 3 (Trino SQL). In CSV exports every column is a string, so cast before math or numeric comparison, e.g. SUM(CAST(metric_value AS BIGINT))
801801
- Reference the Storage Lens metrics table with the {table} placeholder (or its qualified <database>.<table> name). A query that does neither is rejected with a validation_error. Every {table} is replaced; nothing else in the query is changed
802802
- database_name and table_name, if given, must start with a letter and contain only letters, digits and underscores (at most 255 characters); other values are rejected with a validation_error
803803
- Perform aggregations (GROUP BY) when analyzing data across multiple dimensions
@@ -923,7 +923,14 @@ async def storage_lens_run_query(
923923
# Validate user-supplied query is a read-only SELECT
924924
validation_error = validate_query_is_readonly_select(query)
925925
if validation_error:
926-
return format_response('error', {}, validation_error)
926+
return format_response(
927+
'error',
928+
{},
929+
validation_error,
930+
error_type='validation_error',
931+
operation='storage_lens_run_query',
932+
service='Storage Lens',
933+
)
927934

928935
# Get manifest location from args or environment variable
929936
manifest_loc = manifest_location or os.environ.get(ENV_STORAGE_LENS_MANIFEST_LOCATION, '')
@@ -932,6 +939,9 @@ async def storage_lens_run_query(
932939
'error',
933940
{},
934941
f"Missing manifest location. Provide 'manifest_location' parameter or set {ENV_STORAGE_LENS_MANIFEST_LOCATION} environment variable.",
942+
error_type='validation_error',
943+
operation='storage_lens_run_query',
944+
service='Storage Lens',
935945
)
936946

937947
# Get output location from args or environment variable (optional)

‎src/billing-cost-management-mcp-server/tests/tools/test_storage_lens_tools.py‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1443,6 +1443,71 @@ async def test_multi_statement_rejected(self, mock_context):
14431443
mock_instance.query_storage_lens.assert_not_awaited()
14441444

14451445

1446+
# --- Local validation errors carry error_type, operation and service ---
1447+
1448+
1449+
@pytest.mark.asyncio
1450+
@pytest.mark.parametrize(
1451+
'query',
1452+
[
1453+
'DROP TABLE x',
1454+
'INSERT INTO t VALUES (1)',
1455+
'SELECT 1; DROP TABLE t',
1456+
'',
1457+
'SELECT FROM WHERE (',
1458+
],
1459+
)
1460+
async def test_rejected_query_is_validation_error(mock_context, monkeypatch, query):
1461+
"""A query rejected by the read-only check is a validation_error the caller must modify."""
1462+
monkeypatch.setenv('STORAGE_LENS_MANIFEST_LOCATION', 's3://bucket/prefix/')
1463+
stl_mod = _reload_storage_lens_with_identity_decorator()
1464+
with patch.object(stl_mod, 'StorageLensQueryTool') as MockTool:
1465+
res = await stl_mod.storage_lens_run_query(mock_context, query=query) # type: ignore
1466+
MockTool.assert_not_called()
1467+
assert res['status'] == 'error'
1468+
assert res['message'].startswith('Query validation failed')
1469+
assert res['error_type'] == 'validation_error'
1470+
assert res['operation'] == 'storage_lens_run_query'
1471+
assert res['service'] == 'Storage Lens'
1472+
assert add_retry_guidance(res)['retry_strategy'] == 'modify_request'
1473+
1474+
1475+
@pytest.mark.asyncio
1476+
async def test_missing_manifest_location_is_validation_error(mock_context, monkeypatch):
1477+
"""No manifest_location argument and no environment variable is a validation_error."""
1478+
monkeypatch.delenv('STORAGE_LENS_MANIFEST_LOCATION', raising=False)
1479+
stl_mod = _reload_storage_lens_with_identity_decorator()
1480+
res = await stl_mod.storage_lens_run_query(mock_context, query='SELECT 1') # type: ignore
1481+
assert res['status'] == 'error'
1482+
assert 'Missing manifest location' in res['message']
1483+
assert res['error_type'] == 'validation_error'
1484+
assert res['operation'] == 'storage_lens_run_query'
1485+
assert res['service'] == 'Storage Lens'
1486+
assert add_retry_guidance(res)['retry_strategy'] == 'modify_request'
1487+
1488+
1489+
@pytest.mark.asyncio
1490+
async def test_description_states_dialect_and_csv_string_columns():
1491+
"""The served description names the Trino dialect and that CSV export columns are strings."""
1492+
from awslabs.billing_cost_management_mcp_server.tools import storage_lens_tools as stl_mod
1493+
1494+
importlib.reload(stl_mod)
1495+
tool = await stl_mod.storage_lens_server.get_tool('storage-lens')
1496+
assert tool is not None
1497+
description = tool.description or ''
1498+
assert 'Athena engine version 3 (Trino SQL)' in description
1499+
assert 'In CSV exports every column is a string' in description
1500+
assert 'SUM(CAST(metric_value AS BIGINT))' in description
1501+
assert 'standard SQL' not in description
1502+
1503+
1504+
def test_module_docstring_has_no_missing_file_pointer():
1505+
"""The module docstring does not point at a reference file the package does not ship."""
1506+
from awslabs.billing_cost_management_mcp_server.tools import storage_lens_tools as stl_mod
1507+
1508+
assert 'storage_lens_metrics_reference.md' not in (stl_mod.__doc__ or '')
1509+
1510+
14461511
# --- Direct unit tests for validate_query_is_readonly_select ---
14471512

14481513

0 commit comments

Comments
 (0)