Repository: apache/airflow

Issue:
Using TaskGroup without context manager (Graph view visual bug)
**Apache Airflow version**: 2.0.0

**Kubernetes version (if you are using kubernetes)** (use `kubectl version`): n/a

**What happened**:
When I do not use the context manager for the task group and instead call the add function to add the tasks, those tasks show up on the Graph view.
![Screen Shot 2021-03-17 at 2 06 17 PM](https://user-images.githubusercontent.com/5952735/111544849-5939b200-8732-11eb-80dc-89c013aeb083.png) 

However, when I click on the task group item on the Graph UI, it will fix the issue. When I close the task group item, the tasks will not be displayed as expected.
![Screen Shot 2021-03-17 at 2 06 21 PM](https://user-images.githubusercontent.com/5952735/111544848-58a11b80-8732-11eb-928b-3c76207a0107.png)

**What you expected to happen**:
I expected the tasks inside the task group to not display on the Graph view.
![Screen Shot 2021-03-17 at 3 17 34 PM](https://user-images.githubusercontent.com/5952735/111545824-eaf5ef00-8733-11eb-99c2-75b051bfefe1.png)

**How to reproduce it**:
Render this DAG in Airflow

```python
from airflow.models import DAG
from airflow.operators.bash import BashOperator
from airflow.operators.dummy import DummyOperator
from airflow.utils.task_group import TaskGroup
from datetime import datetime
with DAG(dag_id="example_task_group", start_date=datetime(2021, 1, 1), tags=["example"], catchup=False) as dag:
    start = BashOperator(task_id="start", bash_command='echo 1; sleep 10; echo 2;')
    tg = TaskGroup("section_1", tooltip="Tasks for section_1")
    task_1 = DummyOperator(task_id="task_1")
    task_2 = BashOperator(task_id="task_2", bash_command='echo 1')
    task_3 = DummyOperator(task_id="task_3")
    tg.add(task_1)
    tg.add(task_2)
    tg.add(task_3)
```


Discussion:
I was able to reproduce this with the supplied sample DAG.

@yuqian90 any ideas on this? You probably know this graph view task group logic fairly well after your work on #14661.
Hi @ryanahamilton and @wolfier , I can explain why this is happening, but I'll need some suggestions from you and others on how to improve this.
There are two ways of putting a task into a `TaskGroup`:
1. Create the task inside the contextmanager of the `TaskGroup`;
2. Pass the `task_group` as an argument to the constructor of the task.

The reproducing example is doing neither of these. Instead it's calling `TaskGroup.add` **after** the task has already been created. At the time of creating the task, there was no `TaskGroup` contextmanager nor was the `task_group` specified. So the tasks were already added to the root `task_group` of the dag, resulting in this inconsistent behaviour.

To fix this, you just need to use a contextmanager. Or do this:
```python
from airflow.models import DAG
from airflow.operators.bash import BashOperator
from airflow.operators.dummy import DummyOperator
from airflow.utils.task_group import TaskGroup
from datetime import datetime
with DAG(dag_id="example_task_group", start_date=datetime(2021, 1, 1), tags=["example"], catchup=False) as dag:
    start = BashOperator(task_id="start", bash_command='echo 1; sleep 10; echo 2;')
    tg = TaskGroup("section_1", tooltip="Tasks for section_1")
    task_1 = DummyOperator(task_id="task_1", task_group=tg)
    task_2 = BashOperator(task_id="task_2", bash_command='echo 1', task_group=tg)
    task_3 = DummyOperator(task_id="task_3", task_group=tg)
```

Some ideas for improvement, not sure what's best:
1. Add a check somewhere to make sure a task is not added to more than one `TaskGroup`. Where should such a check be?
2. Make `TaskGroup.add` a private method. The reason it currently looks like a public method is that it's called in places such as the constructor of `BaseOperator`. But strictly speaking, users of Airflow should not be encouraged to call `TaskGroup.add` directly because it only does what the name says under strict pre-conditions (i.e. during the construction of the task and not after).
> Add a check somewhere to make sure a task is not added to more than one TaskGroup. Where should such a check be?

We should definitely prevent adding task to more than 1 task group.

Resolving patch:
diff --git a/airflow/exceptions.py b/airflow/exceptions.py
--- a/airflow/exceptions.py
+++ b/airflow/exceptions.py
@@ -185,6 +185,23 @@ class DuplicateTaskIdFound(AirflowException):
     """Raise when a Task with duplicate task_id is defined in the same DAG."""
 
 
+class TaskAlreadyInTaskGroup(AirflowException):
+    """Raise when a Task cannot be added to a TaskGroup since it already belongs to another TaskGroup."""
+
+    def __init__(self, task_id: str, existing_group_id: Optional[str], new_group_id: str) -> None:
+        super().__init__(task_id, new_group_id)
+        self.task_id = task_id
+        self.existing_group_id = existing_group_id
+        self.new_group_id = new_group_id
+
+    def __str__(self) -> str:
+        if self.existing_group_id is None:
+            existing_group = "the DAG's root group"
+        else:
+            existing_group = f"group {self.existing_group_id!r}"
+        return f"cannot add {self.task_id!r} to {self.new_group_id!r} (already in {existing_group})"
+
+
 class SerializationError(AirflowException):
     """A problem occurred when trying to serialize a DAG."""
 
diff --git a/airflow/utils/task_group.py b/airflow/utils/task_group.py
--- a/airflow/utils/task_group.py
+++ b/airflow/utils/task_group.py
@@ -24,7 +24,12 @@
 import weakref
 from typing import TYPE_CHECKING, Any, Dict, Generator, List, Optional, Sequence, Set, Tuple, Union
 
-from airflow.exceptions import AirflowDagCycleException, AirflowException, DuplicateTaskIdFound
+from airflow.exceptions import (
+    AirflowDagCycleException,
+    AirflowException,
+    DuplicateTaskIdFound,
+    TaskAlreadyInTaskGroup,
+)
 from airflow.models.taskmixin import DAGNode, DependencyMixin
 from airflow.serialization.enums import DagAttributeTypes
 from airflow.utils.helpers import validate_group_key
@@ -186,7 +191,16 @@ def __iter__(self):
                 yield child
 
     def add(self, task: DAGNode) -> None:
-        """Add a task to this TaskGroup."""
+        """Add a task to this TaskGroup.
+
+        :meta private:
+        """
+        from airflow.models.abstractoperator import AbstractOperator
+
+        existing_tg = task.task_group
+        if isinstance(task, AbstractOperator) and existing_tg is not None and existing_tg != self:
+            raise TaskAlreadyInTaskGroup(task.node_id, existing_tg.node_id, self.node_id)
+
         # Set the TG first, as setting it might change the return value of node_id!
         task.task_group = weakref.proxy(self)
         key = task.node_id


Test patch:
diff --git a/tests/utils/test_task_group.py b/tests/utils/test_task_group.py
--- a/tests/utils/test_task_group.py
+++ b/tests/utils/test_task_group.py
@@ -20,6 +20,7 @@
 import pytest
 
 from airflow.decorators import dag, task_group as task_group_decorator
+from airflow.exceptions import TaskAlreadyInTaskGroup
 from airflow.models import DAG
 from airflow.models.xcom_arg import XComArg
 from airflow.operators.bash import BashOperator
@@ -1201,3 +1202,24 @@ def nested_topo(group):
         ],
         task6,
     ]
+
+
+def test_add_to_sub_group():
+    with DAG("test_dag", start_date=pendulum.parse("20200101")):
+        tg = TaskGroup("section")
+        task = EmptyOperator(task_id="task")
+        with pytest.raises(TaskAlreadyInTaskGroup) as ctx:
+            tg.add(task)
+
+    assert str(ctx.value) == "cannot add 'task' to 'section' (already in the DAG's root group)"
+
+
+def test_add_to_another_group():
+    with DAG("test_dag", start_date=pendulum.parse("20200101")):
+        tg = TaskGroup("section_1")
+        with TaskGroup("section_2"):
+            task = EmptyOperator(task_id="task")
+        with pytest.raises(TaskAlreadyInTaskGroup) as ctx:
+            tg.add(task)
+
+    assert str(ctx.value) == "cannot add 'section_2.task' to 'section_1' (already in group 'section_2')"