diff --git a/tools/python/airmaps/dags/update_booking.py b/tools/python/airmaps/dags/update_booking.py new file mode 100644 index 00000000000..7018f80a7ac --- /dev/null +++ b/tools/python/airmaps/dags/update_booking.py @@ -0,0 +1,84 @@ +import logging +import os +from datetime import timedelta + +from airflow import DAG +from airflow.operators.python_operator import PythonOperator +from airflow.utils.dates import days_ago + +from airmaps.instruments import settings +from airmaps.instruments import storage +from booking.download_hotels import download + +logger = logging.getLogger("airmaps") + + +DAG = DAG( + "Update_booking", + schedule_interval=timedelta(days=1), + default_args={ + "owner": "MAPS.ME", + "depends_on_past": True, + "start_date": days_ago(0), + "email": settings.EMAILS, + "email_on_failure": True, + "email_on_retry": False, + "retries": 5, + "retry_delay": timedelta(minutes=5), + "priority_weight": 1, + }, +) + + +BOOKING_STORAGE_PATH = f"{settings.STORAGE_PREFIX}/booking/" + +BOOKING_LOCAL_DIR = os.path.join(settings.MAIN_OUT_PATH, "booking") + + +def download_booking(**kwargs): + booking_version = "" + kwargs["ti"].xcom_push(key="booking_version", value=booking_version) + + os.makedirs(BOOKING_LOCAL_DIR, exist_ok=True) + booking_file = os.path.join(BOOKING_LOCAL_DIR, booking_version) + download( + settings.BOOKING_USER, settings.BOOKING_PASSWORD, booking_file, threads_count=4 + ) + + +def publish_booking(**kwargs): + booking_version = kwargs["ti"].xcom_pull(key="booking_version") + booking_file = os.path.join(BOOKING_LOCAL_DIR, booking_version) + storage.wd_publish(booking_file, BOOKING_STORAGE_PATH) + + +def rm_booking_file(**kwargs): + booking_version = kwargs["ti"].xcom_pull(key="booking_version") + booking_file = os.path.join(BOOKING_LOCAL_DIR, booking_version) + os.remove(booking_file) + + +DOWNLOAD_BOOKING_TASK = PythonOperator( + task_id="Download_booking_task", + provide_context=True, + python_callable=download_booking, + dag=DAG, +) + + +PUBLISH_BOOKING_TASK = PythonOperator( + task_id="Publish_booking_task", + provide_context=True, + python_callable=publish_booking, + dag=DAG, +) + +RM_BOOKING_FILE_TASK = PythonOperator( + task_id="Rm_booking_dir_task", + provide_context=True, + python_callable=rm_booking_file, + dag=DAG, +) + + +DOWNLOAD_BOOKING_TASK >> PUBLISH_BOOKING_TASK >> RM_BOOKING_FILE_TASK diff --git a/tools/python/airmaps/instruments/settings.py b/tools/python/airmaps/instruments/settings.py index 45200b476f2..6726ccf1d40 100644 --- a/tools/python/airmaps/instruments/settings.py +++ b/tools/python/airmaps/instruments/settings.py @@ -10,6 +10,10 @@ WD_LOGIN = "" WD_PASSWORD = "" +# Booking settings +BOOKING_USER = "" +BOOKING_PASSWORD= "" + # Common section EMAILS = [] @@ -45,6 +49,13 @@ def init(default_settings_path: AnyStr): WD_LOGIN = cfg.get_opt("Storage", "WD_LOGIN", WD_LOGIN) WD_PASSWORD = cfg.get_opt("Storage", "WD_PASSWORD", WD_PASSWORD) + # Booking section + global BOOKING_USER + global BOOKING_PASSWORD + + BOOKING_USER = cfg.get_opt("Booking", "BOOKING_USER", BOOKING_USER) + BOOKING_PASSWORD = cfg.get_opt("Booking", "BOOKING_PASSWORD", BOOKING_PASSWORD) + # Common section global EMAILS emails = cfg.get_opt("Common", "EMAILS", "") diff --git a/tools/python/booking/__main__.py b/tools/python/booking/__main__.py index cedc7eb8055..5a16c22c127 100644 --- a/tools/python/booking/__main__.py +++ b/tools/python/booking/__main__.py @@ -74,20 +74,20 @@ def main(): with tqdm(disable=not options.verbose) as progress_bar: if options.download_test_dataset: download_test_data( - options.country_code, options.user, options.password, options.output, options.threads_count, + options.country_code, progress_bar, ) else: download( - options.country_code, options.user, options.password, options.output, options.threads_count, + options.country_code, progress_bar, ) diff --git a/tools/python/booking/download_hotels.py b/tools/python/booking/download_hotels.py index 1ebd0b257c4..42fd9b0def0 100755 --- a/tools/python/booking/download_hotels.py +++ b/tools/python/booking/download_hotels.py @@ -207,7 +207,12 @@ def download_hotels_by_country(api, country): def download( - country_code, user, password, path, threads_count, progress_bar=tqdm(disable=True) + user, + password, + path, + threads_count, + country_code=None, + progress_bar=tqdm(disable=True), ): api = BookingApi(user, password, "2.4") list_api = BookingListApi(api) diff --git a/tools/python/booking/download_test_data.py b/tools/python/booking/download_test_data.py index 6acf1017b16..61e1685031a 100755 --- a/tools/python/booking/download_test_data.py +++ b/tools/python/booking/download_test_data.py @@ -70,7 +70,12 @@ def download_hotels_by_country(api, district_names, country): def download_test_data( - country_code, user, password, path, threads_count, progress_bar=tqdm(disable=True) + user, + password, + path, + threads_count, + country_code=None, + progress_bar=tqdm(disable=True), ): logging.info(f"Starting test dataset download.") api = BookingApi(user, password, "2.4")