Add more logs

astronomer · Dec 20, 2024 · a1d63ee · a1d63ee
1 parent fc7b502
commit a1d63ee
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/cosmos/listeners/dag_run_listener.py b/cosmos/listeners/dag_run_listener.py
@@ -19,6 +19,7 @@ class EventStatus:
 
 
 def is_cosmos_dag(dag: DAG) -> bool:
+    logger.info("is_cosmos_dag: {dag.__class__.__module__}")
     if dag.__class__.__module__.startswith("cosmos."):
         return True
     return False
@@ -27,6 +28,7 @@ def is_cosmos_dag(dag: DAG) -> bool:
 def total_cosmos_task_groups(dag: DAG) -> int:
     cosmos_task_groups = 0
     for group_id, task_group in dag.task_group_dict.items():
+        logger.info("total_cosmos_task_groups: {task_group.__class__}")
         if task_group.__class__.__module__.startswith("cosmos."):
             cosmos_task_groups += 1
     return cosmos_task_groups
@@ -35,9 +37,8 @@ def total_cosmos_task_groups(dag: DAG) -> int:
 def total_cosmos_tasks(dag: DAG) -> int:
     cosmos_tasks = 0
     for task in dag.tasks:
-        task_class = type(task)
-        task_module = task_class.__module__
-        if task_module.startswith("cosmos."):
+        logger.info("total_cosmos_task_groups: {task.__class__.__module__}")
+        if task.__class__.__module__.startswith("cosmos."):
             cosmos_tasks += 1
     return cosmos_tasks