def test_verify_labels_yaml(config_dict):
assert(verify_labels_yaml(config_dict['good_labels']) is True)
assert(verify_labels_yaml(config_dict['missing_label_list']) is False)
assert(verify_labels_yaml(config_dict['label_list_wrong_type']) is False)
assert(verify_labels_yaml(config_dict['label_list_len']) is True)
assert(verify_labels_yaml(config_dict['label_list_obj_wrong_type']) is False)
assert(verify_labels_yaml(config_dict['label_list_obj_wrong_type2']) is False)
verify_labels_yaml
# verify_default_yaml
)
# test config file for errors as a requirement for merging
with open(sys.argv[1], "r") as f:
budget_dict = yaml.load(f, Loader=yaml.SafeLoader)
for project_dict in budget_dict['projects']:
if verify_project_yaml(project_dict):
pass
else:
print(f"Error validating {project_dict}.")
sys.exit(1)
for parent_dict in budget_dict['parent_folders']:
if verify_parent_yaml(parent_dict):
pass
else:
print(f"Error validating {parent_dict}.")
sys.exit(1)
for label_dict in budget_dict['labels']:
if verify_labels_yaml(label_dict):
pass
else:
print(f"Error validating {label_dict}.")
sys.exit(1)
print("Config verification succeeded.")
def test_label_list_obj_wrong_type2(label_dict):
label_dict['label_list'] = [{'key': 'value', 'key2': 'value2'}]
assert(verify_labels_yaml(label_dict) is False)
def test_label_list_obj_wrong_type(label_dict):
label_dict['label_list'] = ['key, value']
assert(verify_labels_yaml(label_dict) is False)
def test_label_list_len(label_dict):
label_dict['label_list'].append({'key2': 'value2'})
assert(verify_labels_yaml(label_dict) is True)
def test_label_list_wrong_type(label_dict):
label_dict['label_list'] = 1234
assert(verify_labels_yaml(label_dict) is False)
def test_no_labels(label_dict):
del label_dict['label_list']
assert(verify_labels_yaml(label_dict) is False)
def test_yes_labels(label_dict):
assert(verify_labels_yaml(label_dict) is True)
def main(gcs_bucket, gcs_file_path, billing_account_id, default_pubsub_topic,
mysql_host, mysql_port, mysql_user, mysql_pass, mysql_db, statsd_host,
local_mode, dry_run):
logformat = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
logging.basicConfig(stream=sys.stdout,
level=logging.INFO,
format=logformat)
# TODO - troubleshoot metric uploading with local container, and test it works on GKE as well.
# Remove this line to see WARNINGs
logging.getLogger('datadog.dogstatsd').setLevel(logging.ERROR)
# Load config.yaml file
if local_mode:
with open("config.yaml", "r") as f:
budget_dict = yaml.load(f, Loader=yaml.SafeLoader)
setup_metrics("localhost")
else:
# Load yaml configuration file from GCS
gcs_client = storage.Client()
bucket = gcs_client.bucket(gcs_bucket)
blob = bucket.blob(gcs_file_path)
yaml_string = blob.download_as_string()
budget_dict = yaml.load(yaml_string, Loader=yaml.SafeLoader)
# Setup metrics with configured statsd host
setup_metrics(statsd_host)
# Setup gcp helper
billing_client = billing_budgets_v1beta1.BudgetServiceClient()
resource_manager_client = resource_manager.Client()
gcp = GcpHelper(billing_client, resource_manager_client)
# Setup mysql connection pool
pool = PooledDB(creator=pymysql,
host=mysql_host,
user=mysql_user,
password=mysql_pass,
database=mysql_db,
autocommit=True,
blocking=True,
maxconnections=5)
mysql_conn = pool.connection()
# Iterate over all configured project budgets
for project_dict in budget_dict['projects']:
config_type = PLUTUS_CONFIG_TYPE_PROJECT
if verify_project_yaml(project_dict):
project = ProjectBudget(project_dict, config_type,
billing_account_id, default_pubsub_topic)
budget = gcp.get_and_update_or_create_budget(project)
if budget is not None:
with mysql_conn.cursor() as mysql_cursor:
upsert_budget(mysql_cursor, budget, project.project_id,
config_type, project.alert_emails)
else:
log.error("Project config verification failed.")
metrics.incr("error_count",
tags=[
"type:misconfig",
f"project_id:{project.project_id}",
f"config_type:{config_type}"
])
sys.exit(1)
# Iterate over all configured parent folder id budgets
for parent_dict in budget_dict['parent_folders']:
parent_id = parent_dict['parent_folder_id']
config_type = PLUTUS_CONFIG_TYPE_PARENT
if verify_parent_yaml(parent_dict):
parent_filter = {'parent.id': parent_id}
for p in resource_manager_client.list_projects(parent_filter):
# Overwrite 'project_id' key for each project under parent folder
parent_dict['project_id'] = p.project_id
project = ProjectBudget(parent_dict, config_type,
billing_account_id,
default_pubsub_topic)
existing_project_budget_id = gcp.has_existing_project_budget(
project)
if existing_project_budget_id is None:
# No project one off budget configured for this projectid
budget = gcp.get_and_update_or_create_budget(project)
if budget is not None:
with mysql_conn.cursor() as mysql_cursor:
upsert_budget(mysql_cursor, budget,
project.project_id, config_type,
project.alert_emails)
else:
log.info(f"Skipping creating parent project budget for \
{p.project_id} since exiting plutus project budget found."
)
existing_parent_budget_id = gcp.has_existing_parent_budget(
parent_id, project)
if existing_parent_budget_id is not None:
# We have a configured plutus project budget. Delete parent budget
gcp.delete_budget(existing_parent_budget_id)
else:
log.error("Parent folder config verification failed.")
metrics.incr("error_count",
tags=[
"type:misconfig", f"parent_id:{parent_id}",
f"config_type:{config_type}"
])
sys.exit(1)
# Iterate over all configured label budgets
for label_dict in budget_dict['labels']:
config_type = PLUTUS_CONFIG_TYPE_LABEL
if verify_labels_yaml(label_dict):
labels_filter = {}
for row in label_dict['label_list']:
for key in row:
labels_filter[f"labels.{key}"] = row[key]
# Find projects that match the labels
for p in resource_manager_client.list_projects(
filter_params=labels_filter):
# Overwrite the 'project_id' key for each project that matches labels
label_dict['project_id'] = p.project_id
project = ProjectBudget(label_dict, config_type,
billing_account_id,
default_pubsub_topic)
existing_project_budget_id = gcp.has_existing_project_budget(
project)
if existing_project_budget_id is None:
# No project one off budget configured for this projectid
budget = gcp.get_and_update_or_create_budget(project)
if budget is not None:
with mysql_conn.cursor() as mysql_cursor:
upsert_budget(mysql_cursor, budget,
project.project_id, config_type,
project.alert_emails)
else:
log.info(f"Skipping creating label project budget for \
{p.project_id} since exiting plutus project budget found."
)
existing_labels_budget_id = gcp.has_existing_labels_budget(
project)
if existing_labels_budget_id is not None:
# We have a configured plutus project budget. Delete labels budget
gcp.delete_budget(existing_labels_budget_id)
else:
log.error("Labels config verification failed.")
metrics.incr("error_count",
tags=["type:misconfig", f"config_type:{config_type}"])
sys.exit(1)
# Default logic removed for now because it will create hundreds of budgets
# And we need to decide good thresholds for this
'''
default_dict = budget_dict['default']
config_type = PLUTUS_CONFIG_TYPE_DEFAULT
if verify_default_yaml(default_dict):
# Query for all projects. If no budget exists for project, add a default budget
for p in resource_manager_client.list_projects():
project_id = p.project_id
project_number = gcp.get_project_number(project_id)
if project_number is not None:
budgets = gcp.get_budgets_by_project(project, project_number)
if len(budgets) == 0:
# No budget exists for this project, so we create one
default_dict['project_id'] = project_id
# TODO - add and test
# project = ProjectBudget(default_dict, config_type,
# billing_account_id, default_pubsub_topic)
log.info(f"Creating default budget for plutus-default-{project_id}")
# budget = gcp.create_budget(project)
# metrics.incr("default_budget_created_count", tags=[])
#if budget is not None:
# with mysql_conn.cursor() as mysql_cursor:
# upsert_budget(mysql_cursor, budget, project.project_id,
# config_type, project.alert_emails)
# TODO - delete plutus-default budget if > 1 budget
elif len(budgets) > 1:
# TODO - if one of the budgets is a plutus-default, delete via API and in Mysql
pass
else:
log.error("Default config verification failed.")
metrics.incr("error_count", tags=["type:misconfig", f"config_type:{config_type}"])
sys.exit(1)
'''
log.info("Plutus run complete.")