openmetadata-ingestion 1.5.13.2__py3-none-any.whl → 1.6.0.0__py3-none-any.whl

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

@@ -0,0 +1,22 @@
+from metadata.data_quality.interface.sqlalchemy.unity_catalog.test_suite_interface import (
+    UnityCatalogTestSuiteInterface,
+)
+from metadata.ingestion.source.database.unitycatalog.lineage import (
+    UnitycatalogLineageSource,
+)
+from metadata.ingestion.source.database.unitycatalog.metadata import UnitycatalogSource
+from metadata.ingestion.source.database.unitycatalog.usage import (
+    UnitycatalogUsageSource,
+)
+from metadata.profiler.interface.sqlalchemy.unity_catalog.profiler_interface import (
+    UnityCatalogProfilerInterface,
+)
+from metadata.utils.service_spec.default import DefaultDatabaseSpec
+
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=UnitycatalogSource,
+    lineage_source_class=UnitycatalogLineageSource,
+    usage_source_class=UnitycatalogUsageSource,
+    profiler_class=UnityCatalogProfilerInterface,
+    test_suite_class=UnityCatalogTestSuiteInterface,
+)

metadata/ingestion/source/database/usage_source.py

@@ -152,7 +152,10 @@ class UsageSource(QueryParserSource, ABC):
             except Exception as exc:
                 if query:
                     logger.debug(
-                        f"###### USAGE QUERY #######\n{mask_query(query, self.dialect.value)}\n##########################"
+                        (
+                            f"###### USAGE QUERY #######\n{mask_query(query, self.dialect.value)}"
+                            "\n##########################"
+                        )
                     )
                 logger.debug(traceback.format_exc())
                 logger.error(f"Source usage processing error: {exc}")

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

@@ -0,0 +1,10 @@
+from metadata.ingestion.source.database.vertica.lineage import VerticaLineageSource
+from metadata.ingestion.source.database.vertica.metadata import VerticaSource
+from metadata.ingestion.source.database.vertica.usage import VerticaUsageSource
+from metadata.utils.service_spec.default import DefaultDatabaseSpec
+
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=VerticaSource,
+    lineage_source_class=VerticaLineageSource,
+    usage_source_class=VerticaUsageSource,
+)

metadata/ingestion/source/messaging/kafka/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.messaging.kafka.metadata import KafkaSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=KafkaSource)

metadata/ingestion/source/messaging/kinesis/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.messaging.kinesis.metadata import KinesisSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=KinesisSource)

metadata/ingestion/source/messaging/redpanda/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.messaging.redpanda.metadata import RedpandaSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=RedpandaSource)

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

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.metadata.alationsink.metadata import AlationsinkSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=AlationsinkSource)

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

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.metadata.amundsen.metadata import AmundsenSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=AmundsenSource)

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

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.metadata.atlas.metadata import AtlasSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=AtlasSource)

metadata/ingestion/source/mlmodel/mlflow/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.mlmodel.mlflow.metadata import MlflowSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=MlflowSource)

metadata/ingestion/source/mlmodel/sagemaker/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.mlmodel.sagemaker.metadata import SagemakerSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=SagemakerSource)

metadata/ingestion/source/pipeline/airbyte/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.pipeline.airbyte.metadata import AirbyteSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=AirbyteSource)

metadata/ingestion/source/pipeline/airflow/connection.py

@@ -16,7 +16,9 @@ from functools import partial, singledispatch
 from typing import Optional
 
 from airflow import settings
+from airflow.models.serialized_dag import SerializedDagModel
 from sqlalchemy.engine import Engine
+from sqlalchemy.orm import sessionmaker
 
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
@@ -102,6 +104,18 @@ def get_connection(connection: AirflowConnection) -> Engine:
         raise SourceConnectionException(msg) from exc
 
 
