feat(mongodb): mongodb_mcp TencentBlueKing#16073

Author: cycker

dbm-ui/backend/dbm_aiagent/mcp_tools/constants.py

@@ -33,6 +33,12 @@ class DBMMcpTools(StrStructuredEnum):
     REDIS_QUERY_ALARM = EnumField("redis-query-alarm", "redis-query-alarm")
     REDIS_BILL = EnumField("redis-bill", "redis-bill")
     REDIS_JOB = EnumField("redis-job", "redis-job")
+    MONGODB_QUERY_META = EnumField("mongodb-query-meta", "mongodb-query-meta")
+    MONGODB_QUERY_STATUS = EnumField("mongodb-query-status", "mongodb-query-status")
+    MONGODB_QUERY_LOG = EnumField("mongodb-query-log", "mongodb-query-log")
+    MONGODB_QUERY_ALARM = EnumField("mongodb-query-alarm", "mongodb-query-alarm")
+    MONGODB_BILL = EnumField("mongodb-bill", "mongodb-bill")
+    MONGODB_JOB = EnumField("mongodb-job", "mongodb-job")
 
 
 class DBMMCPTags(StrStructuredEnum):

dbm-ui/backend/dbm_aiagent/mcp_tools/mongodb/__init__.py

@@ -0,0 +1,236 @@
+"""
+TencentBlueKing is pleased to support the open source community by making 蓝鲸智云-DB管理系统(BlueKing-BK-DBM) available.
+Copyright (C) 2017-2023 THL A29 Limited, a Tencent company. All rights reserved.
+Licensed under the MIT License (the "License"); you may not use this file except in compliance with the License.
+You may obtain a copy of the License at https://opensource.org/licenses/MIT
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+"""
+from collections import defaultdict
+from typing import Dict, List
+
+from django.db.models import F
+
+from backend.configuration.constants import DBType
+from backend.configuration.models import DBAdministrator
+from backend.db_meta.enums import ClusterType, MachineType
+from backend.db_meta.models import AppCache, Cluster, ClusterEntry, Machine, ProxyInstance, StorageInstance
+
+
+def list_my_mongodb_bizs(username: str) -> List:
+    res = []
+    for app in AppCache.objects.all():
+        bk_biz_id = app.bk_biz_id
+        if DBAdministrator.objects.filter(
+            bk_biz_id=bk_biz_id, users__0=username, db_type=DBType.MongoDB.value
+        ).exists():
+            res.append({"bk_biz_id": bk_biz_id, "app_name": app.bk_biz_name, "abbr": app.db_app_abbr})
+    return res
+
+
+def mongodb_list_clusters(bk_biz_id: int) -> List:
+    clusters = Cluster.objects.filter(
+        bk_biz_id=bk_biz_id,
+        cluster_type__in=[ClusterType.MongoReplicaSet, ClusterType.MongoShardedCluster],
+    )
+    result = []
+    for c in clusters:
+        mongos_count = c.proxyinstance_set.filter(machine_type=MachineType.MONGOS).count()
+        storage_count = c.storageinstance_set.count()
+        shard_count = 0
+        if c.cluster_type == ClusterType.MongoShardedCluster:
+            shard_count = (
+                c.nosqlstoragesetdtl_set.filter(
+                    instance__machine__machine_type=MachineType.MONGODB
+                )
+                .values("seg_range")
+                .distinct()
+                .count()
+            )
+        elif c.cluster_type == ClusterType.MongoReplicaSet:
+            shard_count = 1
+        result.append(
+            {
+                "cluster_id": c.id,
+                "bk_cloud_id": c.bk_cloud_id,
+                "cluster_type": c.cluster_type,
+                "immute_domain": c.immute_domain,
+                "alias": c.alias,
+                "region": c.region,
+                "mongos_count": mongos_count,
+                "shard_count": shard_count,
+                "storage_count": storage_count,
+                "mongodb_version": c.major_version,
+            }
+        )
+    return result
+
+
+def get_machine_stats(all_machine_ids) -> Dict:
+    machines = Machine.objects.filter(bk_host_id__in=all_machine_ids).select_related("bk_city")
+    machine_distribution = {
+        "total_count": len(all_machine_ids),
+        "by_sub_zone": defaultdict(int),
+        "by_os": defaultdict(int),
+        "by_device_class": defaultdict(int),
+        "spec_summary": defaultdict(int),
+    }
+    for machine in machines:
+        if machine.bk_sub_zone:
+            machine_distribution["by_sub_zone"][machine.bk_sub_zone] += 1
+        if machine.bk_os_name:
+            machine_distribution["by_os"][machine.bk_os_name] += 1
+        if machine.bk_svr_device_cls_name:
+            machine_distribution["by_device_class"][machine.bk_svr_device_cls_name] += 1
+        if machine.spec_id:
+            machine_distribution["spec_summary"][f"spec_{machine.spec_id}"] += 1
+    return machine_distribution
+
+
+def cluster_overview(immute_domain: str) -> Dict:
+    cluster_obj = Cluster.objects.prefetch_related("tags").get(immute_domain=immute_domain)
+    stats = {
+        "bk_cloud_id": cluster_obj.bk_cloud_id,
+        "bk_biz_id": cluster_obj.bk_biz_id,
+        "cluster_id": cluster_obj.id,
+        "immute_domain": cluster_obj.immute_domain,
+        "alias": cluster_obj.alias,
+        "cluster_type": cluster_obj.cluster_type,
+        "major_version": cluster_obj.major_version,
+        "region": cluster_obj.region,
+        "disaster_tolerance_level": cluster_obj.disaster_tolerance_level,
+        "tags": ["{}:{}".format(tag.key, tag.value) for tag in cluster_obj.tags.all()],
+        "cluster_entries": [
+            {"entry_type": ce.cluster_entry_type, "entry_addr": ce.entry}
+            for ce in ClusterEntry.objects.filter(cluster=cluster_obj)
+        ],
+    }
+    storage_instances = (
+        StorageInstance.objects.filter(cluster=cluster_obj)
+        .select_related("machine", "machine__bk_city")
+        .prefetch_related("bind_entry")
+    )
+    proxy_instances = (
+        ProxyInstance.objects.filter(cluster=cluster_obj)
+        .select_related("machine", "machine__bk_city")
+        .prefetch_related("bind_entry")
+    )
+    storage_stats = {
+        "by_role": defaultdict(int),
+        "by_status": defaultdict(int),
+        "by_machine_type": defaultdict(int),
+        "versions": set(),
+        "machines": set(),
+    }
+    for instance in storage_instances:
+        storage_stats["by_role"][instance.instance_role or instance.machine.machine_type] += 1
+        storage_stats["by_status"][instance.status] += 1
+        storage_stats["by_machine_type"][instance.machine_type] += 1
+        if instance.version:
+            storage_stats["versions"].add(instance.version)
+        storage_stats["machines"].add(instance.machine.bk_host_id)
+    storage_machines = get_machine_stats(storage_stats["machines"])
+    stats["storage_instances"] = {
+        "node_count": storage_instances.count(),
+        "by_role": dict(sorted(storage_stats["by_role"].items())),
+        "by_status": dict(sorted(storage_stats["by_status"].items())),
+        "versions": sorted(list(storage_stats["versions"])),
+        "machine_count": len(storage_stats["machines"]),
+        "by_os": dict(sorted(storage_machines["by_os"].items())),
+        "by_sub_zone": dict(sorted(storage_machines["by_sub_zone"].items())),
+        "by_device_class": dict(sorted(storage_machines["by_device_class"].items())),
+    }
+    proxy_stats = {
+        "by_status": defaultdict(int),
+        "by_machine_type": defaultdict(int),
+        "versions": set(),
+        "machines": set(),
+    }
+    for instance in proxy_instances:
+        proxy_stats["by_status"][instance.status] += 1
+        proxy_stats["by_machine_type"][instance.machine_type] += 1
+        if instance.version:
+            proxy_stats["versions"].add(instance.version)
+        proxy_stats["machines"].add(instance.machine.bk_host_id)
+    proxy_machines = get_machine_stats(proxy_stats["machines"])
+    stats["proxy_instances"] = {
+        "node_count": proxy_instances.count(),
+        "by_status": dict(sorted(proxy_stats["by_status"].items())),
+        "versions": sorted(list(proxy_stats["versions"])),
+        "machine_count": len(proxy_stats["machines"]),
+        "by_os": dict(sorted(proxy_machines["by_os"].items())),
+        "by_sub_zone": dict(sorted(proxy_machines["by_sub_zone"].items())),
+        "by_device_class": dict(sorted(proxy_machines["by_device_class"].items())),
+    }
+    return stats
+
+
+def cluster_mongos(immute_domain: str) -> List:
+    """集群 Mongos 列表"""
+    c_obj = Cluster.objects.get(immute_domain=immute_domain)
+    mongos_instances = c_obj.proxyinstance_set.filter(machine_type=MachineType.MONGOS)
+    return [
+        {
+            "address": "{}:{}".format(s.machine.ip, s.port),
+            "status": s.status,
+            "version": s.version or "",
+            "sub_zone": s.machine.bk_sub_zone or "",
+            "cls_name": s.machine.bk_svr_device_cls_name or "",
+        }
+        for s in mongos_instances
+    ]
+
+
+def cluster_shards(immute_domain: str) -> List:
+    """集群分片(Shard)节点列表，按 IP 聚合端口"""
+    c_obj = Cluster.objects.get(immute_domain=immute_domain)
+    storage_objs = c_obj.storageinstance_set.filter(machine_type=MachineType.MONGODB)
+    host_ports = defaultdict(list)
+    for ins in storage_objs:
+        host_ports[ins.machine.ip].append(ins.port)
+    result = []
+    for ip, ports in host_ports.items():
+        m_obj = Machine.objects.filter(ip=ip, bk_cloud_id=c_obj.bk_cloud_id, bk_biz_id=c_obj.bk_biz_id).first()
+        if m_obj:
+            result.append(
+                {
+                    "ip": ip,
+                    "ports": ports,
+                    "sub_zone": m_obj.bk_sub_zone or "",
+                    "cls_name": m_obj.bk_svr_device_cls_name or "",
+                }
+            )
+    return result
+
+
+def list_clusters_by_hosts(hosts: List) -> List[Dict]:
+    cluster_host = []
+    storage_data = (
+        Cluster.objects.filter(storageinstance__machine__ip__in=hosts)
+        .values(
+            "immute_domain",
+            host=F("storageinstance__machine__ip"),
+            instance_role=F("storageinstance__instance_role"),
+        )
+        .distinct()
+    )
+    cluster_host.extend(list(storage_data))
+    proxy_data = (
+        Cluster.objects.filter(proxyinstance__machine__ip__in=hosts)
+        .values(
+            "immute_domain",
+            host=F("proxyinstance__machine__ip"),
+            instance_role=F("proxyinstance__machine_type"),
+        )
+        .distinct()
+    )
+    cluster_host.extend(list(proxy_data))
+    seen = set()
+    unique_results = []
+    for item in cluster_host:
+        key = (item["immute_domain"], item["host"], str(item.get("instance_role", "")))
+        if key not in seen:
+            seen.add(key)
+            unique_results.append(item)
+    return unique_results

