Fix: db2 i series connection fix(no ssl, port in con args) (#25781)

* db2 i series connection fix(no ssl, port in con args)

* column type fixes

* pyformat, add tests

* fixes

Author: harshsoni2024

ingestion/src/metadata/ingestion/source/database/column_type_parser.py

@@ -217,6 +217,18 @@ class ColumnTypeParser:
         "CIDR": "CIDR",
         "INET": "INET",
         "TSRANGE": "DATETIMERANGE",
+        # DB2
+        "XMLVARCHAR": "XML",
+        "XMLCLOB": "XML",
+        "XMLFILE": "XML",
+        "TIMESTMP": "TIMESTAMP",
+        "LONGVARCHAR": "VARCHAR",
+        "GRAPHIC": "CHAR",
+        "VARGRAPHIC": "VARCHAR",
+        "LONGVARGRAPHIC": "VARCHAR",
+        "DBCLOB": "CLOB",
+        "DECFLOAT": "DOUBLE",
+        "CHARACTER": "CHAR",
         # ORACLE
         "BINARY_DOUBLE": "DOUBLE",
         "BINARY_FLOAT": "FLOAT",

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

@@ -12,8 +12,11 @@
 """
 Source connection handler
 """
+import importlib
+import sys
 from pathlib import Path
-from typing import Optional
+from typing import Any, Dict, Optional
+from urllib.parse import quote_plus
 
 from sqlalchemy.engine import Engine
 
@@ -22,6 +25,7 @@
 )
 from metadata.generated.schema.entity.services.connections.database.db2Connection import (
     Db2Connection,
+    Db2Scheme,
 )
 from metadata.generated.schema.entity.services.connections.testConnectionResult import (
     TestConnectionResult,
@@ -30,6 +34,7 @@
     create_generic_db_connection,
     get_connection_args_common,
     get_connection_url_common,
+    get_password_secret,
 )
 from metadata.ingestion.connections.test_connections import test_connection_db_common
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
@@ -44,6 +49,50 @@
 logger = ingestion_logger()
 
 
