TiDB CloudのAPIを使って、Dedicatedクラスタを時間指定で起動・停止してみた

terraform {
  # AWSプロバイダーのバージョン指定
  required_providers {
    aws = {
      source  = "hashicorp/aws"
      version = "~> 6.14.0"
    }
  }
  backend "s3" {
    bucket = "xxxxxxxxxxxx"
    region = "ap-northeast-1"
    key    = "tidb-auto-startstop.tfstate"
  }
}
provider "aws" {
  region = "ap-northeast-1"
}

##EventBridge Sheduler
data "aws_iam_policy_document" "events_assume_role" {
  statement {
    effect = "Allow"
    principals {
      type        = "Service"
      identifiers = ["scheduler.amazonaws.com"]
    }
    actions = ["sts:AssumeRole"]
  }
}
resource "aws_iam_role" "iam_for_events" {
  name               = "eventbridge-tidb-auto-startstop-role"
  assume_role_policy = data.aws_iam_policy_document.events_assume_role.json
}

resource "aws_iam_role_policy_attachment" "iam_for_events_AWSLambdaRole" {
  role       = aws_iam_role.iam_for_events.name
  policy_arn = "arn:aws:iam::aws:policy/service-role/AWSLambdaRole"
}

##Lambda
data "aws_iam_policy_document" "lambda_assume_role" {
  statement {
    effect = "Allow"
    principals {
      type        = "Service"
      identifiers = ["lambda.amazonaws.com"]
    }
    actions = ["sts:AssumeRole"]
  }
}
resource "aws_iam_policy" "lambda" {
  name        = "lambda-tidb-auto-startstop-policy"
  description = "lambda-tidb-auto-startstop-policy"
  policy      = file("./policy/lambda_tidb_api_iam.json")
}
resource "aws_iam_role" "iam_for_lambda" {
  name               = "lambda-tidb-auto-startstop-role"
  assume_role_policy = data.aws_iam_policy_document.lambda_assume_role.json
}

resource "aws_iam_role_policy_attachment" "iam_for_lambda_basic" {
  role       = aws_iam_role.iam_for_lambda.name
  policy_arn = "arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"
}

resource "aws_iam_role_policy_attachment" "iam_for_lambda_custom" {
  role       = aws_iam_role.iam_for_lambda.name
  policy_arn = aws_iam_policy.lambda.arn
}

## Lambda(自動起動停止共通)
data "archive_file" "lambda_tidb" {
  type        = "zip"
  source_dir  = "./lambda"
  output_path = "tidb_api.zip"
}
resource "aws_lambda_function" "lambda_tidb" {
  filename         = "tidb_api.zip"
  function_name    = "tidb-auto-startstop"
  role             = aws_iam_role.iam_for_lambda.arn
  handler          = "tidb_api.lambda_handler"
  source_code_hash = data.archive_file.lambda_tidb.output_base64sha256
  runtime          = "python3.12"
  timeout          = 60
  layers = [
    "arn:aws:lambda:ap-northeast-1:133490724326:layer:AWS-Parameters-and-Secrets-Lambda-Extension:11"
  ]
  environment {
    variables = {
      ProjectID   = var.ProjectID
      ClusterName = var.ClusterName
    }
  }
  lifecycle {
    ignore_changes = [environment]
  }
}

##EventBridge Sheduler
resource "aws_scheduler_schedule" "tidb_start" {
  name        = "tidb-auto-start-schedule"
  description = "TiDB Cloud Auto Start"
  flexible_time_window {
    mode = "OFF"
  }
  schedule_expression          = "cron(0 8 * * ? *)"
  schedule_expression_timezone = "Asia/Tokyo"
  target {
    arn      = aws_lambda_function.lambda_tidb.arn
    role_arn = aws_iam_role.iam_for_events.arn
    input    = <<EOF
        {"pause": "start"}
        EOF
  }
}
resource "aws_scheduler_schedule" "tidb_stop" {
  name        = "tidb-auto-stop-schedule"
  description = "TiDB Cloud Auto Stop"
  flexible_time_window {
    mode = "OFF"
  }
  schedule_expression          = "cron(0 18 * * ? *)"
  schedule_expression_timezone = "Asia/Tokyo"
  target {
    arn      = aws_lambda_function.lambda_tidb.arn
    role_arn = aws_iam_role.iam_for_events.arn
    input    = <<EOF
        {"pause": "stop"}
        EOF
  }
}

import json
import os
import requests
from requests.auth import HTTPDigestAuth

PROJECT_ID = os.environ['ProjectID']
CLUSTER_NAME = os.environ['ClusterName']

def lambda_handler(event, context):
    # SecureStringの取得
    end_point = 'http://localhost:2773'
    path_public_key = "/systemsmanager/parameters/get/?name=/tidb-tools/TiDBAPIPublicKey&withDecryption=true"
    path_private_key = "/systemsmanager/parameters/get/?name=/tidb-tools/TiDBAPIPrivateKey&withDecryption=true"
    url_public_key = end_point + path_public_key
    url_private_key = end_point + path_private_key
    headers = {
        'X-Aws-Parameters-Secrets-Token': os.environ['AWS_SESSION_TOKEN']
    }

    try:
        tidb_public_key = requests.get(url_public_key, headers=headers)
        tidb_public_key.raise_for_status()
        tidb_private_key = requests.get(url_private_key, headers=headers)
        tidb_private_key.raise_for_status()
    except requests.exceptions.RequestException as e:
        raise RuntimeError(f"TiDB CloudのAPIキーの取得に失敗しました: {e}")

    tidb_public_key_value = tidb_public_key.json()['Parameter']['Value']
    tidb_private_key_value = tidb_private_key.json()['Parameter']['Value']

    # 起動停止の分岐
    pause_check = event['pause']
    # pausedかどうかなため、停止がTrue、起動がFalse
    if pause_check == "start":
        is_pause = False
    elif pause_check == "stop":
        is_pause = True
    else:
        raise ValueError("EventBridge Schedulerのpauseフィールドが不正です。startまたはstopを指定してください。")

    # ClusterNameからClusterIDへの変換
    # プロジェクト内のCluster情報の取得
    cluster_info_url = f'https://api.tidbcloud.com/api/v1beta/projects/{PROJECT_ID}/clusters'
    params = {'page': 1, 'page_size': 100}
    try:
        cluster_response = requests.get(cluster_info_url, params=params, auth=HTTPDigestAuth(tidb_public_key_value, tidb_private_key_value))
        cluster_response.raise_for_status()
    except requests.exceptions.RequestException as e:
        raise RuntimeError(f"クラスタ情報の取得に失敗しました: {e}")

    # jsonに変換
    cluster_response_json = cluster_response.json()
    cluster_id = None

    # 指定したClusterNameに一致するIDの取得
    for item in cluster_response_json['items']:
        if item['name'] == CLUSTER_NAME.strip():
            cluster_id = item['id']
            break
    if not cluster_id:
        raise ValueError(f"指定されたクラスタ名: {CLUSTER_NAME} に一致するクラスタが見つかりません。")

    if is_pause:
        status_change_url = f'https://dedicated.tidbapi.com/v1beta1/clusters/{cluster_id}:pauseCluster'
    else:
        status_change_url = f'https://dedicated.tidbapi.com/v1beta1/clusters/{cluster_id}:resumeCluster'

    requests.post(status_change_url, auth=HTTPDigestAuth(tidb_public_key_value, tidb_private_key_value))