+class AirflowPipelineDetailsAccessError(Exception):
+    """
+    Raise when Pipeline information is not retrieved
+    """
+
+
+class AirflowTaskDetailsAccessError(Exception):
+    """
+    Raise when Task detail information is not retrieved
+    """
+
+
 def test_connection(
     metadata: OpenMetadata,
     engine: Engine,
@@ -114,8 +128,37 @@ def test_connection(
     of a metadata workflow or during an Automation Workflow
     """
 
-    test_fn = {"CheckAccess": partial(test_connection_engine_step, engine)}
-
+    session_maker = sessionmaker(bind=engine)
+    session = session_maker()
+
+    def test_pipeline_details_access(session):
+        try:
+            result = session.query(SerializedDagModel).first()
+            return result
+        except Exception as e:
+            raise AirflowPipelineDetailsAccessError(
+                f"Pipeline details access error: {e}"
+            )
+
+    def test_task_detail_access(session):
+        try:
+            json_data_column = (
+                SerializedDagModel._data  # For 2.3.0 onwards # pylint: disable=protected-access
+                if hasattr(SerializedDagModel, "_data")
+                else SerializedDagModel.data  # For 2.2.5 and 2.1.4
+            )
+            result = session.query(json_data_column).first()
+
+            retrieved_tasks = result[0]["dag"]["tasks"]
+            return retrieved_tasks
+        except Exception as e:
+            raise AirflowTaskDetailsAccessError(f"Task details access error : {e}")
+
+    test_fn = {
+        "CheckAccess": partial(test_connection_engine_step, engine),
+        "PipelineDetailsAccess": partial(test_pipeline_details_access, session),
+        "TaskDetailAccess": partial(test_task_detail_access, session),
+    }
     return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,

metadata/ingestion/source/pipeline/airflow/metadata.py

@@ -17,7 +17,7 @@ from datetime import datetime
 from enum import Enum
 from typing import Dict, Iterable, List, Optional, cast
 
-from airflow.models import BaseOperator, DagRun, TaskInstance
+from airflow.models import BaseOperator, DagRun, DagTag, TaskInstance
 from airflow.models.dag import DagModel
 from airflow.models.serialized_dag import SerializedDagModel
 from airflow.serialization.serialized_objects import SerializedDAG
@@ -57,6 +57,7 @@ from metadata.generated.schema.type.entityReferenceList import EntityReferenceLi
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.connections.session import create_and_bind_session
+from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.models.pipeline_status import OMetaPipelineStatus
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.pipeline.airflow.lineage_parser import (
@@ -74,9 +75,12 @@ from metadata.utils import fqn
 from metadata.utils.constants import ENTITY_REFERENCE_TYPE_MAP
 from metadata.utils.helpers import clean_uri, datetime_to_ts
 from metadata.utils.logger import ingestion_logger
+from metadata.utils.tag_utils import get_ometa_tag_and_classification, get_tag_labels
 
 logger = ingestion_logger()
 
+AIRFLOW_TAG_CATEGORY = "AirflowTags"
+
 
 class AirflowTaskStatus(Enum):
     SUCCESS = "success"
@@ -152,6 +156,31 @@ class AirflowSource(PipelineServiceSource):
             return task["__var"]
         return task
 
+    def get_all_tags(self, dag_id: str) -> List[str]:
+        try:
+            tag_query = (
+                self.session.query(DagTag.name)
+                .filter(DagTag.dag_id == dag_id)
+                .distinct()
+                .all()
+            )
+            return [tag[0] for tag in tag_query]
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Could not extract tags details due to {exc}")
+        return []
+
+    def yield_tag(
+        self, pipeline_details: AirflowDagDetails
+    ) -> Iterable[Either[OMetaTagAndClassification]]:
+        yield from get_ometa_tag_and_classification(
+            tags=self.get_all_tags(dag_id=pipeline_details.dag_id),
+            classification_name=AIRFLOW_TAG_CATEGORY,
+            tag_description="Airflow Tag",
+            classification_description="Tags associated with airflow entities.",
+            include_tags=self.source_config.includeTags,
+        )
+
     def get_pipeline_status(self, dag_id: str) -> List[DagRun]:
         """
         Return the DagRuns of given dag
@@ -328,32 +357,42 @@ class AirflowSource(PipelineServiceSource):
             ).filter(
                 DagModel.is_paused == False  # pylint: disable=singleton-comparison
             )
-        for serialized_dag in session_query.yield_per(100):
-            try:
-                data = serialized_dag[1]["dag"]
-                dag = AirflowDagDetails(
-                    dag_id=serialized_dag[0],
-                    fileloc=serialized_dag[2],
-                    data=AirflowDag.model_validate(serialized_dag[1]),
-                    max_active_runs=data.get("max_active_runs", None),
-                    description=data.get("_description", None),
-                    start_date=data.get("start_date", None),
-                    tasks=list(
-                        map(self._extract_serialized_task, data.get("tasks", []))
-                    ),
-                    schedule_interval=get_schedule_interval(data),
-                    owner=self.fetch_dag_owners(data),
-                )
+        limit = 100  # Number of records per batch
+        offset = 0  # Start
+
+        while True:
+            paginated_query = session_query.limit(limit).offset(offset)
+            results = paginated_query.all()
+            if not results:
+                break
+            for serialized_dag in results:
+                try:
+                    data = serialized_dag[1]["dag"]
+                    dag = AirflowDagDetails(
+                        dag_id=serialized_dag[0],
+                        fileloc=serialized_dag[2],
+                        data=AirflowDag.model_validate(serialized_dag[1]),
+                        max_active_runs=data.get("max_active_runs", None),
+                        description=data.get("_description", None),
+                        start_date=data.get("start_date", None),
+                        tasks=list(
+                            map(self._extract_serialized_task, data.get("tasks", []))
+                        ),
+                        schedule_interval=get_schedule_interval(data),
+                        owner=self.fetch_dag_owners(data),
+                    )
 
-                yield dag
-            except ValidationError as err:
-                logger.debug(traceback.format_exc())
-                logger.warning(
-                    f"Error building pydantic model for {serialized_dag} - {err}"
-                )
-            except Exception as err:
-                logger.debug(traceback.format_exc())
-                logger.warning(f"Wild error yielding dag {serialized_dag} - {err}")
+                    yield dag
+                except ValidationError as err:
+                    logger.debug(traceback.format_exc())
+                    logger.warning(
+                        f"Error building pydantic model for {serialized_dag} - {err}"
+                    )
+                except Exception as err:
+                    logger.debug(traceback.format_exc())
+                    logger.warning(f"Wild error yielding dag {serialized_dag} - {err}")
+
+            offset += limit
 
     def fetch_dag_owners(self, data) -> Optional[str]:
         """
@@ -441,7 +480,6 @@ class AirflowSource(PipelineServiceSource):
         try:
             # Airflow uses /dags/dag_id/grid to show pipeline / dag
             source_url = f"{clean_uri(self.service_connection.hostPort)}/dags/{pipeline_details.dag_id}/grid"
-
             pipeline_request = CreatePipelineRequest(
                 name=EntityName(pipeline_details.dag_id),
                 description=Markdown(pipeline_details.description)
@@ -459,6 +497,12 @@ class AirflowSource(PipelineServiceSource):
                 service=FullyQualifiedEntityName(self.context.get().pipeline_service),
                 owners=self.get_owner(pipeline_details.owner),
                 scheduleInterval=pipeline_details.schedule_interval,
+                tags=get_tag_labels(
+                    metadata=self.metadata,
+                    tags=pipeline_details.data.dag.tags,
+                    classification_name=AIRFLOW_TAG_CATEGORY,
+                    include_tags=self.source_config.includeTags,
+                ),
             )
             yield Either(right=pipeline_request)
             self.register_record(pipeline_request=pipeline_request)

metadata/ingestion/source/pipeline/airflow/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.pipeline.airflow.metadata import AirflowSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=AirflowSource)

metadata/ingestion/source/pipeline/dagster/metadata.py

@@ -161,7 +161,9 @@ class DagsterSource(PipelineServiceSource):
                 )
             )
 
-    def yield_tag(self, *_, **__) -> Iterable[Either[OMetaTagAndClassification]]:
+    def yield_tag(
+        self, pipeline_details: DagsterPipeline
+    ) -> Iterable[Either[OMetaTagAndClassification]]:
         yield from get_ometa_tag_and_classification(
             tags=[self.context.get().repository_name],
             classification_name=DAGSTER_TAG_CATEGORY,

metadata/ingestion/source/pipeline/dagster/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.pipeline.dagster.metadata import DagsterSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=DagsterSource)

metadata/ingestion/source/pipeline/databrickspipeline/service_spec.py

@@ -0,0 +1,6 @@
+from metadata.ingestion.source.pipeline.databrickspipeline.metadata import (
+    DatabrickspipelineSource,
+)
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=DatabrickspipelineSource)

metadata/ingestion/source/pipeline/dbtcloud/client.py

@@ -35,6 +35,7 @@ from metadata.utils.helpers import clean_uri
 from metadata.utils.logger import ometa_logger
 
 logger = ometa_logger()
+API_VERSION = "api/v2"
 
 
 class DBTCloudClient:
@@ -44,9 +45,13 @@ class DBTCloudClient:
 
     def __init__(self, config: DBTCloudConnection):
         self.config = config
+
+        self.job_ids = self.config.jobIds
+        self.project_ids = self.config.projectIds
+
         client_config: ClientConfig = ClientConfig(
             base_url=clean_uri(self.config.host),
-            api_version="api/v2",
+            api_version=API_VERSION,
             auth_header=AUTHORIZATION_HEADER,
             auth_token=lambda: (self.config.token.get_secret_value(), 0),
             allow_redirects=True,
@@ -63,27 +68,43 @@ class DBTCloudClient:
         self.client = REST(client_config)
         self.graphql_client = REST(graphql_client_config)
 
-    def test_get_jobs(self) -> Optional[List[DBTJob]]:
+    def _get_jobs(
+        self, job_id: str = None, project_id: str = None
+    ) -> Optional[List[DBTJob]]:
         """
-        test fetch jobs for an account in dbt cloud
+        fetch jobs for an account in dbt cloud
         """
-        job_path = f"{self.config.jobId}/" if self.config.jobId else ""
-        result = self.client.get(f"/accounts/{self.config.accountId}/jobs/{job_path}")
-        job_list = (
-            [DBTJob.model_validate(result["data"])]
-            if self.config.jobId
-            else DBTJobList.model_validate(result).Jobs
-        )
+        job_list = []
+        try:
+            job_path = f"{job_id}/" if job_id else ""
+            project_path = f"?project_id={project_id}" if project_id else ""
+            result = self.client.get(
+                f"/accounts/{self.config.accountId}/jobs/{job_path}{project_path}"
+            )
+            job_list = (
+                [DBTJob.model_validate(result["data"])]
+                if job_id
+                else DBTJobList.model_validate(result).Jobs
+            )
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.error(
+                f"Failed to get job info for project_id: `{project_id}` or job_id: `{job_id}` : {exc}"
+            )
         return job_list
 
-    def test_get_runs(self, job_id: int) -> Optional[List[DBTRun]]:
+    def test_get_jobs(self) -> List[DBTJob]:
+        """
+        test fetch jobs for an account in dbt cloud
+        """
+        job_list = self.client.get(f"/accounts/{self.config.accountId}/jobs/")
+        return DBTJobList.model_validate(job_list).Jobs
+
+    def test_get_runs(self) -> List[DBTRun]:
         """
         test fetch runs for a job in dbt cloud
         """
-        result = self.client.get(
-            f"/accounts/{self.config.accountId}/runs/",
-            data={"job_definition_id": job_id},
-        )
+        result = self.client.get(f"/accounts/{self.config.accountId}/runs/")
         run_list = DBTRunList.model_validate(result).Runs
         return run_list
 
@@ -92,17 +113,29 @@ class DBTCloudClient:
         list jobs for an account in dbt cloud
         """
         try:
-            job_path = f"{self.config.jobId}/" if self.config.jobId else ""
-            result = self.client.get(
-                f"/accounts/{self.config.accountId}/jobs/{job_path}"
-            )
-            if result:
-                job_list = (
-                    [DBTJob.model_validate(result.get("data"))]
-                    if self.config.jobId
-                    else DBTJobList.model_validate(result).Jobs
-                )
-                return job_list
+            jobs = []
+            # case when job_ids are specified and project_ids are not
+            if self.job_ids and not self.project_ids:
+                for job_id in self.job_ids:
+                    jobs.extend(self._get_jobs(job_id=job_id))
+            # case when project_ids are specified or both are specified
+            elif self.project_ids:
+                for project_id in self.project_ids:
+                    results = self._get_jobs(project_id=project_id)
+                    if self.job_ids:
+                        jobs.extend(
+                            [
+                                result
+                                for result in results
+                                if str(result.id) in self.job_ids
+                            ]
+                        )
+                    else:
+                        jobs.extend(results)
+            else:
+                results = self._get_jobs()
+                jobs.extend(results)
+            return jobs
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.error(f"Unable to get job info :{exc}")
@@ -141,6 +174,9 @@ class DBTCloudClient:
 
             if result.get("data") and result["data"].get("job"):
                 model_list = DBTModelList.model_validate(result["data"]["job"]).models
+                logger.debug(
+                    f"Successfully fetched models from dbt for job_id:{job_id} run_id:{run_id}: {model_list}"
+                )
                 return model_list
 
         except Exception as exc:
@@ -150,7 +186,7 @@ class DBTCloudClient:
 
     def get_models_and_seeds_details(self, job_id: int, run_id: int):
         """
-        get model details for a job in dbt cloud for lineage
+        get parent model details for a job in dbt cloud for lineage
         """
         try:
             query_params = {
@@ -163,9 +199,12 @@ class DBTCloudClient:
             if result.get("data") and result["data"].get("job"):
                 result = DBTModelList.model_validate(result["data"]["job"])
                 parents_list = result.models + result.seeds
+                logger.debug(
+                    f"Successfully fetched parent models from dbt for job_id:{job_id} run_id:{run_id}: {parents_list}"
+                )
                 return parents_list
 
         except Exception as exc:
             logger.debug(traceback.format_exc())
-            logger.warning(f"Unable to get model info :{exc}")
+            logger.warning(f"Unable to get parents model info :{exc}")
         return None

metadata/ingestion/source/pipeline/dbtcloud/connection.py

@@ -50,11 +50,9 @@ def test_connection(
     of a metadata workflow or during an Automation Workflow
     """
 
-    job_id = int(service_connection.jobId) if service_connection.jobId else 0
-
     test_fn = {
         "GetJobs": client.test_get_jobs,
-        "GetRuns": partial(client.test_get_runs, job_id=job_id),
+        "GetRuns": partial(client.test_get_runs),
     }
 
     return test_connection_steps(

metadata/ingestion/source/pipeline/dbtcloud/models.py

@@ -35,7 +35,7 @@ class DBTJob(BaseModel):
 
 
 class DBTJobList(BaseModel):
-    Jobs: Optional[List[DBTJob]] = Field([], alias="data")
+    Jobs: List[DBTJob] = Field(alias="data")
 
 
 class DBTRun(BaseModel):

metadata/ingestion/source/pipeline/dbtcloud/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.pipeline.dbtcloud.metadata import DbtcloudSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=DbtcloudSource)

metadata/ingestion/source/pipeline/domopipeline/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.pipeline.domopipeline.metadata import DomopipelineSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=DomopipelineSource)

metadata/ingestion/source/pipeline/fivetran/metadata.py

@@ -275,7 +275,7 @@ class FivetranSource(PipelineServiceSource):
                 )
 
     def get_pipeline_name(self, pipeline_details: FivetranPipelineDetails) -> str:
-        return pipeline_details.pipeline_name
+        return pipeline_details.pipeline_display_name or pipeline_details.pipeline_name
 
     def get_source_url(
         self,

metadata/ingestion/source/pipeline/fivetran/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.pipeline.fivetran.metadata import FivetranSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=FivetranSource)

metadata/ingestion/source/pipeline/flink/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.pipeline.flink.metadata import FlinkSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=FlinkSource)

metadata/ingestion/source/pipeline/gluepipeline/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.pipeline.gluepipeline.metadata import GluepipelineSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=GluepipelineSource)

metadata/ingestion/source/pipeline/kafkaconnect/client.py

@@ -54,22 +54,25 @@ class KafkaConnectClient:
         auth = None
         ssl_verify = config.verifySSL
         if config.KafkaConnectConfig:
-            auth = f"{config.KafkaConnectConfig.username}:{config.KafkaConnectConfig.password}"
+            auth = f"{config.KafkaConnectConfig.username}:{config.KafkaConnectConfig.password.get_secret_value()}"
         self.client = KafkaConnect(url=url, auth=auth, ssl_verify=ssl_verify)
 
     def get_cluster_info(self) -> Optional[dict]:
         """
         Get the version and other details of the Kafka Connect cluster.
         """
-        try:
-            result = self.client.get_cluster_info()
-            return result
-
-        except Exception as exc:
-            logger.debug(traceback.format_exc())
-            logger.error(f"Unable to get cluster info :{exc}")
+        return self.client.get_cluster_info()
 
-        return None
+    def get_connectors_list(
+        self,
+        expand: str = None,
+        pattern: str = None,
+        state: str = None,
+    ) -> dict:
+        """
+        Get the list of connectors from Kafka Connect cluster.
+        """
+        return self.client.list_connectors(expand=expand, pattern=pattern, state=state)
 
     def get_connectors(
         self,
@@ -85,12 +88,8 @@ class KafkaConnectClient:
             pattern (str): Only list connectors that match the regex pattern.
             state (str): Only list connectors that match the state.
         """
-
         try:
-            result = self.client.list_connectors(
-                expand=expand, pattern=pattern, state=state
-            )
-            return result
+            return self.get_connectors_list(expand=expand, pattern=pattern, state=state)
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.error(f"Unable to get connectors list {exc}")

metadata/ingestion/source/pipeline/kafkaconnect/connection.py

@@ -51,7 +51,7 @@ def test_connection(
 
     test_fn = {
         "GetClusterInfo": client.get_cluster_info,
-        "GetPipelines": client.get_connectors,
+        "GetPipelines": client.get_connectors_list,
         "GetPlugins": client.get_connector_plugins,
     }
 

metadata/ingestion/source/pipeline/kafkaconnect/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.pipeline.kafkaconnect.metadata import KafkaconnectSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=KafkaconnectSource)

metadata/ingestion/source/pipeline/nifi/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.pipeline.nifi.metadata import NifiSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=NifiSource)

metadata/ingestion/source/pipeline/openlineage/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.pipeline.openlineage.metadata import OpenlineageSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=OpenlineageSource)