openmetadata-ingestion 1.3.4.0__py3-none-any.whl → 1.4.0.0__py3-none-any.whl

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

@@ -11,6 +11,8 @@
 """
 RedPanda source ingestion
 """
+from typing import Optional
+
 from metadata.generated.schema.entity.services.connections.messaging.redpandaConnection import (
     RedpandaConnection,
 )
@@ -24,7 +26,9 @@ from metadata.ingestion.source.messaging.common_broker_source import CommonBroke
 
 class RedpandaSource(CommonBrokerSource):
     @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: RedpandaConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, RedpandaConnection):

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

@@ -126,7 +126,9 @@ class AmundsenSource(Source):
         self.test_connection()
 
     @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: AmundsenConnection = config.serviceConnection.__root__.config

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

@@ -93,7 +93,9 @@ class AtlasSource(Source):
         self.test_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: AtlasConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, AtlasConnection):

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

@@ -55,7 +55,9 @@ class MlflowSource(MlModelServiceSource):
     """
 
     @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: MlflowConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, MlflowConnection):
@@ -126,7 +128,7 @@ class MlflowSource(MlModelServiceSource):
                 run.data, latest_version.run_id, model.name
             ),
             mlStore=self._get_ml_store(latest_version),
-            service=self.context.mlmodel_service,
+            service=self.context.get().mlmodel_service,
             sourceUrl=source_url,
         )
         yield Either(right=mlmodel_request)

metadata/ingestion/source/mlmodel/mlmodel_service.py

@@ -39,7 +39,7 @@ from metadata.ingestion.models.delete_entity import DeleteEntity
 from metadata.ingestion.models.topology import (
     NodeStage,
     ServiceTopology,
-    TopologyContext,
+    TopologyContextManager,
     TopologyNode,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
@@ -100,7 +100,7 @@ class MlModelServiceSource(TopologyRunnerMixin, Source, ABC):
     service_connection: MlModelConnection.__fields__["config"].type_
 
     topology = MlModelServiceTopology()
-    context = TopologyContext.create(topology)
+    context = TopologyContextManager(topology)
     mlmodel_source_state: Set = set()
 
     def __init__(
@@ -179,7 +179,7 @@ class MlModelServiceSource(TopologyRunnerMixin, Source, ABC):
                 entity_type=MlModel,
                 entity_source_state=self.mlmodel_source_state,
                 mark_deleted_entity=self.source_config.markDeletedMlModels,
-                params={"service": self.context.mlmodel_service},
+                params={"service": self.context.get().mlmodel_service},
             )
 
     def register_record(self, mlmodel_request: CreateMlModelRequest) -> None:

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

@@ -67,7 +67,9 @@ class SagemakerSource(MlModelServiceSource):
         self.sagemaker = self.client
 
     @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: SageMakerConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, SageMakerConnection):
@@ -136,7 +138,7 @@ class SagemakerSource(MlModelServiceSource):
                 name=model.name,
                 algorithm=self._get_algorithm(),  # Setting this to a constant
                 mlStore=self._get_ml_store(model.name),
-                service=self.context.mlmodel_service,
+                service=self.context.get().mlmodel_service,
             )
             yield Either(right=mlmodel_request)
             self.register_record(mlmodel_request=mlmodel_request)
@@ -158,7 +160,10 @@ class SagemakerSource(MlModelServiceSource):
         """
         try:
             model_info = self.sagemaker.describe_model(ModelName=model_name)
