概要

TiDB Cloudの請求情報を確認するにはOwnerかBilling Adminの権限が必要ですが、支払いに使用するクレジットカードを変更できるなど開発チームに気軽に確認してもらうには強すぎる権限になります。

そこでNew RelicのMetric APIを使ってTiDB Cloudの請求情報を取り込み可視化してみました。

構成図

EventBridge Schedulerで定期的にLambdaを実行します。

Lambdaの実装

Pythonで実装します。TiDB Cloud APIから取得したデータをNew Relic Metric APIにPOSTできるよう整形します。

def build_metrics(billing_data, billed_month, unix_timestamp):
    metrics = []
    for project in billing_data["summaryByProject"]["projects"]:
        metrics.append({
            "name": "tidb_cloud.runningTotal",
            "type": "gauge",
            "value": float(project["runningTotal"]),
            "timestamp": unix_timestamp,
            "attributes": {
                "projectName": project["projectName"],
                "billedMonth": billed_month
            }
        })
        metrics.append({
            "name": "tidb_cloud.totalCost",
            "type": "gauge",
            "value": float(project["totalCost"]),
            "timestamp": unix_timestamp,
            "attributes": {
                "projectName": project["projectName"],
                "billedMonth": billed_month
            }
        })
    for charge in billing_data["summaryByProject"]["otherCharges"]:
        metrics.append({
            "name": "tidb_cloud.runningTotal",
            "type": "gauge",
            "value": float(charge["runningTotal"]),
            "timestamp": unix_timestamp,
            "attributes": {
                "projectName": charge["chargeName"],
                "billedMonth": billed_month
            }
        })
        metrics.append({
            "name": "tidb_cloud.totalCost",
            "type": "gauge",
            "value": float(charge["totalCost"]),
            "timestamp": unix_timestamp,
            "attributes": {
                "projectName": charge["chargeName"],
                "billedMonth": billed_month
            }
        })
    return metrics

全体のコード

APIアクセスに必要なキーはParameter StoreやSecrets Managerで管理しているので、保存先だけ環境変数から取り出しています。

lambda_function.py

import os
import json
import requests
import boto3
from requests.auth import HTTPDigestAuth
from datetime import datetime
from zoneinfo import ZoneInfo

# 環境変数から設定情報を取得
tidb_public_key_name = os.environ['TIDB_CLOUD_PUBLIC_KEY_NAME']
tidb_private_key_name = os.environ['TIDB_CLOUD_PRIVATE_KEY_NAME']
nr_secret_name = os.environ['NR_SECRET_NAME']

ssm = boto3.client('ssm')
secretsmanager = boto3.client('secretsmanager')


def lambda_handler(event, context):
    try:
        target_date = datetime.now(ZoneInfo('Asia/Tokyo'))
        unix_timestamp = int(target_date.timestamp())

        # 現在の月
        date_str = str(target_date.date())
        year_month = date_str[0:7]

        # APIエンドポイントのURLを構築
        tidb_api = f"https://billing.tidbapi.com/v1beta1/bills/{year_month}"

        # データ取得
        tidb_public_key = get_tidb_public_key()
        tidb_private_key = get_tidb_private_key()
        tidb_billing = requests.get(tidb_api, auth=HTTPDigestAuth(tidb_public_key, tidb_private_key))

        # レスポンスを処理
        if tidb_billing.status_code == 200:
            print(f"TiDB API call succeeded. Response: {tidb_billing.text}")
            # New Relicにデータを送信
            billing_data = tidb_billing.json()
            billed_month = billing_data["overview"]["billedMonth"]
            metrics = build_metrics(billing_data, billed_month, unix_timestamp)
            metrics_data = json.dumps([{"metrics": metrics}])
            print(metrics_data)

            nr_license_key = get_nr_license_key()
            response = post_to_newrelic(metrics_data, nr_license_key)
            if response.status_code == 202:
                print(f"New Relic API call succeeded. Response: {response.text}")
            else:
                print(f"New Relic API call failed. Status code: {response.status_code}, Error: {response.text}")

        else:
            print(f"TiDB API call failed. Status code: {tidb_billing.status_code}, Error: {tidb_billing.text}")

    except Exception as e:
        print(f"An error occurred: {e}")
        return {
            "statusCode": 500,
            "body": f"An error occurred: {e}"
        }


