Skip to content

Commit

Permalink
feat(tag-filtering): enable tag filtering metrics
Browse files Browse the repository at this point in the history
  • Loading branch information
Pavel committed Nov 9, 2024
1 parent 6966493 commit 2a34ea2
Show file tree
Hide file tree
Showing 3 changed files with 113 additions and 54 deletions.
115 changes: 85 additions & 30 deletions app/exporter.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,15 +12,16 @@

class MetricExporter:
def __init__(
self,
polling_interval_seconds,
metric_name,
aws_access_key,
aws_access_secret,
aws_assumed_role_name,
group_by,
targets,
metric_type,
self,
polling_interval_seconds,
metric_name,
aws_access_key,
aws_access_secret,
aws_assumed_role_name,
group_by,
targets,
metric_type,
tag_filters=None, # Added tag_filters parameter
):
self.polling_interval_seconds = polling_interval_seconds
self.metric_name = metric_name
Expand All @@ -30,21 +31,30 @@ def __init__(
self.aws_assumed_role_name = aws_assumed_role_name
self.group_by = group_by
self.metric_type = metric_type # Store metrics
# we have verified that there is at least one target
self.tag_filters = tag_filters # Store tag filters

# We have verified that there is at least one target
self.labels = set(targets[0].keys())
# for now we only support exporting one type of cost (Usage)

# For now we only support exporting one type of cost (Usage)
self.labels.add("ChargeType")

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(self):
# every time we clear up all the existing labels before setting new ones
# 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:
Expand All @@ -59,33 +69,66 @@ def get_aws_account_session_via_iam_user(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"]

def get_aws_account_session_default(self, account_id):
sts_client = boto3.client(
"sts",
)
sts_client = boto3.client("sts")

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"]

def query_aws_cost_explorer(self, aws_client, group_by):
def query_aws_cost_explorer(self, aws_client, group_by, tag_filters=None):
end_date = datetime.today()
start_date = end_date - relativedelta(days=1)

# Keep the 'groups' code as specified
groups = list()
if group_by["enabled"]:
for group in group_by["groups"]:
groups.append({"Type": group["type"], "Key": group["key"]})

# Build the base filter with RECORD_TYPE
base_filter = {
"Dimensions": {
"Key": "RECORD_TYPE",
"Values": ["Usage"]
}
}

# Include tag filters if provided
if tag_filters:
tag_filter_list = []
for tag_key, tag_values in tag_filters.items():
tag_filter = {
"Tags": {
"Key": tag_key,
"Values": tag_values,
"MatchOptions": ["EQUALS"] # Adjust as needed
}
}
tag_filter_list.append(tag_filter)

# Combine the base filter and tag filters using 'And'
combined_filter = {
"And": [base_filter] + tag_filter_list
}
else:
combined_filter = base_filter

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=combined_filter,
Granularity="DAILY",
Metrics=[self.metric_type], # Use dynamic metrics
GroupBy=groups,
Expand All @@ -106,7 +149,13 @@ def fetch(self, aws_account):
aws_session_token=aws_credentials["SessionToken"],
region_name="us-east-1",
)
cost_response = self.query_aws_cost_explorer(aws_client, self.group_by)

# Pass tag_filters to query_aws_cost_explorer
cost_response = self.query_aws_cost_explorer(
aws_client,
self.group_by,
self.tag_filters, # Include tag filters
)

for result in cost_response:
if not self.group_by["enabled"]:
Expand All @@ -126,19 +175,25 @@ def fetch(self, aws_account):
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"]
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.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.aws_daily_cost_usd.labels(
**aws_account, **group_key_values, ChargeType="Usage"
).set(merged_minor_cost)
10 changes: 10 additions & 0 deletions exporter_config.yaml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,11 @@ metrics:
tag_value: other
# Allowed values for metric type are AmortizedCost, BlendedCost, NetAmortizedCost, NetUnblendedCost, NormalizedUsageAmount, UnblendedCost, and UsageQuantity
metric_type: AmortizedCost
tag_filters:
ProjectName:
- dev-team-1
EnvironmentName:
- dev

- metric_name: aws_daily_cost_usd # change the metric name if needed
group_by:
Expand All @@ -46,6 +51,11 @@ metrics:
tag_value: other
# Allowed values for metric type are AmortizedCost, BlendedCost, NetAmortizedCost, NetUnblendedCost, NormalizedUsageAmount, UnblendedCost, and UsageQuantity
metric_type: AmortizedCost
tag_filters:
ProjectName:
- dev-team-1
EnvironmentName:
- dev

target_aws_accounts:
# here defines a list of target AWS accounts
Expand Down
42 changes: 18 additions & 24 deletions main.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ def validate_configs(config):
]

