openmetadata-ingestion 1.3.3.0rc2__py3-none-any.whl → 1.4.0.0rc1__py3-none-any.whl

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

@@ -18,9 +18,11 @@ from enum import Enum
 from typing import Iterable, List, Optional, cast
 
 from airflow.models import BaseOperator, DagRun, TaskInstance
+from airflow.models.dag import DagModel
 from airflow.models.serialized_dag import SerializedDagModel
 from airflow.serialization.serialized_objects import SerializedDAG
 from pydantic import BaseModel, ValidationError
+from sqlalchemy import join
 from sqlalchemy.orm import Session
 
 from metadata.generated.schema.api.data.createPipeline import CreatePipelineRequest
@@ -111,7 +113,9 @@ class AirflowSource(PipelineServiceSource):
         self._session = None
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata) -> "AirflowSource":
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ) -> "AirflowSource":
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: AirflowConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, AirflowConnection):
@@ -225,7 +229,7 @@ class AirflowSource(PipelineServiceSource):
 
             for dag_run in dag_run_list:
                 if (
-                    dag_run.run_id and self.context.task_names
+                    dag_run.run_id and self.context.get().task_names
                 ):  # Airflow dags can have old task which are turned off/commented out in code
                     tasks = self.get_task_instances(
                         dag_id=dag_run.dag_id,
@@ -245,7 +249,7 @@ class AirflowSource(PipelineServiceSource):
                             ),  # Might be None for running tasks
                         )  # Log link might not be present in all Airflow versions
                         for task in tasks
