openmetadata-ingestion 1.7.5.3__py3-none-any.whl → 1.8.1.0__py3-none-any.whl

metadata/ingestion/source/database/trino/connection.py

@@ -13,7 +13,7 @@
 Source connection handler
 """
 from copy import deepcopy
-from typing import Optional, cast
+from typing import Optional
 from urllib.parse import quote_plus
 
 from requests import Session
@@ -25,7 +25,6 @@ from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
 )
 from metadata.generated.schema.entity.services.connections.database.common import (
-    azureConfig,
     basicAuth,
     jwtAuth,
     noConfigAuthenticationTypes,
@@ -148,11 +147,17 @@ def get_connection(connection: TrinoConnection) -> Engine:
         )
         connection.connectionArguments.root["verify"] = {"verify": connection.verify}
     if hasattr(connection.authType, "azureConfig"):
-        access_token_obj = get_azure_token(connection)
-
+        azure_client = AzureClient(connection.authType.azureConfig).create_client()
+        if not connection.authType.azureConfig.scopes:
+            raise ValueError(
+                "Azure Scopes are missing, please refer https://learn.microsoft.com/en-gb/azure/mysql/flexible-server/how-to-azure-ad#2---retrieve-microsoft-entra-access-token and fetch the resource associated with it, for e.g. https://ossrdbms-aad.database.windows.net/.default"
+            )
+        access_token_obj = azure_client.get_token(
+            *connection.authType.azureConfig.scopes.split(",")
+        )
         if not connection.connectionOptions:
             connection.connectionOptions = init_empty_connection_options()
-        connection.connectionOptions.root["access_token"] = access_token_obj
+        connection.connectionOptions.root["access_token"] = access_token_obj.token
     return create_generic_db_connection(
         connection=connection_copy,
         get_connection_url_fn=get_connection_url,
@@ -160,80 +165,6 @@ def get_connection(connection: TrinoConnection) -> Engine:
     )
 
 
-def get_connection_dict(connection: TrinoConnection) -> dict:
-    """
-    Create connection
-    """
-    url = get_connection(connection).url
-    connection_copy = deepcopy(connection)
-
-    connection_dict = {
-        "driver": url.drivername,
-        "host": url.host,
-        "port": url.port,
-        "user": url.username,
-        "catalog": url.database,
-        "schema": url.query.get("schema"),
-    }
-    connection_dict.update(url.query)
-
-    if connection_copy.proxies:
-        connection_dict["http_session"] = connection_copy.proxies
-
-    if connection_copy.connectionArguments and connection_copy.connectionArguments.root:
-        connection_with_options_secrets(lambda: connection_copy)
-        connection_dict.update(get_connection_args_common(connection_copy))
-
-    if isinstance(connection_copy.authType, basicAuth.BasicAuth):
-        connection_dict["auth"] = {
-            "authType": "basic",
-            "username": connection.username,
-            "password": connection.authType.password.get_secret_value()
-            if connection.authType.password
-            else None,
-        }
-        connection_dict["http_scheme"] = "https"
-
-    elif isinstance(connection_copy.authType, jwtAuth.JwtAuth):
-        connection_dict["auth"] = {
-            "authType": "jwt",
-            "jwt": connection.authType.jwt.get_secret_value(),
-        }
-        connection_dict["http_scheme"] = "https"
-
-    elif hasattr(connection_copy.authType, "azureConfig"):
-        connection_dict["auth"] = {
-            "authType": "jwt",
-            "jwt": get_azure_token(connection_copy),
-        }
-        connection_dict["http_scheme"] = "https"
-
-    elif (
-        connection_copy.authType
-        == noConfigAuthenticationTypes.NoConfigAuthenticationTypes.OAuth2
-    ):
-        connection_dict["auth"] = {"authType": "oauth2"}
-        connection_dict["http_scheme"] = "https"
-
-    return connection_dict
-
-
-def get_azure_token(connection: TrinoConnection) -> str:
-    """
-    Get the azure token for the trino connection
-    """
-    auth_type = cast(azureConfig.AzureConfigurationSource, connection.authType)
-
-    if not auth_type.azureConfig.scopes:
-        raise ValueError(
-            "Azure Scopes are missing, please refer https://learn.microsoft.com/en-gb/azure/mysql/flexible-server/how-to-azure-ad#2---retrieve-microsoft-entra-access-token and fetch the resource associated with it, for e.g. https://ossrdbms-aad.database.windows.net/.default"
-        )
-
-    azure_client = AzureClient(auth_type.azureConfig).create_client()
-
-    return azure_client.get_token(*auth_type.azureConfig.scopes.split(",")).token
-
-
 def test_connection(
     metadata: OpenMetadata,
     engine: Engine,

metadata/ingestion/source/database/trino/profiler/system_tables_profiler.py

@@ -13,7 +13,7 @@ System table profiler
 """
 from datetime import datetime
 from decimal import Decimal
