Fix#9251: DB2 connection config and ingestion update (#9322)

* DB2 connection config and ingestion update * Update ingestion/src/metadata/ingestion/source/database/common_db_source.py Co-authored-by: Ayush Shah <ayush@getcollate.io> * Update ingestion/src/metadata/ingestion/source/database/common_db_source.py Co-authored-by: Ayush Shah <ayush@getcollate.io> * Update bootstrap/sql/com.mysql.cj.jdbc.Driver/v007__create_db_connection_info.sql Co-authored-by: Ayush Shah <ayush@getcollate.io>
2025-12-05 03:54:23 +00:00 · 2022-12-16 07:43:18 +01:00 · 2022-12-16 07:43:18 +01:00 · 819001182f
commit 819001182f
parent a633befe2a
13 changed files with 43 additions and 23 deletions
--- a/bootstrap/sql/com.mysql.cj.jdbc.Driver/v007__create_db_connection_info.sql
+++ b/bootstrap/sql/com.mysql.cj.jdbc.Driver/v007__create_db_connection_info.sql
@ -36,3 +36,10 @@ CREATE TABLE IF NOT EXISTS alert_action_def (
    PRIMARY KEY (id),
    UNIQUE (name)
 );
+
+UPDATE dbservice_entity
+SET json = JSON_INSERT(
+        JSON_REMOVE(json, '$.connection.config.databaseSchema'),
+        '$.connection.config.database',
+        JSON_EXTRACT(json, '$.connection.config.databaseSchema')
+    ) where serviceType in ('Db2');
--- a/bootstrap/sql/org.postgresql.Driver/v007__create_db_connection_info.sql
+++ b/bootstrap/sql/org.postgresql.Driver/v007__create_db_connection_info.sql
@ -29,4 +29,13 @@ CREATE TABLE IF NOT EXISTS alert_action_def (
    json JSONB NOT NULL,
    PRIMARY KEY (id),
    UNIQUE (name)
-);
+);
+
+UPDATE dbservice_entity
+SET json = jsonb_set(json, '{connection,config,database}', json#>'{connection,config,databaseSchema}')
+where serviceType in ('Db2')
+  and json#>'{connection,config,databaseSchema}' is not null;
+
+UPDATE dbservice_entity
+SET json = json::jsonb #- '{connection,config,databaseSchema}'
+where serviceType in ('Db2');
--- a/ingestion/setup.py
+++ b/ingestion/setup.py
@ -150,7 +150,7 @@ plugins: Dict[str, Set[str]] = {
    "okta": {"okta~=2.3.0"},
    "mlflow": {"mlflow-skinny~=1.26.1"},
    "sklearn": {"scikit-learn==1.0.2"},
-    "db2": {"ibm-db-sa==0.3.7"},
+    "db2": {"ibm-db-sa==0.3.8"},
    "clickhouse": {"clickhouse-driver==0.2.5", "clickhouse-sqlalchemy==0.2.3"},
    "databricks": {"sqlalchemy-databricks==0.1.0"},
    "singlestore": {"pymysql>=1.0.2"},
--- a/ingestion/src/metadata/ingestion/source/database/common_db_source.py
+++ b/ingestion/src/metadata/ingestion/source/database/common_db_source.py
@ -175,7 +175,11 @@ class CommonDbSourceService(
                f"Table description error for table [{schema_name}.{table_name}]: {exc}"
            )
        else:
-            description = table_info["text"]
+            if hasattr(table_info, "text"):
+                description = table_info["text"]
+                # DB2 connector does not return a str type
+                if isinstance(description, list):
+                    description = description[0]
        return description

    def get_tables_name_and_type(self) -> Optional[Iterable[Tuple[str, str]]]:
--- a/ingestion/src/metadata/orm_profiler/orm/functions/conn_test.py
+++ b/ingestion/src/metadata/orm_profiler/orm/functions/conn_test.py
@ -47,5 +47,6 @@ def _(*_, **__):


@compiles(ConnTestFn, Dialects.Db2)
+@compiles(ConnTestFn, Dialects.IbmDbSa)
 def _(*_, **__):
    return "SELECT 42 FROM SYSIBM.SYSDUMMY1;"
--- a/ingestion/src/metadata/orm_profiler/orm/registry.py
+++ b/ingestion/src/metadata/orm_profiler/orm/registry.py
@ -60,6 +60,7 @@ class Dialects(Enum):
    AzureSQL = "azuresql"
    SingleStore = "singlestore"
    SQLite = "sqlite"
+    IbmDbSa = "ibm_db_sa"


 # Sometimes we want to skip certain types for computing metrics.
--- a/ingestion/tests/unit/test_source_connection.py
+++ b/ingestion/tests/unit/test_source_connection.py
@ -512,21 +512,13 @@ class SouceConnectionTest(TestCase):
        assert expected_url == get_connection_url(singleStore_conn_obj)

    def test_db2_url(self):
-        # connection arguments without db
-        expected_url = "db2+ibm_db://openmetadata_user:@localhost:50000"
-        db2_conn_obj = Db2Connection(
-            scheme=Db2Scheme.db2_ibm_db,
-            username="openmetadata_user",
-            hostPort="localhost:50000",
-        )
-        assert expected_url == get_connection_url(db2_conn_obj)
-
        # connection arguments with db
-        expected_url = "db2+ibm_db://openmetadata_user:@localhost:50000"
+        expected_url = "db2+ibm_db://openmetadata_user:@localhost:50000/testdb"
        db2_conn_obj = Db2Connection(
+            scheme=Db2Scheme.db2_ibm_db,
            username="openmetadata_user",
            hostPort="localhost:50000",
-            scheme=Db2Scheme.db2_ibm_db,
+            database="testdb",
        )
        assert expected_url == get_connection_url(db2_conn_obj)

@ -720,6 +712,7 @@ class SouceConnectionTest(TestCase):
            hostPort="localhost:443",
            connectionArguments=None,
            scheme=Db2Scheme.db2_ibm_db,
+            database="testdb",
        )
        assert expected_args == get_connection_args(db2_conn_obj)

@ -731,6 +724,7 @@ class SouceConnectionTest(TestCase):
            hostPort="localhost:443",
            connectionArguments={"user": "user-to-be-impersonated"},
            scheme=Db2Scheme.db2_ibm_db,
+            database="testdb",
        )
        assert expected_args == get_connection_args(db2_conn_obj)

