def deploy_dag(): """Deploys a dag to a Cloud Composer environment""" logger.log( logging.INFO, f"Entered deploy_dag -- {API_BASE_PATH_V1}/dag/deploy api POST method") req_data = request.get_json() if not req_data: return {'error': "Empty JSON payload"}, 500 try: api_validator.validate_payload(req_data) if 'project_id' in req_data: project_id, location, composer_environment = api_service.get_gcp_composer_details( req_data) else: project_id, location, composer_environment = api_service.get_gcp_composer_details( None) airflow_dag_bucket_name = api_service.get_dag_bucket( project_id, location, composer_environment) dag_name = req_data['dag_name'] next_actions = { 'trigger': f'{API_BASE_PATH_V1}/dag/trigger/{dag_name}' } if req_data['mode'] == 'GCS': deploy_file = api_service.gcs_download_file( project_id, req_data['bucket_name'], req_data['file_path']) gcs_dag_path = api_service.deploy_dag(project_id, 'GCS', airflow_dag_bucket_name, dag_file=deploy_file) return jsonify(dag_name=dag_name, dag_gcs_path=gcs_dag_path, next_actions=next_actions) if req_data['mode'] == 'GIT': deploy_file = api_service.git_download_file( req_data['git_url'], req_data['repo_name'], req_data['file_path']) git_dag_path = api_service.deploy_dag(project_id, 'GIT', airflow_dag_bucket_name, dag_file=deploy_file) return jsonify(dag_name=dag_name, dag_gcs_path=git_dag_path, next_actions=next_actions) if req_data['mode'] == 'INLINE': gcs_dag_path = api_service.deploy_dag(project_id, 'INLINE', airflow_dag_bucket_name, dag_data=req_data) return jsonify(dag_name=dag_name, dag_gcs_path=gcs_dag_path, next_actions=next_actions) except: return {'error': traceback.print_exc()}, 500
def test_validate_payload_gcs_invalid(): json = { 'project_id': 'mock_project_id', 'location': 'mock_gcp_location', 'composer_environment': 'mock_composer_environment', 'dag_name': 'test_trigger_dag', 'mode': 'GCS', 'file_path': 'dags/dag_workflow_simple.py' } with pytest.raises(ValueError): api_validator.validate_payload(json)
def test_validate_payload_git_invalid(): json = { 'project_id': 'mock_project_id', 'location': 'mock_gcp_location', 'composer_environment': 'mock_composer_environment', 'dag_name': 'test_trigger_dag', 'mode': 'GIT', 'git_url': 'path.to.git.repo', 'file_path': 'mock/path/to/dag.py' } with pytest.raises(ValueError): api_validator.validate_payload(json)
def test_validate_payload_inline_invalid(): json = { 'project_id': 'mock_project_id', 'location': 'mock_gcp_location', 'composer_environment': 'mock_composer_environment', 'dag_name': 'test_trigger_dag', 'mode': 'INLINE', 'kubernetes_pod_operators': [{ 'task_id': 'k8s_pod_operator_example_task_01', 'name': 'k8s_pod_example_01' }] } with pytest.raises(ValueError): api_validator.validate_payload(json)
def validate_dag(): """Validates a dag to ensure that it is error free and compatible with Cloud Composer""" logger.log( logging.INFO, f"Entered validate_dag -- {API_BASE_PATH_V1}/dag/validate api POST method" ) req_data = request.get_json() if not req_data: return {'error': "Empty JSON payload"}, 500 try: api_validator.validate_payload(req_data) if 'project_id' in req_data: project_id, location, composer_environment = api_service.get_gcp_composer_details( req_data) else: project_id, location, composer_environment = api_service.get_gcp_composer_details( None) next_actions = {'deploy': f'{API_BASE_PATH_V1}/dag/deploy'} if req_data['mode'] == 'GCS': deploy_file = api_service.gcs_download_file( project_id, req_data['bucket_name'], req_data['file_path']) validation_json = api_service.validate_dag('GCS', deploy_file) validation_json['next_actions'] = next_actions return jsonify(validation_json) if req_data['mode'] == 'GIT': deploy_file = api_service.git_download_file( req_data['git_url'], req_data['repo_name'], req_data['file_path']) validation_json = api_service.validate_dag('GIT', deploy_file) validation_json['next_actions'] = next_actions return jsonify(validation_json) if req_data['mode'] == 'INLINE': validation_json = api_service.validate_dag('INLINE', req_data) validation_json['next_actions'] = next_actions return jsonify(validation_json) except: return {'error': traceback.print_exc()}, 500
def test_validate_payload_gcs_valid(): json = { 'project_id': 'mock_project_id', 'location': 'mock_gcp_location', 'composer_environment': 'mock_composer_environment', 'dag_name': 'test_trigger_dag', 'mode': 'GCS', 'bucket_name': 'mock_bucket_name', 'file_path': 'dags/dag_workflow_simple.py' } is_valid = api_validator.validate_payload(json) assert is_valid
def test_validate_payload_git_valid(): json = { 'project_id': 'mock_project_id', 'location': 'mock_gcp_location', 'composer_environment': 'mock_composer_environment', 'dag_name': 'test_trigger_dag', 'mode': 'GIT', 'git_url': 'path.to.git.repo', 'repo_name': 'repo-test-dir', 'file_path': 'mock/path/to/dag.py' } is_valid = api_validator.validate_payload(json) assert is_valid
def test_validate_payload_inline_valid(): json = { 'project_id': 'mock_project_id', 'location': 'mock_gcp_location', 'composer_environment': 'mock_composer_environment', 'dag_name': 'test_trigger_dag', 'mode': 'INLINE', 'kubernetes_pod_operators': [{ 'task_id': 'k8s_pod_operator_example_task_01', 'name': 'k8s_pod_example_01', 'image': 'bash' }] } is_valid = api_validator.validate_payload(json) assert is_valid