Migrate Postgres example DAGs to new design #22458 (#24148)

* Migrate Postgres example DAGs to new design #22458 * Fix static checks
apache · Jun 3, 2022 · c60bb9e · c60bb9e
1 parent fb1187d
commit c60bb9e
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 23 deletions.
diff --git a/airflow/providers/postgres/example_dags/__init__.py b/airflow/providers/postgres/example_dags/__init__.py
diff --git a/docs/apache-airflow-providers-postgres/index.rst b/docs/apache-airflow-providers-postgres/index.rst
@@ -39,7 +39,7 @@ Content
     :maxdepth: 1
     :caption: Resources
 
-    Example DAGs <https://github.com/apache/airflow/tree/main/airflow/providers/postgres/example_dags>
+    Example DAGs <https://github.com/apache/airflow/tree/main/tests/system/providers/postgres>
 
 .. toctree::
     :maxdepth: 1

diff --git a/docs/apache-airflow-providers-postgres/operators/postgres_operator_howto_guide.rst b/docs/apache-airflow-providers-postgres/operators/postgres_operator_howto_guide.rst
@@ -42,7 +42,7 @@ Creating a Postgres database table
 
 The code snippets below are based on Airflow-2.0
 
-.. exampleinclude:: /../../airflow/providers/postgres/example_dags/example_postgres.py
+.. exampleinclude:: /../../tests/system/providers/postgres/example_postgres.py
     :language: python
     :start-after: [START postgres_operator_howto_guide]
     :end-before: [END postgres_operator_howto_guide_create_pet_table]
@@ -160,7 +160,7 @@ Passing Server Configuration Parameters into PostgresOperator
 PostgresOperator provides the optional ``runtime_parameters`` attribute which makes it possible to set
 the `server configuration parameter values <https://www.postgresql.org/docs/current/runtime-config-client.html>`_ for the SQL request during runtime.
 
-.. exampleinclude:: /../../airflow/providers/postgres/example_dags/example_postgres.py
+.. exampleinclude:: /../../tests/system/providers/postgres/example_postgres.py
     :language: python
     :start-after: [START postgres_operator_howto_guide_get_birth_date]
     :end-before: [END postgres_operator_howto_guide_get_birth_date]
@@ -171,7 +171,7 @@ The complete Postgres Operator DAG
 
 When we put everything together, our DAG should look like this:
 
-.. exampleinclude:: /../../airflow/providers/postgres/example_dags/example_postgres.py
+.. exampleinclude:: /../../tests/system/providers/postgres/example_postgres.py
     :language: python
     :start-after: [START postgres_operator_howto_guide]
     :end-before: [END postgres_operator_howto_guide]

diff --git a/...postgres/example_dags/example_postgres.py → ...em/providers/postgres/example_postgres.py b/...postgres/example_dags/example_postgres.py → ...em/providers/postgres/example_postgres.py
@@ -14,18 +14,23 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-
-# [START postgres_operator_howto_guide]
 import datetime
+import os
 
 from airflow import DAG
 from airflow.providers.postgres.operators.postgres import PostgresOperator
 
+# [START postgres_operator_howto_guide]
+
+
 # create_pet_table, populate_pet_table, get_all_pets, and get_birth_date are examples of tasks created by
 # instantiating the Postgres Operator
 
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
+DAG_ID = "postgres_operator_dag"
+
 with DAG(
-    dag_id="postgres_operator_dag",
+    dag_id=DAG_ID,
     start_date=datetime.datetime(2020, 2, 2),
     schedule_interval="@once",
     catchup=False,
@@ -72,3 +77,14 @@
 
     create_pet_table >> populate_pet_table >> get_all_pets >> get_birth_date
     # [END postgres_operator_howto_guide]
+
+    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)