From 1e4e75ef6e63ea442189d8853606fc835fe32234 Mon Sep 17 00:00:00 2001
From: lixintao <lixintao@uniontech.com>
Date: Fri, 30 Aug 2024 10:09:28 +0800
Subject: [PATCH] yq add custom api for json backfill

Signed-off-by: lixintao <lixintao@uniontech.com>
---
 src/rtk/_base.py         | 1 +
 src/rtk/_cargo.py        | 4 ++++
 src/rtk/json_backfill.py | 8 +++++---
 src/rtk/remote_runner.py | 5 ++++-
 4 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/src/rtk/_base.py b/src/rtk/_base.py
index dd1b2ff..748fc56 100644
--- a/src/rtk/_base.py
+++ b/src/rtk/_base.py
@@ -72,6 +72,7 @@ class Args(Enum):
     json_backfill_task_id = "json_backfill_task_id"
     json_backfill_user = "json_backfill_user"
     json_backfill_password = "json_backfill_password"
+    json_backfill_custom_api = "json_backfill_custom_api"
 
 
 def transform_app_name(app_name):
diff --git a/src/rtk/_cargo.py b/src/rtk/_cargo.py
index b2fdf35..5472204 100644
--- a/src/rtk/_cargo.py
+++ b/src/rtk/_cargo.py
@@ -189,6 +189,9 @@ def remote_runner(parser, sub_parser_remote):
     sub_parser_remote.add_argument(
         "--json_backfill_password", default="", help="json报告回填的密码"
     )
+    sub_parser_remote.add_argument(
+        "--json_backfill_custom_api", default="", help="json报告回填的自定义api，默认是/api/youqu/yqresult/"
+    )
 
     local_kwargs, args = local_runner(parser, sub_parser_remote)
     from src.rtk._base import Args
@@ -208,6 +211,7 @@ def remote_runner(parser, sub_parser_remote):
         Args.json_backfill_task_id.value: args.json_backfill_task_id,
         Args.json_backfill_user.value: args.json_backfill_user,
         Args.json_backfill_password.value: args.json_backfill_password,
+        Args.json_backfill_custom_api.value: args.json_backfill_custom_api,
     }
     _remote_kwargs = {
         "remote_kwargs": remote_kwargs,
diff --git a/src/rtk/json_backfill.py b/src/rtk/json_backfill.py
index 6fb496c..554facb 100644
--- a/src/rtk/json_backfill.py
+++ b/src/rtk/json_backfill.py
@@ -9,6 +9,7 @@
 import json
 import os
 import re
+from urllib.parse import urljoin
 
 from setting import conf
 from src.rtk.api_client import ApiClient
@@ -17,10 +18,11 @@ from src.rtk.api_client import ApiClient
 class JsonBackfill:
     __author__ = "huangmingqiang@uniontech.com"
 
-    def __init__(self, base_url, username, password):
+    def __init__(self, base_url, username, password, custom_api="api/youqu/yqresult/"):
         self.base_url = base_url
         self.username = username
         self.password = password
+        self.custom_api = custom_api
         self.api = ApiClient(
             base_url=self.base_url,
             username=self.username,
@@ -65,10 +67,10 @@ class JsonBackfill:
                 tpl["longrepr"] = value.get("longrepr")
                 tpl["pm_ip"] = _ip
                 res = self.api.post(
-                    url=f"{self.base_url}/api/youqu/yqresult/",
+                    url=urljoin(self.base_url, self.custom_api),
                     json=tpl
                 )
-                print(res)
+                print(self.custom_api, res)
 
 
 if __name__ == '__main__':
diff --git a/src/rtk/remote_runner.py b/src/rtk/remote_runner.py
index b406c34..3d24e72 100644
--- a/src/rtk/remote_runner.py
+++ b/src/rtk/remote_runner.py
@@ -99,6 +99,7 @@ class RemoteRunner:
             Args.json_backfill_task_id.value: remote_kwargs.get("json_backfill_task_id"),
             Args.json_backfill_user.value: remote_kwargs.get("json_backfill_user"),
             Args.json_backfill_password.value: remote_kwargs.get("json_backfill_password"),
+            Args.json_backfill_custom_api.value: remote_kwargs.get("json_backfill_custom_api") or "api/youqu/yqresult/"
         }
         # 客户端地址
         if "/home/" not in GlobalConfig.ROOT_DIR:
@@ -515,13 +516,15 @@ class RemoteRunner:
             self.default.get(Args.json_backfill_base_url.value),
             self.default.get(Args.json_backfill_task_id.value),
             self.default.get(Args.json_backfill_user.value),
-            self.default.get(Args.json_backfill_password.value)
+            self.default.get(Args.json_backfill_password.value),
+            self.default.get(Args.json_backfill_custom_api.value)
         ]):
             from src.rtk.json_backfill import JsonBackfill
             JsonBackfill(
                 base_url=self.default.get(Args.json_backfill_base_url.value),
                 username=self.default.get(Args.json_backfill_user.value),
                 password=self.default.get(Args.json_backfill_password.value),
+                custom_api=self.default.get(Args.json_backfill_custom_api.value),
             ).remote_backfill(self.server_detail_json_path, self.default.get(Args.json_backfill_task_id.value))
         # 分布式执行的情况下需要汇总结果
         if not self.default.get(Args.parallel.value):
-- 
Gitee