From 1342d845381804bb51e1aadaf7c2c385ead3cb5c Mon Sep 17 00:00:00 2001 From: Tatiana Al-Chueyr Date: Sun, 10 Nov 2024 11:40:54 +0000 Subject: [PATCH 1/3] Fix dbt docs outside of Astro Since Cosmos 1.7.1, users who do not use Astro started getting 'Access is Denied' when trying to access the dbt docs menu item. This was an uninted 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 | 64 +++++++++++++++++++++++++++++++++---- 3 files changed, 75 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..f8b96f09a 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,41 @@ 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 Airflow flask app + 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 +90,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 +347,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")] From 846eb45a3cf55c5af31360357c8b474fe9492f29 Mon Sep 17 00:00:00 2001 From: Tatiana Al-Chueyr Date: Mon, 11 Nov 2024 15:02:31 +0000 Subject: [PATCH 2/3] Fix integration test in Airflow 2.4 --- tests/plugin/test_plugin.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/plugin/test_plugin.py b/tests/plugin/test_plugin.py index f8b96f09a..3f7186b1d 100644 --- a/tests/plugin/test_plugin.py +++ b/tests/plugin/test_plugin.py @@ -68,6 +68,7 @@ def app_within_astro_cloud(module_monkeypatch) -> FlaskClient: appbuilder.sm.check_authorization = lambda *args, **kwargs: True if cosmos.plugin.dbt_docs_view not in appbuilder.baseviews: + app._got_first_request = False # unregister blueprints registered in global Airflow flask app del app.blueprints["DbtDocsView"] keys_to_delete = [view_name for view_name in app.view_functions.keys() if view_name.startswith("DbtDocsView")] From 298c21c4099727f27ec0775ba3b80056d0d5c116 Mon Sep 17 00:00:00 2001 From: Tatiana Al-Chueyr Date: Mon, 11 Nov 2024 15:11:24 +0000 Subject: [PATCH 3/3] Improve comments --- tests/plugin/test_plugin.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/plugin/test_plugin.py b/tests/plugin/test_plugin.py index 3f7186b1d..963df9f70 100644 --- a/tests/plugin/test_plugin.py +++ b/tests/plugin/test_plugin.py @@ -68,8 +68,8 @@ def app_within_astro_cloud(module_monkeypatch) -> FlaskClient: appbuilder.sm.check_authorization = lambda *args, **kwargs: True if cosmos.plugin.dbt_docs_view not in appbuilder.baseviews: - app._got_first_request = False - # unregister blueprints registered in global Airflow flask app + # 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]