awslabs.cost-explorer-mcp-server 0.0.5__py3-none-any.whl → 0.0.6__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -18,7 +18,13 @@ import boto3
18
18
  import os
19
19
  import re
20
20
  import sys
21
- from datetime import datetime
21
+ from awslabs.cost_explorer_mcp_server.constants import (
22
+ VALID_DIMENSIONS,
23
+ VALID_GROUP_BY_DIMENSIONS,
24
+ VALID_GROUP_BY_TYPES,
25
+ VALID_MATCH_OPTIONS,
26
+ )
27
+ from datetime import datetime, timezone
22
28
  from loguru import logger
23
29
  from typing import Any, Dict, Optional, Tuple
24
30
 
@@ -58,6 +64,80 @@ def get_cost_explorer_client():
58
64
  return _cost_explorer_client
59
65
 
60
66
 
67
+ def validate_dimension_key(dimension_key: str) -> Dict[str, Any]:
68
+ """Validate that the dimension key is supported by AWS Cost Explorer.
69
+
70
+ Args:
71
+ dimension_key: The dimension key to validate
72
+
73
+ Returns:
74
+ Empty dictionary if valid, or an error dictionary
75
+ """
76
+ try:
77
+ dimension_upper = dimension_key.upper()
78
+ if dimension_upper not in VALID_DIMENSIONS:
79
+ return {
80
+ 'error': f"Invalid dimension key '{dimension_key}'. Valid dimensions are: {', '.join(VALID_DIMENSIONS)}"
81
+ }
82
+ return {}
83
+ except Exception as e:
84
+ return {'error': f'Error validating dimension key: {str(e)}'}
85
+
86
+
87
+ def get_available_dimension_values(
88
+ key: str, billing_period_start: str, billing_period_end: str
89
+ ) -> Dict[str, Any]:
90
+ """Get available values for a specific dimension."""
91
+ # Validate dimension key first
92
+ dimension_validation = validate_dimension_key(key)
93
+ if 'error' in dimension_validation:
94
+ return dimension_validation
95
+
96
+ # Validate date range (no granularity constraint for dimension values)
97
+ is_valid, error_message = validate_date_range(billing_period_start, billing_period_end)
98
+ if not is_valid:
99
+ return {'error': error_message}
100
+
101
+ try:
102
+ ce = get_cost_explorer_client()
103
+ response = ce.get_dimension_values(
104
+ TimePeriod={'Start': billing_period_start, 'End': billing_period_end},
105
+ Dimension=key.upper(),
106
+ )
107
+ dimension_values = response['DimensionValues']
108
+ values = [value['Value'] for value in dimension_values]
109
+ return {'dimension': key.upper(), 'values': values}
110
+ except Exception as e:
111
+ logger.error(
112
+ f'Error getting dimension values for {key.upper()} ({billing_period_start} to {billing_period_end}): {e}'
113
+ )
114
+ return {'error': str(e)}
115
+
116
+
117
+ def get_available_tag_values(
118
+ tag_key: str, billing_period_start: str, billing_period_end: str
119
+ ) -> Dict[str, Any]:
120
+ """Get available values for a specific tag key."""
121
+ # Validate date range (no granularity constraint for tag values)
122
+ is_valid, error_message = validate_date_range(billing_period_start, billing_period_end)
123
+ if not is_valid:
124
+ return {'error': error_message}
125
+
126
+ try:
127
+ ce = get_cost_explorer_client()
128
+ response = ce.get_tags(
129
+ TimePeriod={'Start': billing_period_start, 'End': billing_period_end},
130
+ TagKey=tag_key,
131
+ )
132
+ tag_values = response['Tags']
133
+ return {'tag_key': tag_key, 'values': tag_values}
134
+ except Exception as e:
135
+ logger.error(
136
+ f'Error getting tag values for {tag_key} ({billing_period_start} to {billing_period_end}): {e}'
137
+ )
138
+ return {'error': str(e)}
139
+
140
+
61
141
  def validate_date_format(date_str: str) -> Tuple[bool, str]:
