TiDB CloudのAPIを使って、Dedicatedクラスタを時間指定で起動・停止してみた
こんにちは、ゲームソリューション部のsoraです。
今回は、TiDB CloudのAPIを使って、Dedicatedクラスタを時間指定で起動・停止してみたことについて書いていきます。
はじめに
今回構築する構成は以下です。
EventBridge Schedulerで毎日決まった時間にLambdaを実行して、指定したTiDB CloudのDedicatedクラスタを起動・停止します。
開発環境で平日の日中帯のみ起動したいといった場合に役に立つと思います。
TiDB CloudのAPIについては、以下の公式ドキュメントを参考に進めていきます。
環境構築
AWS上のリソースをTerraformで構築します。
TiDB Cloudのクラスタに対して、日本時間の8時に起動、18時に停止する設定にしています。
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
}
}
Lambdaのコードは以下です。
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))
TiDB CloudのAPIキーは事前に作成して、Parameter StoreにてSecureStringとして管理しています。
TiDB CloudのDedicatedクラスタの指定については、Lambdaの環境変数としてTiDB CloudのプロジェクトIDとDedicatedクラスタ名を設定しています。
動作確認
EventBridge Schedulerは起動用と停止用の2つを作成しており、それぞれ指定の時間になると動作するようになっています。
TiDB Cloud側を確認すると、指定時間にPausingやResumingとなっており、停止・起動が実行されていることを確認できました。
最後に
今回は、TiDB CloudのAPIを使って、Dedicatedクラスタを時間指定で起動・停止してみたことを記事にしました。
どなたかの参考になれば幸いです。