diff --git a/python-sdk/src/astro/sql/operators/cleanup.py b/python-sdk/src/astro/sql/operators/cleanup.py index 31538a7938..084289496e 100644 --- a/python-sdk/src/astro/sql/operators/cleanup.py +++ b/python-sdk/src/astro/sql/operators/cleanup.py @@ -232,7 +232,14 @@ def resolve_tables_from_tasks( # noqa: C901 and isinstance(task, MappedOperator) and issubclass(task.operator_class, OPERATOR_CLASSES_WITH_TABLE_OUTPUT) ): - for t in task.output.resolve(context): + from airflow.models.xcom_arg import XComArg + + try: + task_output = task.output + except AttributeError: + task_output = XComArg(operator=task) + + for t in task_output.resolve(context): if isinstance(t, BaseTable): res.append(t) except AirflowException: # pragma: no cover