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

metadata/ingestion/source/database/databricks/client.py

@@ -14,14 +14,23 @@ Client to interact with databricks apis
 import json
 import traceback
 from datetime import timedelta
-from typing import Iterable, List
+from typing import Iterable, List, Optional, Tuple, Union
 
 import requests
+from sqlalchemy import text
+from sqlalchemy.engine import Engine
 
 from metadata.generated.schema.entity.services.connections.database.databricksConnection import (
     DatabricksConnection,
 )
+from metadata.generated.schema.entity.services.connections.pipeline.databricksPipelineConnection import (
+    DatabricksPipelineConnection,
+)
 from metadata.ingestion.ometa.client import APIError
+from metadata.ingestion.source.database.databricks.queries import (
+    DATABRICKS_GET_COLUMN_LINEAGE_FOR_JOB,
+    DATABRICKS_GET_TABLE_LINEAGE_FOR_JOB,
+)
 from metadata.utils.constants import QUERY_WITH_DBT, QUERY_WITH_OM_VERSION
 from metadata.utils.helpers import datetime_to_ts
 from metadata.utils.logger import ingestion_logger
@@ -30,6 +39,8 @@ logger = ingestion_logger()
 API_TIMEOUT = 10
 PAGE_SIZE = 100
 QUERIES_PATH = "/sql/history/queries"
+API_VERSION = "/api/2.0"
+JOB_API_VERSION = "/api/2.1"
 
 
 class DatabricksClientException(Exception):
