From d5758b2fb2833f9c03cf87d9ce85410a20b7fb01 Mon Sep 17 00:00:00 2001
From: Onkar Ravgan <onkar.10r@gmail.com>
Date: Sat, 9 Apr 2022 16:30:18 +0530
Subject: [PATCH] Fixed Db2 source config (#3978)

Co-authored-by: Onkar Ravgan <onkarravgan@Onkars-MacBook-Pro.local>
---
 ingestion/examples/workflows/db2.json         | 40 +++++++++----------
 .../src/metadata/ingestion/source/db2.py      | 17 ++++----
 .../src/metadata/utils/source_connections.py  |  4 ++
 3 files changed, 34 insertions(+), 27 deletions(-)

diff --git a/ingestion/examples/workflows/db2.json b/ingestion/examples/workflows/db2.json
index 11deff095a2..adac0a9c05a 100644
--- a/ingestion/examples/workflows/db2.json
+++ b/ingestion/examples/workflows/db2.json
@@ -1,26 +1,26 @@
 {
-    "source": {
-      "type": "db2",
+  "source": {
+    "type": "db2",
+    "serviceName": "local_db2",
+    "serviceConnection": {
       "config": {
-        "username": "db2inst1",
-        "password": "password",
-        "database": "metadata",
-        "service_name": "local_db2",
-        "table_filter_pattern": {
-          "excludes": []
-        }
+        "type": "Db2",
+        "username": "openmetadata_user",
+        "password": "openmetadata_password",
+        "hostPort": "localhost:50000",
+        "database": "custom_database_name"
       }
     },
-    "sink": {
-      "type": "metadata-rest",
-      "config": {}
-    },
-    "metadata_server": {
-      "type": "metadata-server",
-      "config": {
-        "api_endpoint": "http://localhost:8585/api",
-        "auth_provider_type": "no-auth"
-      }
+    "sourceConfig": {"config": {"enableDataProfiler": false}}
+  },
+  "sink": {
+    "type": "metadata-rest",
+    "config": {}
+  },
+  "workflowConfig": {
+    "openMetadataServerConfig": {
+      "hostPort": "http://localhost:8585/api",
+      "authProvider": "no-auth"
     }
   }
-  
\ No newline at end of file
+}
diff --git a/ingestion/src/metadata/ingestion/source/db2.py b/ingestion/src/metadata/ingestion/source/db2.py
index 6acbe87a8c0..0645354fdc1 100644
--- a/ingestion/src/metadata/ingestion/source/db2.py
+++ b/ingestion/src/metadata/ingestion/source/db2.py
@@ -15,8 +15,11 @@ from sqlalchemy.engine import reflection
 from metadata.generated.schema.metadataIngestion.workflow import (
     OpenMetadataServerConfig,
 )
+from metadata.generated.schema.metadataIngestion.workflow import (
+    Source as WorkflowSource,
+)
+from metadata.ingestion.api.source import InvalidSourceException
 from metadata.ingestion.source.sql_source import SQLSource
-from metadata.ingestion.source.sql_source_common import SQLConnectionConfig
 
 
 @reflection.cache
@@ -30,16 +33,16 @@ from metadata.generated.schema.entity.services.connections.database.db2Connectio
 )
 
 
-class Db2Config(DB2Connection, SQLConnectionConfig):
-    def get_connection_url(self):
-        return super().get_connection_url()
-
-
 class Db2Source(SQLSource):
     def __init__(self, config, metadata_config):
         super().__init__(config, metadata_config)
 
     @classmethod
     def create(cls, config_dict, metadata_config: OpenMetadataServerConfig):
-        config = Db2Config.parse_obj(config_dict)
+        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
+        connection: DB2Connection = config.serviceConnection.__root__.config
+        if not isinstance(connection, DB2Connection):
+            raise InvalidSourceException(
+                f"Expected DB2Connection, but got {connection}"
+            )
         return cls(config, metadata_config)
diff --git a/ingestion/src/metadata/utils/source_connections.py b/ingestion/src/metadata/utils/source_connections.py
index 8aa968131f3..4f799f69753 100644
--- a/ingestion/src/metadata/utils/source_connections.py
+++ b/ingestion/src/metadata/utils/source_connections.py
@@ -22,6 +22,9 @@ from metadata.generated.schema.entity.services.connections.database.clickhouseCo
 from metadata.generated.schema.entity.services.connections.database.databricksConnection import (
     DatabricksConnection,
 )
+from metadata.generated.schema.entity.services.connections.database.db2Connection import (
+    DB2Connection,
+)
 from metadata.generated.schema.entity.services.connections.database.hiveConnection import (
     HiveSQLConnection,
 )
@@ -97,6 +100,7 @@ def get_connection_url(connection):
 @get_connection_url.register(ClickhouseConnection)
 @get_connection_url.register(SingleStoreConnection)
 @get_connection_url.register(VerticaConnection)
+@get_connection_url.register(DB2Connection)
 def _(connection):
     return get_connection_url_common(connection)