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