if len(config["target_aws_accounts"]) == 0:
logging.error("There should be at least one target AWS accounts defined in the config!")
logging.error("There should be at least one target AWS account defined in the config!")
sys.exit(1)

labels = config["target_aws_accounts"][0].keys()
Expand All @@ -65,21 +65,23 @@ def validate_configs(config):
sys.exit(1)

for config_metric in config["metrics"]:
group_label_names = set()

if config_metric["group_by"]["enabled"]:
if len(config_metric["group_by"]["groups"]) < 1 or len(config_metric["group_by"]["groups"]) > 2:
logging.error("If group_by is enabled, there should be at least one group, and at most two groups!")
sys.exit(1)
group_label_names = set()

for group in config_metric["group_by"]["groups"]:
if group["label_name"] in group_label_names:
logging.error("Group label names should be unique!")
sys.exit(1)
else:
group_label_names.add(group["label_name"])
if group_label_names and (group_label_names & set(labels)):
logging.error("Some label names in group_by are the same as AWS account labels!")
sys.exit(1)

if group_label_names & set(labels):
logging.error("Some label names in group_by are the same as AWS account labels!")
sys.exit(1)

# Validate metric_type
if config_metric["metric_type"] not in valid_metric_types:
Expand All @@ -88,27 +90,18 @@ def validate_configs(config):
)
sys.exit(1)

for i in range(1, len(config["target_aws_accounts"])):
if labels != config["target_aws_accounts"][i].keys():
logging.error("All the target AWS accounts should have the same set of keys (labels)!")
sys.exit(1)

for config_metric in config["metrics"]:

if config_metric["group_by"]["enabled"]:
if len(config_metric["group_by"]["groups"]) < 1 or len(config_metric["group_by"]["groups"]) > 2:
logging.error("If group_by is enabled, there should be at least one group, and at most two groups!")
# Validate tag_filters if present
if "tag_filters" in config_metric:
tag_filters = config_metric["tag_filters"]
if not isinstance(tag_filters, dict):
logging.error("tag_filters should be a dictionary.")
sys.exit(1)
group_label_names = set()
for group in config_metric["group_by"]["groups"]:
if group["label_name"] in group_label_names:
logging.error("Group label names should be unique!")
for tag_key, tag_values in tag_filters.items():
if not isinstance(tag_values, list):
logging.error(f"Values for tag '{tag_key}' should be a list.")
sys.exit(1)
else:
group_label_names.add(group["label_name"])
if group_label_names and (group_label_names & set(labels)):
logging.error("Some label names in group_by are the same as AWS account labels!")
sys.exit(1)

# No need to repeat the validation loops; they have been consolidated above.


def main(config):
Expand All @@ -124,6 +117,7 @@ def main(config):
metric_name=config_metric["metric_name"],
group_by=config_metric["group_by"],
metric_type=config_metric["metric_type"],
tag_filters=config_metric.get("tag_filters", None), # Added tag_filters parameter
)
app_metrics.run_metrics()
time.sleep(config["polling_interval_seconds"])
Expand Down

0 comments on commit 2a34ea2

Please sign in to comment.