-from typing import Any, Dict, List, Optional, Set, Union
+from typing import Any, Dict, List, Optional, Set, Type, Union
 
 from more_itertools import partition
 from pydantic import field_validator
@@ -25,7 +25,11 @@ from metadata.profiler.interface.sqlalchemy.stored_statistics_profiler import (
     StoredStatisticsSource,
 )
 from metadata.profiler.metrics.core import Metric
-from metadata.profiler.metrics.registry import Metrics
+from metadata.profiler.registry import MetricRegistry
+from metadata.utils.dependency_injector.dependency_injector import (
+    Inject,
+    inject_class_attributes,
+)
 from metadata.utils.logger import profiler_logger
 from metadata.utils.lru_cache import LRU_CACHE_SIZE, LRUCache
 from metadata.utils.ssl_manager import get_ssl_connection
@@ -61,23 +65,28 @@ class TableStats(BaseModel):
     columns: Dict[str, ColumnStats] = {}
 
 
+@inject_class_attributes
 class TrinoStoredStatisticsSource(StoredStatisticsSource):
     """Trino system profile source"""
 
-    metric_stats_map: Dict[Metrics, str] = {
-        Metrics.NULL_RATIO: "nulls_fractions",
-        Metrics.DISTINCT_COUNT: "distinct_values_count",
-        Metrics.ROW_COUNT: "row_count",
-        Metrics.MAX: "high_value",
-        Metrics.MIN: "low_value",
-    }
+    metrics: Inject[Type[MetricRegistry]]
+
+    @classmethod
+    def get_metric_stats_map(cls) -> Dict[MetricRegistry, str]:
+        return {
+            cls.metrics.NULL_RATIO: "nulls_fractions",
+            cls.metrics.DISTINCT_COUNT: "distinct_values_count",
+            cls.metrics.ROW_COUNT: "row_count",
+            cls.metrics.MAX: "high_value",
+            cls.metrics.MIN: "low_value",
+        }
 
-    metric_stats_by_name: Dict[str, str] = {
-        k.name: v for k, v in metric_stats_map.items()
-    }
+    @classmethod
+    def get_metric_stats_by_name(cls) -> Dict[str, str]:
+        return {k.name: v for k, v in cls.get_metric_stats_map().items()}
 
-    def get_statistics_metrics(self) -> Set[Metrics]:
-        return set(self.metric_stats_map.keys())
+    def get_statistics_metrics(self) -> Set[MetricRegistry]:
+        return set(self.get_metric_stats_map().keys())
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
@@ -96,7 +105,7 @@ class TrinoStoredStatisticsSource(StoredStatisticsSource):
                 f"Column {column} not found in table {table_name}. Statistics might be stale or missing."
             )
         result = {
-            m.name(): getattr(column_stats, self.metric_stats_by_name[m.name()])
+            m.name(): getattr(column_stats, self.get_metric_stats_by_name()[m.name()])
             for m in metric
         }
         result.update(self.get_hybrid_statistics(table_stats, column_stats))
