From bb56abd4af825c96651db5fdf0039dba33e5598b Mon Sep 17 00:00:00 2001
From: sushi30 <imri.paran@gmail.com>
Date: Fri, 19 Apr 2024 12:31:07 +0200
Subject: [PATCH] use service connection for nosql adaptor factory

---
 .../src/metadata/profiler/adaptors/factory.py | 29 ++++++++++---------
 .../interface/nosql/profiler_interface.py     | 10 +++++--
 .../tests/integration/profiler/conftest.py    |  2 +-
 3 files changed, 24 insertions(+), 17 deletions(-)

diff --git a/ingestion/src/metadata/profiler/adaptors/factory.py b/ingestion/src/metadata/profiler/adaptors/factory.py
index 273f6646d111..50005158b582 100644
--- a/ingestion/src/metadata/profiler/adaptors/factory.py
+++ b/ingestion/src/metadata/profiler/adaptors/factory.py
@@ -13,8 +13,12 @@
 """
 from typing import Callable
 
-from pymongo import MongoClient
-
+from metadata.generated.schema.entity.services.connections.database.dynamoDBConnection import (
+    DynamoDBConnection,
+)
+from metadata.generated.schema.entity.services.connections.database.mongoDBConnection import (
+    MongoDBConnection,
+)
 from metadata.profiler.adaptors.dynamodb import DynamoDB
 from metadata.profiler.adaptors.mongodb import MongoDB
 from metadata.profiler.adaptors.nosql_adaptor import NoSQLAdaptor
@@ -32,12 +36,6 @@ class NoSQLAdaptorFactory(Factory):
     and can be created using the `construct` method.
     """
 
-    def __init__(self):
-        """
-        Initialize a new instance of NoSQLClientFactory.
-        """
-        super().__init__()
-
     def register(self, interface_type: str, interface_class: NoSQLAdaptorConstructor):
         """
         Register a client type with its constructor.
@@ -56,7 +54,7 @@ def create(self, interface_type: any, *args, **kwargs) -> NoSQLAdaptor:
         Create a client instance of the type of the given source client.
 
         Args:
-            interface_type (any): The source client instance.
+            interface_type (str): The type of the source connection.
 
         Returns:
             NoSQLAdaptor: The created client instance.
@@ -64,12 +62,15 @@ def create(self, interface_type: any, *args, **kwargs) -> NoSQLAdaptor:
         Raises:
             ValueError: If the type of the source client is not registered.
         """
-        client_class = self._interface_type.get(type(interface_type).__name__)
+        client_class = self._interface_type.get(interface_type)
         if not client_class:
-            raise ValueError(f"Unknown NoSQL source: {type(interface_type).__name__}")
-        return client_class(interface_type)
+            raise ValueError(f"Unknown NoSQL source: {interface_type}")
+        return client_class(*args, **kwargs)
 
 
 factory = NoSQLAdaptorFactory()
-factory.register(type(MongoClient).__name__, MongoDB)
-factory.register("dynamodb.ServiceResource", DynamoDB)
+adaptors = {
+    MongoDBConnection.__name__: MongoDB,
+    DynamoDBConnection.__name__: DynamoDB,
+}
+factory.register_many(adaptors)
diff --git a/ingestion/src/metadata/profiler/interface/nosql/profiler_interface.py b/ingestion/src/metadata/profiler/interface/nosql/profiler_interface.py
index 62f6b2b00f36..f90389b9d84c 100644
--- a/ingestion/src/metadata/profiler/interface/nosql/profiler_interface.py
+++ b/ingestion/src/metadata/profiler/interface/nosql/profiler_interface.py
@@ -165,7 +165,10 @@ def _get_sampler(self) -> NoSQLSampler:
         return sampler_factory_.create(
             self.service_connection_config.__class__.__name__,
             table=self.table,
-            client=factory.create(self.connection),
+            client=factory.create(
+                self.service_connection_config.config.__class__.__name__,
+                client=self.connection,
+            ),
             profile_sample_config=self.profile_sample_config,
             partition_details=self.partition_details,
             profile_sample_query=self.profile_query,
@@ -187,7 +190,10 @@ def get_all_metrics(
     ):
         """get all profiler metrics"""
         profile_results = {"table": {}, "columns": defaultdict(dict)}
-        runner = factory.create(self.connection)
+        runner = factory.create(
+            self.service_connection_config.config.__class__.__name__,
+            client=self.connection,
+        )
         metric_list = [
             self.compute_metrics(runner, metric_func) for metric_func in metric_funcs
         ]
diff --git a/ingestion/tests/integration/profiler/conftest.py b/ingestion/tests/integration/profiler/conftest.py
index e92906bf32bf..fb29d005652b 100644
--- a/ingestion/tests/integration/profiler/conftest.py
+++ b/ingestion/tests/integration/profiler/conftest.py
@@ -2,7 +2,7 @@
 
 import boto3
 import pytest
-from integration.integration_base import int_admin_ometa
+from ...integration.integration_base import int_admin_ometa
 from testcontainers.localstack import LocalStackContainer
 
 from metadata.generated.schema.api.services.createDatabaseService import (