From 9b966496b8e10d2947b6146150e03a8c27271b2b Mon Sep 17 00:00:00 2001 From: Long Zhang Date: Thu, 14 Mar 2024 10:21:41 +0100 Subject: [PATCH] fix: metric labels are not cleaned up when new data is fetched (#8) * format the code * fix: metric labels are not cleaned up when new data is fetched --- app/exporter.py | 86 +++++++++++++++++++++++-------------------------- package.json | 6 ++-- 2 files changed, 43 insertions(+), 49 deletions(-) diff --git a/app/exporter.py b/app/exporter.py index edcbad1..78d3ec7 100644 --- a/app/exporter.py +++ b/app/exporter.py @@ -2,17 +2,26 @@ # -*- coding:utf-8 -*- # Filename: exporter.py +import logging import time -import boto3 -import botocore from datetime import datetime + +import boto3 from dateutil.relativedelta import relativedelta from prometheus_client import Gauge -import logging class MetricExporter: - def __init__(self, polling_interval_seconds, metric_name, aws_access_key, aws_access_secret, aws_assumed_role_name, group_by, targets): + def __init__( + self, + polling_interval_seconds, + metric_name, + aws_access_key, + aws_access_secret, + aws_assumed_role_name, + group_by, + targets, + ): self.polling_interval_seconds = polling_interval_seconds self.metric_name = metric_name self.targets = targets @@ -27,14 +36,15 @@ def __init__(self, polling_interval_seconds, metric_name, aws_access_key, aws_ac if group_by["enabled"]: for group in group_by["groups"]: self.labels.add(group["label_name"]) - self.aws_daily_cost_usd = Gauge( - self.metric_name, "Daily cost of an AWS account in USD", self.labels) + self.aws_daily_cost_usd = Gauge(self.metric_name, "Daily cost of an AWS account in USD", self.labels) def run_metrics_loop(self): while True: + # every time we clear up all the existing labels before setting new ones + self.aws_daily_cost_usd .clear() + for aws_account in self.targets: - logging.info("querying cost data for aws account %s" % - aws_account["Publisher"]) + logging.info("querying cost data for aws account %s" % aws_account["Publisher"]) try: self.fetch(aws_account) except Exception as e: @@ -50,8 +60,7 @@ def get_aws_account_session(self, account_id): ) assumed_role_object = sts_client.assume_role( - RoleArn=f"arn:aws:iam::{account_id}:role/{self.aws_assumed_role_name}", - RoleSessionName="AssumeRoleSession1" + RoleArn=f"arn:aws:iam::{account_id}:role/{self.aws_assumed_role_name}", RoleSessionName="AssumeRoleSession1" ) return assumed_role_object["Credentials"] @@ -62,54 +71,37 @@ def query_aws_cost_explorer(self, aws_client, group_by): groups = list() if group_by["enabled"]: for group in group_by["groups"]: - groups.append({ - "Type": group["type"], - "Key": group["key"] - }) + groups.append({"Type": group["type"], "Key": group["key"]}) response = aws_client.get_cost_and_usage( - TimePeriod={ - "Start": start_date.strftime("%Y-%m-%d"), - "End": end_date.strftime("%Y-%m-%d") - }, - Filter={ - "Dimensions": { - "Key": "RECORD_TYPE", - "Values": ["Usage"] - } - }, + TimePeriod={"Start": start_date.strftime("%Y-%m-%d"), "End": end_date.strftime("%Y-%m-%d")}, + Filter={"Dimensions": {"Key": "RECORD_TYPE", "Values": ["Usage"]}}, Granularity="DAILY", - Metrics=[ - "UnblendedCost" - ], - GroupBy=groups + Metrics=["UnblendedCost"], + GroupBy=groups, ) return response["ResultsByTime"] def fetch(self, aws_account): - aws_credentials = self.get_aws_account_session( - aws_account["Publisher"]) + aws_credentials = self.get_aws_account_session(aws_account["Publisher"]) aws_client = boto3.client( "ce", aws_access_key_id=aws_credentials["AccessKeyId"], aws_secret_access_key=aws_credentials["SecretAccessKey"], aws_session_token=aws_credentials["SessionToken"], - region_name="us-east-1" + region_name="us-east-1", ) - cost_response = self.query_aws_cost_explorer( - aws_client, self.group_by) + cost_response = self.query_aws_cost_explorer(aws_client, self.group_by) for result in cost_response: if not self.group_by["enabled"]: cost = float(result["Total"]["UnblendedCost"]["Amount"]) - self.aws_daily_cost_usd.labels( - **aws_account, ChargeType="Usage").set(cost) + self.aws_daily_cost_usd.labels(**aws_account, ChargeType="Usage").set(cost) else: merged_minor_cost = 0 for item in result["Groups"]: - cost = float(item["Metrics"] - ["UnblendedCost"]["Amount"]) + cost = float(item["Metrics"]["UnblendedCost"]["Amount"]) group_key_values = dict() for i in range(len(self.group_by["groups"])): @@ -117,20 +109,22 @@ def fetch(self, aws_account): value = item["Keys"][i].split("$")[1] else: value = item["Keys"][i] - group_key_values.update( - {self.group_by["groups"][i]["label_name"]: value}) + group_key_values.update({self.group_by["groups"][i]["label_name"]: value}) - if self.group_by["merge_minor_cost"]["enabled"] and \ - cost < self.group_by["merge_minor_cost"]["threshold"]: + if ( + self.group_by["merge_minor_cost"]["enabled"] + and cost < self.group_by["merge_minor_cost"]["threshold"] + ): merged_minor_cost += cost else: - self.aws_daily_cost_usd.labels( - **aws_account, **group_key_values, ChargeType="Usage").set(cost) + self.aws_daily_cost_usd.labels(**aws_account, **group_key_values, ChargeType="Usage").set(cost) if merged_minor_cost > 0: group_key_values = dict() for i in range(len(self.group_by["groups"])): group_key_values.update( - {self.group_by["groups"][i]["label_name"]: self.group_by["merge_minor_cost"]["tag_value"]}) - self.aws_daily_cost_usd.labels( - **aws_account, **group_key_values, ChargeType="Usage").set(merged_minor_cost) + {self.group_by["groups"][i]["label_name"]: self.group_by["merge_minor_cost"]["tag_value"]} + ) + self.aws_daily_cost_usd.labels(**aws_account, **group_key_values, ChargeType="Usage").set( + merged_minor_cost + ) diff --git a/package.json b/package.json index 27588ca..5d8ed1f 100644 --- a/package.json +++ b/package.json @@ -1,4 +1,4 @@ { - "name": "aws-cost-exporter", - "version": "v1.0.3" - } + "name": "aws-cost-exporter", + "version": "v1.0.4" +}