Migrate Qubole example DAGs to new design apache#22460

chethanuk · Jun 3, 2022 · 03265d7 · 03265d7
1 parent f294a26
commit 03265d7
Show file tree

Hide file tree

Showing 5 changed files with 40 additions and 33 deletions.
diff --git a/airflow/providers/qubole/example_dags/__init__.py b/airflow/providers/qubole/example_dags/__init__.py
diff --git a/docs/apache-airflow-providers-qubole/index.rst b/docs/apache-airflow-providers-qubole/index.rst
@@ -38,7 +38,7 @@ Content
     :maxdepth: 1
     :caption: Resources
 
-    Example DAGs <https://github.com/apache/airflow/tree/main/airflow/providers/qubole/example_dags>
+    Example DAGs <https://github.com/apache/airflow/tree/main/tests/system/providers/qubole>
     PyPI Repository <https://pypi.org/project/apache-airflow-providers-qubole/>
     Installing from sources <installing-providers-from-sources>
 

diff --git a/docs/apache-airflow-providers-qubole/operators/index.rst b/docs/apache-airflow-providers-qubole/operators/index.rst
@@ -29,4 +29,4 @@ Qubole Operators
 
 .. note::
     You can learn how to use Google Cloud integrations by analyzing the
-    `source code <https://github.com/apache/airflow/tree/main/airflow/providers/qubole/example_dags/>`_ of the particular example DAGs.
+    `source code <https://github.com/apache/airflow/tree/main/tests/system/providers/qubole/>`_ of the particular example DAGs.
diff --git a/docs/apache-airflow-providers-qubole/operators/qubole.rst b/docs/apache-airflow-providers-qubole/operators/qubole.rst
@@ -43,15 +43,15 @@ Run Hive command
 
 To run query that shows all tables you can use
 
-.. exampleinclude:: /../../airflow/providers/qubole/example_dags/example_qubole.py
+.. exampleinclude:: /../../tests/system/providers/qubole/example_qubole.py
     :language: python
     :dedent: 4
     :start-after: [START howto_operator_qubole_run_hive_query]
     :end-before: [END howto_operator_qubole_run_hive_query]
 
 Also you can run script that locates in the bucket by passing path to query file
 
-.. exampleinclude:: /../../airflow/providers/qubole/example_dags/example_qubole.py
+.. exampleinclude:: /../../tests/system/providers/qubole/example_qubole.py
     :language: python
     :dedent: 4
     :start-after: [START howto_operator_qubole_run_hive_script]
@@ -62,7 +62,7 @@ Run Hadoop command
 
 To run jar file in your Hadoop cluster use
 
-.. exampleinclude:: /../../airflow/providers/qubole/example_dags/example_qubole.py
+.. exampleinclude:: /../../tests/system/providers/qubole/example_qubole.py
     :language: python
     :dedent: 4
     :start-after: [START howto_operator_qubole_run_hadoop_jar]
@@ -73,7 +73,7 @@ Run Pig command
 
 To run script script in *Pig Latin* in your Hadoop cluster use
 
-.. exampleinclude:: /../../airflow/providers/qubole/example_dags/example_qubole.py
+.. exampleinclude:: /../../tests/system/providers/qubole/example_qubole.py
     :language: python
     :dedent: 4
     :start-after: [START howto_operator_qubole_run_pig_script]
@@ -84,7 +84,7 @@ Run Shell command
 
 To run Shell-script script use
 
-.. exampleinclude:: /../../airflow/providers/qubole/example_dags/example_qubole.py
+.. exampleinclude:: /../../tests/system/providers/qubole/example_qubole.py
     :language: python
     :dedent: 4
     :start-after: [START howto_operator_qubole_run_shell_script]
@@ -95,7 +95,7 @@ Run Presto command
 
 To run query using Presto use
 
-.. exampleinclude:: /../../airflow/providers/qubole/example_dags/example_qubole.py
+.. exampleinclude:: /../../tests/system/providers/qubole/example_qubole.py
     :language: python
     :dedent: 4
     :start-after: [START howto_operator_qubole_run_presto_query]
