add fetch_notices_by_query DAG

Author: CaptainOfHacks

dags/fetch_notices_by_query.py

@@ -0,0 +1,68 @@
+from airflow.decorators import dag, task
+from airflow.operators.dummy import DummyOperator
+from airflow.operators.python import BranchPythonOperator
+from airflow.utils.trigger_rule import TriggerRule
+from dags import DEFAULT_DAG_ARGUMENTS
+from dags.dags_utils import get_dag_param, push_dag_downstream, pull_dag_upstream
+from dags.operators.DagBatchPipelineOperator import NOTICE_IDS_KEY, TriggerNoticeBatchPipelineOperator
+from dags.pipelines.notice_fetcher_pipelines import notice_fetcher_by_query_pipeline
+from ted_sws.event_manager.adapters.event_log_decorator import event_log
+from ted_sws.event_manager.model.event_message import TechnicalEventMessage, EventMessageMetadata, \
+    EventMessageProcessType
+
+DAG_NAME = "fetch_notices_by_query"
+BATCH_SIZE = 2000
+QUERY_DAG_KEY = "query"
+TRIGGER_COMPLETE_WORKFLOW_DAG_KEY = "trigger_complete_workflow"
+TRIGGER_PARTIAL_WORKFLOW_TASK_ID = "trigger_partial_notice_proc_workflow"
+TRIGGER_COMPLETE_WORKFLOW_TASK_ID = "trigger_complete_notice_proc_workflow"
+CHECK_IF_TRIGGER_COMPLETE_WORKFLOW_TASK_ID = "check_if_trigger_complete_workflow"
+FINISH_FETCH_BY_DATE_TASK_ID = "finish_fetch_by_query"
+
+
+@dag(default_args=DEFAULT_DAG_ARGUMENTS,
+     catchup=False,
+     tags=['fetch'])
+def fetch_notices_by_query():
+    @task
+    @event_log(TechnicalEventMessage(
+        message="fetch_by_query_notice_from_ted",
+        metadata=EventMessageMetadata(
+            process_type=EventMessageProcessType.DAG, process_name=DAG_NAME
+        ))
+    )
+    def fetch_by_query_notice_from_ted():
+        notice_ids = notice_fetcher_by_query_pipeline(query=get_dag_param(key=QUERY_DAG_KEY, raise_error=True))
+        if not notice_ids:
+            raise Exception("No notices has been fetched!")
+        push_dag_downstream(key=NOTICE_IDS_KEY, value=notice_ids)
+
+    trigger_complete_workflow = TriggerNoticeBatchPipelineOperator(task_id=TRIGGER_COMPLETE_WORKFLOW_TASK_ID,
+                                                                   execute_only_one_step=False
+                                                                   )
+    trigger_normalisation_workflow = TriggerNoticeBatchPipelineOperator(
+        task_id=TRIGGER_PARTIAL_WORKFLOW_TASK_ID,
+        batch_size=BATCH_SIZE,
+        execute_only_one_step=True)
+
+    def _branch_selector():
+        trigger_complete_workflow = get_dag_param(key=TRIGGER_COMPLETE_WORKFLOW_DAG_KEY,
+                                                  default_value=True)
+        push_dag_downstream(key=NOTICE_IDS_KEY, value=pull_dag_upstream(key=NOTICE_IDS_KEY))
+        if trigger_complete_workflow:
+            return [TRIGGER_COMPLETE_WORKFLOW_TASK_ID]
+        return [TRIGGER_PARTIAL_WORKFLOW_TASK_ID]
+
+    branch_task = BranchPythonOperator(
+        task_id=CHECK_IF_TRIGGER_COMPLETE_WORKFLOW_TASK_ID,
+        python_callable=_branch_selector,
+    )
+
+    finish_step = DummyOperator(task_id=FINISH_FETCH_BY_DATE_TASK_ID,
+                                trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS)
+
+    fetch_by_query_notice_from_ted() >> branch_task >> [trigger_normalisation_workflow,
+                                                        trigger_complete_workflow] >> finish_step
+
+
+dag = fetch_notices_by_query()

dags/pipelines/notice_fetcher_pipelines.py

@@ -1,6 +1,7 @@
 from datetime import datetime, timedelta
 from typing import List
 
+
 def notice_fetcher_by_date_pipeline(date_wild_card: str = None) -> List[str]:
     from pymongo import MongoClient
     from ted_sws import config
@@ -21,3 +22,18 @@ def notice_fetcher_by_date_pipeline(date_wild_card: str = None) -> List[str]:
                                                     notice_fetched_date=notice_publication_date)
 
     return notice_ids
+
+
+def notice_fetcher_by_query_pipeline(query: str = None) -> List[str]:
+    from pymongo import MongoClient
+    from ted_sws import config
+    from ted_sws.data_manager.adapters.notice_repository import NoticeRepository
+    from ted_sws.notice_fetcher.adapters.ted_api import TedAPIAdapter, TedRequestAPI
+    from ted_sws.notice_fetcher.services.notice_fetcher import NoticeFetcher
+
+    ted_api_query = {"q": query}
+    mongodb_client = MongoClient(config.MONGO_DB_AUTH_URL)
+    notice_ids = NoticeFetcher(notice_repository=NoticeRepository(mongodb_client=mongodb_client),
+                               ted_api_adapter=TedAPIAdapter(
+                                   request_api=TedRequestAPI())).fetch_notices_by_query(query=ted_api_query)
+    return notice_ids