perf: k8s tree api (#9169)

Co-authored-by: feng <1304903146@qq.com>

apps/assets/api/mixin.py

@@ -61,6 +61,7 @@ class SerializeToTreeNodeMixin:
                 'meta': {
                     'type': 'asset',
                     'data': {
+                        'platform_type': asset.platform.type,
                         'org_name': asset.org_name,
                         'sftp': asset.platform_id in sftp_enabled_platform,
                     },

apps/assets/utils/__init__.py

@@ -0,0 +1,2 @@
+from .k8s import *
+from .node import *

apps/assets/utils/k8s.py

@@ -0,0 +1,192 @@
+# -*- coding: utf-8 -*-
+from urllib3.exceptions import MaxRetryError
+from urllib.parse import urlencode, parse_qsl
+
+from kubernetes import client
+from kubernetes.client import api_client
+from kubernetes.client.api import core_v1_api
+from kubernetes.client.exceptions import ApiException
+
+from rest_framework.generics import get_object_or_404
+
+from common.utils import get_logger
+from common.tree import TreeNode
+from assets.models import Account, Asset
+
+from ..const import CloudTypes, Category
+
+logger = get_logger(__file__)
+
+
+class KubernetesClient:
+    def __init__(self, url, token):
+        self.url = url
+        self.token = token
+
+    def get_api(self):
+        configuration = client.Configuration()
+        configuration.host = self.url
+        configuration.verify_ssl = False
+        configuration.api_key = {"authorization": "Bearer " + self.token}
+        c = api_client.ApiClient(configuration=configuration)
+        api = core_v1_api.CoreV1Api(c)
+        return api
+
+    def get_namespace_list(self):
+        api = self.get_api()
+        namespace_list = []
+        for ns in api.list_namespace().items:
+            namespace_list.append(ns.metadata.name)
+        return namespace_list
+
+    def get_services(self):
+        api = self.get_api()
+        ret = api.list_service_for_all_namespaces(watch=False)
+        for i in ret.items:
+            print("%s \t%s \t%s \t%s \t%s \n" % (
+                i.kind, i.metadata.namespace, i.metadata.name, i.spec.cluster_ip, i.spec.ports))
+
+    def get_pod_info(self, namespace, pod):
+        api = self.get_api()
+        resp = api.read_namespaced_pod(namespace=namespace, name=pod)
+        return resp
+
+    def get_pod_logs(self, namespace, pod):
+        api = self.get_api()
+        log_content = api.read_namespaced_pod_log(pod, namespace, pretty=True, tail_lines=200)
+        return log_content
+
+    def get_pods(self):
+        api = self.get_api()
+        try:
+            ret = api.list_pod_for_all_namespaces(watch=False, _request_timeout=(3, 3))
+        except MaxRetryError:
+            logger.warning('Kubernetes connection timed out')
+            return
+        except ApiException as e:
+            if e.status == 401:
+                logger.warning('Kubernetes User not authenticated')
+            else:
+                logger.warning(e)
+            return
+        data = {}
+        for i in ret.items:
+            namespace = i.metadata.namespace
+            pod_info = {
+                'pod_name': i.metadata.name,
+                'containers': [j.name for j in i.spec.containers]
+            }
+            if namespace in data:
+                data[namespace].append(pod_info)
+            else:
+                data[namespace] = [pod_info, ]
+        return data
+
+    @staticmethod
+    def get_kubernetes_data(app_id, username):
+        asset = get_object_or_404(Asset, id=app_id)
+        account = get_object_or_404(Account, asset=asset, username=username)
+        k8s = KubernetesClient(asset.address, account.secret)
+        return k8s.get_pods()
+
+
+class KubernetesTree:
+    def __init__(self, tree_id):
+        self.tree_id = str(tree_id)
+
+    @staticmethod
+    def create_tree_id(pid, tp, v):
+        i = dict(parse_qsl(pid))
+        i[tp] = v
+        tree_id = urlencode(i)
+        return tree_id
+
+    def as_tree_node(self, app):
+        pid = app.create_app_tree_pid(self.tree_id)
+        app_id = str(app.id)
+        node = self.create_tree_node(
+            app_id, pid, app.name, 'k8s'
+        )
+        return node
+
+    def as_asset_tree_node(self, asset):
+        i = urlencode({'asset_id': self.tree_id})
+        node = self.create_tree_node(
+            i, str(asset.id), str(asset), 'asset',
+        )
+        return node
+
+    def as_account_tree_node(self, account, parent_info):
+        username = account.username
+        name = f'{account.name}({account.username})'
+        pid = urlencode({'asset_id': self.tree_id})
+        i = self.create_tree_id(pid, 'account', username)
+        parent_info.update({'account': username})
+        node = self.create_tree_node(
+            i, pid, name, 'account', icon='user-tie'
+        )
+        return node
+
+    def as_namespace_pod_tree_node(self, name, tp, counts=0, is_container=False):
+        i = self.create_tree_id(self.tree_id, tp, name)
+        name = name if is_container else f'{name}({counts})'
+        node = self.create_tree_node(
+            i, self.tree_id, name, tp, icon='cloud', is_container=is_container
+        )
+        return node
+
+    @staticmethod
+    def create_tree_node(id_, pid, name, identity, icon='', is_container=False):
+        node = {
+            'id': id_,
+            'name': name,
+            'title': name,
+            'pId': pid,
+            'isParent': not is_container,
+            'open': False,
+            'iconSkin': icon,
+            'meta': {
+                'type': 'k8s',
+                'data': {
+                    'category': Category.CLOUD,
+                    'type': CloudTypes.K8S,
+                    'identity': identity
+                }
+            }
+        }
+        return node
+
+    def async_tree_node(self, parent_info):
+        pod_name = parent_info.get('pod')
+        asset_id = parent_info.get('asset_id')
+        namespace = parent_info.get('namespace')
+        account_username = parent_info.get('account')
+
+        tree = []
+        data = KubernetesClient.get_kubernetes_data(asset_id, account_username)
+        if not data:
+            return tree
+
+        if pod_name:
+            for container in next(
+                    filter(
+                        lambda x: x['pod_name'] == pod_name, data[namespace]
+                    )
+            )['containers']:
+                container_node = self.as_namespace_pod_tree_node(
+                    container, 'container', is_container=True
+                )
+                tree.append(container_node)
+        elif namespace:
+            for pod in data[namespace]:
+                pod_nodes = self.as_namespace_pod_tree_node(
+                    pod['pod_name'], 'pod', len(pod['containers'])
+                )
+                tree.append(pod_nodes)
+        elif account_username:
+            for namespace, pods in data.items():
+                namespace_node = self.as_namespace_pod_tree_node(
+                    namespace, 'namespace', len(pods)
+                )
+                tree.append(namespace_node)
+        return tree

apps/assets/utils/node.py

@@ -7,8 +7,8 @@ from common.struct import Stack
 from common.db.models import output_as_string
 from orgs.utils import ensure_in_real_or_default_org, current_org
 
-from .locks import NodeTreeUpdateLock
-from .models import Node, Asset
+from ..locks import NodeTreeUpdateLock
+from ..models import Node, Asset
 
 logger = get_logger(__file__)
 

apps/perms/api/user_permission/tree/node_with_asset.py

@@ -1,26 +1,31 @@
 import abc
+from urllib.parse import parse_qsl
 
 from django.conf import settings
 from django.db.models import F, Value, CharField
 from rest_framework.generics import ListAPIView
+from rest_framework.request import Request
 from rest_framework.response import Response
+from rest_framework.generics import get_object_or_404
 
 from assets.models import Asset
+from assets.utils import KubernetesTree
 from assets.api import SerializeToTreeNodeMixin
 from perms.hands import Node
 from perms.models import PermNode
 from perms.utils.user_permission import (
     UserGrantedNodesQueryUtils, UserGrantedAssetsQueryUtils,
 )
+from perms.utils import PermAccountUtil
 from perms.utils.permission import AssetPermissionUtil
 from common.utils import get_object_or_none, lazyproperty
 from common.utils.common import timeit
 
-
 from ..mixin import SelfOrPKUserMixin
 from .mixin import RebuildTreeMixin
 
 __all__ = [
+    'UserGrantedK8sAsTreeApi',
     'UserPermedNodesWithAssetsAsTreeApi',
     'UserPermedNodeChildrenWithAssetsAsTreeApi'
 ]
@@ -123,3 +128,41 @@ class UserPermedNodeChildrenWithAssetsAsTreeApi(BaseUserNodeWithAssetAsTreeApi):
     def node_key_for_serializer_assets(self):
         return self.query_node_key
 
+
+class UserGrantedK8sAsTreeApi(
+    SelfOrPKUserMixin,
+    ListAPIView
+):
+    """ 用户授权的K8s树 """
+
+    @staticmethod
+    def asset(asset_id):
+        kwargs = {'id': asset_id, 'is_active': True}
+        asset = get_object_or_404(Asset, **kwargs)
+        return asset
+
+    def list(self, request: Request, *args, **kwargs):
+        tree_id = request.query_params.get('tree_id')
+        key = request.query_params.get('key', {})
+
+        tree = []
+        util = PermAccountUtil()
+        parent_info = dict(parse_qsl(key))
+        account_username = parent_info.get('account')
+
+        asset_id = parent_info.get('asset_id')
+        asset_id = tree_id if not asset_id else asset_id
+
+        if tree_id and not account_username:
+            asset = self.asset(asset_id)
+            accounts = util.get_permed_accounts_for_user(self.user, asset)
+            asset_node = KubernetesTree(tree_id).as_asset_tree_node(asset)
+            tree.append(asset_node)
+            for account in accounts:
+                account_node = KubernetesTree(tree_id).as_account_tree_node(
+                    account, parent_info,
+                )
+                tree.append(account_node)
+        else:
+            tree = KubernetesTree(key).async_tree_node(parent_info)
+        return Response(data=tree)

apps/perms/urls/user_permission.py

@@ -34,9 +34,13 @@ user_permission_urlpatterns = [
     path('<str:user>/nodes/children-with-assets/tree/',
          api.UserPermedNodeChildrenWithAssetsAsTreeApi.as_view(),
          name='user-node-children-with-assets-as-tree'),
+
+    path('<str:user>/nodes/children-with-k8s/tree/',
+         api.UserGrantedK8sAsTreeApi.as_view(),
+         name='user-nodes-children-with-k8s-as-tree'),
+
     path('<str:user>/nodes-with-assets/tree/', api.UserPermedNodesWithAssetsAsTreeApi.as_view(),
          name='user-nodes-with-assets-as-tree'),
-
     # accounts
     path('<str:user>/assets/<uuid:asset_id>/accounts/', api.UserPermedAssetAccountsApi.as_view(),
          name='user-permed-asset-accounts'),