def _sync_perm_for_dag(self, dag, session: Optional[Session] = None):
"""Sync DAG specific permissions, if necessary"""
from flask_appbuilder.security.sqla import models as sqla_models
from airflow.security.permissions import DAG_PERMS, resource_name_for_dag
def needs_perm_views(dag_id: str) -> bool:
dag_resource_name = resource_name_for_dag(dag_id)
for permission_name in DAG_PERMS:
if not (
session.query(sqla_models.PermissionView)
.join(sqla_models.Permission)
.join(sqla_models.ViewMenu)
.filter(sqla_models.Permission.name == permission_name)
.filter(sqla_models.ViewMenu.name == dag_resource_name)
.one_or_none()
):
return True
return False
if dag.access_control or needs_perm_views(dag.dag_id):
self.log.debug("Syncing DAG permissions: %s to the DB", dag.dag_id)
from airflow.www.security import ApplessAirflowSecurityManager
security_manager = ApplessAirflowSecurityManager(session=session)
security_manager.sync_perm_for_dag(dag.dag_id, dag.access_control)
def _serialize_dag_capturing_errors(dag, session):
"""
Try to serialize the dag to the DB, but make a note of any errors.
We can't place them directly in import_errors, as this may be retried, and work the next time
"""
if dag.is_subdag:
return []
try:
# We cant use bulk_write_to_db as we want to capture each error individually
dag_was_updated = SerializedDagModel.write_dag(
dag,
min_update_interval=settings.
MIN_SERIALIZED_DAG_UPDATE_INTERVAL,
session=session,
)
if dag_was_updated:
self.log.debug("Syncing DAG permissions: %s to the DB",
dag.dag_id)
from airflow.www.security import ApplessAirflowSecurityManager
security_manager = ApplessAirflowSecurityManager(
session=session)
security_manager.sync_perm_for_dag(dag.dag_id,
dag.access_control)
return []
except OperationalError:
raise
except Exception: # pylint: disable=broad-except
return [(dag.fileloc,
traceback.format_exc(
limit=-self.dagbag_import_error_traceback_depth))]
def _sync_perm_for_dag(self, dag, session: Session = None):
"""Sync DAG specific permissions, if necessary"""
from airflow.security.permissions import DAG_ACTIONS, resource_name_for_dag
from airflow.www.fab_security.sqla.models import Action, Permission, Resource
root_dag_id = dag.parent_dag.dag_id if dag.parent_dag else dag.dag_id
def needs_perms(dag_id: str) -> bool:
dag_resource_name = resource_name_for_dag(dag_id)
for permission_name in DAG_ACTIONS:
if not (
session.query(Permission)
.join(Action)
.join(Resource)
.filter(Action.name == permission_name)
.filter(Resource.name == dag_resource_name)
.one_or_none()
):
return True
return False
if dag.access_control or needs_perms(root_dag_id):
self.log.debug("Syncing DAG permissions: %s to the DB", root_dag_id)
from airflow.www.security import ApplessAirflowSecurityManager
security_manager = ApplessAirflowSecurityManager(session=session)
security_manager.sync_perm_for_dag(root_dag_id, dag.access_control)
def test_sync_perm_for_dag(self, mock_security_manager):
"""
Test that dagbag._sync_perm_for_dag will call ApplessAirflowSecurityManager.sync_perm_for_dag
when DAG specific perm views don't exist already or the DAG has access_control set.
"""
delete_dag_specific_permissions()
with create_session() as session:
security_manager = ApplessAirflowSecurityManager(session)
mock_sync_perm_for_dag = mock_security_manager.return_value.sync_perm_for_dag
mock_sync_perm_for_dag.side_effect = security_manager.sync_perm_for_dag
dagbag = DagBag(
dag_folder=os.path.join(TEST_DAGS_FOLDER,
"test_example_bash_operator.py"),
include_examples=False,
)
dag = dagbag.dags["test_example_bash_operator"]
def _sync_perms():
mock_sync_perm_for_dag.reset_mock()
dagbag._sync_perm_for_dag(dag, session=session)
# permviews dont exist
_sync_perms()
mock_sync_perm_for_dag.assert_called_once_with(
"test_example_bash_operator", None)
# permviews now exist
_sync_perms()
mock_sync_perm_for_dag.assert_not_called()
# Always sync if we have access_control
dag.access_control = {"Public": {"can_read"}}
_sync_perms()
mock_sync_perm_for_dag.assert_called_once_with(
"test_example_bash_operator", {"Public": {"can_read"}})