feat: add IBM DB2 support (#10306)

### What problem does this PR solve? issue：#5617 change：add IBM DB2 support in ExeSQL ### Type of change - [x] New Feature (non-breaking change which adds functionality)
2025-11-30 08:56:59 +00:00 · 2025-09-26 14:55:19 +08:00 · 2025-09-26 14:55:19 +08:00 · bd94b5dfb5
commit bd94b5dfb5
parent ef59c5bab9
15 changed files with 136 additions and 30 deletions
--- a/agent/component/agent_with_tools.py
+++ b/agent/component/agent_with_tools.py
@ -137,7 +137,7 @@ class Agent(LLM, ToolBase):
            res.update(cpn.get_input_form())
        return res

-    @timeout(os.environ.get("COMPONENT_EXEC_TIMEOUT", 20*60))
+    @timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 20*60)))
    def _invoke(self, **kwargs):
        if kwargs.get("user_prompt"):
            usr_pmt = ""
--- a/agent/component/base.py
+++ b/agent/component/base.py
@ -431,7 +431,7 @@ class ComponentBase(ABC):
        self.set_output("_elapsed_time", time.perf_counter() - self.output("_created_time"))
        return self.output()

-    @timeout(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10*60))
+    @timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10*60)))
    def _invoke(self, **kwargs):
        raise NotImplementedError()

--- a/agent/component/invoke.py
+++ b/agent/component/invoke.py
@ -53,7 +53,7 @@ class InvokeParam(ComponentParamBase):
 class Invoke(ComponentBase, ABC):
    component_name = "Invoke"

-    @timeout(os.environ.get("COMPONENT_EXEC_TIMEOUT", 3))
+    @timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 3)))
    def _invoke(self, **kwargs):
        args = {}
        for para in self._param.variables:
--- a/agent/component/llm.py
+++ b/agent/component/llm.py
@ -201,7 +201,7 @@ class LLM(ComponentBase):
            for txt in self.chat_mdl.chat_streamly(msg[0]["content"], msg[1:], self._param.gen_conf(), images=self.imgs, **kwargs):
                yield delta(txt)

-    @timeout(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10*60))
+    @timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10*60)))
    def _invoke(self, **kwargs):
        def clean_formated_answer(ans: str) -> str:
            ans = re.sub(r"^.*</think>", "", ans, flags=re.DOTALL)
--- a/agent/component/message.py
+++ b/agent/component/message.py
@ -127,7 +127,7 @@ class Message(ComponentBase):
        ]
        return any([re.search(p, content) for p in patt])

-    @timeout(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10*60))
+    @timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10*60)))
    def _invoke(self, **kwargs):
        rand_cnt = random.choice(self._param.content)
        if self._param.stream and not self._is_jinjia2(rand_cnt):
--- a/agent/component/switch.py
+++ b/agent/component/switch.py
@ -61,7 +61,7 @@ class SwitchParam(ComponentParamBase):
 class Switch(ComponentBase, ABC):
    component_name = "Switch"

-    @timeout(os.environ.get("COMPONENT_EXEC_TIMEOUT", 3))
+    @timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 3)))
    def _invoke(self, **kwargs):
        for cond in self._param.conditions:
            res = []
--- a/agent/tools/code_exec.py
+++ b/agent/tools/code_exec.py
@ -156,7 +156,7 @@ class CodeExec(ToolBase, ABC):
            self.set_output("_ERROR", "construct code request error: " + str(e))

        try:
-            resp = requests.post(url=f"http://{settings.SANDBOX_HOST}:9385/run", json=code_req, timeout=os.environ.get("COMPONENT_EXEC_TIMEOUT", 10*60))
+            resp = requests.post(url=f"http://{settings.SANDBOX_HOST}:9385/run", json=code_req, timeout=int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10*60)))
            logging.info(f"http://{settings.SANDBOX_HOST}:9385/run,  code_req: {code_req}, resp.status_code {resp.status_code}:")
            if resp.status_code != 200:
                resp.raise_for_status()
