From c71f00a254be334efb45f09ebfda326508a98124 Mon Sep 17 00:00:00 2001
From: Ephraim Anierobi <splendidzigy24@gmail.com>
Date: Tue, 5 Oct 2021 13:41:23 +0100
Subject: [PATCH 1/2] Move docker decorator example dag to docker provider

This example dag errors out during startup when we set `AIRFLOW__CORE__EXAMPLE_DAGS=True` and docker
provider is not installed.

separate the examples

remove doc text since it's not used in documentation
---
 .../tutorial_taskflow_api_etl_virtualenv.py   | 81 +++++++++++++++++++
 ...rial_taskflow_api_etl_docker_virtualenv.py |  4 +-
 docs/apache-airflow/tutorial_taskflow_api.rst |  4 +-
 3 files changed, 85 insertions(+), 4 deletions(-)
 create mode 100644 airflow/example_dags/tutorial_taskflow_api_etl_virtualenv.py
 rename airflow/{ => providers/docker}/example_dags/tutorial_taskflow_api_etl_docker_virtualenv.py (96%)

diff --git a/airflow/example_dags/tutorial_taskflow_api_etl_virtualenv.py b/airflow/example_dags/tutorial_taskflow_api_etl_virtualenv.py
new file mode 100644
index 0000000000000..09aefcb50805e
--- /dev/null
+++ b/airflow/example_dags/tutorial_taskflow_api_etl_virtualenv.py
@@ -0,0 +1,81 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+
+from datetime import datetime
+
+from airflow.decorators import dag, task
+
+
+@dag(schedule_interval=None, start_date=datetime(2021, 1, 1), catchup=False, tags=['example'])
+def tutorial_taskflow_api_etl_virtualenv():
+    """
+    ### TaskFlow API example using virtualenv
+    This is a simple ETL data pipeline example which demonstrates the use of
+    the TaskFlow API using three simple tasks for Extract, Transform, and Load.
+    """
+
+    @task.virtualenv(
+        use_dill=True,
+        system_site_packages=False,
+        requirements=['funcsigs'],
+    )
+    def extract():
+        """
+        #### Extract task
+        A simple Extract task to get data ready for the rest of the data
+        pipeline. In this case, getting data is simulated by reading from a
+        hardcoded JSON string.
+        """
+        import json
+
+        data_string = '{"1001": 301.27, "1002": 433.21, "1003": 502.22}'
+
+        order_data_dict = json.loads(data_string)
+        return order_data_dict
+
+    @task(multiple_outputs=True)
+    def transform(order_data_dict: dict):
+        """
+        #### Transform task
+        A simple Transform task which takes in the collection of order data and
+        computes the total order value.
+        """
+        total_order_value = 0
+
+        for value in order_data_dict.values():
+            total_order_value += value
+
+        return {"total_order_value": total_order_value}
+
+    @task()
+    def load(total_order_value: float):
+        """
+        #### Load task
+        A simple Load task which takes in the result of the Transform task and
+        instead of saving it to end user review, just prints it out.
+        """
+
+        print(f"Total order value is: {total_order_value:.2f}")
+
+    order_data = extract()
+    order_summary = transform(order_data)
+    load(order_summary["total_order_value"])
+
+
+tutorial_etl_dag = tutorial_taskflow_api_etl_virtualenv()
diff --git a/airflow/example_dags/tutorial_taskflow_api_etl_docker_virtualenv.py b/airflow/providers/docker/example_dags/tutorial_taskflow_api_etl_docker_virtualenv.py
similarity index 96%
rename from airflow/example_dags/tutorial_taskflow_api_etl_docker_virtualenv.py
rename to airflow/providers/docker/example_dags/tutorial_taskflow_api_etl_docker_virtualenv.py
index c89ea9b358ec2..caab91d47b2b0 100644
--- a/airflow/example_dags/tutorial_taskflow_api_etl_docker_virtualenv.py
+++ b/airflow/providers/docker/example_dags/tutorial_taskflow_api_etl_docker_virtualenv.py
@@ -28,7 +28,7 @@
 
 # [START instantiate_dag]
 @dag(schedule_interval=None, start_date=datetime(2021, 1, 1), catchup=False, tags=['example'])
