feng-tao commented on a change in pull request #6564: [AIRFLOW-5911] Simplify 
lineage API and improve robustness
URL: https://github.com/apache/airflow/pull/6564#discussion_r347732639
 
 

 ##########
 File path: airflow/lineage/__init__.py
 ##########
 @@ -92,51 +116,46 @@ def prepare_lineage(func):
     * "auto" -> picks up any outlets from direct upstream tasks that have 
outlets defined, as such that
       if A -> B -> C and B does not have outlets but A does, these are 
provided as inlets.
     * "list of task_ids" -> picks up outlets from the upstream task_ids
-    * "list of datasets" -> manually defined list of DataSet
+    * "list of datasets" -> manually defined list of data
 
     """
     @wraps(func)
     def wrapper(self, context, *args, **kwargs):
         self.log.debug("Preparing lineage inlets and outlets")
 
-        task_ids = set(self._inlets['task_ids']).intersection(  # 
pylint:disable=protected-access
-            self.get_flat_relative_ids(upstream=True)
-        )
-        if task_ids:
-            inlets = self.xcom_pull(context,
-                                    task_ids=task_ids,
-                                    dag_id=self.dag_id,
-                                    key=PIPELINE_OUTLETS)
-            inlets = [item for sublist in inlets if sublist for item in 
sublist]
-            inlets = [DataSet.map_type(i['typeName'])(data=i['attributes'])
-                      for i in inlets]
-            self.inlets.extend(inlets)
+        if isinstance(self._inlets, str):
+            self._inlets = [self._inlets, ]
+
+        if isinstance(self._inlets, list):
+            task_ids = set(
+                filter(lambda x: isinstance(x, str) and x.lower() != AUTO, 
self._inlets)
+            ).union(
+                map(lambda op: op.task_id,
+                    filter(lambda op: isinstance(op, Operator), self._inlets))
+            ).intersection(self.get_flat_relative_ids(upstream=True))
 
-        if self._inlets['auto']:  # pylint:disable=protected-access
-            # dont append twice
-            task_ids = set(self._inlets['task_ids']).symmetric_difference(  # 
pylint:disable=protected-access
-                self.upstream_task_ids
-            )
-            inlets = self.xcom_pull(context,
-                                    task_ids=task_ids,
-                                    dag_id=self.dag_id,
-                                    key=PIPELINE_OUTLETS)
-            inlets = [item for sublist in inlets if sublist for item in 
sublist]
-            inlets = [DataSet.map_type(i['typeName'])(data=i['attributes'])
-                      for i in inlets]
+            if AUTO.upper() in self._inlets or AUTO.lower() in self._inlets:
+                task_ids = 
task_ids.union(task_ids.symmetric_difference(self.upstream_task_ids))
+
+            inlets = self.xcom_pull(context, task_ids=task_ids,
+                                    dag_id=self.dag_id, key=PIPELINE_OUTLETS)
+            inlets = [_get_instance(item) for sublist in inlets if sublist for 
item in sublist]
             self.inlets.extend(inlets)
 
-        if self._inlets['datasets']:  # pylint:disable=protected-access
-            self.inlets.extend(self._inlets['datasets'])  # 
pylint:disable=protected-access
+            self.inlets.extend([_get_instance(_as_rendered_dict(i, context))
 
 Review comment:
   if we put antlar to parse sql, will it put in the render_template field?

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
[email protected]


With regards,
Apache Git Services

Reply via email to