--- a/agent/tools/exesql.py
+++ b/agent/tools/exesql.py
@ -21,6 +21,7 @@ import pandas as pd
 import pymysql
 import psycopg2
 import pyodbc
+import ibm_db
 from agent.tools.base import ToolParamBase, ToolBase, ToolMeta
 from api.utils.api_utils import timeout

@ -53,7 +54,7 @@ class ExeSQLParam(ToolParamBase):
        self.max_records = 1024

    def check(self):
-        self.check_valid_value(self.db_type, "Choose DB type", ['mysql', 'postgres', 'mariadb', 'mssql'])
+        self.check_valid_value(self.db_type, "Choose DB type", ['mysql', 'postgres', 'mariadb', 'mssql', 'IBM DB2'])
        self.check_empty(self.database, "Database name")
        self.check_empty(self.username, "database username")
        self.check_empty(self.host, "IP Address")
@ -123,6 +124,52 @@ class ExeSQL(ToolBase, ABC):
                    r'PWD=' + self._param.password
            )
            db = pyodbc.connect(conn_str)
+        elif self._param.db_type == 'IBM DB2':
+            conn_str = (
+                f"DATABASE={self._param.database};"
+                f"HOSTNAME={self._param.host};"
+                f"PORT={self._param.port};"
+                f"PROTOCOL=TCPIP;"
+                f"UID={self._param.username};"
+                f"PWD={self._param.password};"
+            )
+            try:
+                conn = ibm_db.connect(conn_str, "", "")
+            except Exception as e:
+                raise Exception("Database Connection Failed! \n" + str(e))
+
+            sql_res = []
+            formalized_content = []
+            for single_sql in sqls:
+                single_sql = single_sql.replace("```", "").strip()
+                if not single_sql:
+                    continue
+                single_sql = re.sub(r"\[ID:[0-9]+\]", "", single_sql)
+
+                stmt = ibm_db.exec_immediate(conn, single_sql)
+                rows = []
+                row = ibm_db.fetch_assoc(stmt)
+                while row and len(rows) < self._param.max_records:
+                    rows.append(row)
+                    row = ibm_db.fetch_assoc(stmt)
+
+                if not rows:
+                    sql_res.append({"content": "No record in the database!"})
+                    continue
+
+                df = pd.DataFrame(rows)
+                for col in df.columns:
+                    if pd.api.types.is_datetime64_any_dtype(df[col]):
+                        df[col] = df[col].dt.strftime("%Y-%m-%d")
+
+                sql_res.append(convert_decimals(df.to_dict(orient="records")))
+                formalized_content.append(df.to_markdown(index=False, floatfmt=".6f"))
+
+            ibm_db.close(conn)
+
+            self.set_output("json", sql_res)
+            self.set_output("formalized_content", "\n\n".join(formalized_content))
+            return self.output("formalized_content")
        try:
            cursor = db.cursor()
        except Exception as e:
--- a/api/apps/canvas_app.py
+++ b/api/apps/canvas_app.py
@ -348,6 +348,22 @@ def test_db_connect():
            cursor = db.cursor()
            cursor.execute("SELECT 1")
            cursor.close()
+        elif req["db_type"] == 'IBM DB2':
+            import ibm_db
+            conn_str = (
+                f"DATABASE={req['database']};"
+                f"HOSTNAME={req['host']};"
+                f"PORT={req['port']};"
+                f"PROTOCOL=TCPIP;"
+                f"UID={req['username']};"
+                f"PWD={req['password']};"
+            )
+            logging.info(conn_str)
+            conn = ibm_db.connect(conn_str, "", "")
+            stmt = ibm_db.exec_immediate(conn, "SELECT 1 FROM sysibm.sysdummy1")
+            ibm_db.fetch_assoc(stmt)
+            ibm_db.close(conn)
+            return get_json_result(data="Database Connection Successful!")
        else:
            return server_error_response("Unsupported database type.")
        if req["db_type"] != 'mssql':
--- a/pyproject.toml
+++ b/pyproject.toml
@ -132,6 +132,7 @@ dependencies = [
    "litellm>=1.74.15.post1",
    "flask-mail>=0.10.0",
    "lark>=1.2.2",
+    "ibm-db>=3.2.7",
 ]

 [project.optional-dependencies]
