From 942f7135c701700bc319673139c536973352a7b0 Mon Sep 17 00:00:00 2001 From: diodonfrost Date: Fri, 9 Jun 2023 00:12:25 +0200 Subject: [PATCH] feat(terraform): add documentdb scheduler This commit introduces a new feature to our Terraform configuration, enabling the addition of a Documentdb scheduler module. This module allows for the automated stopping or starting of Documentdb clusters based on predefined tags. By leveraging this scheduler, we can effectively manage the usage and costs associated with our Documentdb clusters, ensuring they are active only when needed and maximizing resource efficiency. --- main.tf | 1 + package/scheduler/documentdb_handler.py | 74 +++++++++++++++++++++++++ package/scheduler/exceptions.py | 43 ++++++++++++++ package/scheduler/main.py | 2 + variables.tf | 7 ++- 5 files changed, 126 insertions(+), 1 deletion(-) create mode 100644 package/scheduler/documentdb_handler.py diff --git a/main.tf b/main.tf index d5b54dbe..c8da9b36 100644 --- a/main.tf +++ b/main.tf @@ -274,6 +274,7 @@ resource "aws_lambda_function" "this" { SCHEDULE_ACTION = var.schedule_action TAG_KEY = local.scheduler_tag["key"] TAG_VALUE = local.scheduler_tag["value"] + DOCUMENTDB_SCHEDULE = tostring(var.documentdb_schedule) EC2_SCHEDULE = tostring(var.ec2_schedule) ECS_SCHEDULE = tostring(var.ecs_schedule) RDS_SCHEDULE = tostring(var.rds_schedule) diff --git a/package/scheduler/documentdb_handler.py b/package/scheduler/documentdb_handler.py new file mode 100644 index 00000000..ed1c9125 --- /dev/null +++ b/package/scheduler/documentdb_handler.py @@ -0,0 +1,74 @@ +# -*- coding: utf-8 -*- + +"""documentdb instances scheduler.""" + +from typing import Dict, List + +import boto3 + +from botocore.exceptions import ClientError + +from .exceptions import documentdb_exception +from .filter_resources_by_tags import FilterByTags + + +class DocumentDBScheduler: + """documentdb scheduler.""" + + def __init__(self, region_name=None) -> None: + """Initialize documentdb scheduler.""" + if region_name: + self.documentdb = boto3.client("docdb", region_name=region_name) + else: + self.documentdb = boto3.client("docdb") + self.tag_api = FilterByTags(region_name=region_name) + + def stop(self, aws_tags: List[Dict]) -> None: + """Aws documentdb cluster stop function. + + Stop documentdb clusters with defined tags. + + :param list[map] aws_tags: + Aws tags to use for filter resources. + For example: + [ + { + 'Key': 'string', + 'Values': [ + 'string', + ] + } + ] + """ + for cluster_arn in self.tag_api.get_resources("rds:cluster", aws_tags): + cluster_id = cluster_arn.split(":")[-1] + try: + self.documentdb.stop_db_cluster(DBClusterIdentifier=cluster_id) + print(f"Stop documentdb cluster {cluster_id}") + except ClientError as exc: + documentdb_exception("documentdb cluster", cluster_id, exc) + + def start(self, aws_tags: List[Dict]) -> None: + """Aws documentdb cluster start function. + + Start documentdb clusters with defined tags. + + :param list[map] aws_tags: + Aws tags to use for filter resources. + For example: + [ + { + 'Key': 'string', + 'Values': [ + 'string', + ] + } + ] + """ + for cluster_arn in self.tag_api.get_resources("rds:cluster", aws_tags): + cluster_id = cluster_arn.split(":")[-1] + try: + self.documentdb.start_db_cluster(DBClusterIdentifier=cluster_id) + print(f"Start documentdb cluster {cluster_id}") + except ClientError as exc: + documentdb_exception("documentdb cluster", cluster_id, exc) diff --git a/package/scheduler/exceptions.py b/package/scheduler/exceptions.py index 2427f53c..3e7f745c 100644 --- a/package/scheduler/exceptions.py +++ b/package/scheduler/exceptions.py @@ -5,6 +5,49 @@ import logging +def documentdb_exception(resource_name: str, resource_id: str, exception) -> None: + """Exception raised during execution of documentdb scheduler. + + Log instance, spot instance and autoscaling groups exceptions + on the specific aws resources. + + :param str resource_name: + Aws resource name + :param str resource_id: + Aws resource id + :param str exception: + Human readable string describing the exception + """ + info_codes = ["InvalidDBClusterStateFault"] + warning_codes = [ + "InvalidDBClusterStateFault", + "DBClusterNotFoundFault", + "DBClusterParameterGroupNotFound", + ] + + if exception.response["Error"]["Code"] in info_codes: + logging.info( + "%s %s: %s", + resource_name, + resource_id, + exception, + ) + elif exception.response["Error"]["Code"] in warning_codes: + logging.warning( + "%s %s: %s", + resource_name, + resource_id, + exception, + ) + else: + logging.error( + "Unexpected error on %s %s: %s", + resource_name, + resource_id, + exception, + ) + + def ec2_exception(resource_name: str, resource_id: str, exception) -> None: """Exception raised during execution of ec2 scheduler. diff --git a/package/scheduler/main.py b/package/scheduler/main.py index bfe7a028..1202752e 100644 --- a/package/scheduler/main.py +++ b/package/scheduler/main.py @@ -6,6 +6,7 @@ from .autoscaling_handler import AutoscalingScheduler from .cloudwatch_handler import CloudWatchAlarmScheduler +from .documentdb_handler import DocumentDBScheduler from .ecs_handler import EcsScheduler from .instance_handler import InstanceScheduler from .rds_handler import RdsScheduler @@ -34,6 +35,7 @@ def lambda_handler(event, context): _strategy = {} _strategy[AutoscalingScheduler] = os.getenv("AUTOSCALING_SCHEDULE") + _strategy[DocumentDBScheduler] = os.getenv("DOCUMENTDB_SCHEDULE") _strategy[InstanceScheduler] = os.getenv("EC2_SCHEDULE") _strategy[EcsScheduler] = os.getenv("ECS_SCHEDULE") _strategy[RdsScheduler] = os.getenv("RDS_SCHEDULE") diff --git a/variables.tf b/variables.tf index 8e9fbf21..0a2a6505 100644 --- a/variables.tf +++ b/variables.tf @@ -67,13 +67,18 @@ variable "ec2_schedule" { default = false } +variable "documentdb_schedule" { + description = "Enable scheduling on documentdb resources" + type = bool + default = false +} + variable "ecs_schedule" { description = "Enable scheduling on ecs services" type = bool default = false } - variable "rds_schedule" { description = "Enable scheduling on rds resources" type = any