Add databricks workflow system to integration tests (#1565)

We want to add the Databricks workflow system to our integration tests so that we can be notified in case of any failure during RC testing. This PR copies the Databricks workflow system test from the Airflow repo and integrates it with our integration tests.

.circleci/integration-tests/Dockerfile.astro_cloud

@@ -87,6 +87,7 @@ RUN mkdir -p ${AIRFLOW_HOME}/dags
 COPY . .
 RUN cp -r example_* ${AIRFLOW_HOME}/dags
 RUN cp master_dag.py  ${AIRFLOW_HOME}/dags/
+COPY databricks  ${AIRFLOW_HOME}/dags/
 RUN cp astronomer_migration_dag.py  ${AIRFLOW_HOME}/dags/
 RUN cp nuke-config.yml  ${AIRFLOW_HOME}/dags/
 # we have an issue with connexion==3.0.0, so for now pinning previous stable version

.circleci/integration-tests/master_dag.py

@@ -310,7 +310,10 @@ def prepare_dag_dependency(task_info, execution_time):
     chain(*kubernetes_trigger_tasks)
 
     # Databricks DAG
-    databricks_task_info = [{"databricks_dag": "example_async_databricks"}]
+    databricks_task_info = [
+        {"databricks_dag": "example_async_databricks"},
+        {"databricks_workflow_dag": "example_databricks_workflow"},
+    ]
     databricks_trigger_tasks, ids = prepare_dag_dependency(databricks_task_info, "{{ ds }}")
     dag_run_ids.extend(ids)
     chain(*databricks_trigger_tasks)

.circleci/integration-tests/script.sh

@@ -35,7 +35,7 @@ function clean() {
   [[ -f ${SCRIPT_PATH}/setup.cfg ]] && rm "${SCRIPT_PATH}"/setup.cfg
   [[ -f ${SCRIPT_PATH}/packages.txt ]] && rm "${SCRIPT_PATH}"/packages.txt
   [[ -f ${SCRIPT_PATH}/requirements.txt ]] && rm "${SCRIPT_PATH}"/requirements.txt
-  find  . -name "example_*" -exec rm {} \;
+  find . -name "example_*" ! -name "example_databricks_workflow.py" -exec rm {} \;
 }
 
 if [ "$1" == "-h" ]; then

.github/scripts/refactor_dag.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+import argparse
+
+
+def remove_lines_after(file_path, target_line):
+    """
+    Remove all lines in a file after the first occurrence of a specified target line.
+
+    This function reads a file and rewrites it, keeping only the lines up to and including
+    the first occurrence of the target line. All lines following the target line are removed.
+
+    Args:
+        file_path (str): The path to the file to be modified.
+        target_line (str): The line after which all content in the file should be removed.
+    """
+    with open(file_path) as input_file:
+        lines = input_file.readlines()
+
+    with open(file_path, "w") as output_file:
+        for line in lines:
+            if target_line in line:
+                break
+            output_file.write(line)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("file_path", help="file path", type=str)
+    args = parser.parse_args()
+    target_line = "from tests.system.utils.watcher import watcher"
+    remove_lines_after(args.file_path, target_line)

.github/workflows/reuse-wf-deploy-to-astro-cloud.yaml

@@ -43,6 +43,15 @@ jobs:
         id: get_git_revision
         run: echo "git_rev=$(git rev-parse HEAD)" >> $GITHUB_OUTPUT
 
+
+      - name: Clone airflow and copy system tests
+        working-directory: .circleci/integration-tests
+        run: |
+            git clone https://github.com/apache/airflow.git
+            cp -r airflow/tests/system/providers/databricks .
+            python ../../.github/scripts/refactor_dag.py databricks/example_databricks_workflow.py
+            cat databricks/example_databricks_workflow.py
+
       - name: deploy
         working-directory: .circleci/integration-tests
         run: |

dev/Dockerfile

@@ -14,7 +14,6 @@ RUN apt-get install -y --no-install-recommends \
 COPY setup.cfg ${AIRFLOW_HOME}/astronomer_providers/setup.cfg
 COPY pyproject.toml ${AIRFLOW_HOME}/astronomer_providers/pyproject.toml
 
-
 RUN pip install -e "${AIRFLOW_HOME}/astronomer_providers[all,tests,mypy]"
 RUN pip install apache-airflow-providers-slack