issue-21370: db2 custom driver installation (#21638)

* db2 custom driver installation

* pylint changes

* typo fix

(cherry picked from commit 4a3b6f4934877d4ec196448227cf3a57204886c4)

ingestion/setup.py

@@ -207,7 +207,7 @@ plugins: Dict[str, Set[str]] = {
         VERSIONS["azure-storage-blob"],
         VERSIONS["azure-identity"],
     },
-    "db2": {"ibm-db-sa~=0.4.1", "ibm-db>=2.0.0"},
+    "db2": {"ibm-db-sa~=0.4.1", "ibm-db>=3.2.6"},
     "db2-ibmi": {"sqlalchemy-ibmi~=0.9.3"},
     "databricks": {
         VERSIONS["sqlalchemy-databricks"],

ingestion/src/metadata/ingestion/source/database/db2/connection.py

@@ -33,13 +33,28 @@ from metadata.ingestion.connections.builders import (
 )
 from metadata.ingestion.connections.test_connections import test_connection_db_common
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.database.db2.utils import (
+    check_clidriver_version,
+    install_clidriver,
+)
 from metadata.utils.constants import THREE_MIN, UTF_8
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
 
 
 def get_connection(connection: Db2Connection) -> Engine:
     """
     Create connection
     """
+    # Install ibm_db with specific version
+    clidriver_version = connection.clidriverVersion
+
+    if clidriver_version:
+        clidriver_version = check_clidriver_version(clidriver_version)
+        if clidriver_version:
+            install_clidriver(clidriver_version.value)
+
     # prepare license
     # pylint: disable=import-outside-toplevel
     if connection.license and connection.licenseFileName:

ingestion/src/metadata/ingestion/source/database/db2/utils.py

@@ -12,9 +12,33 @@
 """
 Module to define overriden dialect methods
 """
+from enum import Enum
+
 from sqlalchemy import and_, join, sql
 from sqlalchemy.engine import reflection
 
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+BASE_CLIDRIVER_URL = (
+    "https://public.dhe.ibm.com/ibmdl/export/pub/software/data/db2/drivers/odbc_cli"
+)
+
+
+class DB2CLIDriverVersions(Enum):
+    """
+    Enum for the DB2 CLI Driver versions
+    """
+
+    V11_1_4 = "11.1.4"
+    V11_5_4 = "11.5.4"
+    V11_5_5 = "11.5.5"
+    V11_5_6 = "11.5.6"
+    V11_5_8 = "11.5.8"
+    V11_5_9 = "11.5.9"
+    V12_1_0 = "12.1.0"
+
 
 @reflection.cache
 def get_unique_constraints(
@@ -61,3 +85,109 @@ def get_unique_constraints(
                 }
             )
     return unique_consts
+
+
+def check_clidriver_version(clidriver_version: str):
+    """
+    Check if the CLI Driver version is valid
+    """
+    if clidriver_version not in [v.value for v in DB2CLIDriverVersions]:
+        logger.warning(f"Invalid CLI Driver version provided: {clidriver_version}")
+        return None
+    return DB2CLIDriverVersions(clidriver_version)
+
+
+# pylint: disable=too-many-statements,too-many-branches
+def install_clidriver(clidriver_version: str) -> None:
+    """
+    Install the CLI Driver for DB2
+    """
+    # pylint: disable=import-outside-toplevel
+    import os
+    import platform
+    import subprocess
+    import sys
+    from urllib.request import URLError, urlopen
+
+    import pkg_resources
+
+    clidriver_version = f"v{clidriver_version}"
+    system = platform.system().lower()
+    is_64bits = platform.architecture()[0] == "64bit"
+    clidriver_url = None
+    default_clidriver_url = None
+
+    def is_valid_url(url: str) -> bool:
+        """Check if the URL is valid and accessible"""
+        try:
+            with urlopen(url) as _:
+                return True
+        except URLError:
+            return False
+
+    if system == "darwin":  # macOS
+        machine = platform.machine().lower()
+        if machine == "arm64":  # Apple Silicon
+            default_clidriver_url = f"{BASE_CLIDRIVER_URL}/macarm64_odbc_cli.tar.gz"
+            clidriver_url = f"{BASE_CLIDRIVER_URL}/macarm64_odbc_cli.tar.gz"
+        elif machine == "x86_64":  # Intel
+            default_clidriver_url = f"{BASE_CLIDRIVER_URL}/macos64_odbc_cli.tar.gz"
+            clidriver_url = (
+                f"{BASE_CLIDRIVER_URL}/{str(clidriver_version)}/macos64_odbc_cli.tar.gz"
+            )
+    elif system == "linux":
+        if is_64bits:
+            default_clidriver_url = f"{BASE_CLIDRIVER_URL}/linuxx64_odbc_cli.tar.gz"
+            clidriver_url = f"{BASE_CLIDRIVER_URL}/{str(clidriver_version)}/linuxx64_odbc_cli.tar.gz"
+        else:
+            default_clidriver_url = f"{BASE_CLIDRIVER_URL}/linuxia32_odbc_cli.tar.gz"
+            clidriver_url = f"{BASE_CLIDRIVER_URL}/{str(clidriver_version)}/linuxia32_odbc_cli.tar.gz"
+    elif system == "windows":
+        if is_64bits:
+            default_clidriver_url = f"{BASE_CLIDRIVER_URL}/ntx64_odbc_cli.zip"
+            clidriver_url = (
+                f"{BASE_CLIDRIVER_URL}/{str(clidriver_version)}/ntx64_odbc_cli.zip"
+            )
+        else:
+            default_clidriver_url = f"{BASE_CLIDRIVER_URL}/nt32_odbc_cli.zip"
+            clidriver_url = (
+                f"{BASE_CLIDRIVER_URL}/{str(clidriver_version)}/nt32_odbc_cli.zip"
+            )
+    else:
+        logger.error(
+            f"Unsupported operating system for db2 driver installation: {system}"
+        )
+        return None
+
+    # set env variables for CLIDRIVER_VERSION and IBM_DB_INSTALLER_URL
+    os.environ["CLIDRIVER_VERSION"] = clidriver_version
+    if is_valid_url(clidriver_url):
+        os.environ["IBM_DB_INSTALLER_URL"] = clidriver_url
+    else:
+        os.environ["IBM_DB_INSTALLER_URL"] = default_clidriver_url
+    logger.info(f"Set IBM_DB_INSTALLER_URL to {os.environ['IBM_DB_INSTALLER_URL']}")
+    logger.info(f"Set CLIDRIVER_VERSION to {os.environ['CLIDRIVER_VERSION']}")
+    # Uninstall ibm_db if it is already installed
+    try:
+        pkg_resources.get_distribution("ibm_db")
+        # If we get here, ibm_db is installed, so uninstall it first
+        subprocess.check_call(
+            [sys.executable, "-m", "pip", "uninstall", "-y", "ibm_db"]
+        )
+    except pkg_resources.DistributionNotFound:
+        # ibm_db is not installed, proceed with installation
+        pass
+    # Install ibm_db with specific flags
+    subprocess.check_call(
+        [
+            sys.executable,
+            "-m",
+            "pip",
+            "install",
+            "ibm_db~=3.2.6",
+            "--no-binary",
+            ":all:",
+            "--no-cache-dir",
+        ]
+    )
+    return None

openmetadata-spec/src/main/resources/json/schema/entity/services/connections/database/db2Connection.json

@@ -63,6 +63,11 @@
       "description": "License to connect to DB2.",
       "type": "string"
     },
+    "clidriverVersion": {
+      "title": "CLI Driver Version",
+      "description": "CLI Driver version to connect to DB2. If not provided, the latest version will be used.",
+      "type": "string"
+    },
     "connectionOptions": {
       "title": "Connection Options",
       "$ref": "../connectionBasicType.json#/definitions/connectionOptions"

openmetadata-ui/src/main/resources/ui/src/generated/entity/services/connections/database/db2Connection.ts

@@ -14,6 +14,10 @@
  * Db2 Connection Config
  */
 export interface Db2Connection {
+    /**
+     * CLI Driver version to connect to DB2. If not provided, the latest version will be used.
+     */
+    clidriverVersion?:    string;
     connectionArguments?: { [key: string]: any };
     connectionOptions?:   { [key: string]: string };
     /**