示例#1
0
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)
示例#5
0
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