Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

DatabricksPlugin - Fix dag view redirect URL by using url_for redirect #41040

Merged
merged 3 commits into from
Jul 26, 2024
Merged

DatabricksPlugin - Fix dag view redirect URL by using url_for redirect #41040

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
0cd22f2
Fix dag view redirect URL by using url_for redirect
pankajkoti Jul 26, 2024
80370e3
Add a quick test
pankajkoti Jul 26, 2024
dfa26ac
Mark test as db test
pankajkoti Jul 26, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 3 additions & 5 deletions airflow/providers/databricks/plugins/databricks_workflow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,6 @@
from flask_appbuilder.api import expose
from packaging.version import Version

from airflow.configuration import conf
from airflow.exceptions import AirflowException, TaskInstanceNotFound
from airflow.models import BaseOperator, BaseOperatorLink
from airflow.models.dag import DAG, clear_task_instances
Expand Down Expand Up @@ -413,8 +412,7 @@ class RepairDatabricksTasks(AirflowBaseView, LoggingMixin):
@expose("/repair_databricks_job/<string:dag_id>/<string:run_id>", methods=("GET",))
@get_auth_decorator()
def repair(self, dag_id: str, run_id: str):
view = conf.get("webserver", "dag_default_view")
return_url = self._get_return_url(dag_id, view)
return_url = self._get_return_url(dag_id, run_id)

tasks_to_repair = request.values.get("tasks_to_repair")
self.log.info("Tasks to repair: %s", tasks_to_repair)
Expand Down Expand Up @@ -450,8 +448,8 @@ def repair(self, dag_id: str, run_id: str):
return redirect(return_url)

@staticmethod
def _get_return_url(dag_id: str, view) -> str:
return f"/dags/{dag_id}/{view}"
def _get_return_url(dag_id: str, run_id: str) -> str:
return url_for("Airflow.grid", dag_id=dag_id, dag_run_id=run_id)
pankajkoti marked this conversation as resolved.
Show resolved Hide resolved


repair_databricks_view = RepairDatabricksTasks()
Expand Down
12 changes: 12 additions & 0 deletions tests/providers/databricks/plugins/test_databricks_workflow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
from unittest.mock import MagicMock, Mock, patch

import pytest
from flask import url_for

from airflow.exceptions import AirflowException
from airflow.models.dagrun import DagRun
Expand Down Expand Up @@ -144,6 +145,17 @@ def test_get_task_instance(app):
assert result == dag_run


def test_get_return_url_dag_id_run_id(app):
dag_id = "example_dag"
run_id = "example_run"

expected_url = url_for("Airflow.grid", dag_id=dag_id, dag_run_id=run_id)

with app.app_context():
actual_url = RepairDatabricksTasks._get_return_url(dag_id, run_id)
assert actual_url == expected_url, f"Expected {expected_url}, got {actual_url}"


@pytest.mark.db_test
def test_workflow_job_run_link(app):
with app.app_context():
Expand Down
Loading