62
142
  """Validate that a date string is in YYYY-MM-DD format and is a valid date.
63
143
 
@@ -141,55 +221,6 @@ def validate_date_range(
141
221
  return True, ''
142
222
 
143
223
 
144
- def get_dimension_values(
145
- key: str, billing_period_start: str, billing_period_end: str
146
- ) -> Dict[str, Any]:
147
- """Get available values for a specific dimension."""
148
- # Validate date range (no granularity constraint for dimension values)
149
- is_valid, error_message = validate_date_range(billing_period_start, billing_period_end)
150
- if not is_valid:
151
- return {'error': error_message}
152
-
153
- try:
154
- ce = get_cost_explorer_client()
155
- response = ce.get_dimension_values(
156
- TimePeriod={'Start': billing_period_start, 'End': billing_period_end},
157
- Dimension=key.upper(),
158
- )
159
- dimension_values = response['DimensionValues']
160
- values = [value['Value'] for value in dimension_values]
161
- return {'dimension': key.upper(), 'values': values}
162
- except Exception as e:
163
- logger.error(
164
- f'Error getting dimension values for {key.upper()} ({billing_period_start} to {billing_period_end}): {e}'
165
- )
166
- return {'error': str(e)}
167
-
168
-
169
- def get_tag_values(
170
- tag_key: str, billing_period_start: str, billing_period_end: str
171
- ) -> Dict[str, Any]:
172
- """Get available values for a specific tag key."""
173
- # Validate date range (no granularity constraint for tag values)
174
- is_valid, error_message = validate_date_range(billing_period_start, billing_period_end)
175
- if not is_valid:
176
- return {'error': error_message}
177
-
178
- try:
179
- ce = get_cost_explorer_client()
180
- response = ce.get_tags(
181
- TimePeriod={'Start': billing_period_start, 'End': billing_period_end},
182
- TagKey=tag_key,
183
- )
184
- tag_values = response['Tags']
185
- return {'tag_key': tag_key, 'values': tag_values}
186
- except Exception as e:
187
- logger.error(
188
- f'Error getting tag values for {tag_key} ({billing_period_start} to {billing_period_end}): {e}'
189
- )
190
- return {'error': str(e)}
191
-
192
-
193
224
  def validate_match_options(match_options: list, filter_type: str) -> Dict[str, Any]:
194
225
  """Validate MatchOptions based on filter type.
195
226
 
