Update worker_single_notice_process_orchestrator.py

Author: CaptainOfHacks

dags/worker_single_notice_process_orchestrator.py

@@ -1,43 +1,43 @@
+from airflow.decorators import dag
+from airflow.operators.python import get_current_context, BranchPythonOperator, PythonOperator
 from airflow.utils.trigger_rule import TriggerRule
 from pymongo import MongoClient
 
+from dags import DEFAULT_DAG_ARGUMENTS
 from dags.dags_utils import pull_dag_upstream, push_dag_downstream
 from ted_sws import config
 from ted_sws.core.model.manifestation import METSManifestation
-from ted_sws.core.model.notice import NoticeStatus
+from ted_sws.core.model.notice import NoticeStatus, Notice
 from ted_sws.data_manager.adapters.mapping_suite_repository import MappingSuiteRepositoryMongoDB
 from ted_sws.data_manager.adapters.notice_repository import NoticeRepository
-from ted_sws.data_sampler.services.notice_xml_indexer import index_notice_by_id
-from ted_sws.notice_metadata_processor.services.metadata_normalizer import normalise_notice_by_id
-
-from airflow.decorators import dag
-from airflow.operators.python import get_current_context, BranchPythonOperator, PythonOperator
-
-from dags import DEFAULT_DAG_ARGUMENTS
-from ted_sws.notice_metadata_processor.services.notice_eligibility import notice_eligibility_checker_by_id
-from ted_sws.notice_packager.services.notice_packager import create_notice_package
-from ted_sws.notice_transformer.adapters.rml_mapper import RMLMapper
-from ted_sws.notice_transformer.services.notice_transformer import transform_notice_by_id
-from ted_sws.notice_validator.services.shacl_test_suite_runner import validate_notice_by_id_with_shacl_suite
-from ted_sws.notice_validator.services.sparql_test_suite_runner import validate_notice_by_id_with_sparql_suite
-from ted_sws.event_manager.adapters.event_logger import EventLogger
+from ted_sws.data_sampler.services.notice_xml_indexer import index_notice
 from ted_sws.event_manager.adapters.event_log_decorator import event_log
+from ted_sws.event_manager.adapters.event_logger import EventLogger
 from ted_sws.event_manager.model.event_message import NoticeEventMessage, EventMessageProcessType, EventMessageMetadata, \
     TechnicalEventMessage
 from ted_sws.event_manager.services.logger_from_context import get_logger_from_dag_context, \
     handle_event_message_metadata_dag_context, get_task_id_from_dag_context
+from ted_sws.notice_metadata_processor.services.metadata_normalizer import normalise_notice
+from ted_sws.notice_metadata_processor.services.notice_eligibility import notice_eligibility_checker
+from ted_sws.notice_packager.services.notice_packager import create_notice_package
+from ted_sws.notice_transformer.adapters.rml_mapper import RMLMapper
+from ted_sws.notice_transformer.services.notice_transformer import transform_notice
+from ted_sws.notice_validator.services.shacl_test_suite_runner import validate_notice_with_shacl_suite
+from ted_sws.notice_validator.services.sparql_test_suite_runner import validate_notice_with_sparql_suite
 
 NOTICE_ID = "notice_id"
 MAPPING_SUITE_ID = "mapping_suite_id"
-DAG_NAME = "old_worker_single_notice_process_orchestrator"
+NOTICE_OBJECT = "notice_object"
+MAPPING_SUITE_OBJECT = "mapping_suite_object"
+DAG_NAME = "worker_single_notice_process_orchestrator"
 
 
 @dag(default_args=DEFAULT_DAG_ARGUMENTS,
      schedule_interval=None,
      max_active_runs=128,
      concurrency=128,
      tags=['worker', 'pipeline'])
-def old_worker_single_notice_process_orchestrator():
+def worker_single_notice_process_orchestrator():
     """
 
     :return:
@@ -57,8 +57,10 @@ def _index_notice_xml_content():
         """
         notice_id = pull_dag_upstream(NOTICE_ID)
         mongodb_client = MongoClient(config.MONGO_DB_AUTH_URL)
