From 9b20b67039b296187cf1d44024836fbe30cb9204 Mon Sep 17 00:00:00 2001
From: Aaron3S <chenyang@fit2cloud.com>
Date: Tue, 23 Apr 2024 18:59:39 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8D=E6=89=A7=E8=A1=8C?=
 =?UTF-8?q?=E5=BF=AB=E6=8D=B7=E5=91=BD=E4=BB=A4=E6=97=B6=20local=5Fconnect?=
 =?UTF-8?q?ion=20=E6=B2=A1=E6=9C=89=E8=A2=AB=E6=AD=A3=E7=A1=AE=E8=AE=BE?=
 =?UTF-8?q?=E7=BD=AE?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/ops/ansible/runner.py | 14 +++++++++++---
 apps/ops/models/job.py     |  1 +
 receptor                   |  1 -
 3 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/apps/ops/ansible/runner.py b/apps/ops/ansible/runner.py
index 6af2fffb4..dbbdbdcc3 100644
--- a/apps/ops/ansible/runner.py
+++ b/apps/ops/ansible/runner.py
@@ -14,8 +14,10 @@ __all__ = ['AdHocRunner', 'PlaybookRunner', 'SuperPlaybookRunner', 'UploadFileRu
 
 class AdHocRunner:
     cmd_modules_choices = ('shell', 'raw', 'command', 'script', 'win_shell')
+    need_local_connection_modules_choices = ("mysql", "postgresql", "sqlserver", "huawei")
 
-    def __init__(self, inventory, module, module_args='', pattern='*', project_dir='/tmp/', extra_vars=None,
+    def __init__(self, inventory, job_module, module, module_args='', pattern='*', project_dir='/tmp/',
+                 extra_vars=None,
                  dry_run=False, timeout=-1):
         if extra_vars is None:
             extra_vars = {}
@@ -23,6 +25,7 @@ class AdHocRunner:
         self.inventory = inventory
         self.pattern = pattern
         self.module = module
+        self.job_module = job_module
         self.module_args = module_args
         self.project_dir = project_dir
         self.cb = DefaultCallback()
@@ -30,8 +33,7 @@ class AdHocRunner:
         self.extra_vars = extra_vars
         self.dry_run = dry_run
         self.timeout = timeout
-        # enable local connection
-        self.extra_vars.update({"LOCAL_CONNECTION_ENABLED": "1"})
+        self.envs = {}
 
     def check_module(self):
         if self.module not in self.cmd_modules_choices:
@@ -40,8 +42,13 @@ class AdHocRunner:
             raise CommandInBlackListException(
                 "Command is rejected by black list: {}".format(self.module_args.split()[0]))
 
+    def set_local_connection(self):
+        if self.job_module in self.need_local_connection_modules_choices:
+            self.envs.update({"LOCAL_CONNECTION_ENABLED": "1"})
+
     def run(self, verbosity=0, **kwargs):
         self.check_module()
+        self.set_local_connection()
         verbosity = get_ansible_log_verbosity(verbosity)
 
         if not os.path.exists(self.project_dir):
@@ -53,6 +60,7 @@ class AdHocRunner:
         interface.run(
             timeout=self.timeout if self.timeout > 0 else None,
             extravars=self.extra_vars,
+            envvars=self.envs,
             host_pattern=self.pattern,
             private_data_dir=self.project_dir,
             inventory=self.inventory,
diff --git a/apps/ops/models/job.py b/apps/ops/models/job.py
index 15724ec3d..9e7fcd511 100644
--- a/apps/ops/models/job.py
+++ b/apps/ops/models/job.py
@@ -334,6 +334,7 @@ class JobExecution(JMSOrgBaseModel):
 
             runner = AdHocRunner(
                 self.inventory_path,
+                self.job.module,
                 module,
                 timeout=self.current_job.timeout,
                 module_args=args,
diff --git a/receptor b/receptor
index d32f40889..1efc8cba4 100755
--- a/receptor
+++ b/receptor
@@ -75,7 +75,6 @@ class ReceptorService:
                 print("\n- PID file is corrupted, starting Receptor...")
                 os.remove(self.pid_file)
 
-        os.environ.update({'LOCAL_CONNECTION_ENABLED': '1'})
         os.environ.setdefault('ANSIBLE_LIBRARY', DEFAULT_ANSIBLE_MODULES_DIR)
         os.environ.update({'PYTHONPATH': APPS_DIR})
         process = subprocess.Popen(self.receptor_command)