-            return MlStore(imageRepository=model_info["PrimaryContainer"]["Image"])
+            storage = model_info.get("PrimaryContainer", {}).get("ModelDataUrl")
+            image_repository = model_info.get("PrimaryContainer", {}).get("Image")
+            if image_repository or storage:
+                return MlStore(storage=storage, imageRepository=image_repository)
         except ValidationError as err:
             logger.debug(traceback.format_exc())
             logger.warning(

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

@@ -12,7 +12,7 @@
 Airbyte source to extract metadata
 """
 
-from typing import Iterable
+from typing import Iterable, Optional
 
 from pydantic import BaseModel
 
@@ -75,7 +75,9 @@ class AirbyteSource(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: AirbyteConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, AirbyteConnection):
@@ -116,7 +118,7 @@ class AirbyteSource(PipelineServiceSource):
             tasks=self.get_connections_jobs(
                 pipeline_details.connection, connection_url
             ),
-            service=self.context.pipeline_service,
+            service=self.context.get().pipeline_service,
         )
         yield Either(right=pipeline_request)
         self.register_record(pipeline_request=pipeline_request)
@@ -171,8 +173,8 @@ class AirbyteSource(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(
@@ -235,8 +237,8 @@ class AirbyteSource(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/airflow/lineage_parser.py

@@ -77,10 +77,12 @@ from pydantic import BaseModel
 
 from metadata.generated.schema.entity.data.table import Table
 from metadata.ingestion.ometa.models import T
+from metadata.utils.constants import ENTITY_REFERENCE_CLASS_MAP
 from metadata.utils.deprecation import deprecated
 from metadata.utils.importer import import_from_module
 
 logger = logging.getLogger("airflow.task")
+XLET_KEYS = {"entity", "fqn", "key"}
 
 
 class XLetsMode(Enum):
@@ -213,7 +215,7 @@ def dictionary_lineage_annotation(xlet: dict) -> Dict[str, List[OMEntity]]:
     """
     Handle OM specific inlet/outlet information. E.g.,
 
-    ```
+    ```python
     BashOperator(
         task_id="print_date",
         bash_command="date",
@@ -224,12 +226,63 @@ def dictionary_lineage_annotation(xlet: dict) -> Dict[str, List[OMEntity]]:
         },
     )
     ```
+
+    The Serialized DAG for the old lineage style will look like:
+    ```
+    "inlets": [
+          {
+            "__var": {
+              "tables": [
+                "sample_data.ecommerce_db.shopify.raw_order"
+              ]
+            },
+            "__type": "dict"
+          }
+        ],
+    ```
+    With the new lineage style where we annotate tasks' lineage as:
+    ```python
+    BashOperator(
+        task_id="print_date",
+        bash_command="date",
+        inlets=[
+            {"entity": "container", "fqn": "s3_storage_sample.departments", "key": "test"},
+        ],
+    )
+    ```,
+    the Serialized DAG looks like
+    ```
+    "inlets": [
+          {
+            "__var": {
+              "fqn": "s3_storage_sample.departments",
+              "key": "test",
+              "entity": "container"
+            },
+            "__type": "dict"
+          }
+        ],
+    ```
+    To validate if we are on the first on latter case, we can try to parse the available dict to an OMEntity.
     """
     xlet_dict = xlet
     # This is how the Serialized DAG is giving us the info from _inlets & _outlets
     if isinstance(xlet_dict, dict) and xlet_dict.get("__var"):
         xlet_dict = xlet_dict["__var"]
 
+    # Check if we can properly build the OMEntity class
+    if XLET_KEYS.issubset(xlet_dict.keys()):
+        return {
+            xlet_dict["key"]: [
+                OMEntity(
+                    entity=ENTITY_REFERENCE_CLASS_MAP[xlet_dict["entity"]],
+                    fqn=xlet_dict["fqn"],
+                    key=xlet_dict["key"],
+                )
+            ]
+        }
+
+    # Otherwise, fall to the old lineage style
     return {
         key: [
             # We will convert the old dict lineage method into Tables

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/dagster/queries.py

@@ -14,62 +14,30 @@ GraphQL queries used during ingestion
 """
 
 DAGSTER_PIPELINE_DETAILS_GRAPHQL = """
-            query PipelineRuns{
-  repositoriesOrError{
-    __typename 
-    ... on RepositoryConnection{
-      nodes{
-        id,
-        name,
+query PipelineRuns {
+  repositoriesOrError {
+    __typename
+    ... on RepositoryConnection {
+      nodes {
+        id
+        name
         location {
           id
           name
         }
-        pipelines{
-          id,
-          name,
-          description,
-          runs{
-            id,
-            runId,
-            status,
-            startTime,
-            stats{
-                ... on RunStatsSnapshot {
-                        startTime
-                        endTime
-                        stepsFailed
-                    }
-                }
-            executionPlan{
-              steps{
-                inputs{
-                  name,
-                  dependsOn{
-                    inputs{
-                      name
-                    },
-                    outputs{
-                      name
-                    }
-                  }
-                },
-                outputs{
-                  name
-                }
-              }
-
-            }
-          }
+        pipelines {
+          id
+          name
+          description
         }
       }
     }
-
   }
 }
 """
 
-GRAPHQL_RUNS_QUERY = """query SidebarOpGraphsQuery($selector: PipelineSelector!, $handleID: String!) {
+GRAPHQL_RUNS_QUERY = """
+query SidebarOpGraphsQuery($selector: PipelineSelector!, $handleID: String!) {
   pipelineOrError(params: $selector) {
     __typename
     ... on Pipeline {
@@ -100,41 +68,44 @@ GRAPHQL_RUNS_QUERY = """query SidebarOpGraphsQuery($selector: PipelineSelector!,
 }
 """
 
-TEST_QUERY_GRAPHQL = """query Pipeline {
-            pipelineRunsOrError {
-                __typename
-                ... on PipelineRuns {
-                results {
-                    assetSelection {
-                    path
-                    }
-                    runId
-                    status
-                }
-                }
-            }
-            }"""
-
-GRAPHQL_QUERY_FOR_JOBS = """query PipelineRuns($selector: GraphSelector!){
-                graphOrError(selector: $selector){
-                    ... on Graph{
-                    id,
-                    name,
-                    description,
-                        solidHandles{
-                            handleID,
-                            solid{
-                                name,
-                                inputs{
-                                    dependsOn {
-                                        solid {
-                                            name
-                                        }
-                                    }
-                                },
-                            }
-                        },
+TEST_QUERY_GRAPHQL = """
+query Pipeline {
+  pipelineRunsOrError(limit: 1) {
+    __typename
+    ... on PipelineRuns {
+      results {
+        assetSelection {
+          path
+        }
+        runId
+        status
+      }
+    }
+  }
+}
+"""
 
-                    }
-                }
-                }"""
+GRAPHQL_QUERY_FOR_JOBS = """
+query PipelineRuns($selector: GraphSelector!) {
+  graphOrError(selector: $selector) {
+    ... on Graph {
+      id
+      name
+      description
+      solidHandles {
+        handleID
+        solid {
+          name
+          inputs {
+            dependsOn {
+              solid {
+                name
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+}
+"""