+def _get_ibmi_connection_url(connection: Db2Connection) -> str:
+    """
+    Build connection URL for ibmi scheme.
+
+    sqlalchemy-ibmi validates URL options strictly and does NOT support:
+    - Port in the URL (host:port format) — causes validation error
+    - SSL parameters like SECURITY=SSL — SSL is negotiated via port (9471)
+
+    Port is passed separately via connect_args.
+    """
+    hostname = connection.hostPort.split(":")[0]
+
+    url = f"{connection.scheme.value}://"
+    if connection.username:
+        url += f"{quote_plus(connection.username)}"
+        password = get_password_secret(connection)
+        url += f":{quote_plus(password.get_secret_value())}"
+        url += "@"
+    url += hostname
+    if connection.database:
+        url += f"/{connection.database}"
+    return url
+
+
+def _get_ibmi_connection_args(connection: Db2Connection) -> Dict[str, Any]:
+    """
+    Build connection args for ibmi scheme.
+
+    Port must be passed via connect_args since sqlalchemy-ibmi
+    rejects it in the URL.
+    """
+    args = get_connection_args_common(connection)
+    host_port = connection.hostPort
+    if ":" in host_port:
+        port_str = host_port.split(":")[1]
+        try:
+            args["port"] = int(port_str)
+        except ValueError:
+            raise ValueError(
+                f"Invalid port in hostPort '{host_port}'. Expected format: 'hostname:port'"
+            )
+    return args
+
+
 def get_connection(connection: Db2Connection) -> Engine:
     """
     Create connection
@@ -55,22 +104,43 @@ def get_connection(connection: Db2Connection) -> Engine:
         clidriver_version = check_clidriver_version(clidriver_version)
         if clidriver_version:
             install_clidriver(clidriver_version.value)
+            # Invalidate cached clidriver module so the next import
+            # picks up the freshly installed path
+            sys.modules.pop("clidriver", None)
 
     # prepare license
     # pylint: disable=import-outside-toplevel
     if connection.license and connection.licenseFileName:
         import clidriver
 
+        if clidriver_version:
+            importlib.reload(clidriver)
+
+        license_dir = Path(clidriver.__path__[0], "license")
+        license_dir.mkdir(parents=True, exist_ok=True)
+
         with open(
-            Path(clidriver.__path__[0], "license", connection.licenseFileName),
+            license_dir / connection.licenseFileName,
             "w",
             encoding=UTF_8,
         ) as file:
             file.write(connection.license.encode(UTF_8).decode("unicode-escape"))
 
-    ssl_manager = check_ssl_and_init(connection)
-    if ssl_manager:
-        connection = ssl_manager.setup_ssl(connection)
+    is_ibmi = connection.scheme == Db2Scheme.ibmi
+
+    # SSL setup only for db2+ibm_db scheme.
+    # ibmi scheme negotiates SSL via port (9471), not connection params.
+    if not is_ibmi:
+        ssl_manager = check_ssl_and_init(connection)
+        if ssl_manager:
+            connection = ssl_manager.setup_ssl(connection)
+
+    if is_ibmi:
+        return create_generic_db_connection(
+            connection=connection,
+            get_connection_url_fn=_get_ibmi_connection_url,
+            get_connection_args_fn=_get_ibmi_connection_args,
+        )
 
     return create_generic_db_connection(
         connection=connection,

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

@@ -12,11 +12,9 @@
 import traceback
 from typing import Iterable, Optional
 
-from ibm_db_sa.base import ischema_names
-from ibm_db_sa.reflection import DB2Reflector, OS390Reflector
 from sqlalchemy.engine.reflection import Inspector
 from sqlalchemy.engine.row import LegacyRow
-from sqlalchemy.sql.sqltypes import BOOLEAN
+from sqlalchemy.sql.sqltypes import BINARY, BOOLEAN, VARBINARY
 
 from metadata.generated.schema.entity.services.connections.database.db2Connection import (
     Db2Connection,
@@ -26,18 +24,39 @@
 )
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.database.column_type_parser import create_sqlalchemy_type
 from metadata.ingestion.source.database.common_db_source import CommonDbSourceService
-from metadata.ingestion.source.database.db2.utils import get_unique_constraints
+from metadata.ingestion.source.database.db2.utils import (
+    get_columns_os390,
+    get_unique_constraints,
+)
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
 
+# ibm_db_sa is only required for the db2+ibm_db scheme (z/OS, LUW).
+# The ibmi scheme uses sqlalchemy-ibmi which does not depend on ibm_db_sa.
+try:
+    from ibm_db_sa.base import ischema_names
+    from ibm_db_sa.reflection import DB2Reflector, OS390Reflector
 
-ischema_names.update({"BOOLEAN": BOOLEAN})
-
-
-DB2Reflector.get_unique_constraints = get_unique_constraints
-OS390Reflector.get_unique_constraints = get_unique_constraints
+    ischema_names.update(
+        {
+            "BOOLEAN": BOOLEAN,
+            "BINARY": BINARY,
+            "VARBINARY": VARBINARY,
+            "DECFLOAT": create_sqlalchemy_type("DECFLOAT"),
+            "ROWID": create_sqlalchemy_type("ROWID"),
+            "XMLVARCHAR": create_sqlalchemy_type("XMLVARCHAR"),
+            "XMLCLOB": create_sqlalchemy_type("XMLCLOB"),
+            "XMLFILE": create_sqlalchemy_type("XMLFILE"),
+        }
+    )
+    DB2Reflector.get_unique_constraints = get_unique_constraints
+    OS390Reflector.get_unique_constraints = get_unique_constraints
+    OS390Reflector.get_columns = get_columns_os390
+except ImportError:
+    logger.debug("ibm_db_sa not installed - db2+ibm_db scheme unavailable")
 
 
 class Db2Source(CommonDbSourceService):

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

@@ -16,6 +16,7 @@
 
 from sqlalchemy import and_, join, sql
 from sqlalchemy.engine import reflection
+from sqlalchemy.sql import sqltypes as sa_types
 
 from metadata.utils.logger import ingestion_logger
 
@@ -40,6 +41,67 @@ class DB2CLIDriverVersions(Enum):
     V12_1_0 = "12.1.0"
 
 
+@reflection.cache
+def get_columns_os390(
+    self, connection, table_name, schema=None, **kw
+):  # pylint: disable=unused-argument
+    """Override OS390Reflector.get_columns to handle empty/unrecognized types
+    gracefully instead of emitting SAWarnings."""
+    current_schema = self.denormalize_name(schema or self.default_schema_name)
+    table_name = self.denormalize_name(table_name)
+    syscols = self.sys_columns
+
+    query = (
+        sql.select(
+            syscols.c.colname,
+            syscols.c.typename,
+            syscols.c.defaultval,
+            syscols.c.nullable,
+            syscols.c.length,
+            syscols.c.scale,
+            syscols.c.generated,
+            syscols.c.remark,
+        )
+        .where(
+            and_(
+                syscols.c.tabschema == current_schema,
+                syscols.c.tabname == table_name,
+            )
+        )
+        .order_by(syscols.c.colno)
+    )
+    sa_columns = []
+    for r in connection.execute(query):
+        coltype = r[1].strip().upper() if r[1] else ""
+        if coltype in ["DECIMAL", "NUMERIC"]:
+            coltype = self.ischema_names.get(coltype)(int(r[4]), int(r[5]))
+        elif coltype in ["CHARACTER", "CHAR", "VARCHAR", "GRAPHIC", "VARGRAPHIC"]:
+            coltype = self.ischema_names.get(coltype)(int(r[4]))
+        elif coltype and coltype in self.ischema_names:
+            coltype = self.ischema_names[coltype]
+        else:
+            if not coltype:
+                logger.warning(f"Empty type for column '{r[0]}' - ingesting as UNKNOWN")
+            else:
+                logger.warning(
+                    f"Did not recognize type '{coltype}' of column '{r[0]}'"
+                    " - ingesting as UNKNOWN"
+                )
+            coltype = sa_types.NULLTYPE
+
+        sa_columns.append(
+            {
+                "name": self.normalize_name(r[0]),
+                "type": coltype,
+                "nullable": r[3] == "Y",
+                "default": r[2] or None,
+                "autoincrement": r[6] not in (" ", None),
+                "comment": r[7] or None,
+            }
+        )
+    return sa_columns
+
+
 @reflection.cache
 def get_unique_constraints(
     self, connection, table_name, schema=None, **kw