def get_code(dag_id: str) -> str: """Return python code of a given dag_id. :param dag_id: DAG id :return: code of the DAG """ dag = check_and_get_dag(dag_id=dag_id) try: return DagCode.get_code_by_fileloc(dag.fileloc) except (OSError, DagCodeNotFound) as exception: error_message = f"Error {str(exception)} while reading Dag id {dag_id} Code" raise AirflowException(error_message, exception)
def test_code_can_be_read_when_no_access_to_file(self): """ Test that code can be retrieved from DB when you do not have access to Code file. Source Code should atleast exist in one of DB or File. """ example_dag = make_example_dags(example_dags_module).get('example_bash_operator') example_dag.sync_to_db() # Mock that there is no access to the Dag File with patch('airflow.models.dagcode.open_maybe_zipped') as mock_open: mock_open.side_effect = FileNotFoundError dag_code = DagCode.get_code_by_fileloc(example_dag.fileloc) for test_string in ['example_bash_operator', 'also_run_this', 'run_this_last']: assert test_string in dag_code