From 05a8401c51715c13c676f8a94d27a9159c18465b Mon Sep 17 00:00:00 2001
From: Mayur Singal <39544459+ulixius9@users.noreply.github.com>
Date: Tue, 10 May 2022 18:33:06 +0530
Subject: [PATCH] Fix #4831: Added multiple db support for mssql (#4836)

---
 .../src/metadata/ingestion/source/mssql.py    | 38 +++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/ingestion/src/metadata/ingestion/source/mssql.py b/ingestion/src/metadata/ingestion/source/mssql.py
index 98b086e26ea..b1b95fdb584 100644
--- a/ingestion/src/metadata/ingestion/source/mssql.py
+++ b/ingestion/src/metadata/ingestion/source/mssql.py
@@ -10,6 +10,11 @@
 #  limitations under the License.
 """MSSQL source module"""
 
+from typing import Iterable
+
+from sqlalchemy.engine.reflection import Inspector
+from sqlalchemy.inspection import inspect
+
 from metadata.generated.schema.entity.services.connections.database.mssqlConnection import (
     MssqlConnection,
 )
@@ -24,6 +29,11 @@ from metadata.generated.schema.metadataIngestion.workflow import (
 )
 from metadata.ingestion.api.source import InvalidSourceException
 from metadata.ingestion.source.sql_source import SQLSource
+from metadata.utils.connections import get_connection
+from metadata.utils.filters import filter_by_database
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
 
 
 class MssqlSource(SQLSource):
@@ -49,3 +59,31 @@ class MssqlSource(SQLSource):
                 "select top 50 * from [{}].[{}]"
             )
         return cls(config, metadata_config)
+
+    def get_databases(self) -> Iterable[Inspector]:
+        if self.service_connection.database:
+            yield from super().get_databases()
+        else:
+            query = "SELECT name FROM master.sys.databases order by name;"
+            results = self.connection.execute(query)
+            db_list = []
+
+            for res in results:
+                db_list.append(list(res))
+
+            for row in db_list:
+                try:
+                    if filter_by_database(
+                        self.source_config.databaseFilterPattern, database_name=row[0]
+                    ):
+                        self.status.filter(row[0], "Database pattern not allowed")
+                        continue
+                    logger.info(f"Ingesting from database: {row[0]}")
+                    self.service_connection.database = row[0]
+                    self.engine = get_connection(
+                        self.config.serviceConnection.__root__.config
+                    )
+                    self.engine.connect()
+                    yield inspect(self.engine)
+                except Exception as err:
+                    logger.error(f"Failed to Connect: {row[0]} due to error {err}")