@@ -108,7 +117,7 @@ class TrinoStoredStatisticsSource(StoredStatisticsSource):
     ) -> dict:
         table_stats = self._get_cached_stats(schema, table_name)
         return {
-            m.name(): getattr(table_stats, self.metric_stats_by_name[m.name()])
+            m.name(): getattr(table_stats, self.get_metric_stats_by_name()[m.name()])
             for m in metric
         }
 
@@ -159,7 +168,7 @@ class TrinoStoredStatisticsSource(StoredStatisticsSource):
     ) -> Dict[str, Any]:
         return {
             # trino stats are in fractions, so we need to convert them to counts (unlike our default profiler)
-            Metrics.NULL_COUNT.name: (
+            self.metrics.NULL_COUNT.name: (
                 int(table_stats.row_count * column_stats.nulls_fraction)
                 if None not in [table_stats.row_count, column_stats.nulls_fraction]
                 else None

metadata/ingestion/source/database/unitycatalog/connection.py

@@ -36,6 +36,12 @@ from metadata.ingestion.connections.test_connections import test_connection_step
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.database.unitycatalog.client import UnityCatalogClient
 from metadata.ingestion.source.database.unitycatalog.models import DatabricksTable
+from metadata.ingestion.source.database.unitycatalog.queries import (
+    UNITY_CATALOG_GET_ALL_SCHEMA_TAGS,
+    UNITY_CATALOG_GET_ALL_TABLE_COLUMNS_TAGS,
+    UNITY_CATALOG_GET_ALL_TABLE_TAGS,
+    UNITY_CATALOG_GET_CATALOGS_TAGS,
+)
 from metadata.utils.constants import THREE_MIN
 from metadata.utils.db_utils import get_host_from_host_port
 from metadata.utils.logger import ingestion_logger
@@ -111,6 +117,32 @@ def test_connection(
                 table_obj.name = table.name
                 break
 
+    def get_tags(
+        service_connection: UnityCatalogConnection, table_obj: DatabricksTable
+    ):
+        engine = get_sqlalchemy_connection(service_connection)
+        with engine.connect() as connection:
+            connection.execute(
+                UNITY_CATALOG_GET_CATALOGS_TAGS.format(
+                    database=table_obj.catalog_name
+                ).replace(";", " limit 1;")
+            )
+            connection.execute(
+                UNITY_CATALOG_GET_ALL_SCHEMA_TAGS.format(
+                    database=table_obj.catalog_name
+                ).replace(";", " limit 1;")
+            )
+            connection.execute(
+                UNITY_CATALOG_GET_ALL_TABLE_TAGS.format(
+                    database=table_obj.catalog_name, schema=table_obj.schema_name
+                ).replace(";", " limit 1;")
+            )
+            connection.execute(
+                UNITY_CATALOG_GET_ALL_TABLE_COLUMNS_TAGS.format(
+                    database=table_obj.catalog_name, schema=table_obj.schema_name
+                ).replace(";", " limit 1;")
+            )
+
     test_fn = {
         "CheckAccess": connection.catalogs.list,
         "GetDatabases": partial(get_catalogs, connection, table_obj),
@@ -118,6 +150,7 @@ def test_connection(
         "GetTables": partial(get_tables, connection, table_obj),
         "GetViews": partial(get_tables, connection, table_obj),
         "GetQueries": client.test_query_api_access,
+        "GetTags": partial(get_tags, service_connection, table_obj),
     }
 
     return test_connection_steps(

metadata/ingestion/source/database/unitycatalog/lineage.py

@@ -33,12 +33,9 @@ from metadata.generated.schema.type.entityLineage import Source as LineageSource
 from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException, Source
-from metadata.ingestion.connections.test_connections import (
-    raise_test_connection_exception,
-)
 from metadata.ingestion.lineage.sql_lineage import get_column_fqn
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.connections import get_test_connection_fn
+from metadata.ingestion.source.connections import test_connection_common
 from metadata.ingestion.source.database.unitycatalog.client import UnityCatalogClient
 from metadata.ingestion.source.database.unitycatalog.connection import get_connection
 from metadata.ingestion.source.database.unitycatalog.models import LineageTableStreams
@@ -220,8 +217,6 @@ class UnitycatalogLineageSource(Source):
                     )
 
     def test_connection(self) -> None:
-        test_connection_fn = get_test_connection_fn(self.service_connection)
-        result = test_connection_fn(
+        test_connection_common(
             self.metadata, self.connection_obj, self.service_connection
         )
-        raise_test_connection_exception(result)

metadata/ingestion/source/database/unitycatalog/metadata.py

@@ -65,20 +65,33 @@ from metadata.ingestion.source.database.external_table_lineage_mixin import (
 from metadata.ingestion.source.database.multi_db_source import MultiDBSource
 from metadata.ingestion.source.database.stored_procedures_mixin import QueryByProcedure
 from metadata.ingestion.source.database.unitycatalog.client import UnityCatalogClient
-from metadata.ingestion.source.database.unitycatalog.connection import get_connection
+from metadata.ingestion.source.database.unitycatalog.connection import (
+    get_connection,
+    get_sqlalchemy_connection,
+)
 from metadata.ingestion.source.database.unitycatalog.models import (
     ColumnJson,
     ElementType,
     ForeignConstrains,
     Type,
 )
+from metadata.ingestion.source.database.unitycatalog.queries import (
+    UNITY_CATALOG_GET_ALL_SCHEMA_TAGS,
+    UNITY_CATALOG_GET_ALL_TABLE_COLUMNS_TAGS,
+    UNITY_CATALOG_GET_ALL_TABLE_TAGS,
+    UNITY_CATALOG_GET_CATALOGS_TAGS,
+)
 from metadata.utils import fqn
 from metadata.utils.filters import filter_by_database, filter_by_schema, filter_by_table
 from metadata.utils.helpers import retry_with_docker_host
 from metadata.utils.logger import ingestion_logger
+from metadata.utils.tag_utils import get_ometa_tag_and_classification
 
 logger = ingestion_logger()
 
+UNITY_CATALOG_TAG = "UNITY CATALOG TAG"
+UNITY_CATALOG_TAG_CLASSIFICATION = "UNITY CATALOG TAG CLASSIFICATION"
+
 
 class UnitycatalogSource(
     ExternalTableLineageMixin, DatabaseServiceSource, MultiDBSource
@@ -179,10 +192,14 @@ class UnitycatalogSource(
         From topology.
         Prepare a database request and pass it to the sink
         """
+        catalog = self.client.catalogs.get(database_name)
         yield Either(
             right=CreateDatabaseRequest(
                 name=database_name,
                 service=self.context.get().database_service,
+                owners=self.get_owner_ref(catalog.owner),
+                description=catalog.comment,
+                tags=self.get_database_tag_labels(database_name),
             )
         )
 
@@ -227,6 +244,9 @@ class UnitycatalogSource(
         From topology.
         Prepare a database schema request and pass it to the sink
         """
+        schema = self.client.schemas.get(
+            full_name=f"{self.context.get().database}.{schema_name}"
+        )
         yield Either(
             right=CreateDatabaseSchemaRequest(
                 name=EntityName(schema_name),
@@ -238,6 +258,9 @@ class UnitycatalogSource(
                         database_name=self.context.get().database,
                     )
                 ),
+                description=schema.comment,
+                owners=self.get_owner_ref(schema.owner),
+                tags=self.get_schema_tag_labels(schema_name),
             )
         )
 
@@ -315,7 +338,7 @@ class UnitycatalogSource(
                 (db_name, schema_name, table_name)
             ] = table.storage_location
         try:
-            columns = list(self.get_columns(table.columns))
+            columns = list(self.get_columns(table_name, table.columns))
             (
                 primary_constraints,
                 foreign_constraints,
@@ -341,6 +364,7 @@ class UnitycatalogSource(
                     )
                 ),
                 owners=self.get_owner_ref(table.owner),
+                tags=self.get_tag_labels(table_name),
             )
             yield Either(right=table_request)
 
@@ -477,11 +501,12 @@ class UnitycatalogSource(
                 f"Unable to add description to complex datatypes for column [{column.name}]: {exc}"
             )
 
-    def get_columns(self, column_data: List[ColumnInfo]) -> Iterable[Column]:
+    def get_columns(
+        self, table_name: str, column_data: List[ColumnInfo]
+    ) -> Iterable[Column]:
         """
         process table regular columns info
         """
-
         for column in column_data:
             parsed_string = {}
             if column.type_text:
@@ -500,6 +525,9 @@ class UnitycatalogSource(
             parsed_string["dataLength"] = parsed_string.get("dataLength", 1)
             if column.comment:
                 parsed_string["description"] = Markdown(column.comment)
+            parsed_string["tags"] = self.get_column_tag_labels(
+                table_name=table_name, column={"name": column.name}
+            )
             parsed_column = Column(**parsed_string)
             self.add_complex_datatype_descriptions(
                 column=parsed_column,
@@ -507,10 +535,97 @@ class UnitycatalogSource(
             )
             yield parsed_column
 
+    def yield_database_tag(
+        self, database_name: str
+    ) -> Iterable[Either[OMetaTagAndClassification]]:
+        """Get Unity Catalog database/catalog tags using SQL query"""
+        query_tag_fqn_builder_mapping = (
+            (
+                UNITY_CATALOG_GET_CATALOGS_TAGS.format(database=database_name),
+                lambda tag: [self.context.get().database_service, database_name],
+            ),
+            (
+                UNITY_CATALOG_GET_ALL_SCHEMA_TAGS.format(database=database_name),
+                lambda tag: [
+                    self.context.get().database_service,
+                    database_name,
+                    tag.schema_name,
+                ],
+            ),
+        )
+        try:
+            with get_sqlalchemy_connection(
+                self.service_connection
+            ).connect() as connection:
+                for query, tag_fqn_builder in query_tag_fqn_builder_mapping:
+                    for tag in connection.execute(query):
+                        yield from get_ometa_tag_and_classification(
+                            tag_fqn=FullyQualifiedEntityName(
+                                fqn._build(*tag_fqn_builder(tag))
+                            ),  # pylint: disable=protected-access
+                            tags=[tag.tag_value],
+                            classification_name=tag.tag_name,
+                            tag_description=UNITY_CATALOG_TAG,
+                            classification_description=UNITY_CATALOG_TAG_CLASSIFICATION,
+                            metadata=self.metadata,
+                            system_tags=True,
+                        )
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Error getting tags for catalog/schema {database_name}: {exc}"
+            )
+
     def yield_tag(
         self, schema_name: str
     ) -> Iterable[Either[OMetaTagAndClassification]]:
-        """No tags being processed"""
+        """Get Unity Catalog schema tags using SQL query"""
+        database = self.context.get().database
+        query_tag_fqn_builder_mapping = (
+            (
+                UNITY_CATALOG_GET_ALL_TABLE_TAGS.format(
+                    database=database, schema=schema_name
+                ),
+                lambda tag: [
+                    self.context.get().database_service,
+                    database,
+                    schema_name,
+                    tag.table_name,
+                ],
+            ),
+            (
+                UNITY_CATALOG_GET_ALL_TABLE_COLUMNS_TAGS.format(
+                    database=database, schema=schema_name
+                ),
+                lambda tag: [
+                    self.context.get().database_service,
+                    database,
+                    schema_name,
+                    tag.table_name,
+                    tag.column_name,
+                ],
+            ),
+        )
+        try:
+            with get_sqlalchemy_connection(
+                self.service_connection
+            ).connect() as connection:
+                for query, tag_fqn_builder in query_tag_fqn_builder_mapping:
+                    for tag in connection.execute(query):
+                        yield from get_ometa_tag_and_classification(
+                            tag_fqn=FullyQualifiedEntityName(
+                                fqn._build(*tag_fqn_builder(tag))
+                            ),  # pylint: disable=protected-access
+                            tags=[tag.tag_value],
+                            classification_name=tag.tag_name,
+                            tag_description=UNITY_CATALOG_TAG,
+                            classification_description=UNITY_CATALOG_TAG_CLASSIFICATION,
+                            metadata=self.metadata,
+                            system_tags=True,
+                        )
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Error getting tags for schema {schema_name}: {exc}")
 
     def get_stored_procedures(self) -> Iterable[Any]:
         """Not implemented"""
@@ -527,24 +642,22 @@ class UnitycatalogSource(
         """Nothing to close"""
 
     # pylint: disable=arguments-renamed
-    def get_owner_ref(
-        self, table_owner: Optional[str]
-    ) -> Optional[EntityReferenceList]:
+    def get_owner_ref(self, owner: Optional[str]) -> Optional[EntityReferenceList]:
         """
         Method to process the table owners
         """
         if self.source_config.includeOwners is False:
             return None
         try:
-            if not table_owner or not isinstance(table_owner, str):
+            if not owner or not isinstance(owner, str):
                 return None
-            owner_ref = self.metadata.get_reference_by_email(email=table_owner)
+            owner_ref = self.metadata.get_reference_by_email(email=owner)
             if owner_ref:
                 return owner_ref
-            table_name = table_owner.split("@")[0]
-            owner_ref = self.metadata.get_reference_by_name(name=table_name)
+            owner_name = owner.split("@")[0]
+            owner_ref = self.metadata.get_reference_by_name(name=owner_name)
             return owner_ref
         except Exception as exc:
             logger.debug(traceback.format_exc())
-            logger.warning(f"Error processing owner {table_owner}: {exc}")
+            logger.warning(f"Error processing owner {owner}: {exc}")
         return None

metadata/ingestion/source/database/unitycatalog/queries.py

@@ -0,0 +1,19 @@
+"""
+SQL queries for Unity Catalog
+"""
+
+UNITY_CATALOG_GET_CATALOGS_TAGS = """
+SELECT * FROM `{database}`.information_schema.catalog_tags;
+"""
+
+UNITY_CATALOG_GET_ALL_SCHEMA_TAGS = """
+SELECT * FROM `{database}`.information_schema.schema_tags;
+"""
+
+UNITY_CATALOG_GET_ALL_TABLE_TAGS = """
+SELECT * FROM `{database}`.information_schema.table_tags WHERE schema_name = '{schema}';
+"""
+
+UNITY_CATALOG_GET_ALL_TABLE_COLUMNS_TAGS = """
+SELECT * FROM `{database}`.information_schema.column_tags WHERE schema_name = '{schema}';
+"""

metadata/ingestion/source/messaging/messaging_service.py

@@ -34,9 +34,6 @@ from metadata.ingestion.api.delete import delete_entity_from_source
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import Source
 from metadata.ingestion.api.topology_runner import TopologyRunnerMixin
-from metadata.ingestion.connections.test_connections import (
-    raise_test_connection_exception,
-)
 from metadata.ingestion.models.delete_entity import DeleteEntity
 from metadata.ingestion.models.topology import (
     NodeStage,
@@ -45,7 +42,7 @@ from metadata.ingestion.models.topology import (
     TopologyNode,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.connections import get_connection, get_test_connection_fn
+from metadata.ingestion.source.connections import get_connection, test_connection_common
 from metadata.utils import fqn
 from metadata.utils.filters import filter_by_topic
 from metadata.utils.helpers import retry_with_docker_host
@@ -202,11 +199,9 @@ class MessagingServiceSource(TopologyRunnerMixin, Source, ABC):
         """By default, nothing to prepare"""
 
     def test_connection(self) -> None:
-        test_connection_fn = get_test_connection_fn(self.service_connection)
-        result = test_connection_fn(
+        test_connection_common(
             self.metadata, self.connection_obj, self.service_connection
         )
-        raise_test_connection_exception(result)
 
     def mark_topics_as_deleted(self) -> Iterable[Either[DeleteEntity]]:
         """Method to mark the topics as deleted"""

metadata/ingestion/source/metadata/alationsink/metadata.py

@@ -33,12 +33,9 @@ from metadata.generated.schema.metadataIngestion.workflow import (
 )
 from metadata.ingestion.api.models import Either, Entity
 from metadata.ingestion.api.steps import InvalidSourceException, Source
-from metadata.ingestion.connections.test_connections import (
-    raise_test_connection_exception,
-)
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.ometa.utils import model_str
-from metadata.ingestion.source.connections import get_connection, get_test_connection_fn
+from metadata.ingestion.source.connections import get_connection, test_connection_common
 from metadata.ingestion.source.metadata.alationsink.client import AlationSinkClient
 from metadata.ingestion.source.metadata.alationsink.constants import (
     SERVICE_TYPE_MAPPER,
@@ -475,8 +472,6 @@ class AlationsinkSource(Source):
         """Not required to implement"""
 
     def test_connection(self) -> None:
-        test_connection_fn = get_test_connection_fn(self.service_connection)
-        result = test_connection_fn(
+        test_connection_common(
             self.metadata, self.alation_sink_client, self.service_connection
         )
-        raise_test_connection_exception(result)

metadata/ingestion/source/metadata/amundsen/metadata.py

@@ -56,13 +56,10 @@ from metadata.generated.schema.type.entityReferenceList import EntityReferenceLi
 from metadata.ingestion.api.common import Entity
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException, Source
-from metadata.ingestion.connections.test_connections import (
-    raise_test_connection_exception,
-)
 from metadata.ingestion.models.user import OMetaUserProfile
 from metadata.ingestion.ometa.client_utils import get_chart_entities_from_id
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.connections import get_connection, get_test_connection_fn
+from metadata.ingestion.source.connections import get_connection, test_connection_common
 from metadata.ingestion.source.database.column_type_parser import ColumnTypeParser
 from metadata.ingestion.source.metadata.amundsen.queries import (
     NEO4J_AMUNDSEN_DASHBOARD_QUERY,
@@ -465,8 +462,6 @@ class AmundsenSource(Source):
         return None
 
     def test_connection(self) -> None:
-        test_connection_fn = get_test_connection_fn(self.service_connection)
-        result = test_connection_fn(
+        test_connection_common(
             self.metadata, self.connection_obj, self.service_connection
         )
-        raise_test_connection_exception(result)

metadata/ingestion/source/metadata/atlas/metadata.py

@@ -41,12 +41,9 @@ from metadata.generated.schema.type.entityLineage import EntitiesEdge
 from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.ingestion.api.models import Either, Entity, StackTraceError
 from metadata.ingestion.api.steps import InvalidSourceException, Source
-from metadata.ingestion.connections.test_connections import (
-    raise_test_connection_exception,
-)
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.connections import get_connection, get_test_connection_fn
+from metadata.ingestion.source.connections import get_connection, test_connection_common
 from metadata.ingestion.source.database.column_type_parser import ColumnTypeParser
 from metadata.ingestion.source.metadata.atlas.client import AtlasClient
 from metadata.utils import fqn
@@ -480,8 +477,6 @@ class AtlasSource(Source):
         return None
 
     def test_connection(self) -> None:
-        test_connection_fn = get_test_connection_fn(self.service_connection)
-        result = test_connection_fn(
+        test_connection_common(
             self.metadata, self.connection_obj, self.service_connection
         )
-        raise_test_connection_exception(result)