Skip to content

Commit

Permalink
use service connection for nosql adaptor factory
Browse files Browse the repository at this point in the history
  • Loading branch information
sushi30 committed Apr 19, 2024
1 parent 23b3f30 commit bb56abd
Show file tree
Hide file tree
Showing 3 changed files with 24 additions and 17 deletions.
29 changes: 15 additions & 14 deletions ingestion/src/metadata/profiler/adaptors/factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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.
Expand All @@ -56,20 +54,23 @@ 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.
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)
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
]
Expand Down
2 changes: 1 addition & 1 deletion ingestion/tests/integration/profiler/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down

0 comments on commit bb56abd

Please sign in to comment.