@@ -43,30 +54,58 @@ class DatabricksClient:
     DatabricksClient creates a Databricks connection based on DatabricksCredentials.
     """
 
-    def __init__(self, config: DatabricksConnection):
+    def __init__(
+        self,
+        config: Union[DatabricksConnection, DatabricksPipelineConnection],
+        engine: Optional[Engine] = None,
+    ):
         self.config = config
         base_url, *_ = self.config.hostPort.split(":")
-        api_version = "/api/2.0"
-        job_api_version = "/api/2.1"
         auth_token = self.config.token.get_secret_value()
-        self.base_url = f"https://{base_url}{api_version}"
+        self.base_url = f"https://{base_url}{API_VERSION}"
         self.base_query_url = f"{self.base_url}{QUERIES_PATH}"
-        self.base_job_url = f"https://{base_url}{job_api_version}/jobs"
+        self.base_job_url = f"https://{base_url}{JOB_API_VERSION}/jobs"
         self.jobs_list_url = f"{self.base_job_url}/list"
         self.jobs_run_list_url = f"{self.base_job_url}/runs/list"
         self.headers = {
             "Authorization": f"Bearer {auth_token}",
             "Content-Type": "application/json",
         }
+        self.api_timeout = self.config.connectionTimeout or 120
+        self.job_table_lineage: dict[str, list[dict[str, str]]] = {}
+        self.job_column_lineage: dict[
+            str, dict[Tuple[str, str], list[Tuple[str, str]]]
+        ] = {}
+        self.engine = engine
         self.client = requests
 
     def test_query_api_access(self) -> None:
         res = self.client.get(
-            self.base_query_url, headers=self.headers, timeout=API_TIMEOUT
+            self.base_query_url, headers=self.headers, timeout=self.api_timeout
         )
         if res.status_code != 200:
             raise APIError(res.json)
 
+    def test_lineage_query(self) -> None:
+        try:
+            with self.engine.connect() as connection:
+                test_table_lineage = connection.execute(
+                    text(DATABRICKS_GET_TABLE_LINEAGE_FOR_JOB + " LIMIT 1")
+                )
+                test_column_lineage = connection.execute(
+                    text(DATABRICKS_GET_COLUMN_LINEAGE_FOR_JOB + " LIMIT 1")
+                )
+                # Check if queries executed successfully by fetching results
+                table_result = test_table_lineage.fetchone()
+                column_result = test_column_lineage.fetchone()
+                logger.info("Lineage queries executed successfully")
+        except Exception as exc:
+            logger.debug(f"Error testing lineage queries: {traceback.format_exc()}")
+            raise DatabricksClientException(
+                f"Failed to test lineage queries Make sure you have access"
+                "to the tables table_lineage and column_lineage: {exc}"
+            )
+
     def _run_query_paginator(self, data, result, end_time, response):
         while True:
             if response:
@@ -85,7 +124,7 @@ class DatabricksClient:
                     self.base_query_url,
                     data=json.dumps(data),
                     headers=self.headers,
-                    timeout=API_TIMEOUT,
+                    timeout=self.api_timeout,
                 ).json()
                 yield from response.get("res") or []
 
@@ -117,7 +156,7 @@ class DatabricksClient:
                         self.base_query_url,
                         data=json.dumps(data),
                         headers=self.headers,
-                        timeout=API_TIMEOUT,
+                        timeout=self.api_timeout,
                     ).json()
 
                     result = response.get("res") or []
@@ -145,7 +184,7 @@ class DatabricksClient:
             self.jobs_list_url,
             data=json.dumps(data),
             headers=self.headers,
-            timeout=API_TIMEOUT,
+            timeout=self.api_timeout,
         )
         if response.status_code != 200:
             raise DatabricksClientException(response.text)
@@ -154,6 +193,7 @@ class DatabricksClient:
         """
         Method returns List all the created jobs in a Databricks Workspace
         """
+        self.cache_lineage()
         try:
             iteration_count = 1
             data = {"limit": PAGE_SIZE, "expand_tasks": True, "offset": 0}
@@ -162,7 +202,7 @@ class DatabricksClient:
                 self.jobs_list_url,
                 data=json.dumps(data),
                 headers=self.headers,
-                timeout=API_TIMEOUT,
+                timeout=self.api_timeout,
             ).json()
 
             yield from response.get("jobs") or []
@@ -174,7 +214,7 @@ class DatabricksClient:
                     self.jobs_list_url,
                     data=json.dumps(data),
                     headers=self.headers,
-                    timeout=API_TIMEOUT,
+                    timeout=self.api_timeout,
                 ).json()
                 iteration_count += 1
                 yield from response.get("jobs") or []
@@ -200,7 +240,7 @@ class DatabricksClient:
                 self.jobs_run_list_url,
                 params=params,
                 headers=self.headers,
-                timeout=API_TIMEOUT,
+                timeout=self.api_timeout,
             ).json()
 
             yield from response.get("runs") or []
@@ -212,7 +252,7 @@ class DatabricksClient:
                     self.jobs_run_list_url,
                     params=params,
                     headers=self.headers,
-                    timeout=API_TIMEOUT,
+                    timeout=self.api_timeout,
                 ).json()
 
                 yield from response.get("runs") or []
@@ -220,3 +260,98 @@ class DatabricksClient:
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.error(exc)
+
+    def get_table_lineage(self, job_id: str) -> List[dict[str, str]]:
+        """
+        Method returns table lineage for a job by the specified job_id
+        """
+        try:
+            return self.job_table_lineage.get(str(job_id))
+        except Exception as exc:
+            logger.debug(
+                f"Error getting table lineage for job {job_id} due to {traceback.format_exc()}"
+            )
+            logger.error(exc)
+        return []
+
+    def get_column_lineage(
+        self, job_id: str, TableKey: Tuple[str, str]
+    ) -> List[Tuple[str, str]]:
+        """
+        Method returns column lineage for a job by the specified job_id and table key
+        """
+        try:
+            return self.job_column_lineage.get(str(job_id), {}).get(TableKey)
+        except Exception as exc:
+            logger.debug(
+                f"Error getting column lineage for table {TableKey} due to {traceback.format_exc()}"
+            )
+            logger.error(exc)
+        return []
+
+    def run_lineage_query(self, query: str) -> List[dict]:
+        """
+        Method runs a lineage query and returns the result
+        """
+        try:
+            with self.engine.connect() as connection:
+                result = connection.execute(text(query))
+                return result
+
+        except Exception as exc:
+            logger.debug(f"Error caching table lineage due to {traceback.format_exc()}")
+            logger.error(exc)
+        return []
+
+    def cache_lineage(self):
+        """
+        Method caches table and column lineage for a job by the specified job_id
+        """
+        logger.info(f"Caching table lineage")
+        table_lineage = self.run_lineage_query(DATABRICKS_GET_TABLE_LINEAGE_FOR_JOB)
+        if table_lineage:
+            for row in table_lineage:
+                try:
+                    if row.job_id not in self.job_table_lineage:
+                        self.job_table_lineage[row.job_id] = []
+                    self.job_table_lineage[row.job_id].append(
+                        {
+                            "source_table_full_name": row.source_table_full_name,
+                            "target_table_full_name": row.target_table_full_name,
+                        }
+                    )
+                except Exception as exc:
+                    logger.debug(
+                        f"Error parsing row: {row} due to {traceback.format_exc()}"
+                    )
+                    continue
+
+        # Not every job has column lineage, so we need to check if the job exists in the column_lineage table
+        # we will cache the column lineage for jobs that have column lineage
+        logger.info("Caching column lineage")
+        column_lineage = self.run_lineage_query(DATABRICKS_GET_COLUMN_LINEAGE_FOR_JOB)
+        if column_lineage:
+            for row in column_lineage:
+                try:
+                    table_key = (
+                        row.source_table_full_name,
+                        row.target_table_full_name,
+                    )
+                    column_pair = (
+                        row.source_column_name,
+                        row.target_column_name,
+                    )
+
+                    if row.job_id not in self.job_column_lineage:
+                        self.job_column_lineage[row.job_id] = {}
+
+                    if table_key not in self.job_column_lineage[row.job_id]:
+                        self.job_column_lineage[row.job_id][table_key] = []
+
+                    self.job_column_lineage[row.job_id][table_key].append(column_pair)
+
+                except Exception as exc:
+                    logger.debug(
+                        f"Error parsing row: {row} due to {traceback.format_exc()}"
+                    )
+                    continue

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

@@ -86,3 +86,29 @@ DATABRICKS_GET_COLUMN_TAGS = textwrap.dedent(
 )
 
 DATABRICKS_DDL = "SHOW CREATE TABLE `{table_name}`"
+
+DATABRICKS_GET_TABLE_LINEAGE_FOR_JOB = textwrap.dedent(
+    """
+    SELECT DISTINCT
+        entity_id as job_id,
+        entity_type as job_type,
+        source_table_full_name as source_table_full_name,
+        target_table_full_name as target_table_full_name
+    FROM system.access.table_lineage
+    WHERE entity_type ILIKE 'job'
+    """
+)
+
+DATABRICKS_GET_COLUMN_LINEAGE_FOR_JOB = textwrap.dedent(
+    """
+    SELECT DISTINCT
+        entity_id as job_id,
+        entity_type as job_type,
+        source_table_full_name as source_table_full_name,
+        source_column_name as source_column_name,
+        target_table_full_name as target_table_full_name,
+        target_column_name as target_column_name
+    FROM system.access.column_lineage
+    WHERE entity_type ILIKE 'job'
+    """
+)

metadata/ingestion/source/database/datalake/metadata.py

@@ -248,7 +248,7 @@ class DatalakeSource(DatabaseServiceSource):
         if self.source_config.includeTables:
             for key_name in self.client.get_table_names(bucket_name, prefix):
                 table_name = self.standardize_table_name(bucket_name, key_name)
-
+                logger.info(f"Processing table: {table_name}")
                 if self.filter_dl_table(table_name):
                     continue
 

metadata/ingestion/source/database/dbt/metadata.py

@@ -83,7 +83,6 @@ from metadata.ingestion.source.database.dbt.dbt_service import (
 )
 from metadata.ingestion.source.database.dbt.dbt_utils import (
     check_ephemeral_node,
-    check_or_create_test_suite,
     create_test_case_parameter_definitions,
     create_test_case_parameter_values,
     generate_entity_link,
@@ -1024,9 +1023,6 @@ class DbtSource(DbtServiceSource):
                 logger.debug(f"Processing DBT Test Case for node: {manifest_node.name}")
                 entity_link_list = generate_entity_link(dbt_test)
                 for entity_link_str in entity_link_list:
-                    test_suite = check_or_create_test_suite(
-                        self.metadata, entity_link_str
-                    )
                     table_fqn = get_table_fqn(entity_link_str)
                     logger.debug(f"Table fqn found: {table_fqn}")
                     source_elements = table_fqn.split(fqn.FQN_SEPARATOR)
@@ -1056,7 +1052,6 @@ class DbtSource(DbtServiceSource):
                                     manifest_node.name
                                 ),
                                 entityLink=entity_link_str,
-                                testSuite=test_suite.fullyQualifiedName,
                                 parameterValues=create_test_case_parameter_values(
                                     dbt_test
                                 ),

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

@@ -20,6 +20,9 @@ from metadata.clients.azure_client import AzureClient
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
 )
+from metadata.generated.schema.entity.services.connections.database.common.azureConfig import (
+    AzureConfigurationSource,
+)
 from metadata.generated.schema.entity.services.connections.database.common.basicAuth import (
     BasicAuth,
 )
@@ -34,6 +37,7 @@ from metadata.ingestion.connections.builders import (
     get_connection_args_common,
     get_connection_url_common,
 )
+from metadata.ingestion.connections.connection import BaseConnection
 from metadata.ingestion.connections.test_connections import (
     test_connection_db_schema_sources,
 )
@@ -45,51 +49,54 @@ from metadata.ingestion.source.database.mysql.queries import (
 from metadata.utils.constants import THREE_MIN
 
 
-def get_connection(connection: MysqlConnection) -> Engine:
-    """
-    Create connection
-    """
-    if hasattr(connection.authType, "azureConfig"):
-        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-micr"
-                "osoft-entra-access-token and fetch the resource associated with it,"
-                " for e.g. https://ossrdbms-aad.database.windows.net/.default"
+class MySQLConnection(BaseConnection[MysqlConnection, Engine]):
+    def get_client(self) -> Engine:
+        """
+        Return the SQLAlchemy Engine for MySQL.
+        """
+        connection = self.service_connection
+
+        if isinstance(connection.authType, AzureConfigurationSource):
+            if not connection.authType.azureConfig:
+                raise ValueError("Azure Config is missing")
+            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-micr"
+                    "osoft-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(",")
             )
-        access_token_obj = azure_client.get_token(
-            *connection.authType.azureConfig.scopes.split(",")
+            connection.authType = BasicAuth(password=access_token_obj.token)  # type: ignore
+        return create_generic_db_connection(
+            connection=connection,
+            get_connection_url_fn=get_connection_url_common,
+            get_connection_args_fn=get_connection_args_common,
         )
-        connection.authType = BasicAuth(password=access_token_obj.token)
-    return create_generic_db_connection(
-        connection=connection,
-        get_connection_url_fn=get_connection_url_common,
-        get_connection_args_fn=get_connection_args_common,
-    )
 
-
-def test_connection(
-    metadata: OpenMetadata,
-    engine: Engine,
-    service_connection: MysqlConnection,
-    automation_workflow: Optional[AutomationWorkflow] = None,
-    timeout_seconds: Optional[int] = THREE_MIN,
-) -> TestConnectionResult:
-    """
-    Test connection. This can be executed either as part
-    of a metadata workflow or during an Automation Workflow
-    """
-    queries = {
-        "GetQueries": MYSQL_TEST_GET_QUERIES
-        if not service_connection.useSlowLogs
-        else MYSQL_TEST_GET_QUERIES_SLOW_LOGS,
-    }
-    return test_connection_db_schema_sources(
-        metadata=metadata,
-        engine=engine,
-        service_connection=service_connection,
-        automation_workflow=automation_workflow,
-        timeout_seconds=timeout_seconds,
-        queries=queries,
-    )
+    def test_connection(
+        self,
+        metadata: OpenMetadata,
+        automation_workflow: Optional[AutomationWorkflow] = None,
+        timeout_seconds: Optional[int] = THREE_MIN,
+    ) -> TestConnectionResult:
+        """
+        Test connection. This can be executed either as part
+        of a metadata workflow or during an Automation Workflow
+        """
+        queries = {
+            "GetQueries": MYSQL_TEST_GET_QUERIES
+            if not self.service_connection.useSlowLogs
+            else MYSQL_TEST_GET_QUERIES_SLOW_LOGS,
+        }
+        return test_connection_db_schema_sources(
+            metadata=metadata,
+            engine=self.get_client(),
+            service_connection=self.service_connection,
+            automation_workflow=automation_workflow,
+            timeout_seconds=timeout_seconds,
+            queries=queries,
+        )

metadata/ingestion/source/database/mysql/service_spec.py

@@ -1,3 +1,4 @@
+from metadata.ingestion.source.database.mysql.connection import MySQLConnection
 from metadata.ingestion.source.database.mysql.lineage import MysqlLineageSource
 from metadata.ingestion.source.database.mysql.metadata import MysqlSource
 from metadata.ingestion.source.database.mysql.usage import MysqlUsageSource
@@ -7,4 +8,5 @@ ServiceSpec = DefaultDatabaseSpec(
     metadata_source_class=MysqlSource,
     lineage_source_class=MysqlLineageSource,
     usage_source_class=MysqlUsageSource,
+    connection_class=MySQLConnection,
 )

metadata/ingestion/source/database/postgres/metadata.py

@@ -40,7 +40,11 @@ from metadata.generated.schema.entity.services.ingestionPipelines.status import
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
-from metadata.generated.schema.type.basic import EntityName, FullyQualifiedEntityName
+from metadata.generated.schema.type.basic import (
+    EntityName,
+    FullyQualifiedEntityName,
+    Markdown,
+)
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
@@ -300,9 +304,15 @@ class PostgresSource(CommonDbSourceService, MultiDBSource):
         """List stored procedures"""
         if self.source_config.includeStoredProcedures:
             yield from self._get_stored_procedures_internal(
-                POSTGRES_GET_STORED_PROCEDURES
+                POSTGRES_GET_STORED_PROCEDURES.format(
+                    schema_name=self.context.get().database_schema
+                )
+            )
+            yield from self._get_stored_procedures_internal(
+                POSTGRES_GET_FUNCTIONS.format(
+                    schema_name=self.context.get().database_schema
+                )
             )
-            yield from self._get_stored_procedures_internal(POSTGRES_GET_FUNCTIONS)
 
     def yield_stored_procedure(
         self, stored_procedure
@@ -311,7 +321,9 @@ class PostgresSource(CommonDbSourceService, MultiDBSource):
         try:
             stored_procedure_request = CreateStoredProcedureRequest(
                 name=EntityName(stored_procedure.name),
-                description=None,
+                description=Markdown(stored_procedure.description)
+                if stored_procedure.description
+                else None,
                 storedProcedureCode=StoredProcedureCode(
                     language=STORED_PROC_LANGUAGE_MAP.get(stored_procedure.language),
                     code=stored_procedure.definition,

metadata/ingestion/source/database/postgres/models.py

@@ -24,3 +24,4 @@ class PostgresStoredProcedure(BaseModel):
     definition: str
     language: Optional[str] = None
     procedure_type: Optional[str] = Field(None, alias="procedure_type")
+    description: Optional[str] = Field(None, alias="description")

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

@@ -225,10 +225,12 @@ POSTGRES_GET_STORED_PROCEDURES = """
         proargtypes AS argument_types,
         prorettype::regtype AS return_type,
         prosrc AS definition,