-        index_notice_by_id(notice_id=notice_id, mongodb_client=mongodb_client)
-        push_dag_downstream(NOTICE_ID, notice_id)
+        notice_repository = NoticeRepository(mongodb_client=mongodb_client)
+        notice = notice_repository.get(reference=notice_id)
+        notice = index_notice(notice=notice)
+        push_dag_downstream(NOTICE_OBJECT, notice)
 
     @event_log(is_loggable=False)
     def _normalise_notice_metadata(**context_args):
@@ -70,19 +72,16 @@ def _normalise_notice_metadata(**context_args):
         event_message: NoticeEventMessage = NoticeEventMessage()
         event_message.start_record()
 
-        notice_id = pull_dag_upstream(NOTICE_ID)
-        mongodb_client = MongoClient(config.MONGO_DB_AUTH_URL)
-        notice_repository = NoticeRepository(mongodb_client=mongodb_client)
-        normalised_notice = normalise_notice_by_id(notice_id=notice_id, notice_repository=notice_repository)
-        notice_repository.update(notice=normalised_notice)
-        push_dag_downstream(NOTICE_ID, notice_id)
+        notice: Notice = pull_dag_upstream(NOTICE_OBJECT)
+        normalised_notice = normalise_notice(notice=notice)
+        push_dag_downstream(NOTICE_OBJECT, normalised_notice)
 
         context = get_current_context()
 
         handle_event_message_metadata_dag_context(event_message, context)
-        event_message.notice_id = notice_id
+        event_message.notice_id = notice.ted_id
         event_message.domain_action = get_task_id_from_dag_context(context)
-        event_message.message = f"Normalising notice({notice_id}) metadata"
+        event_message.message = f"Normalising notice({notice.ted_id}) metadata"
         event_message.end_record()
         event_logger.info(event_message)
 
@@ -92,20 +91,18 @@ def _check_eligibility_for_transformation(**context_args):
         event_message: NoticeEventMessage = NoticeEventMessage()
         event_message.start_record()
 
-        notice_id = pull_dag_upstream(NOTICE_ID)
+        notice = pull_dag_upstream(NOTICE_OBJECT)
         mongodb_client = MongoClient(config.MONGO_DB_AUTH_URL)
         notice_repository = NoticeRepository(mongodb_client=mongodb_client)
         mapping_suite_repository = MappingSuiteRepositoryMongoDB(mongodb_client=mongodb_client)
-        result = notice_eligibility_checker_by_id(notice_id=notice_id,
-                                                  notice_repository=notice_repository,
-                                                  mapping_suite_repository=mapping_suite_repository)
+        result = notice_eligibility_checker(notice=notice, mapping_suite_repository=mapping_suite_repository)
+        notice_repository.update(notice=notice)
         mapping_suite_id = None
+        notice_id = notice.ted_id
         if result:
             notice_id, mapping_suite_id = result
-
             push_dag_downstream(MAPPING_SUITE_ID, mapping_suite_id)
         push_dag_downstream(NOTICE_ID, notice_id)
-
         context = get_current_context()
 
         handle_event_message_metadata_dag_context(event_message, context)
@@ -123,16 +120,12 @@ def _preprocess_xml_manifestation(**context_args):
         event_message: NoticeEventMessage = NoticeEventMessage()
         event_message.start_record()
 
-        notice_id = pull_dag_upstream(NOTICE_ID)
+        notice = pull_dag_upstream(NOTICE_OBJECT)
         mapping_suite_id = pull_dag_upstream(MAPPING_SUITE_ID)
-        mongodb_client = MongoClient(config.MONGO_DB_AUTH_URL)
-        notice_repository = NoticeRepository(mongodb_client=mongodb_client)
-        notice = notice_repository.get(reference=notice_id)
         notice.update_status_to(new_status=NoticeStatus.PREPROCESSED_FOR_TRANSFORMATION)
-        notice_repository.update(notice=notice)
-        push_dag_downstream(NOTICE_ID, notice_id)
+        push_dag_downstream(NOTICE_OBJECT, notice)
         push_dag_downstream(MAPPING_SUITE_ID, mapping_suite_id)
-
+        notice_id = notice.ted_id
         context = get_current_context()
 
         handle_event_message_metadata_dag_context(event_message, context)
@@ -148,19 +141,16 @@ def _transform_notice(**context_args):
         event_logger: EventLogger = get_logger_from_dag_context(context_args)
         event_message: NoticeEventMessage = NoticeEventMessage()
         event_message.start_record()