-                        if task.task_id in self.context.task_names
+                        if task.task_id in self.context.get().task_names
                     ]
 
                     pipeline_status = PipelineStatus(
@@ -258,8 +262,8 @@ class AirflowSource(PipelineServiceSource):
                     pipeline_fqn = fqn.build(
                         metadata=self.metadata,
                         entity_type=Pipeline,
-                        service_name=self.context.pipeline_service,
-                        pipeline_name=self.context.pipeline,
+                        service_name=self.context.get().pipeline_service,
+                        pipeline_name=self.context.get().pipeline,
                     )
                     yield Either(
                         right=OMetaPipelineStatus(
@@ -289,11 +293,23 @@ class AirflowSource(PipelineServiceSource):
             if hasattr(SerializedDagModel, "_data")
             else SerializedDagModel.data  # For 2.2.5 and 2.1.4
         )
-        for serialized_dag in self.session.query(
+
+        session_query = self.session.query(
             SerializedDagModel.dag_id,
             json_data_column,
             SerializedDagModel.fileloc,
-        ).yield_per(100):
+        )
+        if not self.source_config.includeUnDeployedPipelines:
+            session_query = session_query.select_from(
+                join(
+                    SerializedDagModel,
+                    DagModel,
+                    SerializedDagModel.dag_id == DagModel.dag_id,
+                )
+            ).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(
@@ -417,17 +433,17 @@ class AirflowSource(PipelineServiceSource):
                 tasks=self.get_tasks_from_dag(
                     pipeline_details, self.service_connection.hostPort
                 ),
-                service=self.context.pipeline_service,
+                service=self.context.get().pipeline_service,
                 owner=self.get_owner(pipeline_details.owner),
                 scheduleInterval=pipeline_details.schedule_interval,
             )
             yield Either(right=pipeline_request)
             self.register_record(pipeline_request=pipeline_request)
-            self.context.task_names = {
+            self.context.get().task_names = {
                 task.name for task in pipeline_request.tasks or []
             }
         except TypeError as err:
-            self.context.task_names = set()
+            self.context.get().task_names = set()
             yield Either(
                 left=StackTraceError(
                     name=pipeline_details.dag_id,
@@ -439,7 +455,7 @@ class AirflowSource(PipelineServiceSource):
                 )
             )
         except ValidationError as err:
-            self.context.task_names = set()
+            self.context.get().task_names = set()
             yield Either(
                 left=StackTraceError(
                     name=pipeline_details.dag_id,
@@ -449,7 +465,7 @@ class AirflowSource(PipelineServiceSource):
             )
 
         except Exception as err:
-            self.context.task_names = set()
+            self.context.get().task_names = set()
             yield Either(
                 left=StackTraceError(
                     name=pipeline_details.dag_id,
@@ -472,8 +488,8 @@ class AirflowSource(PipelineServiceSource):
         pipeline_fqn = fqn.build(
             metadata=self.metadata,
             entity_type=Pipeline,
-            service_name=self.context.pipeline_service,
-            pipeline_name=self.context.pipeline,
+            service_name=self.context.get().pipeline_service,
+            pipeline_name=self.context.get().pipeline,
         )
         pipeline_entity = self.metadata.get_by_name(entity=Pipeline, fqn=pipeline_fqn)
         if not pipeline_entity:

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

@@ -68,7 +68,9 @@ class DagsterSource(PipelineServiceSource):
     """
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: DagsterConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, DagsterConnection):
@@ -91,8 +93,8 @@ class DagsterSource(PipelineServiceSource):
         """Method to collect all the tasks from dagster and return it in a task list"""
         jobs = self.client.get_jobs(
             pipeline_name=pipeline_name,
-            repository_name=self.context.repository_name,
-            repository_location=self.context.repository_location,
+            repository_name=self.context.get().repository_name,
+            repository_location=self.context.get().repository_location,
         )
         task_list: List[Task] = []
         if jobs:
@@ -126,10 +128,10 @@ class DagsterSource(PipelineServiceSource):
                 displayName=pipeline_details.name,
                 description=pipeline_details.description,
                 tasks=self._get_task_list(pipeline_name=pipeline_details.name),
-                service=self.context.pipeline_service,
+                service=self.context.get().pipeline_service,
                 tags=get_tag_labels(
                     metadata=self.metadata,
-                    tags=[self.context.repository_name],
+                    tags=[self.context.get().repository_name],
                     classification_name=DAGSTER_TAG_CATEGORY,
                     include_tags=self.source_config.includeTags,
                 ),
@@ -150,7 +152,7 @@ class DagsterSource(PipelineServiceSource):
 
     def yield_tag(self, *_, **__) -> Iterable[Either[OMetaTagAndClassification]]:
         yield from get_ometa_tag_and_classification(
-            tags=[self.context.repository_name],
+            tags=[self.context.get().repository_name],
             classification_name=DAGSTER_TAG_CATEGORY,
             tag_description="Dagster Tag",
             classification_description="Tags associated with dagster entities",
@@ -187,8 +189,8 @@ class DagsterSource(PipelineServiceSource):
             pipeline_fqn = fqn.build(
                 metadata=self.metadata,
                 entity_type=Pipeline,
-                service_name=self.context.pipeline_service,
-                pipeline_name=self.context.pipeline,
+                service_name=self.context.get().pipeline_service,
+                pipeline_name=self.context.get().pipeline,
             )
             pipeline_status_yield = OMetaPipelineStatus(
                 pipeline_fqn=pipeline_fqn,
@@ -211,8 +213,8 @@ class DagsterSource(PipelineServiceSource):
         pipeline_fqn = fqn.build(
             metadata=self.metadata,
             entity_type=Pipeline,
-            service_name=self.context.pipeline_service,
-            pipeline_name=self.context.pipeline,
+            service_name=self.context.get().pipeline_service,
+            pipeline_name=self.context.get().pipeline,
         )
         pipeline_entity = self.metadata.get_by_name(
             entity=Pipeline, fqn=pipeline_fqn, fields=["tasks"]
@@ -222,8 +224,8 @@ class DagsterSource(PipelineServiceSource):
                 runs = self.client.get_task_runs(
                     task.name,
                     pipeline_name=pipeline_details.name,
-                    repository_name=self.context.repository_name,
-                    repository_location=self.context.repository_location,
+                    repository_name=self.context.get().repository_name,
+                    repository_location=self.context.get().repository_location,
                 )
                 for run in runs.solidHandle.stepStats.nodes or []:
                     yield from self._get_task_status(run=run, task_name=task.name)
@@ -248,8 +250,8 @@ class DagsterSource(PipelineServiceSource):
         try:
             results = self.client.get_run_list()
             for result in results:
-                self.context.repository_location = result.location.name
-                self.context.repository_name = result.name
+                self.context.get().repository_location = result.location.name
+                self.context.get().repository_name = result.name
                 for job in result.pipelines or []:
                     yield job
         except Exception as exc:
@@ -272,7 +274,7 @@ class DagsterSource(PipelineServiceSource):
         try:
             url = (
                 f"{clean_uri(self.service_connection.host)}/locations/"
-                f"{self.context.repository_location}/jobs/{pipeline_name}/"
+                f"{self.context.get().repository_location}/jobs/{pipeline_name}/"
             )
             if task_name:
                 url = f"{url}{task_name}"

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

@@ -14,7 +14,7 @@ Databricks pipeline source to extract metadata
 """
 
 import traceback
-from typing import Any, Iterable, List
+from typing import Any, Iterable, List, Optional
 
 from pydantic import ValidationError
 
@@ -68,7 +68,9 @@ class DatabrickspipelineSource(PipelineServiceSource):
     """
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         """Create class instance"""
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: DatabricksPipelineConnection = (
@@ -91,14 +93,14 @@ class DatabrickspipelineSource(PipelineServiceSource):
         self, pipeline_details: Any
     ) -> Iterable[Either[CreatePipelineRequest]]:
         """Method to Get Pipeline Entity"""
-        self.context.job_id_list = []
+        self.context.get().job_id_list = []
         try:
             pipeline_request = CreatePipelineRequest(
                 name=pipeline_details["job_id"],
                 displayName=pipeline_details["settings"].get("name"),
                 description=pipeline_details["settings"].get("name"),
                 tasks=self.get_tasks(pipeline_details),
-                service=self.context.pipeline_service,
+                service=self.context.get().pipeline_service,
             )
             yield Either(right=pipeline_request)
             self.register_record(pipeline_request=pipeline_request)
@@ -133,7 +135,7 @@ class DatabrickspipelineSource(PipelineServiceSource):
 
     def get_tasks(self, pipeline_details: dict) -> List[Task]:
         task_list = []
-        self.context.append(key="job_id_list", value=pipeline_details["job_id"])
+        self.context.get().append(key="job_id_list", value=pipeline_details["job_id"])
 
         downstream_tasks = self.get_downstream_tasks(
             pipeline_details["settings"].get("tasks")
@@ -192,7 +194,7 @@ class DatabrickspipelineSource(PipelineServiceSource):
         return dependent_tasks
 
     def yield_pipeline_status(self, pipeline_details) -> Iterable[OMetaPipelineStatus]:
-        for job_id in self.context.job_id_list:
+        for job_id in self.context.get().job_id_list:
             try:
                 runs = self.client.get_job_runs(job_id=job_id)
                 for attempt in runs or []:
@@ -227,8 +229,8 @@ class DatabrickspipelineSource(PipelineServiceSource):
                         pipeline_fqn = fqn.build(
                             metadata=self.metadata,
                             entity_type=Pipeline,
-                            service_name=self.context.pipeline_service,
-                            pipeline_name=self.context.pipeline,
+                            service_name=self.context.get().pipeline_service,
+                            pipeline_name=self.context.get().pipeline,
                         )
                         yield Either(
                             right=OMetaPipelineStatus(

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

@@ -59,7 +59,9 @@ class DomopipelineSource(PipelineServiceSource):
     """
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config = WorkflowSource.parse_obj(config_dict)
         connection: DomoPipelineConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, DomoPipelineConnection):
@@ -94,7 +96,7 @@ class DomopipelineSource(PipelineServiceSource):
                 displayName=pipeline_details.get("name"),
                 description=pipeline_details.get("description", ""),
                 tasks=[task],
-                service=self.context.pipeline_service,
+                service=self.context.get().pipeline_service,
                 startDate=pipeline_details.get("created"),
                 sourceUrl=source_url,
             )
@@ -164,8 +166,8 @@ class DomopipelineSource(PipelineServiceSource):
                 pipeline_fqn = fqn.build(
                     metadata=self.metadata,
                     entity_type=Pipeline,
-                    service_name=self.context.pipeline_service,
-                    pipeline_name=self.context.pipeline,
+                    service_name=self.context.get().pipeline_service,
+                    pipeline_name=self.context.get().pipeline,
                 )
                 yield Either(
                     right=OMetaPipelineStatus(

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

@@ -67,7 +67,9 @@ class FivetranSource(PipelineServiceSource):
     """
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: FivetranConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, FivetranConnection):
@@ -97,7 +99,7 @@ class FivetranSource(PipelineServiceSource):
             name=pipeline_details.pipeline_name,
             displayName=pipeline_details.pipeline_display_name,
             tasks=self.get_connections_jobs(pipeline_details),
-            service=self.context.pipeline_service,
+            service=self.context.get().pipeline_service,
             sourceUrl=self.get_source_url(
                 connector_id=pipeline_details.source.get("id"),
                 group_id=pipeline_details.group.get("id"),
@@ -163,8 +165,8 @@ class FivetranSource(PipelineServiceSource):
                 pipeline_fqn = fqn.build(
                     metadata=self.metadata,
                     entity_type=Pipeline,
-                    service_name=self.context.pipeline_service,
-                    pipeline_name=self.context.pipeline,
+                    service_name=self.context.get().pipeline_service,
+                    pipeline_name=self.context.get().pipeline,
                 )
                 pipeline_entity = self.metadata.get_by_name(
                     entity=Pipeline, fqn=pipeline_fqn

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

@@ -14,7 +14,7 @@ Glue pipeline source to extract metadata
 """
 
 import traceback
-from typing import Any, Iterable, List
+from typing import Any, Iterable, List, Optional
 
 from metadata.generated.schema.api.data.createPipeline import CreatePipelineRequest
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
@@ -72,7 +72,9 @@ class GluepipelineSource(PipelineServiceSource):
         self.glue = self.connection
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: GluePipelineConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, GluePipelineConnection):
@@ -103,7 +105,7 @@ class GluepipelineSource(PipelineServiceSource):
             name=pipeline_details[NAME],
             displayName=pipeline_details[NAME],
             tasks=self.get_tasks(pipeline_details),
-            service=self.context.pipeline_service,
+            service=self.context.get().pipeline_service,
             sourceUrl=source_url,
         )
         yield Either(right=pipeline_request)
@@ -172,8 +174,8 @@ class GluepipelineSource(PipelineServiceSource):
                     pipeline_fqn = fqn.build(
                         metadata=self.metadata,
                         entity_type=Pipeline,
-                        service_name=self.context.pipeline_service,
-                        pipeline_name=self.context.pipeline,
+                        service_name=self.context.get().pipeline_service,
+                        pipeline_name=self.context.get().pipeline,
                     )
                     yield Either(
                         right=OMetaPipelineStatus(

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

@@ -81,7 +81,9 @@ class NifiSource(PipelineServiceSource):
     """
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: NifiConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, NifiConnection):
@@ -142,7 +144,7 @@ class NifiSource(PipelineServiceSource):
             displayName=pipeline_details.name,
             sourceUrl=f"{clean_uri(self.service_connection.hostPort)}{pipeline_details.uri}",
             tasks=self._get_tasks_from_details(pipeline_details),
-            service=self.context.pipeline_service,
+            service=self.context.get().pipeline_service,
         )
         yield Either(right=pipeline_request)
         self.register_record(pipeline_request=pipeline_request)

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

@@ -0,0 +1,88 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Source connection handler
+"""
+from typing import Optional
+
+from confluent_kafka import Consumer as KafkaConsumer
+from confluent_kafka import TopicPartition
+
+from metadata.generated.schema.entity.automations.workflow import (
+    Workflow as AutomationWorkflow,
+)
+from metadata.generated.schema.entity.services.connections.pipeline.openLineageConnection import (
+    OpenLineageConnection,
+)
+from metadata.generated.schema.entity.services.connections.pipeline.openLineageConnection import (
+    SecurityProtocol as KafkaSecProtocol,
+)
+from metadata.ingestion.connections.test_connections import (
+    SourceConnectionException,
+    test_connection_steps,
+)
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+
+
+def get_connection(connection: OpenLineageConnection) -> KafkaConsumer:
+    """
+    Create connection
+    """
+    try:
+        config = {
+            "bootstrap.servers": connection.brokersUrl,
+            "group.id": connection.consumerGroupName,
+            "auto.offset.reset": connection.consumerOffsets.value,
+        }
+        if connection.securityProtocol.value == KafkaSecProtocol.SSL.value:
+            config.update(
+                {
+                    "security.protocol": connection.securityProtocol.value,
+                    "ssl.ca.location": connection.sslConfig.__root__.caCertificate,
+                    "ssl.certificate.location": connection.sslConfig.__root__.sslCertificate,
+                    "ssl.key.location": connection.sslConfig.__root__.sslKey,
+                }
+            )
+
+        kafka_consumer = KafkaConsumer(config)
+        kafka_consumer.subscribe([connection.topicName])
+
+        return kafka_consumer
+    except Exception as exc:
+        msg = f"Unknown error connecting with {connection}: {exc}."
+        raise SourceConnectionException(msg)
+
+
+def test_connection(
+    metadata: OpenMetadata,
+    client: KafkaConsumer,
+    service_connection: OpenLineageConnection,
+    automation_workflow: Optional[AutomationWorkflow] = None,
+) -> None:
+    """
+    Test connection. This can be executed either as part
+    of a metadata workflow or during an Automation Workflow
+    """
+
+    def custom_executor():
+        _ = client.get_watermark_offsets(
+            TopicPartition(service_connection.topicName, 0)
+        )
+
+    test_fn = {"GetWatermarkOffsets": custom_executor}
+
+    test_connection_steps(
+        metadata=metadata,
+        test_fn=test_fn,
+        service_type=service_connection.type.value,
+        automation_workflow=automation_workflow,
+    )