diff --git a/hic_aws_costing_tools/aws_costs.py b/hic_aws_costing_tools/aws_costs.py index d25a312..0841a9b 100755 --- a/hic_aws_costing_tools/aws_costs.py +++ b/hic_aws_costing_tools/aws_costs.py @@ -6,16 +6,19 @@ DEFAULT_COST_TYPE = "UnblendedCost" DEFAULT_GRANULARITY = "MONTHLY" +# Previously we excluded these types by default. Now we just include Usage instead. DEFAULT_EXCLUDE_RECORD_TYPES = [ - "Credit", - "Refund", - "Tax", - # These two aren't documented on - # https://docs.aws.amazon.com/awsaccountbilling/latest/aboutv2/manage-cost-categories.html#cost-categories-terms - # but were confirmed in AWS Support ticket 171570162800825 - "Enterprise Discount Program Discount", - "Solution Provider Program Discount", + # "Credit", + # "Refund", + # "Tax", + # # These two aren't documented on + # # https://docs.aws.amazon.com/awsaccountbilling/latest/aboutv2/manage-cost-categories.html#cost-categories-terms + # # but were confirmed in AWS Support ticket 171570162800825 + # "Enterprise Discount Program Discount", + # "Solution Provider Program Discount", ] +# https://docs.aws.amazon.com/awsaccountbilling/latest/aboutv2/manage-cost-categories.html#cost-categories-terms +DEFAULT_INCLUDE_RECORD_TYPES = ["Usage"] EXPECTED_UNIT = "USD" @@ -53,8 +56,62 @@ def _get_group_by(ce, time_period, dimension): raise ValueError(f"Invalid dimension: {dimension}") +def _get_filter(regions, exclude_types, include_types): + filter_count = 0 + region_filter = None + exclude_filter = None + include_filter = None + filter = None + + if regions: + filter_count += 1 + region_filter = dict( + Dimensions={ + "Key": "REGION", + "Values": regions, + } + ) + if exclude_types: + filter_count += 1 + exclude_filter = dict( + Not=dict( + Dimensions={ + "Key": "RECORD_TYPE", + "Values": exclude_types, + } + ) + ) + if include_types: + filter_count += 1 + include_filter = dict( + Dimensions={ + "Key": "RECORD_TYPE", + "Values": include_types, + } + ) + + if filter_count > 1: + filter = dict(And=[]) + for f in [region_filter, exclude_filter, include_filter]: + if f: + if filter: + filter["And"].append(f) + else: + filter = f + + return filter + + def costs_for_regions( - *, time_period, granularity, regions, session, group1, group2, exclude_types + *, + time_period, + granularity, + regions, + session, + group1, + group2, + exclude_types, + include_types, ): if session: ce = session.client("ce") @@ -73,28 +130,9 @@ def costs_for_regions( TimePeriod=time_period, ) - exclude_record_types = dict( - Not=dict( - Dimensions={ - "Key": "RECORD_TYPE", - "Values": exclude_types, - } - ) - ) - if regions: - kwargs["Filter"] = dict( - And=[ - exclude_record_types, - dict( - Dimensions={ - "Key": "REGION", - "Values": regions, - } - ), - ] - ) - else: - kwargs["Filter"] = exclude_record_types + filter = _get_filter(regions, exclude_types, include_types) + if filter: + kwargs["Filter"] = filter while not r or "NextPageToken" in r: # print(f"get_cost_and_usage({kwargs})") @@ -212,6 +250,7 @@ def get_raw_cost_data( group1, group2, exclude_types, + include_types, apply_value_mappings, ): session = None @@ -235,6 +274,7 @@ def get_raw_cost_data( group1=group1, group2=group2, exclude_types=exclude_types, + include_types=include_types, ) if apply_value_mappings: @@ -280,6 +320,7 @@ def create_costs_message( group1, group2, exclude_types, + include_types, output, ): results, all_values1, all_values2, value_map1, value_map2 = get_raw_cost_data( @@ -290,6 +331,7 @@ def create_costs_message( group1=group1, group2=group2, exclude_types=exclude_types, + include_types=include_types, apply_value_mappings=True, ) @@ -347,6 +389,7 @@ def create_costs_plain_output( group1, group2, exclude_types, + include_types, output, ): results, all_values1, all_values2, value_map1, value_map2 = get_raw_cost_data( @@ -357,6 +400,7 @@ def create_costs_plain_output( group1=group1, group2=group2, exclude_types=exclude_types, + include_types=include_types, apply_value_mappings=True, ) diff --git a/hic_aws_costing_tools/main.py b/hic_aws_costing_tools/main.py index c5e1197..53e2a10 100644 --- a/hic_aws_costing_tools/main.py +++ b/hic_aws_costing_tools/main.py @@ -4,6 +4,7 @@ DEFAULT_COST_TYPE, DEFAULT_EXCLUDE_RECORD_TYPES, DEFAULT_GRANULARITY, + DEFAULT_INCLUDE_RECORD_TYPES, create_costs_message, create_costs_plain_output, get_time_period, @@ -51,6 +52,12 @@ def main(): default=DEFAULT_EXCLUDE_RECORD_TYPES, help=f"Exclude these record types (default {DEFAULT_EXCLUDE_RECORD_TYPES})", ) + parser.add_argument( + "--include-types", + nargs="*", + default=DEFAULT_INCLUDE_RECORD_TYPES, + help=f"Include these record types (default {DEFAULT_INCLUDE_RECORD_TYPES})", + ) parser.add_argument( "--output", choices=["auto", "summary", "full", "csv", "flat"], @@ -71,6 +78,7 @@ def main(): group1=args.group1, group2=args.group2, exclude_types=args.exclude_types, + include_types=args.include_types, output=args.output, ) else: @@ -84,6 +92,7 @@ def main(): group1=args.group1, group2=args.group2, exclude_types=args.exclude_types, + include_types=args.include_types, output=args.output, ) print(title) diff --git a/tests/test_costbot.py b/tests/test_costbot.py index 529b8bc..252fed7 100644 --- a/tests/test_costbot.py +++ b/tests/test_costbot.py @@ -85,6 +85,91 @@ def side_effect(*args, **kwargs): assert value_map == expected_value_map +@pytest.mark.parametrize( + "regions,exclude,include,expected", + [ + (["mars"], [], [], {"Dimensions": {"Key": "REGION", "Values": ["mars"]}}), + ( + [], + ["exclude"], + [], + {"Not": {"Dimensions": {"Key": "RECORD_TYPE", "Values": ["exclude"]}}}, + ), + ( + [], + [], + ["include"], + {"Dimensions": {"Key": "RECORD_TYPE", "Values": ["include"]}}, + ), + ( + ["mars"], + ["exclude"], + [], + { + "And": [ + {"Dimensions": {"Key": "REGION", "Values": ["mars"]}}, + { + "Not": { + "Dimensions": {"Key": "RECORD_TYPE", "Values": ["exclude"]} + } + }, + ] + }, + ), + ( + ["mars", "jupiter"], + [], + ["include", "i2"], + { + "And": [ + {"Dimensions": {"Key": "REGION", "Values": ["mars", "jupiter"]}}, + {"Dimensions": {"Key": "RECORD_TYPE", "Values": ["include", "i2"]}}, + ] + }, + ), + ( + [], + ["exclude", "e2"], + ["include"], + { + "And": [ + { + "Not": { + "Dimensions": { + "Key": "RECORD_TYPE", + "Values": ["exclude", "e2"], + } + } + }, + {"Dimensions": {"Key": "RECORD_TYPE", "Values": ["include"]}}, + ] + }, + ), + ( + ["mars", "jupiter"], + ["exclude", "e2"], + ["include", "i2"], + { + "And": [ + {"Dimensions": {"Key": "REGION", "Values": ["mars", "jupiter"]}}, + { + "Not": { + "Dimensions": { + "Key": "RECORD_TYPE", + "Values": ["exclude", "e2"], + } + } + }, + {"Dimensions": {"Key": "RECORD_TYPE", "Values": ["include", "i2"]}}, + ] + }, + ), + ], +) +def test_get_filter(regions, exclude, include, expected): + assert aws_costs._get_filter(regions, exclude, include) == expected + + @pytest.mark.parametrize("scenario", ["dummy-services", "dummy-proj"]) def test_costs_for_regions(mocker, scenario): group1 = "accountname" @@ -134,6 +219,7 @@ def side_effect(*args, **kwargs): group1=group1, group2=group2, exclude_types=["Credit", "Refund"], + include_types=[], ) assert all_values1 == {"000000000001", "000000000002"}