-
-        notice_id = pull_dag_upstream(NOTICE_ID)
+        notice = pull_dag_upstream(NOTICE_OBJECT)
         mapping_suite_id = pull_dag_upstream(MAPPING_SUITE_ID)
         mongodb_client = MongoClient(config.MONGO_DB_AUTH_URL)
-        notice_repository = NoticeRepository(mongodb_client=mongodb_client)
         mapping_suite_repository = MappingSuiteRepositoryMongoDB(mongodb_client=mongodb_client)
+        mapping_suite = mapping_suite_repository.get(reference=mapping_suite_id)
         rml_mapper = RMLMapper(rml_mapper_path=config.RML_MAPPER_PATH)
-        transform_notice_by_id(notice_id=notice_id, mapping_suite_id=mapping_suite_id,
-                               notice_repository=notice_repository, mapping_suite_repository=mapping_suite_repository,
-                               rml_mapper=rml_mapper
-                               )
-        push_dag_downstream(NOTICE_ID, notice_id)
-        push_dag_downstream(MAPPING_SUITE_ID, mapping_suite_id)
+        result_notice = transform_notice(notice=notice, mapping_suite=mapping_suite, rml_mapper=rml_mapper)
+        push_dag_downstream(NOTICE_OBJECT, result_notice)
+        push_dag_downstream(MAPPING_SUITE_OBJECT, mapping_suite)
+        notice_id = result_notice.ted_id
 
         context = get_current_context()
 
@@ -178,23 +168,19 @@ def _resolve_entities_in_the_rdf_manifestation(**context_args):
         event_message: NoticeEventMessage = NoticeEventMessage()
         event_message.start_record()
 
-        notice_id = pull_dag_upstream(NOTICE_ID)
-        mapping_suite_id = pull_dag_upstream(MAPPING_SUITE_ID)
-        mongodb_client = MongoClient(config.MONGO_DB_AUTH_URL)
-        notice_repository = NoticeRepository(mongodb_client=mongodb_client)
-        notice = notice_repository.get(reference=notice_id)
+        notice = pull_dag_upstream(NOTICE_OBJECT)
+        mapping_suite = pull_dag_upstream(MAPPING_SUITE_OBJECT)
         notice.set_distilled_rdf_manifestation(distilled_rdf_manifestation=notice.rdf_manifestation.copy())
-        notice_repository.update(notice=notice)
-        push_dag_downstream(NOTICE_ID, notice_id)
-        push_dag_downstream(MAPPING_SUITE_ID, mapping_suite_id)
-
+        push_dag_downstream(NOTICE_OBJECT, notice)
+        push_dag_downstream(MAPPING_SUITE_OBJECT, mapping_suite)
+        notice_id = notice.ted_id
         context = get_current_context()
 
         handle_event_message_metadata_dag_context(event_message, context)
         event_message.notice_id = notice_id
         event_message.domain_action = get_task_id_from_dag_context(context)
         event_message.message = f"Resolving notice({notice_id}) entities in the RDF manifestation"
-        event_message.metadata.process_context["mapping_suite_id"] = mapping_suite_id
+        event_message.metadata.process_context["mapping_suite_id"] = mapping_suite.identifier
         event_message.end_record()
         event_logger.info(event_message)
 
@@ -204,28 +190,19 @@ def _validate_transformed_rdf_manifestation(**context_args):
         event_message: NoticeEventMessage = NoticeEventMessage()
         event_message.start_record()
 
