SQLServer采集data和log的体积推送到prometheus

# -*-coding:utf-8 -*-

import datetime
import json
import os
import sys

import requests

import pymssql
from prometheus_client import CollectorRegistry, Gauge, push_to_gateway


def get_db_list(host, port, db, user, passwd):
    conn = pymssql.connect(
        host=host, port=port, user=user, password=passwd, database=db, charset="utf8"
    )
    cursor = conn.cursor()
    sql = """select database_id,name,state from sys.databases where [state] not in(1,2,3,4,6) and name not in ('master','model','tempdb','msdb') ORDER BY name asc"""

    cursor.execute(sql)
    res = cursor.fetchall()

    db_list = []

    for i in res:
        db_list.append(i[1])

    return db_list


def push_metrics(host, port, db, user, passwd, db_list):
    for db in db_list:
        # print(db)
        conn = pymssql.connect(
            host=host,
            port=port,
            user=user,
            password=passwd,
            database='master',
            charset="utf8",
        )
        cursor = conn.cursor()

        detail_sql = "USE [" + str(db) + "];" +  """
            SELECT
                DB_NAME() AS [Database Name]
                ,SUM(size * 8.0 / 1024) AS [Total_Size]
                ,SUM(CASE WHEN type_desc = 'ROWS' THEN size * 8.0 / 1024 ELSE 0 END) AS [Data Size]
                ,SUM(CASE WHEN type_desc = 'LOG' THEN size * 8.0 / 1024 ELSE 0 END) AS [Log Size]
                ,@@SERVERNAME as instance
            FROM sys.master_files 
            WHERE type in (0,1) AND state = 0 AND database_id = DB_ID('""" + str(db) + "');"

        try:
            cursor.execute(detail_sql)
        except:
            continue

        res = cursor.fetchall()

        for row in res:            
            Database_Name = str(row[0])
            Total_Size = str(row[1])
            Data_Size = str(row[2])
            Log_Size = str(row[3])
            instance = str(row[4])

            print(f"节点 {instance}, 库名 {Database_Name}, Total_Size: {Total_Size}, Data_Size: {Data_Size}, Log_Size: {Log_Size} ")

            # 把指标格式化成prometheus识别的格式
            labels = ["Metrics"]

            registry = CollectorRegistry()

            db_metrics = Gauge("Data_Size", "MSSQL基础指标", labels, registry=registry)
            db_metrics.labels("Data_Size").set(Data_Size)

            db_metrics = Gauge("Log_Size", "MSSQL基础指标", labels, registry=registry)
            db_metrics.labels("Log_Size").set(Log_Size)


            push_to_gateway(
                gateway='http://192.168.31.181:9091',
                job="MSSQLBasicMetrics",
                # 注意这里的grouping_key是关键，写的有问题会导致不同instance的指标覆盖
                grouping_key={
                    "instance": instance,
                    "Database_Name":Database_Name
                },
                registry=registry,
                timeout=5,
            )
            

if __name__ == "__main__":
    port = "1433"
    db = "master"
    user = "sa"
    passwd = "Abcd1234"

    host_list = ["192.168.31.181","192.168.31.17"]

    for host in host_list:
        print("当前解析的是 ---- ",host)
        db_list = get_db_list(host, port, db, user, passwd)

        push_metrics(host, port, db, user, passwd, db_list)