--- a/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/database/azureSQLConnection.json
+++ b/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/database/azureSQLConnection.json
@ -76,7 +76,7 @@
      "$ref": "../connectionBasicType.json#/definitions/supportsProfiler"
    },
    "supportsDatabase": {
-      "title": "Supports Profiler",
+      "title": "Supports Database",
      "$ref": "../connectionBasicType.json#/definitions/supportsDatabase"
    }
  },
--- a/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/database/db2Connection.json
+++ b/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/database/db2Connection.json
@ -48,9 +48,9 @@
      "description": "Host and port of the DB2 service.",
      "type": "string"
    },
-    "databaseSchema": {
-      "title": "databaseSchema",
-      "description": "databaseSchema of the data source. This is optional parameter, if you would like to restrict the metadata reading to a single databaseSchema. When left blank, OpenMetadata Ingestion attempts to scan all the databaseSchema.",
+    "database": {
+      "title": "database",
+      "description": "Database of the data source.",
      "type": "string"
    },
    "connectionOptions": {
@ -69,11 +69,15 @@
      "title": "Supports Profiler",
      "$ref": "../connectionBasicType.json#/definitions/supportsProfiler"
    },
+    "supportsDatabase": {
+      "title": "Supports Database",
+      "$ref": "../connectionBasicType.json#/definitions/supportsDatabase"
+    },
    "supportsQueryComment": {
      "title": "Supports Query Comment",
      "$ref": "../connectionBasicType.json#/definitions/supportsQueryComment"
    }
  },
  "additionalProperties": false,
-  "required": ["hostPort", "username"]
+  "required": ["hostPort", "username", "database"]
 }
--- a/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/database/postgresConnection.json
+++ b/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/database/postgresConnection.json
@ -76,7 +76,7 @@
      "$ref": "../connectionBasicType.json#/definitions/supportsProfiler"
    },
    "supportsDatabase": {
-      "title": "Supports Profiler",
+      "title": "Supports Database",
      "$ref": "../connectionBasicType.json#/definitions/supportsDatabase"
    },
    "supportsQueryComment": {
--- a/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/database/prestoConnection.json
+++ b/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/database/prestoConnection.json
@ -75,7 +75,7 @@
      "$ref": "../connectionBasicType.json#/definitions/supportsProfiler"
    },
    "supportsDatabase": {
-      "title": "Supports Profiler",
+      "title": "Supports Database",
      "$ref": "../connectionBasicType.json#/definitions/supportsDatabase"
    },
    "supportsQueryComment": {
--- a/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/database/redshiftConnection.json
+++ b/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/database/redshiftConnection.json
@ -82,7 +82,7 @@
      "$ref": "../connectionBasicType.json#/definitions/supportsProfiler"
    },
    "supportsDatabase": {
-      "title": "Supports Profiler",
+      "title": "Supports Database",
      "$ref": "../connectionBasicType.json#/definitions/supportsDatabase"
    },
    "supportsQueryComment": {
--- a/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/database/snowflakeConnection.json
+++ b/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/database/snowflakeConnection.json
@ -103,7 +103,7 @@
      "$ref": "../connectionBasicType.json#/definitions/supportsProfiler"
    },
    "supportsDatabase": {
-      "title": "Supports Profiler",
+      "title": "Supports Database",
      "$ref": "../connectionBasicType.json#/definitions/supportsDatabase"
    },
    "supportsQueryComment": {