From 2b21f4cbcb82161e6e79c4b023c8e19b5285ba05 Mon Sep 17 00:00:00 2001 From: Tatiana Al-Chueyr Date: Tue, 12 Nov 2024 09:14:08 +0000 Subject: [PATCH] Fix access to the `dbt docs` menu item outside of Astro cloud (#1312) Since Cosmos 1.7.1, users who do not use Astro Cloud started facing 'Access is Denied' when accessing the dbt docs menu item. This was an unintended side effect of the fix for Astro users to be able to see that menu item, introduced in #1280 Closes: #1309 --- cosmos/plugin/__init__.py | 40 +++++++++-------------- cosmos/settings.py | 3 ++ tests/plugin/test_plugin.py | 65 +++++++++++++++++++++++++++++++++---- 3 files changed, 76 insertions(+), 32 deletions(-) diff --git a/cosmos/plugin/__init__.py b/cosmos/plugin/__init__.py index 8ca93926c..5997a5fe3 100644 --- a/cosmos/plugin/__init__.py +++ b/cosmos/plugin/__init__.py @@ -10,7 +10,17 @@ from flask import abort, url_for from flask_appbuilder import AppBuilder, expose -from cosmos.settings import dbt_docs_conn_id, dbt_docs_dir, dbt_docs_index_file_name +from cosmos.settings import dbt_docs_conn_id, dbt_docs_dir, dbt_docs_index_file_name, in_astro_cloud + +if in_astro_cloud: + MENU_ACCESS_PERMISSIONS = [ + (permissions.ACTION_CAN_ACCESS_MENU, "Custom Menu"), + (permissions.ACTION_CAN_READ, permissions.RESOURCE_WEBSITE), + ] +else: + MENU_ACCESS_PERMISSIONS = [ + (permissions.ACTION_CAN_READ, permissions.RESOURCE_WEBSITE), + ] def bucket_and_key(path: str) -> Tuple[str, str]: @@ -201,24 +211,14 @@ def create_blueprint( return super().create_blueprint(appbuilder, endpoint=endpoint, static_folder=self.static_folder) # type: ignore[no-any-return] @expose("/dbt_docs") # type: ignore[misc] - @has_access( - [ - (permissions.ACTION_CAN_ACCESS_MENU, "Custom Menu"), - (permissions.ACTION_CAN_READ, permissions.RESOURCE_WEBSITE), - ] - ) + @has_access(MENU_ACCESS_PERMISSIONS) def dbt_docs(self) -> str: if dbt_docs_dir is None: return self.render_template("dbt_docs_not_set_up.html") # type: ignore[no-any-return,no-untyped-call] return self.render_template("dbt_docs.html") # type: ignore[no-any-return,no-untyped-call] @expose("/dbt_docs_index.html") # type: ignore[misc] - @has_access( - [ - (permissions.ACTION_CAN_ACCESS_MENU, "Custom Menu"), - (permissions.ACTION_CAN_READ, permissions.RESOURCE_WEBSITE), - ] - ) + @has_access(MENU_ACCESS_PERMISSIONS) def dbt_docs_index(self) -> Tuple[str, int, Dict[str, Any]]: if dbt_docs_dir is None: abort(404) @@ -233,12 +233,7 @@ def dbt_docs_index(self) -> Tuple[str, int, Dict[str, Any]]: return html, 200, {"Content-Security-Policy": "frame-ancestors 'self'"} @expose("/catalog.json") # type: ignore[misc] - @has_access( - [ - (permissions.ACTION_CAN_ACCESS_MENU, "Custom Menu"), - (permissions.ACTION_CAN_READ, permissions.RESOURCE_WEBSITE), - ] - ) + @has_access(MENU_ACCESS_PERMISSIONS) def catalog(self) -> Tuple[str, int, Dict[str, Any]]: if dbt_docs_dir is None: abort(404) @@ -250,12 +245,7 @@ def catalog(self) -> Tuple[str, int, Dict[str, Any]]: return data, 200, {"Content-Type": "application/json"} @expose("/manifest.json") # type: ignore[misc] - @has_access( - [ - (permissions.ACTION_CAN_ACCESS_MENU, "Custom Menu"), - (permissions.ACTION_CAN_READ, permissions.RESOURCE_WEBSITE), - ] - ) + @has_access(MENU_ACCESS_PERMISSIONS) def manifest(self) -> Tuple[str, int, Dict[str, Any]]: if dbt_docs_dir is None: abort(404) diff --git a/cosmos/settings.py b/cosmos/settings.py index 7bcf04bb9..5b24321c8 100644 --- a/cosmos/settings.py +++ b/cosmos/settings.py @@ -43,3 +43,6 @@ LINEAGE_NAMESPACE = os.getenv("OPENLINEAGE_NAMESPACE", DEFAULT_OPENLINEAGE_NAMESPACE) AIRFLOW_IO_AVAILABLE = Version(airflow_version) >= Version("2.8.0") + +# The following environment variable is populated in Astro Cloud +in_astro_cloud = os.getenv("ASTRONOMER_ENVIRONMENT") == "cloud" diff --git a/tests/plugin/test_plugin.py b/tests/plugin/test_plugin.py index e8812b56a..963df9f70 100644 --- a/tests/plugin/test_plugin.py +++ b/tests/plugin/test_plugin.py @@ -13,19 +13,22 @@ jinja2.Markup = markupsafe.Markup jinja2.escape = markupsafe.escape +import importlib import sys from importlib.util import find_spec from unittest.mock import MagicMock, PropertyMock, mock_open, patch import pytest +from _pytest.monkeypatch import MonkeyPatch from airflow.utils.db import initdb, resetdb from airflow.www.app import cached_app from airflow.www.extensions.init_appbuilder import AirflowAppBuilder from flask.testing import FlaskClient +import cosmos import cosmos.plugin +import cosmos.settings from cosmos.plugin import ( - dbt_docs_view, iframe_script, open_azure_file, open_file, @@ -43,6 +46,42 @@ def _get_text_from_response(response) -> str: return response.text +@pytest.fixture(scope="module") +def module_monkeypatch(): + mp = MonkeyPatch() + yield mp + mp.undo() + + +@pytest.fixture(scope="module") +def app_within_astro_cloud(module_monkeypatch) -> FlaskClient: + module_monkeypatch.setenv("ASTRONOMER_ENVIRONMENT", "cloud") + importlib.reload(cosmos.settings) + importlib.reload(cosmos.plugin) + importlib.reload(cosmos) + initdb() + + cached_app._cached_app = None + app = cached_app(testing=True) + appbuilder: AirflowAppBuilder = app.extensions["appbuilder"] + + appbuilder.sm.check_authorization = lambda *args, **kwargs: True + + if cosmos.plugin.dbt_docs_view not in appbuilder.baseviews: + # unregister blueprints registered in global context + app._got_first_request = False # Necessary for Airflow 2.4, Flask==2.2.2 & Flask-AppBuilder==4.1.3 + del app.blueprints["DbtDocsView"] + keys_to_delete = [view_name for view_name in app.view_functions.keys() if view_name.startswith("DbtDocsView")] + [app.view_functions.pop(view_name) for view_name in keys_to_delete] + + appbuilder._check_and_init(cosmos.plugin.dbt_docs_view) + appbuilder.register_blueprint(cosmos.plugin.dbt_docs_view) + + yield app.test_client() + + resetdb() + + @pytest.fixture(scope="module") def app() -> FlaskClient: initdb() @@ -52,9 +91,9 @@ def app() -> FlaskClient: appbuilder.sm.check_authorization = lambda *args, **kwargs: True - if dbt_docs_view not in appbuilder.baseviews: - appbuilder._check_and_init(dbt_docs_view) - appbuilder.register_blueprint(dbt_docs_view) + if cosmos.plugin.dbt_docs_view not in appbuilder.baseviews: + appbuilder._check_and_init(cosmos.plugin.dbt_docs_view) + appbuilder.register_blueprint(cosmos.plugin.dbt_docs_view) yield app.test_client() @@ -309,8 +348,20 @@ def test_open_file_local(mock_file): @pytest.mark.parametrize( "url_path", ["/cosmos/dbt_docs", "/cosmos/dbt_docs_index.html", "/cosmos/catalog.json", "/cosmos/manifest.json"] ) -def test_has_access_with_permissions(url_path, app): - dbt_docs_view.appbuilder.sm.check_authorization = MagicMock() - mock_check_auth = dbt_docs_view.appbuilder.sm.check_authorization +def test_has_access_with_permissions_outside_astro_does_not_include_custom_menu(url_path, app): + cosmos.plugin.dbt_docs_view.appbuilder.sm.check_authorization = MagicMock() + mock_check_auth = cosmos.plugin.dbt_docs_view.appbuilder.sm.check_authorization + app.get(url_path) + assert mock_check_auth.call_args[0][0] == [("can_read", "Website")] + + +@pytest.mark.integration +@pytest.mark.parametrize( + "url_path", ["/cosmos/dbt_docs", "/cosmos/dbt_docs_index.html", "/cosmos/catalog.json", "/cosmos/manifest.json"] +) +def test_has_access_with_permissions_in_astro_must_include_custom_menu(url_path, app_within_astro_cloud): + app = app_within_astro_cloud + cosmos.plugin.dbt_docs_view.appbuilder.sm.check_authorization = MagicMock() + mock_check_auth = cosmos.plugin.dbt_docs_view.appbuilder.sm.check_authorization app.get(url_path) assert mock_check_auth.call_args[0][0] == [("menu_access", "Custom Menu"), ("can_read", "Website")]