dbm-ui/backend/dbm_aiagent/mcp_tools/mongodb/impl/__init__.py

@@ -0,0 +1,21 @@
+# -*- coding: utf-8 -*-
+"""
+TencentBlueKing is pleased to support the open source community by making 蓝鲸智云-DB管理系统(BlueKing-BK-DBM) available.
+Copyright (C) 2017-2023 THL A29 Limited, a Tencent company. All rights reserved.
+Licensed under the MIT License (the "License"); you may not use this file except in compliance with the License.
+You may obtain a copy of the License at https://opensource.org/licenses/MIT
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+"""
+from itertools import chain
+
+from backend.db_meta.enums import ClusterType
+from backend.db_meta.models import Cluster
+from backend.dbm_aiagent.mcp_tools.redis.impl.get_source_access_impl import generate_cluster_query_report
+
+
+def generate_mongodb_cluster_query_report(job_log_resp, cluster_domain: str, cluster_all_ips: list):
+    """复用 Redis 访问来源报告生成逻辑（StorageInstance/ProxyInstance 通用）"""
+    tcp_report = generate_cluster_query_report(job_log_resp, cluster_domain, cluster_all_ips)
+    return tcp_report

dbm-ui/backend/dbm_aiagent/mcp_tools/mongodb/impl/cluster_meta.py

