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_r350942010
##########
File path: airflow/lineage/__init__.py
##########
@@ -92,52 +120,53 @@ 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
"""
+ # pylint:disable=protected-access
@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 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]
- self.inlets.extend(inlets)
-
- if self._inlets['datasets']: # pylint:disable=protected-access
- self.inlets.extend(self._inlets['datasets']) #
pylint:disable=protected-access
-
- # outlets
- if self._outlets['datasets']: # pylint:disable=protected-access
- self.outlets.extend(self._outlets['datasets']) #
pylint:disable=protected-access
+ if isinstance(self._inlets, (str, Operator)) or attr.has(self._inlets):
+ self._inlets = [self._inlets, ]
- self.log.debug("inlets: %s, outlets: %s", self.inlets, self.outlets)
+ if self._inlets and isinstance(self._inlets, list):
+ # get task_ids that are specified as parameter and make sure they
are upstream
+ 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))
+
+ # pick up unique direct upstream task_ids if AUTO is specified
+ 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))
- for dataset in chain(self.inlets, self.outlets):
- dataset.set_context(context)
+ _inlets = self.xcom_pull(context, task_ids=task_ids,
+ dag_id=self.dag_id, key=PIPELINE_OUTLETS)
+ # re-instantiate and render the obtained inlets
+ _inlets = [_get_instance(structure(item, Metadata))
+ for sublist in _inlets if sublist for item in sublist]
+ _inlets.extend([_render_object(i, context)
+ for i in self._inlets if attr.has(i)])
+
+ self.inlets.extend(_inlets)
+
+ elif self._inlets:
+ raise AttributeError("inlets is not a list, operator, string or
attr annotated object")
Review comment:
not sure if you have mentioned before, but for hive operator, we would like
the outlet to be the s3 file which most of the time embedded in the hql. Is
there a way to figure out the lineage?
----------------------------------------------------------------
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
