Repository: apache/airflow

Issue:
airflow tasks run --ship-dag not able to generate pickeled dag
**Apache Airflow version**: 2.0.1


**Kubernetes version (if you are using kubernetes)** (use `kubectl version`): No

**Environment**:

- **Cloud provider or hardware configuration**: local machine
- **OS** (e.g. from /etc/os-release): 18.04.5 LTS (Bionic Beaver)
- **Kernel** (e.g. `uname -a`): wsl2

**What happened**:
Getting Pickled_id: None

```
root@12c7fd58e084:/opt/airflow# airflow tasks run example_bash_operator runme_0 now --ship-dag --interactive
[2021-05-09 13:11:33,247] {dagbag.py:487} INFO - Filling up the DagBag from /files/dags
Running <TaskInstance: example_bash_operator.runme_0 2021-05-09T13:11:31.788923+00:00 [None]> on host 12c7fd58e084
Pickled dag <DAG: example_bash_operator> as pickle_id: None
Sending to executor.
[2021-05-09 13:11:34,722] {base_executor.py:82} INFO - Adding to queue: ['airflow', 'tasks', 'run', 'example_bash_operator', 'runme_0', '2021-05-09T13:11:31.788923+00:00', '--local', '--pool', 'default_pool', '--subdir', '/opt/airflow/airflow/example_dags/example_bash_operator.py']
[2021-05-09 13:11:34,756] {local_executor.py:81} INFO - QueuedLocalWorker running ['airflow', 'tasks', 'run', 'example_bash_operator', 'runme_0', '2021-05-09T13:11:31.788923+00:00', '--local', '--pool', 'default_pool', '--subdir', '/opt/airflow/airflow/example_dags/example_bash_operator.py']
[2021-05-09 13:11:34,757] {local_executor.py:386} INFO - Shutting down LocalExecutor; waiting for running tasks to finish.  Signal again if you don't want to wait.
[2021-05-09 13:11:34,817] {dagbag.py:487} INFO - Filling up the DagBag from /opt/airflow/airflow/example_dags/example_bash_operator.py
Running <TaskInstance: example_bash_operator.runme_0 2021-05-09T13:11:31.788923+00:00 [None]> on host 12c7fd58e084
```

**What you expected to happen**:

Pickled_id should get generated 

```
Pickled dag <DAG: example_bash_operator> as pickle_id: None
```

**How to reproduce it**:
run below command from command line in airflow environment
```
airflow tasks run example_bash_operator runme_0 now --ship-dag --interactive
```
**Would like to submit PR for this issue**: YES




Discussion:
Thanks for opening your first issue here! Be sure to follow the issue template!

From my initial findings, we are not calling `session.commit()` after creating pickled dag.
https://github.com/apache/airflow/blob/master/airflow/cli/commands/task_command.py#L82

Please let me know if my findings are valid, I will create Pull Request with Unit Tests 😄 
Thanks
@vikramcse The session context manager takes care of committing the session, the issue might be something else

https://github.com/apache/airflow/blob/3711a29e5e253ede4ab12efe5990e83240e8a9a9/airflow/utils/session.py#L32
@bhavaniravi Appreciate for the reply

> @vikramcse The session context manager takes care of committing the session, the issue might be something else

If we look at the other functionality where `DagPickle` object is used, we can see that the `session.commit()` is used over there
https://github.com/apache/airflow/blob/3711a29e5e253ede4ab12efe5990e83240e8a9a9/airflow/jobs/backfill_job.py#L782-L785
I think @vikramcse is correct on the root cause. `with create_session()` only commits at the end of the block, *after* we accessed `pickle.id`, so it does not handle this.
@uranusjr Thanks for checking the issue.
I am writing the unit tests for the `--ship-dag` run tasks flag, which returns a pickle_id like 
```
Pickled dag <DAG: example_python_operator> as pickle_id: 2
```

and the test

```python
    def test_cli_run_with_ship_dag_and_interactive_flag(self):
        """
        tasks run should return valid pickle_id
        """
        with redirect_stdout(io.StringIO()) as stdout:
            task_command.task_run(
                self.parser.parse_args(
                    [
                        'tasks',
                        'run',
                        'example_python_operator',
                        'print_the_context',
                        DEFAULT_DATE.isoformat(),
                        '--ship-dag',
                        '--interactive',
                    ]
                )
            )

        output = stdout.getvalue()
        print(output)
```

In this case the pickle_id is a dynamic field. It will be helpful if you give me some example for testing the dynamically generating pickle_id

Resolving patch:
diff --git a/airflow/cli/commands/task_command.py b/airflow/cli/commands/task_command.py
--- a/airflow/cli/commands/task_command.py
+++ b/airflow/cli/commands/task_command.py
@@ -100,9 +100,9 @@ def _run_task_by_executor(args, dag, ti):
             with create_session() as session:
                 pickle = DagPickle(dag)
                 session.add(pickle)
-                pickle_id = pickle.id
-                # TODO: This should be written to a log
-                print(f'Pickled dag {dag} as pickle_id: {pickle_id}')
+            pickle_id = pickle.id
+            # TODO: This should be written to a log
+            print(f'Pickled dag {dag} as pickle_id: {pickle_id}')
         except Exception as e:
             print('Could not pickle the DAG')
             print(e)


Test patch: