Migrate Snowflake system tests to new design #22434 (#24151)

* Migrate Snowflake system tests to new design #22434 * Fix flake8
apache · Jun 3, 2022 · c2f10a4 · c2f10a4
1 parent c60bb9e
commit c2f10a4
Show file tree

Hide file tree

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

diff --git a/docs/apache-airflow-providers-snowflake/operators/s3_to_snowflake.rst b/docs/apache-airflow-providers-snowflake/operators/s3_to_snowflake.rst
@@ -36,7 +36,7 @@ a file format (see `docs <https://docs.snowflake.com/en/sql-reference/sql/create
 
 An example usage of the S3ToSnowflakeOperator is as follows:
 
-.. exampleinclude:: /../../airflow/providers/snowflake/example_dags/example_snowflake.py
+.. exampleinclude:: /../../tests/system/providers/snowflake/example_snowflake.py
     :language: python
     :start-after: [START howto_operator_s3_to_snowflake]
     :end-before: [END howto_operator_s3_to_snowflake]
diff --git a/docs/apache-airflow-providers-snowflake/operators/snowflake.rst b/docs/apache-airflow-providers-snowflake/operators/snowflake.rst
@@ -47,7 +47,7 @@ the connection metadata is structured as follows:
 
 An example usage of the SnowflakeOperator is as follows:
 
-.. exampleinclude:: /../../airflow/providers/snowflake/example_dags/example_snowflake.py
+.. exampleinclude:: /../../tests/system/providers/snowflake/example_snowflake.py
     :language: python
     :start-after: [START howto_operator_snowflake]
     :end-before: [END howto_operator_snowflake]

diff --git a/docs/apache-airflow-providers-snowflake/operators/snowflake_to_slack.rst b/docs/apache-airflow-providers-snowflake/operators/snowflake_to_slack.rst
@@ -46,7 +46,7 @@ and contain the resulting dataset (e.g. ASCII formatted dataframe).
 
 An example usage of the SnowflakeToSlackOperator is as follows:
 
-.. exampleinclude:: /../../airflow/providers/snowflake/example_dags/example_snowflake.py
+.. exampleinclude:: /../../tests/system/providers/snowflake/example_snowflake.py
     :language: python
     :start-after: [START howto_operator_snowflake_to_slack]
     :end-before: [END howto_operator_snowflake_to_slack]
diff --git a/...owflake/example_dags/example_snowflake.py → .../providers/snowflake/example_snowflake.py b/...owflake/example_dags/example_snowflake.py → .../providers/snowflake/example_snowflake.py
@@ -18,6 +18,7 @@
 """
 Example use of Snowflake related operators.
 """
+import os
 from datetime import datetime
 
 from airflow import DAG
@@ -47,89 +48,91 @@
 SNOWFLAKE_SLACK_MESSAGE = (
     "Results in an ASCII table:\n```{{ results_df | tabulate(tablefmt='pretty', headers='keys') }}```"
 )
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
+DAG_ID = "example_snowflake"
 
 # [START howto_operator_snowflake]
 
-dag = DAG(
-    'example_snowflake',
+with DAG(
+    DAG_ID,
     start_date=datetime(2021, 1, 1),
     default_args={'snowflake_conn_id': SNOWFLAKE_CONN_ID},
     tags=['example'],
     catchup=False,
-)
-
-
-snowflake_op_sql_str = SnowflakeOperator(
-    task_id='snowflake_op_sql_str',
-    dag=dag,
-    sql=CREATE_TABLE_SQL_STRING,
-    warehouse=SNOWFLAKE_WAREHOUSE,
-    database=SNOWFLAKE_DATABASE,
-    schema=SNOWFLAKE_SCHEMA,
-    role=SNOWFLAKE_ROLE,
-)
-
-snowflake_op_with_params = SnowflakeOperator(
-    task_id='snowflake_op_with_params',
-    dag=dag,
-    sql=SQL_INSERT_STATEMENT,
-    parameters={"id": 56},
-    warehouse=SNOWFLAKE_WAREHOUSE,
-    database=SNOWFLAKE_DATABASE,
-    schema=SNOWFLAKE_SCHEMA,
-    role=SNOWFLAKE_ROLE,
-)
-
-snowflake_op_sql_list = SnowflakeOperator(task_id='snowflake_op_sql_list', dag=dag, sql=SQL_LIST)
-
-snowflake_op_sql_multiple_stmts = SnowflakeOperator(
-    task_id='snowflake_op_sql_multiple_stmts',
-    dag=dag,
-    sql=SQL_MULTIPLE_STMTS,
-)
-
-snowflake_op_template_file = SnowflakeOperator(
-    task_id='snowflake_op_template_file',
-    dag=dag,
-    sql='/path/to/sql/<filename>.sql',
-)
-
-# [END howto_operator_snowflake]
-
-# [START howto_operator_s3_to_snowflake]
-
-copy_into_table = S3ToSnowflakeOperator(
-    task_id='copy_into_table',
-    s3_keys=[S3_FILE_PATH],
-    table=SNOWFLAKE_SAMPLE_TABLE,
-    schema=SNOWFLAKE_SCHEMA,
-    stage=SNOWFLAKE_STAGE,
-    file_format="(type = 'CSV',field_delimiter = ';')",
-    dag=dag,
-)
-
-# [END howto_operator_s3_to_snowflake]
-
-# [START howto_operator_snowflake_to_slack]
-
-slack_report = SnowflakeToSlackOperator(
-    task_id="slack_report",
-    sql=SNOWFLAKE_SLACK_SQL,
-    slack_message=SNOWFLAKE_SLACK_MESSAGE,
-    slack_conn_id=SLACK_CONN_ID,
-    dag=dag,
-)
-
-# [END howto_operator_snowflake_to_slack]
-
-(
-    snowflake_op_sql_str
-    >> [
-        snowflake_op_with_params,
-        snowflake_op_sql_list,
-        snowflake_op_template_file,
-        copy_into_table,
-        snowflake_op_sql_multiple_stmts,
-    ]
-    >> slack_report
-)
+) as dag:
+    # [START snowflake_example_dag]
+    snowflake_op_sql_str = SnowflakeOperator(
+        task_id='snowflake_op_sql_str',
+        sql=CREATE_TABLE_SQL_STRING,
+        warehouse=SNOWFLAKE_WAREHOUSE,
+        database=SNOWFLAKE_DATABASE,
+        schema=SNOWFLAKE_SCHEMA,
+        role=SNOWFLAKE_ROLE,
+    )
+
+    snowflake_op_with_params = SnowflakeOperator(
+        task_id='snowflake_op_with_params',
+        sql=SQL_INSERT_STATEMENT,
+        parameters={"id": 56},
+        warehouse=SNOWFLAKE_WAREHOUSE,
+        database=SNOWFLAKE_DATABASE,
+        schema=SNOWFLAKE_SCHEMA,
+        role=SNOWFLAKE_ROLE,
+    )
+
+    snowflake_op_sql_list = SnowflakeOperator(task_id='snowflake_op_sql_list', sql=SQL_LIST)
+
+    snowflake_op_sql_multiple_stmts = SnowflakeOperator(
+        task_id='snowflake_op_sql_multiple_stmts',
+        sql=SQL_MULTIPLE_STMTS,
+    )
+
+    snowflake_op_template_file = SnowflakeOperator(
+        task_id='snowflake_op_template_file',
+        sql='/path/to/sql/<filename>.sql',
+    )
+
+    # [END howto_operator_snowflake]
+
+    # [START howto_operator_s3_to_snowflake]
+
+    copy_into_table = S3ToSnowflakeOperator(
+        task_id='copy_into_table',
+        s3_keys=[S3_FILE_PATH],
+        table=SNOWFLAKE_SAMPLE_TABLE,
+        schema=SNOWFLAKE_SCHEMA,
+        stage=SNOWFLAKE_STAGE,
+        file_format="(type = 'CSV',field_delimiter = ';')",
+    )
+
+    # [END howto_operator_s3_to_snowflake]
+
+    # [START howto_operator_snowflake_to_slack]
+
+    slack_report = SnowflakeToSlackOperator(
+        task_id="slack_report",
+        sql=SNOWFLAKE_SLACK_SQL,
+        slack_message=SNOWFLAKE_SLACK_MESSAGE,
+        slack_conn_id=SLACK_CONN_ID,
+    )
+
+    # [END howto_operator_snowflake_to_slack]
+
+    (
+        snowflake_op_sql_str
+        >> [
+            snowflake_op_with_params,
+            snowflake_op_sql_list,
+            snowflake_op_template_file,
+            copy_into_table,
+            snowflake_op_sql_multiple_stmts,
+        ]
+        >> slack_report
+    )
+    # [END snowflake_example_dag]
+
+
+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)