From 2d25da149db079febd7be03c1b85e17ee2a275bb Mon Sep 17 00:00:00 2001
From: Mayur Singal <39544459+ulixius9@users.noreply.github.com>
Date: Tue, 8 Feb 2022 00:19:59 +0530
Subject: [PATCH] Fix #2626: add db2 connector (#2661)

Co-authored-by: Mayur SIngal <mayursingal@Mayurs-MacBook-Pro.local>
---
 .../entity/services/databaseService.json      |  6 ++-
 ingestion/examples/workflows/db2.json         | 26 +++++++++++
 ingestion/setup.py                            |  1 +
 .../src/metadata/ingestion/source/db2.py      | 45 +++++++++++++++++++
 4 files changed, 77 insertions(+), 1 deletion(-)
 create mode 100644 ingestion/examples/workflows/db2.json
 create mode 100644 ingestion/src/metadata/ingestion/source/db2.py

diff --git a/catalog-rest-service/src/main/resources/json/schema/entity/services/databaseService.json b/catalog-rest-service/src/main/resources/json/schema/entity/services/databaseService.json
index 938ea3542c8..cb77a1cc4ea 100644
--- a/catalog-rest-service/src/main/resources/json/schema/entity/services/databaseService.json
+++ b/catalog-rest-service/src/main/resources/json/schema/entity/services/databaseService.json
@@ -23,7 +23,8 @@
         "Vertica",
         "Glue",
         "MariaDB",
-        "Druid"
+        "Druid",
+        "Db2"
       ],
       "javaEnums": [
         {
@@ -70,6 +71,9 @@
         },
         {
           "name": "Druid"
+        },
+        {
+          "name": "Db2"
         }
       ]
     },
diff --git a/ingestion/examples/workflows/db2.json b/ingestion/examples/workflows/db2.json
new file mode 100644
index 00000000000..11deff095a2
--- /dev/null
+++ b/ingestion/examples/workflows/db2.json
@@ -0,0 +1,26 @@
+{
+    "source": {
+      "type": "db2",
+      "config": {
+        "username": "db2inst1",
+        "password": "password",
+        "database": "metadata",
+        "service_name": "local_db2",
+        "table_filter_pattern": {
+          "excludes": []
+        }
+      }
+    },
+    "sink": {
+      "type": "metadata-rest",
+      "config": {}
+    },
+    "metadata_server": {
+      "type": "metadata-server",
+      "config": {
+        "api_endpoint": "http://localhost:8585/api",
+        "auth_provider_type": "no-auth"
+      }
+    }
+  }
+  
\ No newline at end of file
diff --git a/ingestion/setup.py b/ingestion/setup.py
index d23fa52ba8d..46b4d1cfa7a 100644
--- a/ingestion/setup.py
+++ b/ingestion/setup.py
@@ -112,6 +112,7 @@ plugins: Dict[str, Set[str]] = {
     "okta": {"okta~=2.3.0"},
     "mlflow": {"mlflow-skinny~=1.22.0"},
     "sklearn": {"scikit-learn==1.0.2"},
+    "db2":{"ibm-db-sa==0.3.7"},
 }
 dev = {
     "boto3==1.20.14",
diff --git a/ingestion/src/metadata/ingestion/source/db2.py b/ingestion/src/metadata/ingestion/source/db2.py
new file mode 100644
index 00000000000..b1f32455793
--- /dev/null
+++ b/ingestion/src/metadata/ingestion/source/db2.py
@@ -0,0 +1,45 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from ibm_db_sa.base import DB2Dialect
+from sqlalchemy.engine import reflection
+
+from metadata.ingestion.ometa.openmetadata_rest import MetadataServerConfig
+from metadata.ingestion.source.sql_source import SQLSource
+from metadata.ingestion.source.sql_source_common import SQLConnectionConfig
+
+
+@reflection.cache
+def get_pk_constraint(self, bind, table_name, schema=None, **kw):
+    return {"constrained_columns": [], "name": "undefined"}
+
+
+DB2Dialect.get_pk_constraint = get_pk_constraint
+
+
+class Db2Config(SQLConnectionConfig):
+    host_port = "localhost:50000"
+    scheme = "db2+ibm_db"
+    service_type = "Db2"
+
+    def get_connection_url(self):
+        return super().get_connection_url()
+
+
+class Db2Source(SQLSource):
+    def __init__(self, config, metadata_config, ctx):
+        super().__init__(config, metadata_config, ctx)
+
+    @classmethod
+    def create(cls, config_dict, metadata_config_dict, ctx):
+        config = Db2Config.parse_obj(config_dict)
+        metadata_config = MetadataServerConfig.parse_obj(metadata_config_dict)
+        return cls(config, metadata_config, ctx)