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")]