def build_metrics(billing_data, billed_month, unix_timestamp):
    metrics = []
    for project in billing_data["summaryByProject"]["projects"]:
        metrics.append({
            "name": "tidb_cloud.runningTotal",
            "type": "gauge",
            "value": float(project["runningTotal"]),
            "timestamp": unix_timestamp,
            "attributes": {
                "projectName": project["projectName"],
                "billedMonth": billed_month
            }
        })
        metrics.append({
            "name": "tidb_cloud.totalCost",
            "type": "gauge",
            "value": float(project["totalCost"]),
            "timestamp": unix_timestamp,
            "attributes": {
                "projectName": project["projectName"],
                "billedMonth": billed_month
            }
        })
    for charge in billing_data["summaryByProject"]["otherCharges"]:
        metrics.append({
            "name": "tidb_cloud.runningTotal",
            "type": "gauge",
            "value": float(charge["runningTotal"]),
            "timestamp": unix_timestamp,
            "attributes": {
                "projectName": charge["chargeName"],
                "billedMonth": billed_month
            }
        })
        metrics.append({
            "name": "tidb_cloud.totalCost",
            "type": "gauge",
            "value": float(charge["totalCost"]),
            "timestamp": unix_timestamp,
            "attributes": {
                "projectName": charge["chargeName"],
                "billedMonth": billed_month
            }
        })
    return metrics


def post_to_newrelic(metrics_data, nr_license_key):
    nr_metric_api = "https://metric-api.newrelic.com/metric/v1"
    headers = {
        "Content-Type": "application/json",
        "Api-Key": nr_license_key
    }
    response = requests.post(nr_metric_api, headers=headers, data=metrics_data)
    return response


def get_nr_license_key():
    nr_secret_param = secretsmanager.get_secret_value(
        SecretId=nr_secret_name
    )
    nr_license_key = nr_secret_param["SecretString"]
    return nr_license_key


def get_tidb_public_key():
    tidb_public_key_param = ssm.get_parameter(
        Name=tidb_public_key_name,
        WithDecryption=False
    )
    tidb_public_key = tidb_public_key_param['Parameter']['Value']
    return tidb_public_key


def get_tidb_private_key():
    tidb_private_key_param = ssm.get_parameter(
        Name=tidb_private_key_name,
        WithDecryption=True
    )
    tidb_private_key = tidb_private_key_param['Parameter']['Value']
    return tidb_private_key

Serverlessでデプロイします。requestsを使うので requirements.txt を配置します。

.
├── lambda_function.py
├── requirements.txt
└── serverless.yml

serverless.yml

service: TiDB-Cloud-Billing
frameworkVersion: '3'

custom:
  NR_SECRET_NAME:
    prod: NewRelicLicenseKeySecret-123456ABCDEF

provider:
  name: aws
  stage: ${opt:stage, 'prod'}
  runtime: python3.11
  region: ap-northeast-1
  iamRoleStatements:
    - Effect: Allow
      Action:
        - secretsmanager:GetSecretValue
        - ssm:GetParameter
        - kms:Decrypt
      Resource:
        - "*"

functions:
  TiDB-Cloud-Billing:
    handler: lambda_function.lambda_handler
    name: ${self:provider.stage}-tidb-cloud-billing
    memorySize: 128
    timeout: 10
    events:
      - schedule:
          method: scheduler
          rate:
            - cron(0 * * * ? *)
          timezone: Asia/Tokyo
    environment:
      TIDB_CLOUD_PUBLIC_KEY_NAME: /${self:provider.stage}/tidb/api/public-key
      TIDB_CLOUD_PRIVATE_KEY_NAME: /${self:provider.stage}/tidb/api/private-key
      NR_SECRET_NAME: ${self:custom.NR_SECRET_NAME.${self:provider.stage}}

plugins:
  - serverless-python-requirements

sls plugin install -n serverless-python-requirements
sls deploy

New Relicで可視化する

しばらく待って、取り込んだデータをNew Relicでクエリしてみます。

SELECT
    latest(`tidb_cloud.totalCost`)
FROM
    Metric FACET projectName SINCE 1 WEEK AGO TIMESERIES

いい感じに可視化できたので、いつでも振り返れるようにダッシュボード化しておきます。

まとめ

初めてNew Relic Metric APIを使ってみましたが、簡単にカスタムメトリクスを作成することができました。New Relicの活用の幅が広がりそうです。

Discussion