-        notice_id = pull_dag_upstream(NOTICE_ID)
-        mapping_suite_id = pull_dag_upstream(MAPPING_SUITE_ID)
-        mongodb_client = MongoClient(config.MONGO_DB_AUTH_URL)
-        notice_repository = NoticeRepository(mongodb_client=mongodb_client)
-        mapping_suite_repository = MappingSuiteRepositoryMongoDB(mongodb_client=mongodb_client)
-        validate_notice_by_id_with_sparql_suite(notice_id=notice_id, mapping_suite_identifier=mapping_suite_id,
-                                                notice_repository=notice_repository,
-                                                mapping_suite_repository=mapping_suite_repository)
-        validate_notice_by_id_with_shacl_suite(notice_id=notice_id, mapping_suite_identifier=mapping_suite_id,
-                                               notice_repository=notice_repository,
-                                               mapping_suite_repository=mapping_suite_repository
-                                               )
-        push_dag_downstream(NOTICE_ID, notice_id)
-        push_dag_downstream(MAPPING_SUITE_ID, mapping_suite_id)
-
+        notice = pull_dag_upstream(NOTICE_OBJECT)
+        mapping_suite = pull_dag_upstream(MAPPING_SUITE_OBJECT)
+        validate_notice_with_sparql_suite(notice=notice, mapping_suite_package=mapping_suite)
+        validate_notice_with_shacl_suite(notice=notice, mapping_suite_package=mapping_suite)
+        notice_id = notice.ted_id
+        push_dag_downstream(NOTICE_OBJECT, notice)
         context = get_current_context()
 
         handle_event_message_metadata_dag_context(event_message, context)
         event_message.notice_id = notice_id
         event_message.domain_action = get_task_id_from_dag_context(context)
         event_message.message = f"Validating notice({notice_id}) transformed RDF manifestation"
-        event_message.metadata.process_context["mapping_suite_id"] = mapping_suite_id
+        event_message.metadata.process_context["mapping_suite_id"] = mapping_suite.identifier
         event_message.end_record()
         event_logger.info(event_message)
 
@@ -235,12 +212,12 @@ def _check_eligibility_for_packing_by_validation_report(**context_args):
         event_message: NoticeEventMessage = NoticeEventMessage()
         event_message.start_record()
 
-        notice_id = pull_dag_upstream(NOTICE_ID)
+        notice = pull_dag_upstream(NOTICE_OBJECT)
         mongodb_client = MongoClient(config.MONGO_DB_AUTH_URL)
         notice_repository = NoticeRepository(mongodb_client=mongodb_client)
-        notice = notice_repository.get(reference=notice_id)
         notice.set_is_eligible_for_packaging(eligibility=True)
         notice_repository.update(notice=notice)
+        notice_id = notice.ted_id
         push_dag_downstream(NOTICE_ID, notice_id)
 
         context = get_current_context()
@@ -302,7 +279,7 @@ def _check_notice_state_before_transform():
         mongodb_client = MongoClient(config.MONGO_DB_AUTH_URL)
         notice_repository = NoticeRepository(mongodb_client=mongodb_client)
         notice = notice_repository.get(reference=notice_id)
-        push_dag_downstream(NOTICE_ID, notice_id)
+        push_dag_downstream(NOTICE_OBJECT, notice)
         if notice.status == NoticeStatus.ELIGIBLE_FOR_TRANSFORMATION:
             mapping_suite_id = pull_dag_upstream(MAPPING_SUITE_ID)
             push_dag_downstream(MAPPING_SUITE_ID, mapping_suite_id)
@@ -360,7 +337,7 @@ def _check_notice_state_before_notice_successfully_processed():
         trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS
     )
 
-    transform_notice = PythonOperator(
+    notice_transform = PythonOperator(
         task_id="transform_notice",
         python_callable=_transform_notice,
         trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS
@@ -442,7 +419,7 @@ def _check_notice_state_before_notice_successfully_processed():
 
     index_notice_xml_content >> normalise_notice_metadata >> check_eligibility_for_transformation >> check_notice_state_before_transform >> [
         preprocess_xml_manifestation, fail_on_state]
-    preprocess_xml_manifestation >> transform_notice >> resolve_entities_in_the_rdf_manifestation >> validate_transformed_rdf_manifestation >> check_eligibility_for_packing_by_validation_report
+    preprocess_xml_manifestation >> notice_transform >> resolve_entities_in_the_rdf_manifestation >> validate_transformed_rdf_manifestation >> check_eligibility_for_packing_by_validation_report
     check_eligibility_for_packing_by_validation_report >> check_notice_state_before_generate_mets_package >> [
         generate_mets_package, fail_on_state]
     generate_mets_package >> check_package_integrity_by_package_structure >> check_notice_state_before_publish_notice_in_cellar >> [
@@ -473,4 +450,4 @@ def _get_task_run():
                     check_notice_state_before_generate_mets_package, check_notice_state_before_publish_notice_in_cellar]
 
 
-dag = old_worker_single_notice_process_orchestrator()
+dag = worker_single_notice_process_orchestrator()