-        'StoredProcedure' as procedure_type
+        'StoredProcedure' as procedure_type,
+        obj_description(pg_proc.oid, 'pg_proc') AS description
     FROM pg_proc
     JOIN pg_namespace ON pg_proc.pronamespace = pg_namespace.oid
-    WHERE prokind = 'p';
+    WHERE prokind = 'p'
+    and pg_namespace.nspname = '{schema_name}';
 """
 
 POSTGRES_GET_FUNCTIONS = """
@@ -238,11 +240,12 @@ SELECT
     proargtypes AS argument_types,
     prorettype :: regtype AS return_type,
     prosrc AS definition,
-    'Function' as procedure_type
+    'Function' as procedure_type,
+    obj_description(pg_proc.oid, 'pg_proc') AS description
 FROM
     pg_proc
     JOIN pg_namespace ON pg_proc.pronamespace = pg_namespace.oid
 WHERE
     prokind = 'f'
-    and pg_namespace.nspname NOT IN ('pg_catalog', 'information_schema');
+    and pg_namespace.nspname = '{schema_name}';
 """

metadata/ingestion/source/database/query_parser_source.py

@@ -20,13 +20,10 @@ from metadata.generated.schema.metadataIngestion.workflow import (
 )
 from metadata.generated.schema.type.tableQuery import TableQuery
 from metadata.ingestion.api.steps import Source
-from metadata.ingestion.connections.test_connections import (
-    raise_test_connection_exception,
-)
 from metadata.ingestion.lineage.masker import masked_query_cache
 from metadata.ingestion.lineage.models import ConnectionTypeDialectMapper
 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.utils.helpers import get_start_and_end, retry_with_docker_host
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.ssl_manager import get_ssl_connection
@@ -134,6 +131,4 @@ class QueryParserSource(Source, ABC):
         masked_query_cache.clear()
 
     def test_connection(self) -> None:
-        test_connection_fn = get_test_connection_fn(self.service_connection)
-        result = test_connection_fn(self.metadata, self.engine, self.service_connection)
-        raise_test_connection_exception(result)
+        test_connection_common(self.metadata, self.engine, self.service_connection)

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

@@ -40,7 +40,7 @@ from metadata.ingestion.connections.test_connections import (
     test_query,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.connections import kill_active_connections
+from metadata.ingestion.source.connections_utils import kill_active_connections
 from metadata.ingestion.source.database.redshift.queries import (
     REDSHIFT_GET_ALL_RELATIONS,
     REDSHIFT_GET_DATABASE_NAMES,

metadata/ingestion/source/database/sample_data.py

@@ -890,6 +890,7 @@ class SampleDataSource(
                 tags=table["tags"],
                 schemaDefinition=table.get("schemaDefinition"),
                 sourceUrl=table.get("sourceUrl"),
+                tablePartition=table.get("tablePartition"),
             )
 
             yield Either(right=table_and_db)
@@ -1602,7 +1603,6 @@ class SampleDataSource(
                         description=test_case["description"],
                         testDefinition=test_case["testDefinitionName"],
                         entityLink=test_case["entityLink"],
-                        testSuite=suite.fullyQualifiedName.root,
                         parameterValues=[
                             TestCaseParameterValue(**param_values)
                             for param_values in test_case["parameterValues"]

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

@@ -29,11 +29,8 @@ from metadata.generated.schema.metadataIngestion.workflow import (
 )
 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.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.saphana.cdata_parser import (
     ParsedLineage,
     parse_registry,
@@ -158,6 +155,4 @@ class SaphanaLineageSource(Source):
             )
 
     def test_connection(self) -> None:
-        test_connection_fn = get_test_connection_fn(self.service_connection)
-        result = test_connection_fn(self.metadata, self.engine, self.service_connection)
-        raise_test_connection_exception(result)
+        test_connection_common(self.metadata, self.engine, self.service_connection)

metadata/ingestion/source/database/sas/metadata.py

@@ -70,12 +70,9 @@ from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.ingestion.api.common import Entity
 from metadata.ingestion.api.models import Either, StackTraceError
 from metadata.ingestion.api.steps import InvalidSourceException
-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.database.database_service import DatabaseServiceSource
 from metadata.ingestion.source.database.sas.client import SASClient
@@ -900,8 +897,6 @@ class SasSource(
         pass
 
     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)