witsml_parser_dag.py

#  Copyright 2020 Google LLC
#  Copyright 2020 EPAM Systems
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

"""Energistics XML ingest."""

import os
from datetime import timedelta

import airflow
import yaml
from airflow import DAG
from airflow.contrib.operators.kubernetes_pod_operator import KubernetesPodOperator
from osdu_airflow.backward_compatibility.default_args import update_default_args
from osdu_manifest.operators.ensure_manifest_integrity import EnsureManifestIntegrityOperator
from osdu_manifest.operators.process_manifest_r3 import ProcessManifestOperatorR3
from osdu_manifest.operators.validate_manifest_schema import ValidateManifestSchemaOperator
from osdu_manifest.operators.update_status import UpdateStatusOperator

PROCESS_SINGLE_MANIFEST_FILE = "process_single_manifest_file_task"
PROCESS_BATCH_MANIFEST_FILE = "batch_upload"
ENSURE_INTEGRITY_TASK = "provide_manifest_integrity_task"
SINGLE_MANIFEST_FILE_FIRST_OPERATOR = "validate_manifest_schema_task"

with open(os.environ["OSDU_API_CONFIG"]) as f:
    k8s_configs = yaml.safe_load(f)

energistics_k8s_configs = k8s_configs["energistics"]["energistics_witsml_parser_k8s_operator"]

default_args = {
    "name": "Energistics_xml_ingest",
    "owner": "energistics",
    "start_date": airflow.utils.dates.days_ago(0),
    "retries": 0,
    "retry_delay": timedelta(minutes=50),
    "trigger_rule": "none_failed",

}

default_args = update_default_args(default_args)

dag = DAG(
    "Energistics_xml_ingest",
    default_args=default_args,
    description="witsml parser dag",
    schedule_interval=None,
    dagrun_timeout=timedelta(minutes=60)
)

update_status_running_op = UpdateStatusOperator(
    task_id="update_status_running_task",
    dag=dag
)

update_status_finished_op = UpdateStatusOperator(
    task_id="update_status_finished_task",
    dag=dag,
    trigger_rule="all_done",
)

process_energistics_op = KubernetesPodOperator(
    namespace="default",
    task_id="witsml_parser_task",
    name="witsml_parser_task",
    do_xcom_push=True,
    cmds=[
        "sh",
        "-c",
        "mkdir -p /airflow/xcom/; "
        "python main.py"
        " --context '{{ dag_run.conf['execution_context'] | tojson }}'"
        " --file_service {{ var.value.core__service__file__host }}"
        " --out /airflow/xcom/return.json",
    ],
    is_delete_operator_pod=True,
    image_pull_policy="Always",
    get_logs=True,
    startup_timeout_seconds=300,
    dag=dag,
    **energistics_k8s_configs
)

validate_schema_operator = ValidateManifestSchemaOperator(
    task_id="validate_manifest_schema_task",
    dag=dag,
    previous_task_id=process_energistics_op.task_id,
    trigger_rule="none_failed_or_skipped"
)

ensure_integrity_op = EnsureManifestIntegrityOperator(
    task_id=ENSURE_INTEGRITY_TASK,
    dag=dag,
    previous_task_id=validate_schema_operator.task_id,
    trigger_rule="none_failed_or_skipped"
)

process_single_manifest_file = ProcessManifestOperatorR3(
    task_id=PROCESS_SINGLE_MANIFEST_FILE,
    dag=dag,
    previous_task_id=ensure_integrity_op.task_id,
    trigger_rule="none_failed_or_skipped"
)

update_status_running_op >> \
process_energistics_op >> \
validate_schema_operator >> \
ensure_integrity_op >> \
process_single_manifest_file >> \
update_status_finished_op