Feature/update logic mssql lineage

metadata-ingestion/src/datahub/ingestion/source/sql/mssql/job_models.py

@@ -1,7 +1,11 @@
 from dataclasses import dataclass, field
 from typing import Dict, List, Optional, Union
 
-from datahub.emitter.mce_builder import make_data_flow_urn, make_data_job_urn
+from datahub.emitter.mce_builder import (
+    make_data_flow_urn,
+    make_data_job_urn,
+    make_dataset_urn_with_platform_instance,
+)
 from datahub.metadata.schema_classes import (
     DataFlowInfoClass,
     DataJobInfoClass,
@@ -11,18 +15,61 @@
 
 @dataclass
 class ProcedureDependency:
+    flow_id: str
+    env: str
+    server: Optional[str]
     db: str
     schema: str
     name: str
     type: str
-    env: str
-    server: Optional[str]
+    incoming: int
+    outgoing: int
     source: str = "mssql"
 
 
 @dataclass
 class ProcedureLineageStream:
-    dependencies: List[ProcedureDependency]
+    dependencies: List[ProcedureDependency] = field(default_factory=list)
+
+    @property
+    def get_input_datasets(self) -> List[str]:
+        return [
+            make_dataset_urn_with_platform_instance(
+                platform=dep.source,
+                name=".".join([dep.db, dep.schema, dep.name]),
+                platform_instance=dep.server,
+                env=dep.env,
+            )
+            for dep in self.dependencies
+            if dep.type in ("U ", "V ") and dep.incoming
+        ]
+
+    @property
+    def get_output_datasets(self) -> List[str]:
+        return [
+            make_dataset_urn_with_platform_instance(
+                platform=dep.source,
+                name=".".join([dep.db, dep.schema, dep.name]),
+                platform_instance=dep.server,
+                env=dep.env,
+            )
+            for dep in self.dependencies
+            if dep.type in ("U ", "V ") and dep.outgoing
+        ]
+
+    @property
+    def get_input_datajobs(self) -> List[str]:
+        return [
+            make_data_job_urn(
+                orchestrator=dep.source,
+                flow_id=dep.flow_id,
+                job_id=dep.name,
+                cluster=dep.env,
+                platform_instance=dep.server,
+            )
+            for dep in self.dependencies
+            if dep.type in ("P ",) and dep.incoming
+        ]
 
     @property
     def as_property(self) -> Dict[str, str]:
@@ -65,6 +112,7 @@ class MSSQLProceduresContainer:
     env: str
     source: str = "mssql"
     type: str = "JOB"
+    schema: str = ""
 
     @property
     def formatted_name(self) -> str:
@@ -183,15 +231,15 @@ def valued_properties(self) -> Dict[str, str]:
         return self.job_properties
 
     @property
-    def as_datajob_input_output_aspect(self) -> DataJobInputOutputClass:
+    def get_datajob_input_output_aspect(self) -> DataJobInputOutputClass:
         return DataJobInputOutputClass(
             inputDatasets=sorted(self.incoming),
             outputDatasets=sorted(self.outgoing),
             inputDatajobs=sorted(self.input_jobs),
         )
 
     @property
-    def as_datajob_info_aspect(self) -> DataJobInfoClass:
+    def get_datajob_info_aspect(self) -> DataJobInfoClass:
         return DataJobInfoClass(
             name=self.entity.full_name,
             type=self.entity.full_type,
@@ -229,7 +277,7 @@ def urn(self) -> str:
         )
 
     @property
-    def as_dataflow_info_aspect(self) -> DataFlowInfoClass:
+    def get_dataflow_info_aspect(self) -> DataFlowInfoClass:
         return DataFlowInfoClass(
             name=self.entity.formatted_name,
             customProperties=self.flow_properties,