-def tutorial_taskflow_api_etl_virtualenv():
+def tutorial_taskflow_api_etl_docker_virtualenv():
     """
     ### TaskFlow API Tutorial Documentation
     This is a simple ETL data pipeline example which demonstrates the use of
@@ -99,7 +99,7 @@ def load(total_order_value: float):
 
 
 # [START dag_invocation]
-tutorial_etl_dag = tutorial_taskflow_api_etl_virtualenv()
+tutorial_etl_dag = tutorial_taskflow_api_etl_docker_virtualenv()
 # [END dag_invocation]
 
 # [END tutorial]
diff --git a/docs/apache-airflow/tutorial_taskflow_api.rst b/docs/apache-airflow/tutorial_taskflow_api.rst
index ea5579ed2945f..ce15e4a66e048 100644
--- a/docs/apache-airflow/tutorial_taskflow_api.rst
+++ b/docs/apache-airflow/tutorial_taskflow_api.rst
@@ -175,7 +175,7 @@ image must have a working Python installed and take in a bash command as the ``c
 
 Below is an example of using the ``@task.docker`` decorator to run a python task.
 
-.. exampleinclude:: /../../airflow/example_dags/tutorial_taskflow_api_etl_docker_virtualenv.py
+.. exampleinclude:: /../../airflow/providers/docker/example_dags/tutorial_taskflow_api_etl_docker_virtualenv.py
     :language: python
     :dedent: 4
     :start-after: [START transform_docker]
@@ -199,7 +199,7 @@ environment on the same machine, you can use the ``@task.virtualenv`` decorator
 decorator will allow you to create a new virtualenv with custom libraries and even a different
 Python version to run your function.
 
-.. exampleinclude:: /../../airflow/example_dags/tutorial_taskflow_api_etl_docker_virtualenv.py
+.. exampleinclude:: /../../airflow/providers/docker/example_dags/tutorial_taskflow_api_etl_docker_virtualenv.py
     :language: python
     :dedent: 4
     :start-after: [START extract_virtualenv]

From 4d708cdddad27123592841b909b00bb0bfde7e03 Mon Sep 17 00:00:00 2001
From: Ephraim Anierobi <splendidzigy24@gmail.com>
Date: Wed, 6 Oct 2021 16:28:35 +0100
Subject: [PATCH 2/2] Add try/except for the example dag

---
 .../tutorial_taskflow_api_etl_docker_virtualenv.py   | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/airflow/providers/docker/example_dags/tutorial_taskflow_api_etl_docker_virtualenv.py b/airflow/providers/docker/example_dags/tutorial_taskflow_api_etl_docker_virtualenv.py
index caab91d47b2b0..1961c30b73f8a 100644
--- a/airflow/providers/docker/example_dags/tutorial_taskflow_api_etl_docker_virtualenv.py
+++ b/airflow/providers/docker/example_dags/tutorial_taskflow_api_etl_docker_virtualenv.py
@@ -98,8 +98,14 @@ def load(total_order_value: float):
     # [END main_flow]
 
 
-# [START dag_invocation]
-tutorial_etl_dag = tutorial_taskflow_api_etl_docker_virtualenv()
-# [END dag_invocation]
+# The try/except here is because Airflow versions less than 2.2.0 doesn't support
+# @task.docker decorator and we use this dag in CI test. Thus, in order not to
+# break the CI test, we added this try/except here.
+try:
+    # [START dag_invocation]
+    tutorial_etl_dag = tutorial_taskflow_api_etl_docker_virtualenv()
+    # [END dag_invocation]
+except AttributeError:
+    pass
 
 # [END tutorial]
