Skip to content

Commit

Permalink
feat(terraform): add documentdb scheduler
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
diodonfrost committed Jun 8, 2023
1 parent 7f8a896 commit 942f713
Show file tree
Hide file tree
Showing 5 changed files with 126 additions and 1 deletion.
1 change: 1 addition & 0 deletions main.tf
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
74 changes: 74 additions & 0 deletions package/scheduler/documentdb_handler.py
Original file line number Diff line number Diff line change
@@ -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)
43 changes: 43 additions & 0 deletions package/scheduler/exceptions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
2 changes: 2 additions & 0 deletions package/scheduler/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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")
Expand Down
7 changes: 6 additions & 1 deletion variables.tf
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 942f713

Please sign in to comment.