ferruzzi commented on code in PR #70635:
URL: https://github.com/apache/airflow/pull/70635#discussion_r3693687555
##########
providers/amazon/src/airflow/providers/amazon/aws/log/cloudwatch_task_handler.py:
##########
@@ -271,8 +274,98 @@ def stream(self, relative_path: str, ti: RuntimeTI) ->
StreamingLogResponse:
except Exception as e:
messages.append(str(e))
+ def _read_trigger_stream(name: str) -> RawLogStream:
+ return (
+ self._parse_log_event_as_dumped_json(event) for event in
self.get_cloudwatch_logs(name, ti)
+ )
+
+ extra_messages, extra_logs = self.merge_trigger_logs(relative_path,
ti, _read_trigger_stream)
+ messages.extend(extra_messages)
+ logs.extend(extra_logs)
+
return messages, logs
+ def get_trigger_stream_names(self, relative_path: str) -> list[str]:
+ """
+ Return the names of any trigger log streams associated with
``relative_path``.
+
+ A task instance can be deferred and resumed more than once -- even by
different
+ triggerer processes -- so there may be more than one ``.trigger.<job
id>.log``
+ stream for a single attempt. Sorted numerically by job id so multiple
deferrals of
+ the same attempt read back in the order they actually happened (job
ids are assigned
+ from a monotonically increasing DB sequence); anything that doesn't
parse as an
+ integer job id sorts after, by name, rather than raising.
+ """
+ prefix = f"{relative_path.replace(':', '_')}.trigger."
+ streams = self.hook.describe_log_streams(log_group=self.log_group,
log_stream_name_prefix=prefix)
+ names = [name for s in streams if (name := s.get("logStreamName"))]
+
+ def _sort_key(name: str) -> tuple[int, int | str]:
+ tail = name[len(prefix) :]
+ job_id_str = tail[: -len(".log")] if tail.endswith(".log") else
tail
+ if job_id_str.isdigit():
+ return (0, int(job_id_str))
+ return (1, name)
+
+ return sorted(names, key=_sort_key)
+
+ def merge_trigger_logs(
+ self,
+ relative_path: str,
+ ti: RuntimeTI,
+ read_stream: Callable[[str], T],
+ ) -> tuple[list[str], list[T]]:
+ """
+ Discover trigger log streams for ``relative_path`` and read each with
``read_stream``.
+
+ Shared by :meth:`stream` and the legacy
:meth:`CloudwatchTaskHandler._read_remote_logs`
+ so both read/format trigger-log content in their own way (raw
JSON-line generators vs.
+ joined strings) while sharing the discovery/ordering/exclusion logic.
See :meth:`stream`
+ for why this lookup exists at all -- CloudWatch has no glob()
equivalent to find trigger
+ streams the way the local-file reader does.
+
+ While the task instance is actively DEFERRED, the stream for the
*currently active*
+ triggerer job is excluded: the UI already tails that one live from the
triggerer over
+ HTTP (``FileTaskHandler._read_from_logs_server``), so re-reading it
here would just be
+ redundant, and on every UI poll during a long-running deferral.
Streams from any
+ *earlier* deferral of the same attempt (a task deferred, resumed, and
deferred again)
+ are still included -- the live HTTP tail only ever covers the current
triggerer job, so
+ without this those earlier logs would otherwise be invisible while
re-deferred.
+
+ :param relative_path: The task's own (base) log stream name.
+ :param ti: The task instance being read for.
+ :param read_stream: Called with each trigger stream's name; return its
content in
+ whatever form the caller needs (e.g. a generator of parsed lines,
or a joined str).
+ :return: (extra messages to append, one result per successfully-read
trigger stream).
+ """
+ messages: list[str] = []
+ results: list[T] = []
+ try:
+ trigger_stream_names = self.get_trigger_stream_names(relative_path)
+ except Exception as e:
+ messages.append(f"Could not list trigger log streams for
{relative_path}: {e}")
+ return messages, results
+
+ if getattr(ti, "state", None) == TaskInstanceState.DEFERRED:
+ current_job_id = getattr(getattr(ti, "triggerer_job", None), "id",
None)
+ if current_job_id is not None:
+ live_suffix = f".trigger.{current_job_id}.log"
+ trigger_stream_names = [
+ name for name in trigger_stream_names if not
name.endswith(live_suffix)
+ ]
Review Comment:
So many stacked `getattrs` here makes me think there should be a better way.
Can you try this, and see if it works? It may need a little tweaking, but
should be cleaner, I think:
```python
if (ti.state == TaskInstanceState.DEFERRED) and (job := getattr(ti,
"triggerer_job", None)):
live_stream = f"{relative_path.replace(':', '_')}.trigger.{job.id}.log"
trigger_stream_names = [name for name in trigger_stream_names if name !=
live_stream]
```
Plus, hardcoding the prefix assembly here makes me think we should likely
move that builder into a shared helper at some point, but that can be done in a
different PR.
--
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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