@@ -200,13 +231,11 @@ def validate_match_options(match_options: list, filter_type: str) -> Dict[str, A
200
231
  Returns:
201
232
  Empty dictionary if valid, or an error dictionary
202
233
  """
203
- if filter_type == 'Dimensions':
204
- valid_options = ['EQUALS', 'CASE_SENSITIVE']
205
- elif filter_type in ['Tags', 'CostCategories']:
206
- valid_options = ['EQUALS', 'ABSENT', 'CASE_SENSITIVE']
207
- else:
234
+ if filter_type not in VALID_MATCH_OPTIONS:
208
235
  return {'error': f'Unknown filter type: {filter_type}'}
209
236
 
237
+ valid_options = VALID_MATCH_OPTIONS[filter_type]
238
+
210
239
  for option in match_options:
211
240
  if option not in valid_options:
212
241
  return {
@@ -253,7 +282,7 @@ def validate_expression(
253
282
 
254
283
  dimension_key = dimension['Key']
255
284
  dimension_values = dimension['Values']
256
- valid_values_response = get_dimension_values(
285
+ valid_values_response = get_available_dimension_values(
257
286
  dimension_key, billing_period_start, billing_period_end
258
287
  )
259
288
  if 'error' in valid_values_response:
@@ -277,7 +306,7 @@ def validate_expression(
277
306
 
278
307
  tag_key = tag['Key']
279
308
  tag_values = tag['Values']
280
- valid_tag_values_response = get_tag_values(
309
+ valid_tag_values_response = get_available_tag_values(
281
310
  tag_key, billing_period_start, billing_period_end
282
311
  )
283
312
  if 'error' in valid_tag_values_response:
@@ -379,11 +408,304 @@ def validate_group_by(group_by: Optional[Dict[str, Any]]) -> Dict[str, Any]:
379
408
  ):
380
409
  return {'error': 'group_by must be a dictionary with "Type" and "Key" keys.'}
381
410
 
382
- if group_by['Type'].upper() not in ['DIMENSION', 'TAG', 'COST_CATEGORY']:
411
+ group_type = group_by['Type'].upper()
412
+ group_key = group_by['Key']
413
+
414
+ if group_type not in VALID_GROUP_BY_TYPES:
383
415
  return {
384
- 'error': 'Invalid group Type. Valid types are DIMENSION, TAG, and COST_CATEGORY.'
416
+ 'error': f'Invalid group Type: {group_type}. Valid types are {", ".join(VALID_GROUP_BY_TYPES)}.'
385
417
  }
386
418
 
419
+ # Validate dimension key if type is DIMENSION
420
+ if group_type == 'DIMENSION':
421
+ dimension_upper = group_key.upper()
422
+ if dimension_upper not in VALID_GROUP_BY_DIMENSIONS:
423
+ return {
424
+ 'error': f'Invalid dimension key for GROUP BY: {group_key}. Valid values for the DIMENSION type are {", ".join(VALID_GROUP_BY_DIMENSIONS)}.'
425
+ }
426
+
387
427
  return {}
388
428
  except Exception as e:
389
429
  return {'error': f'Error validating group_by: {str(e)}'}
430
+
431
+
432
+ def validate_forecast_date_range(
433
+ start_date: str, end_date: str, granularity: str = 'MONTHLY'
434
+ ) -> Tuple[bool, str]:
435
+ """Validate that forecast dates meet AWS Cost Explorer requirements.
436
+
437
+ Args:
438
+ start_date: The forecast start date string in YYYY-MM-DD format
439
+ end_date: The forecast end date string in YYYY-MM-DD format
440
+ granularity: The granularity for the forecast (DAILY or MONTHLY)
441
+
442
+ Returns:
443
+ Tuple of (is_valid, error_message)
444
+ """
445
+ # First validate basic date format and range
446
+ is_valid, error = validate_date_range(start_date, end_date)
447
+ if not is_valid:
448
+ return False, error
449
+
450
+ today = datetime.now(timezone.utc).date()
451
+ start_dt = datetime.strptime(start_date, '%Y-%m-%d').date()
452
+ end_dt = datetime.strptime(end_date, '%Y-%m-%d').date()
453
+
454
+ # AWS requires start date to be equal to or no later than current date
455
+ if start_dt > today:
456
+ return (
457
+ False,
458
+ f"Forecast start date '{start_date}' must be equal to or no later than the current date ({today})",
459
+ )
460
+
461
+ # End date must be in the future
462
+ if end_dt <= today:
463
+ return False, f"Forecast end date '{end_date}' must be in the future (after {today})"
464
+
465
+ # AWS Cost Explorer forecast granularity-specific limits
466
+ date_diff = (end_dt - start_dt).days
467
+
468
+ if granularity.upper() == 'DAILY':
469
+ # DAILY forecasts support maximum 3 months (approximately 93 days)
470
+ if date_diff > 93:
471
+ return (
472
+ False,
473
+ f'DAILY granularity supports a maximum of 3 months (93 days). Current range is {date_diff} days ({start_date} to {end_date}). Please use a shorter date range or MONTHLY granularity.',
474
+ )
475
+ elif granularity.upper() == 'MONTHLY':
476
+ # MONTHLY forecasts support maximum 12 months
477
+ max_forecast_date = datetime.now(timezone.utc).date().replace(year=today.year + 1)
478
+ if end_dt > max_forecast_date:
479
+ return (
480
+ False,
481
+ f"MONTHLY granularity supports a maximum of 12 months in the future. Forecast end date '{end_date}' exceeds the limit (max: {max_forecast_date}).",
482
+ )
483
+
484
+ return True, ''
485
+
486
+
487
+ def validate_comparison_date_range(start_date: str, end_date: str) -> Tuple[bool, str]:
488
+ """Validate that comparison dates meet AWS Cost Explorer comparison API requirements.
489
+
490
+ Args:
491
+ start_date: The start date string in YYYY-MM-DD format
492
+ end_date: The end date string in YYYY-MM-DD format
493
+
494
+ Returns:
495
+ Tuple of (is_valid, error_message)
496
+ """
497
+ # First validate basic date format and range
498
+ is_valid, error = validate_date_range(start_date, end_date)
499
+ if not is_valid:
500
+ return False, error
501
+
502
+ today = datetime.now(timezone.utc).date()
503
+ start_dt = datetime.strptime(start_date, '%Y-%m-%d').date()
504
+ end_dt = datetime.strptime(end_date, '%Y-%m-%d').date()
505
+
506
+ # AWS requires start date to be equal to or no later than current date
507
+ if start_dt > today:
508
+ return (
509
+ False,
510
+ f"Comparison start date '{start_date}' must be equal to or no later than the current date ({today})",
511
+ )
512
+
513
+ # Must start on the first day of a month
514
+ if start_dt.day != 1:
515
+ return (
516
+ False,
517
+ f"Comparison start date '{start_date}' must be the first day of a month (e.g., 2025-01-01)",
518
+ )
519
+
520
+ # Must end on the first day of a month (exclusive end date)
521
+ if end_dt.day != 1:
522
+ return (
523
+ False,
524
+ f"Comparison end date '{end_date}' must be the first day of a month (e.g., 2025-02-01)",
525
+ )
526
+
527
+ # Comparison periods can only go up to the last complete month
528
+ # Calculate the first day of current month (last complete month boundary)
529
+ current_month_start = today.replace(day=1)
530
+ # The comparison period (start_date) cannot be in the current month or future
531
+ if start_dt >= current_month_start:
532
+ # Calculate last complete month for user guidance
533
+ if current_month_start.month == 1:
534
+ last_complete_month = current_month_start.replace(
535
+ year=current_month_start.year - 1, month=12
536
+ )
537
+ else:
538
+ last_complete_month = current_month_start.replace(month=current_month_start.month - 1)
539
+ return (
540
+ False,
541
+ f'Comparison periods can only include complete months. Current month ({current_month_start.strftime("%Y-%m")}) is not complete yet. Latest allowed start date: {last_complete_month.strftime("%Y-%m-%d")}',
542
+ )
543
+
544
+ # Must be exactly one month duration
545
+ # Calculate expected end date (first day of next month)
546
+ if start_dt.month == 12:
547
+ expected_end = start_dt.replace(year=start_dt.year + 1, month=1)
548
+ else:
549
+ expected_end = start_dt.replace(month=start_dt.month + 1)
550
+
551
+ if end_dt != expected_end:
552
+ return (
553
+ False,
554
+ f"Comparison period must be exactly one month. For start date '{start_date}', end date should be '{expected_end.strftime('%Y-%m-%d')}'",
555
+ )
556
+
557
+ # Check 13-month lookback limit (38 months if multi-year enabled, but we'll use 13 as conservative)
558
+ thirteen_months_ago = today.replace(day=1)
559
+ for _ in range(13):
560
+ if thirteen_months_ago.month == 1:
561
+ thirteen_months_ago = thirteen_months_ago.replace(
562
+ year=thirteen_months_ago.year - 1, month=12
563
+ )
564
+ else:
565
+ thirteen_months_ago = thirteen_months_ago.replace(month=thirteen_months_ago.month - 1)
566
+
567
+ if start_dt < thirteen_months_ago:
568
+ return (
569
+ False,
570
+ f"Comparison start date '{start_date}' cannot be more than 13 months ago (earliest: {thirteen_months_ago.strftime('%Y-%m-%d')})",
571
+ )
572
+
573
+ return True, ''
574
+
575
+
576
+ def extract_group_key_from_complex_selector(
577
+ selector: Dict[str, Any], group_by: Dict[str, str]
578
+ ) -> str:
579
+ """Extract group key from complex CostSelector structures dynamically.
580
+
581
+ Args:
582
+ selector: The CostSelector dictionary from API response
583
+ group_by: The GroupBy dictionary with Type and Key
584
+
585
+ Returns:
586
+ String representing the group key
587
+ """
588
+ group_type = group_by.get('Type', '').upper()
589
+ group_key = group_by.get('Key', '')
590
+
591
+ def search_for_group_key(sel_part):
592
+ """Recursively search for the group key in any part of the selector."""
593
+ if isinstance(sel_part, dict):
594
+ # Check if this is the structure we're looking for
595
+ if group_type == 'DIMENSION' and 'Dimensions' in sel_part:
596
+ dim_info = sel_part['Dimensions']
597
+ if dim_info.get('Key') == group_key and 'Values' in dim_info:
598
+ values = dim_info['Values']
599
+ return values[0] if values and values[0] else f'No {group_key}'
600
+
601
+ elif group_type == 'TAG' and 'Tags' in sel_part:
602
+ tag_info = sel_part['Tags']
603
+ if tag_info.get('Key') == group_key and 'Values' in tag_info:
604
+ values = tag_info['Values']
605
+ return values[0] if values and values[0] else f'No {group_key}'
606
+
607
+ elif group_type == 'COST_CATEGORY' and 'CostCategories' in sel_part:
608
+ cc_info = sel_part['CostCategories']
609
+ if cc_info.get('Key') == group_key and 'Values' in cc_info:
610
+ values = cc_info['Values']
611
+ return values[0] if values and values[0] else f'No {group_key}'
612
+
613
+ # Recursively search in nested structures
614
+ for key, value in sel_part.items():
615
+ if key in ['And', 'Or'] and isinstance(value, list):
616
+ for item in value:
617
+ result = search_for_group_key(item)
618
+ if result:
619
+ return result
620
+ elif key == 'Not' and isinstance(value, dict):
621
+ result = search_for_group_key(value)
622
+ if result:
623
+ return result
624
+
625
+ return None
626
+
627
+ result = search_for_group_key(selector)
628
+ return result if result else 'Unknown'
629
+
630
+
631
+ def extract_usage_context_from_selector(selector: Dict[str, Any]) -> Dict[str, str]:
632
+ """Extract all available context from complex selectors dynamically.
633
+
634
+ Args:
635
+ selector: The CostSelector dictionary from API response
636
+
637
+ Returns:
638
+ Dictionary with all available context information
639
+ """
640
+ context = {}
641
+
642
+ def extract_from_structure(sel_part):
643
+ """Recursively extract context from any part of the selector."""
644
+ if isinstance(sel_part, dict):
645
+ # Extract from Dimensions
646
+ if 'Dimensions' in sel_part:
647
+ dim_info = sel_part['Dimensions']
648
+ key = dim_info.get('Key', '')
649
+ values = dim_info.get('Values', [])
650
+ if values and values[0]: # Skip empty values
651
+ context[key.lower()] = values[0]
652
+
653
+ # Extract from Tags
654
+ if 'Tags' in sel_part:
655
+ tag_info = sel_part['Tags']
656
+ tag_key = tag_info.get('Key', '')
657
+ values = tag_info.get('Values', [])
658
+ if values and values[0]:
659
+ context[f'tag_{tag_key.lower()}'] = values[0]
660
+
661
+ # Extract from CostCategories
662
+ if 'CostCategories' in sel_part:
663
+ cc_info = sel_part['CostCategories']
664
+ cc_key = cc_info.get('Key', '')
665
+ values = cc_info.get('Values', [])
666
+ if values and values[0]:
667
+ context[f'category_{cc_key.lower()}'] = values[0]
668
+
669
+ # Recursively process nested structures
670
+ for key, value in sel_part.items():
671
+ if key in ['And', 'Or'] and isinstance(value, list):
672
+ for item in value:
673
+ extract_from_structure(item)
674
+ elif key == 'Not' and isinstance(value, dict):
675
+ extract_from_structure(value)
676
+
677
+ extract_from_structure(selector)
678
+ return context
679
+
680
+
681
+ def create_detailed_group_key(
682
+ group_key: str, context: Dict[str, str], group_by: Dict[str, str]
683
+ ) -> str:
684
+ """Create a detailed group key that includes relevant context.
685
+
686
+ Since AWS always includes SERVICE and USAGE_TYPE, we can use them for context.
687
+
688
+ Args:
689
+ group_key: The primary group key extracted from the selector
690
+ context: Additional context from the selector
691
+ group_by: The GroupBy dictionary with Type and Key
692
+
693
+ Returns:
694
+ Enhanced group key with context
695
+ """
696
+ # Get the always-present context
697
+ service = context.get('service', '')
698
+ usage_type = context.get('usage_type', '')
699
+
700
+ # Create a meaningful key based on what's available
701
+ parts = [group_key]
702
+
703
+ # Add service context if it's not the group key itself
704
+ if service and group_by.get('Key') != 'SERVICE':
705
+ parts.append(service)
706
+
707
+ # Add usage type in parentheses for specificity
708
+ if usage_type:
709
+ return f'{" - ".join(parts)} ({usage_type})'
710
+
711
+ return ' - '.join(parts)
@@ -0,0 +1,88 @@
1
+ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ """Cost Explorer MCP server implementation.
16
+
17
+ Metadata tools for Cost Explorer MCP Server.
18
+ """
19
+
20
+ import os
21
+ import sys
22
+ from awslabs.cost_explorer_mcp_server.helpers import (
23
+ get_available_dimension_values,
24
+ get_available_tag_values,
25
+ )
26
+ from awslabs.cost_explorer_mcp_server.models import DateRange, DimensionKey
27
+ from loguru import logger
28
+ from mcp.server.fastmcp import Context
29
+ from pydantic import Field
30
+ from typing import Any, Dict
31
+
32
+
33
+ # Configure Loguru logging
34
+ logger.remove()
35
+ logger.add(sys.stderr, level=os.getenv('FASTMCP_LOG_LEVEL', 'WARNING'))
36
+
37
+
38
+ async def get_dimension_values(
39
+ ctx: Context, date_range: DateRange, dimension: DimensionKey
40
+ ) -> Dict[str, Any]:
41
+ """Retrieve available dimension values for AWS Cost Explorer.
42
+
43
+ This tool retrieves all available and valid values for a specified dimension (e.g., SERVICE, REGION)
44
+ over a period of time. This is useful for validating filter values or exploring available options
45
+ for cost analysis.
46
+
47
+ Args:
48
+ ctx: MCP context
49
+ date_range: The billing period start and end dates in YYYY-MM-DD format
50
+ dimension: The dimension key to retrieve values for (e.g., SERVICE, REGION, LINKED_ACCOUNT)
51
+
52
+ Returns:
53
+ Dictionary containing the dimension name and list of available values
54
+ """
55
+ try:
56
+ response = get_available_dimension_values(
57
+ dimension.dimension_key, date_range.start_date, date_range.end_date
58
+ )
59
+ return response
60
+ except Exception as e:
61
+ logger.error(f'Error getting dimension values for {dimension.dimension_key}: {e}')
62
+ return {'error': f'Error getting dimension values: {str(e)}'}
63
+
64
+
65
+ async def get_tag_values(
66
+ ctx: Context,
67
+ date_range: DateRange,
68
+ tag_key: str = Field(..., description='The tag key to retrieve values for'),
69
+ ) -> Dict[str, Any]:
70
+ """Retrieve available tag values for AWS Cost Explorer.
71
+
72
+ This tool retrieves all available values for a specified tag key over a period of time.
73
+ This is useful for validating tag filter values or exploring available tag options for cost analysis.
74
+
75
+ Args:
76
+ ctx: MCP context
77
+ date_range: The billing period start and end dates in YYYY-MM-DD format
78
+ tag_key: The tag key to retrieve values for
79
+
80
+ Returns:
81
+ Dictionary containing the tag key and list of available values
82
+ """
83
+ try:
84
+ response = get_available_tag_values(tag_key, date_range.start_date, date_range.end_date)
85
+ return response
86
+ except Exception as e:
87
+ logger.error(f'Error getting tag values for {tag_key}: {e}')
88
+ return {'error': f'Error getting tag values: {str(e)}'}
@@ -0,0 +1,70 @@
1
+ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ from awslabs.cost_explorer_mcp_server.constants import VALID_DIMENSIONS
16
+ from awslabs.cost_explorer_mcp_server.helpers import validate_date_format, validate_date_range
17
+ from pydantic import BaseModel, Field, field_validator
18
+
19
+
20
+ """Data models and validation logic for Cost Explorer MCP Server.
21
+ """
22
+
23
+
24
+ class DateRange(BaseModel):
25
+ """Date range model for cost queries."""
26
+
27
+ start_date: str = Field(
28
+ description='The start date of the billing period in YYYY-MM-DD format. Defaults to last month, if not provided.'
29
+ )
30
+ end_date: str = Field(description='The end date of the billing period in YYYY-MM-DD format.')
31
+
32
+ @field_validator('start_date', 'end_date')
33
+ @classmethod
34
+ def validate_individual_dates(cls, v):
35
+ """Validate that individual dates are in YYYY-MM-DD format and are valid dates."""
36
+ is_valid, error = validate_date_format(v)
37
+ if not is_valid:
38
+ raise ValueError(error)
39
+ return v
40
+
41
+ def model_post_init(self, __context):
42
+ """Validate the date range after both dates are set."""
43
+ is_valid, error = validate_date_range(self.start_date, self.end_date)
44
+ if not is_valid:
45
+ raise ValueError(error)
46
+
47
+ def validate_with_granularity(self, granularity: str):
48
+ """Validate the date range with granularity-specific constraints."""
49
+ is_valid, error = validate_date_range(self.start_date, self.end_date, granularity)
50
+ if not is_valid:
51
+ raise ValueError(error)
52
+
53
+
54
+ class DimensionKey(BaseModel):
55
+ """Dimension key model."""
56
+
57
+ dimension_key: str = Field(
58
+ description=f'The name of the dimension to retrieve values for. Valid values are {", ".join(VALID_DIMENSIONS)}.'
59
+ )
60
+
61
+ @field_validator('dimension_key')
62
+ @classmethod
63
+ def validate_dimension_key(cls, v):
64
+ """Validate that the dimension key is supported by AWS Cost Explorer."""
65
+ dimension_upper = v.upper()
66
+ if dimension_upper not in VALID_DIMENSIONS:
67
+ raise ValueError(
68
+ f"Invalid dimension key '{v}'. Valid dimensions are: {', '.join(VALID_DIMENSIONS)}"
69
+ )
70
+ return v