@@ -0,0 +1,85 @@
+"""
+TencentBlueKing is pleased to support the open source community by making 蓝鲸智云-DB管理系统(BlueKing-BK-DBM) available.
+Copyright (C) 2017-2023 THL A29 Limited, a Tencent company. All rights reserved.
+Licensed under the MIT License (the "License"); you may not use this file except in compliance with the License.
+You may obtain a copy of the License at https://opensource.org/licenses/MIT
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+"""
+import time
+from datetime import datetime
+from typing import Dict, List, Optional, Union
+
+from backend import env
+from backend.components import BKMonitorV3Api
+from backend.dbm_aiagent.mcp_tools.mongodb.tools.comm_tools import parse_time2_long
+
+
+def get_cluster_alarms(
+    immute_domain: Optional[str] = None,
+    start_time: Optional[Union[int, str, datetime]] = None,
+    end_time: Optional[Union[int, str, datetime]] = None,
+) -> Dict:
+    alarms = get_alarms_flat(immute_domain=immute_domain, start_time=start_time, end_time=end_time)
+    by_alert = alarms.get(immute_domain, {})
+    cluster_alarms = [{"alert_name": k, "alert_detail": v} for k, v in by_alert.items()]
+    total = sum(len(v) for v in by_alert.values())
+    return {"total_alarms": total, "alarm_detail": cluster_alarms}
+
+
+def get_alarms_flat(
+    appid: Optional[Union[str, int]] = None,
+    immute_domain: Optional[str] = None,
+    start_time: Optional[Union[int, str, datetime]] = None,
+    end_time: Optional[Union[int, str, datetime]] = None,
+    n_hour: Optional[int] = None,
+    limit: int = 5000,
+) -> Dict[str, List]:
+    """获取告警列表，支持按业务或集群筛选。标签使用 DBM_MONGODB。"""
+    if not appid and not immute_domain:
+        raise ValueError("必须指定 appid 或 immute_domain 中的至少一个")
+    if start_time is not None and end_time is not None:
+        start_timestamp = parse_time2_long(start_time)
+        end_timestamp = parse_time2_long(end_time)
+    else:
+        n_hour = n_hour or 24
+        end_timestamp = int(time.time())
+        start_timestamp = end_timestamp - n_hour * 60 * 60
+    query_conditions = ['labels: "DBM_MONGODB"']
+    if appid:
+        query_conditions.append(f'tags.appid : "{appid}"')
+    if immute_domain:
+        query_conditions.append(f'tags.cluster_domain : "{immute_domain}"')
+    query_string = " AND ".join(query_conditions)
+    data = BKMonitorV3Api.search_alert(
+        {
+            "bk_biz_ids": [env.DBA_APP_BK_BIZ_ID],
+            "start_time": start_timestamp,
+            "end_time": end_timestamp,
+            "offset": 0,
+            "limit": limit,
+            "query_string": query_string,
+        }
+    )
+    # 与 Redis 一致：alerts 列表，按 domain -> alert_name -> list 聚合
+    result = {}
+    for alt in data.get("alerts", []):
+        alarm = {
+            "alert_name": alt.get("alert_name", ""),
+            "description": alt.get("description", ""),
+            "begin_time": alt.get("begin_time", 0),
+            "target_key": alt.get("target_key", ""),
+        }
+        for tag in alt.get("tags", []):
+            alarm[tag["key"]] = tag["value"]
+        immute_domain = alarm.get("cluster_domain", "unknown")
+        alert_name = alarm.get("alert_name", "unknown")
+        alarm.pop("cluster_domain", None)
+        alarm.pop("appid", None)
+        if immute_domain not in result:
+            result[immute_domain] = {}
+        if alert_name not in result[immute_domain]:
+            result[immute_domain][alert_name] = []
+        result[immute_domain][alert_name].append(alarm)
+    return result