From b9107e8b319ad156a8c74dbfe5b0d562c5b3a622 Mon Sep 17 00:00:00 2001
From: "J.C. Zhong" <exzjc@hotmail.com>
Date: Fri, 23 Aug 2024 01:28:12 +0000
Subject: [PATCH 1/2] chore: add parameters for the task of  disable scheduled
 docs

---
 querybook/server/lib/utils/utils.py           |  4 ++
 .../server/tasks/disable_scheduled_docs.py    | 53 ++++++++++++-------
 2 files changed, 38 insertions(+), 19 deletions(-)

diff --git a/querybook/server/lib/utils/utils.py b/querybook/server/lib/utils/utils.py
index ec966607d..3e1072473 100644
--- a/querybook/server/lib/utils/utils.py
+++ b/querybook/server/lib/utils/utils.py
@@ -160,3 +160,7 @@ def get_default_args(func):
         for k, v in signature.parameters.items()
         if v.default is not inspect.Parameter.empty
     }
+
+
+def str_to_bool(value: str | bool):
+    return value.lower() in (True, "yes", "true", "t", "1")
diff --git a/querybook/server/tasks/disable_scheduled_docs.py b/querybook/server/tasks/disable_scheduled_docs.py
index ab1a35198..d216fcacd 100644
--- a/querybook/server/tasks/disable_scheduled_docs.py
+++ b/querybook/server/tasks/disable_scheduled_docs.py
@@ -2,31 +2,25 @@
 from datetime import datetime, timedelta
 from enum import Enum
 from typing import Optional
-from app.flask_app import celery
 
-from models.schedule import (
-    TaskSchedule,
-)
 from app.db import with_session
-from env import QuerybookSettings
+from app.flask_app import celery
 from const.data_doc import DataCellType
 from const.impression import ImpressionItemType
-from const.schedule import TaskRunStatus, ScheduleTaskType
+from const.schedule import ScheduleTaskType, TaskRunStatus
+from env import QuerybookSettings
+from lib.logger import get_logger
 from lib.notify.utils import notify_user
-from lib.scheduled_datadoc.legacy import convert_if_legacy_datadoc_schedule
-from logic.schedule import (
-    DATADOC_SCHEDULE_PREFIX,
-    update_task_schedule,
-    with_task_logging,
-    get_task_run_records,
-    delete_task_schedule,
-)
 from lib.query_analysis.lineage import get_table_statement_type
+from lib.scheduled_datadoc.legacy import convert_if_legacy_datadoc_schedule
+from lib.utils.utils import str_to_bool
 from logic.datadoc import get_data_doc_by_id
-from logic.user import get_user_by_id
 from logic.impression import get_viewers_count_by_item_after_date
-from lib.logger import get_logger
-
+from logic.schedule import (DATADOC_SCHEDULE_PREFIX, delete_task_schedule,
+                            get_task_run_records, update_task_schedule,
+                            with_task_logging)
+from logic.user import get_user_by_id
+from models.schedule import TaskSchedule
 
 logger = get_logger(__name__)
 
@@ -271,8 +265,29 @@ def disable_deactivated_scheduled_docs(
 
 @celery.task(bind=True)
 @with_task_logging()
-def disable_scheduled_docs(self):
-    tasks_to_disable = disable_deactivated_scheduled_docs()
+def disable_scheduled_docs(
+    self,
+    notifier: str = None,
+    disable_if_inactive_owner=True,
+    disable_if_failed_for_n_runs=5,
+    disable_if_no_impression_for_n_days=30,
+    skip_if_no_impression_but_non_select=False,
+    skip_if_no_impression_with_export=False,
+):
+    disable_config = DisableConfig(
+        disable_if_inactive_owner=disable_if_inactive_owner,
+        disable_if_failed_for_n_runs=disable_if_failed_for_n_runs,
+        disable_if_no_impression_for_n_days=disable_if_no_impression_for_n_days,
+        skip_if_no_impression_but_non_select=str_to_bool(
+            skip_if_no_impression_but_non_select
+        ),
+        skip_if_no_impression_with_export=str_to_bool(
+            skip_if_no_impression_with_export
+        ),
+    )
+    tasks_to_disable = disable_deactivated_scheduled_docs(
+        notifier=notifier, disable_config=disable_config
+    )
     if len(tasks_to_disable) == 0:
         logger.info("No scheduled docs disabled.")
     else:

From d4fbc40ac8269256bcc934e787fc53bab2150e83 Mon Sep 17 00:00:00 2001
From: "J.C. Zhong" <exzjc@hotmail.com>
Date: Fri, 23 Aug 2024 01:58:08 +0000
Subject: [PATCH 2/2] fix linter

---
 querybook/server/lib/utils/utils.py              |  2 +-
 querybook/server/tasks/disable_scheduled_docs.py | 10 +++++++---
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/querybook/server/lib/utils/utils.py b/querybook/server/lib/utils/utils.py
index 3e1072473..dad878ba2 100644
--- a/querybook/server/lib/utils/utils.py
+++ b/querybook/server/lib/utils/utils.py
@@ -162,5 +162,5 @@ def get_default_args(func):
     }
 
 
-def str_to_bool(value: str | bool):
+def str_to_bool(value: Union[str, bool]):
     return value.lower() in (True, "yes", "true", "t", "1")
diff --git a/querybook/server/tasks/disable_scheduled_docs.py b/querybook/server/tasks/disable_scheduled_docs.py
index d216fcacd..e5709c367 100644
--- a/querybook/server/tasks/disable_scheduled_docs.py
+++ b/querybook/server/tasks/disable_scheduled_docs.py
@@ -16,9 +16,13 @@
 from lib.utils.utils import str_to_bool
 from logic.datadoc import get_data_doc_by_id
 from logic.impression import get_viewers_count_by_item_after_date
-from logic.schedule import (DATADOC_SCHEDULE_PREFIX, delete_task_schedule,
-                            get_task_run_records, update_task_schedule,
-                            with_task_logging)
+from logic.schedule import (
+    DATADOC_SCHEDULE_PREFIX,
+    delete_task_schedule,
+    get_task_run_records,
+    update_task_schedule,
+    with_task_logging,
+)
 from logic.user import get_user_by_id
 from models.schedule import TaskSchedule