#!/usr/bin/env python3
# Alertmanager告警入库 兼容Mariadb + Python3.6
from flask import Flask, request, jsonify
import pymysql
from datetime import datetime
import logging

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
app = Flask(__name__)

# 数据库配置，字符集改为Mariadb支持的 utf8mb4_unicode_ci
DB_CONF = {
    "host": "127.0.0.1",
    "port": 3306,
    "user": "root",
    "password": "hp93000",
    "database": "alert_db",
    "charset": "utf8mb4",
    "cursorclass": pymysql.cursors.DictCursor
}

def get_db_connection():
    try:
        conn = pymysql.connect(**DB_CONF)
        # 手动指定兼容Mariadb的排序规则
        with conn.cursor() as cur:
            cur.execute("SET NAMES utf8mb4 COLLATE utf8mb4_unicode_ci;")
        return conn
    except Exception as e:
        logging.error(f"数据库连接失败：{e}")
        return None

@app.route("/webhook/alert", methods=["POST"])
def alert_handle():
    conn = get_db_connection()
    if not conn:
        return jsonify({"code": 500, "msg": "数据库连接失败"}), 500
    cursor = conn.cursor()

    try:
        data = request.get_json()
        alert_arr = data.get("alerts", [])
        insert_sql = """
        INSERT INTO alert_log (
            alertname, job, instance, port, severity, summary, description,
            starts_at, ends_at, status
        ) VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
        """
        batch = []
        for alert in alert_arr:
            lab = alert.get("labels", {})
            ann = alert.get("annotations", {})
            stat = alert.get("status")
            start_time = datetime.fromisoformat(alert["startsAt"].replace("Z", ""))
            end_time = None
            if alert.get("endsAt"):
                end_time = datetime.fromisoformat(alert["endsAt"].replace("Z", ""))

            row_data = (
                lab.get("alertname", ""),
                lab.get("job", ""),
                lab.get("instance", ""),
                lab.get("port", ""),
                lab.get("severity", ""),
                ann.get("summary", ""),
                ann.get("description", ""),
                start_time,
                end_time,
                stat
            )
            batch.append(row_data)

        if batch:
            cursor.executemany(insert_sql, batch)
            conn.commit()
            logging.info(f"成功入库 {len(batch)} 条告警")
        return jsonify({"code": 0, "msg": "success"}), 200

    except Exception as err:
        logging.error(f"入库异常：{str(err)}")
        conn.rollback()
        return jsonify({"code": 500, "msg": str(err)}), 500
    finally:
        cursor.close()
        conn.close()

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8090, debug=False)
