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