@ -157,6 +158,9 @@ test = [
    "requests-toolbelt>=1.0.0",
 ]

+[[tool.uv.index]]
+url = "https://mirrors.aliyun.com/pypi/simple"
+
 [tool.setuptools]
 packages = [
    'agent',
@ -170,9 +174,6 @@ packages = [
    'sdk.python.ragflow_sdk',
 ]

-[[tool.uv.index]]
-url = "https://mirrors.aliyun.com/pypi/simple"
-
 [tool.ruff]
 line-length = 200
 exclude = [".venv", "rag/svr/discord_svr.py"]
--- a/rag/flow/base.py
+++ b/rag/flow/base.py
@ -58,6 +58,6 @@ class ProcessBase(ComponentBase):
        self.set_output("_elapsed_time", time.perf_counter() - self.output("_created_time"))
        return self.output()

-    @timeout(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10 * 60))
+    @timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10 * 60)))
    async def _invoke(self, **kwargs):
        raise NotImplementedError()
--- a/uv.lock
+++ b/uv.lock
@ -1,4 +1,5 @@
 version = 1
+revision = 3
 requires-python = ">=3.10, <3.13"
 resolution-markers = [
    "python_full_version >= '3.12' and sys_platform == 'darwin'",
@ -2503,6 +2504,32 @@ wheels = [
    { url = "https://mirrors.aliyun.com/pypi/packages/34/87/7940713f929d0280cff1bde207479cb588a0d3a4dd49a0e2e69bfff46363/hyppo-0.4.0-py3-none-any.whl", hash = "sha256:4e75565b8deb601485cd7bc1b5c3f44e6ddf329136fc81e65d011f9b4e95132f" },
 ]

+[[package]]
+name = "ibm-db"
+version = "3.2.7"
+source = { registry = "https://mirrors.aliyun.com/pypi/simple" }
+sdist = { url = "https://mirrors.aliyun.com/pypi/packages/dd/a5/182413f7fe28ee7a67d8be5afa1139ccc60cfb5c13c0e9be81e2205bddbb/ibm_db-3.2.7.tar.gz", hash = "sha256:b3c3b4550364a43bf1daa4519b668e6e00e7c3935291f8c444c4ec989417e861" }
+wheels = [
+    { url = "https://mirrors.aliyun.com/pypi/packages/50/cd/a6549ed35424875f07ea89dbd44093b7d9dc4a03d9e29e56c5167bd7d568/ibm_db-3.2.7-cp310-cp310-macosx_10_15_x86_64.whl", hash = "sha256:7c2b451ffe67be602e93d94b2d2042dd051ec0757cfd6e4d7344cb594f2d3508" },
+    { url = "https://mirrors.aliyun.com/pypi/packages/50/5a/be83503ec6ef9b2a47175b38b1099595a8d06237ac3fcc82d967b834672a/ibm_db-3.2.7-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:9a1b139a9c21ff7216aac83ba29dceb6c8a9df3d6aee44ff1fe845cb60d3caed" },
+    { url = "https://mirrors.aliyun.com/pypi/packages/cc/fa/379405785d27d10110992ee17f150c8ef1ee0c3eadca2d1451c8c03ff075/ibm_db-3.2.7-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1e5e60297b4680cc566caa67f513aa68883ef48b0c612028a38883620807b09c" },
+    { url = "https://mirrors.aliyun.com/pypi/packages/52/e1/c18aee07888666b3249b4f54d3cb67ae5041600b5fc3ed281817e868eaa8/ibm_db-3.2.7-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bf1c30e67e9e573e33524c393a1426e0dffa2da34ba42a0ec510e0f75766976f" },
+    { url = "https://mirrors.aliyun.com/pypi/packages/26/01/f80a407192351aa304206504d6b15ccbb061f45dc9fc3e37c8c00c7e222b/ibm_db-3.2.7-cp310-cp310-win32.whl", hash = "sha256:171014c2caa0419055943ff3badae5118cc3a191360f03b80c8366ef374d5c28" },
+    { url = "https://mirrors.aliyun.com/pypi/packages/98/1b/29f98a0d4d9896635d7b7fa53a51f8753214f0deed23ac7987d726299b12/ibm_db-3.2.7-cp310-cp310-win_amd64.whl", hash = "sha256:3425c158a65dd43e4b09dc968c18042a656ed6ef2e1db0164f032e97681823b7" },
+    { url = "https://mirrors.aliyun.com/pypi/packages/c8/04/18e42549f498569db7a437453db51ae3d61105ad4da7b1fe64e9e59aedee/ibm_db-3.2.7-cp311-cp311-macosx_10_15_x86_64.whl", hash = "sha256:ba493d07e1845b8b1169ad27ace92f0ff540cc9a623f2753b8c68dc66c59d7df" },
+    { url = "https://mirrors.aliyun.com/pypi/packages/c5/f3/85c8963ee8047183c435059abaf40771d40c6e9268ca32d66be0b66b7a6c/ibm_db-3.2.7-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:abed0a7c644b9ddf2c49bf5c0938f936f0b2dffd1703c9819440021be141716e" },
+    { url = "https://mirrors.aliyun.com/pypi/packages/6d/fb/8d2ff4b6bcd6b05d13af855bedc47e597430a6c3372e0ab7579659cad9bb/ibm_db-3.2.7-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1cabd3d3e8c879ef60d91e1fe1356cf8603f8b4b69cc7dda39d4a8698a055044" },
+    { url = "https://mirrors.aliyun.com/pypi/packages/77/26/c43e02bb4cf62b1e93bf617440f5da6d3888935ef09d4fbd86fe07f3f920/ibm_db-3.2.7-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:aab5dceec45d69b0bbd333be66597dbaedf663c6c56a0fbd6196ecd1836e4095" },
+    { url = "https://mirrors.aliyun.com/pypi/packages/05/d0/787d7a9864d3782238ca3b14a4484c642931e1363a760b0828c9ee26ad58/ibm_db-3.2.7-cp311-cp311-win32.whl", hash = "sha256:16272ad07912051d9ab5cbe3a9e2d3d888365d071334f9620d8e0b2ed69ee4f9" },
+    { url = "https://mirrors.aliyun.com/pypi/packages/1e/7c/998c663d0a65984c76c2c2c8d01cdbdd174bd817359e0e4024b9b316a9cc/ibm_db-3.2.7-cp311-cp311-win_amd64.whl", hash = "sha256:4b479e92b6954ab7f65c9d247a65fb0cde6a48899f71a8881b58023c0ace1f49" },
+    { url = "https://mirrors.aliyun.com/pypi/packages/4d/0f/048679ca8516b73f3547c64f968b9654181d79f6cd2706914f37c2486da3/ibm_db-3.2.7-cp312-cp312-macosx_10_15_x86_64.whl", hash = "sha256:24e8a538475997f20569f221247808507b63349df51119fe9b2f8e48a0bf6f9b" },
+    { url = "https://mirrors.aliyun.com/pypi/packages/f0/36/4cc64c70ebc74b2765005cd5357df18512c358cc6f5e87c6b0e70cff4053/ibm_db-3.2.7-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:24a53fb8e3c200bf2a55095f1ae4c065f2136f8be87ca1db89a874bd82d88ea5" },
+    { url = "https://mirrors.aliyun.com/pypi/packages/09/d7/ae63f1257736c5e6a06cd2be133b4bc38f68893504f046b4c850144b65cd/ibm_db-3.2.7-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:91f68be7bd0d2940023da43d0a94f196fe267ca825df7874b8174583c8678ea0" },
+    { url = "https://mirrors.aliyun.com/pypi/packages/d9/24/af465d93f549b0dbc944f256cdb2b470574018285e1478b6c50305a609ac/ibm_db-3.2.7-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5d39fe5001c078f2b824d1805ca9737060203a00a9dd9a8fe4b6f6b32b271cb5" },
+    { url = "https://mirrors.aliyun.com/pypi/packages/48/04/c512eed2c701e8762e90f000fd1a768dd749ffb070a62b8cd92722c16327/ibm_db-3.2.7-cp312-cp312-win32.whl", hash = "sha256:20388753f52050e07e845b74146dbbe3f892dcfdfb015638e8f57c2fb2e056b8" },
+    { url = "https://mirrors.aliyun.com/pypi/packages/07/cc/ae978e6d020f3b17f2e68cc2c60fe8381fffd1271608e871b5b4ee6434f2/ibm_db-3.2.7-cp312-cp312-win_amd64.whl", hash = "sha256:6e507ddf93b8406b0b88ff6bf07658a3100ce98cb1e735e5ec8e0a56e30ea856" },
+]
+
 [[package]]
 name = "idna"
 version = "3.10"
@ -5323,6 +5350,7 @@ dependencies = [
    { name = "html-text" },
    { name = "httpx", extra = ["socks"] },
    { name = "huggingface-hub" },
+    { name = "ibm-db" },
    { name = "infinity-emb" },
    { name = "infinity-sdk" },
    { name = "itsdangerous" },
@ -5479,6 +5507,7 @@ requires-dist = [
    { name = "html-text", specifier = "==0.6.2" },
    { name = "httpx", extras = ["socks"], specifier = "==0.27.2" },
    { name = "huggingface-hub", specifier = ">=0.25.0,<0.26.0" },
+    { name = "ibm-db", specifier = ">=3.2.7" },
    { name = "infinity-emb", specifier = ">=0.0.66,<0.0.67" },
    { name = "infinity-sdk", specifier = "==0.6.0.dev5" },
    { name = "itsdangerous", specifier = "==2.1.2" },
@ -5563,6 +5592,7 @@ requires-dist = [
    { name = "yfinance", specifier = "==0.2.65" },
    { name = "zhipuai", specifier = "==2.0.1" },
 ]
+provides-extras = ["full"]

 [package.metadata.requires-dev]
 test = [
--- a/web/src/pages/agent/options.ts
+++ b/web/src/pages/agent/options.ts
@ -2133,12 +2133,16 @@ export const QWeatherTimePeriodOptions = [
  '30d',
 ];

-export const ExeSQLOptions = ['mysql', 'postgres', 'mariadb', 'mssql'].map(
-  (x) => ({
-    label: upperFirst(x),
-    value: x,
-  }),
-);
+export const ExeSQLOptions = [
+  'mysql',
+  'postgres',
+  'mariadb',
+  'mssql',
+  'IBM DB2',
+].map((x) => ({
+  label: upperFirst(x),
+  value: x,
+}));

 export const WenCaiQueryTypeOptions = [
  'stock',
--- a/web/src/pages/data-flow/options.ts
+++ b/web/src/pages/data-flow/options.ts
@ -2133,12 +2133,16 @@ export const QWeatherTimePeriodOptions = [
  '30d',
 ];

-export const ExeSQLOptions = ['mysql', 'postgres', 'mariadb', 'mssql'].map(
-  (x) => ({
-    label: upperFirst(x),
-    value: x,
-  }),
-);
+export const ExeSQLOptions = [
+  'mysql',
+  'postgres',
+  'mariadb',
+  'mssql',
+  'IBM DB2',
+].map((x) => ({
+  label: upperFirst(x),
+  value: x,
+}));

 export const WenCaiQueryTypeOptions = [
  'stock',
--- a/web/src/pages/flow/constant.tsx
+++ b/web/src/pages/flow/constant.tsx
@ -2911,12 +2911,16 @@ export const QWeatherTimePeriodOptions = [
  '30d',
 ];

-export const ExeSQLOptions = ['mysql', 'postgres', 'mariadb', 'mssql'].map(
-  (x) => ({
-    label: upperFirst(x),
-    value: x,
-  }),
-);
+export const ExeSQLOptions = [
+  'mysql',
+  'postgres',
+  'mariadb',
+  'mssql',
+  'IBM DB2',
+].map((x) => ({
+  label: upperFirst(x),
+  value: x,
+}));

 export const SwitchElseTo = 'end_cpn_id';