@@ -106,23 +106,23 @@ Run DB commands
 
 To run query as `DbTap <https://docs.qubole.com/en/latest/rest-api/dbtap_api/index.html>`_ use
 
-.. exampleinclude:: /../../airflow/providers/qubole/example_dags/example_qubole.py
+.. exampleinclude:: /../../tests/system/providers/qubole/example_qubole.py
     :language: python
     :dedent: 4
     :start-after: [START howto_operator_qubole_run_db_tap_query]
     :end-before: [END howto_operator_qubole_run_db_tap_query]
 
 To run DB export command use
 
-.. exampleinclude:: /../../airflow/providers/qubole/example_dags/example_qubole.py
+.. exampleinclude:: /../../tests/system/providers/qubole/example_qubole.py
     :language: python
     :dedent: 4
     :start-after: [START howto_operator_qubole_run_db_export]
     :end-before: [END howto_operator_qubole_run_db_export]
 
 To run DB import command use
 
-.. exampleinclude:: /../../airflow/providers/qubole/example_dags/example_qubole.py
+.. exampleinclude:: /../../tests/system/providers/qubole/example_qubole.py
     :language: python
     :dedent: 4
     :start-after: [START howto_operator_qubole_run_db_import]
@@ -133,7 +133,7 @@ Run Spark commands
 
 To run Scala script as a Spark job use
 
-.. exampleinclude:: /../../airflow/providers/qubole/example_dags/example_qubole.py
+.. exampleinclude:: /../../tests/system/providers/qubole/example_qubole.py
     :language: python
     :dedent: 4
     :start-after: [START howto_operator_qubole_run_spark_scala]
@@ -153,7 +153,7 @@ File or directory existence
 
 To wait for file or directory existence in cluster use
 
-.. exampleinclude:: /../../airflow/providers/qubole/example_dags/example_qubole.py
+.. exampleinclude:: /../../tests/system/providers/qubole/example_qubole.py
     :language: python
     :dedent: 4
     :start-after: [START howto_sensor_qubole_run_file_sensor]
@@ -173,7 +173,7 @@ Partition existence
 
 To wait for table partition existence in cluster use
 
-.. exampleinclude:: /../../airflow/providers/qubole/example_dags/example_qubole.py
+.. exampleinclude:: /../../tests/system/providers/qubole/example_qubole.py
     :language: python
     :dedent: 4
     :start-after: [START howto_sensor_qubole_run_partition_sensor]

diff --git a/...ers/qubole/example_dags/example_qubole.py → ...system/providers/qubole/example_qubole.py b/...ers/qubole/example_dags/example_qubole.py → ...system/providers/qubole/example_qubole.py
@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import os
 import filecmp
 import random
 import textwrap
@@ -35,8 +36,12 @@
 
 START_DATE = datetime(2021, 1, 1)
 
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
+DAG_ID = "example_qubole_operator"
+DAG2_ID = "example_qubole_sensor"
+
 with DAG(
-    dag_id='example_qubole_operator',
+    dag_id=DAG_ID,
     schedule_interval=None,
     start_date=START_DATE,
     tags=['example'],
@@ -218,8 +223,14 @@ def main(args: Array[String]) {
 
     branching >> db_import >> spark_cmd >> join
 
+    from tests.system.utils.watcher import watcher
+
+    # This test needs watcher in order to properly mark success/failure
+    # when "tearDown" task with trigger rule is part of the DAG
+    list(dag.tasks) >> watcher()
+
 with DAG(
-    dag_id='example_qubole_sensor',
+    dag_id=DAG2_ID,
     schedule_interval=None,
     start_date=START_DATE,
     tags=['example'],
@@ -268,3 +279,16 @@ def main(args: Array[String]) {
     # [END howto_sensor_qubole_run_partition_sensor]
 
     check_s3_file >> check_hive_partition
+
+    from tests.system.utils.watcher import watcher
+
+    # This test needs watcher in order to properly mark success/failure
+    # when "tearDown" task with trigger rule is part of the DAG
+    list(dag.tasks) >> watcher()
+
+
+from tests.system.utils import get_test_run  # noqa: E402
+
+# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest)
+test_run = get_test_run(dag)
+test_run_2 = get_test_run(dag2)