Browse Source

HUE-3797 [scheduler] Basics to support cron scheduled queries

Romain 6 years ago
parent
commit
65277b686b

+ 28 - 14
desktop/core/src/desktop/lib/scheduler/lib/beat.py

@@ -15,7 +15,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from django_celery_beat.models import PeriodicTask, IntervalSchedule
+from django_celery_beat.models import PeriodicTask, CrontabSchedule, IntervalSchedule
 
 from desktop.lib.scheduler.lib.api import Api
 
@@ -23,16 +23,30 @@ from desktop.lib.scheduler.lib.api import Api
 class CeleryBeatApi(Api):
 
   def submit_schedule(self, request, coordinator, mapping):
-    schedule, created = IntervalSchedule.objects.get_or_create(
-      every=10,
-      period=IntervalSchedule.SECONDS,
-    )
-
-    task, created = PeriodicTask.objects.get_or_create(
-      interval=schedule,
-      name='Scheduled query N',
-      task='notebook.tasks.run_sync_query',
-    )
-
-    task.enabled = True
-    task.save()
+    is_cron = True
+
+    if True:
+      schedule, created = CrontabSchedule.objects.get_or_create(
+        minute='*',
+        hour='*',
+        day_of_week='*',
+        day_of_month='*',
+        month_of_year='*'
+      )
+
+      task, created = PeriodicTask.objects.get_or_create(
+        crontab=schedule,
+        name='Scheduled query N',
+        task='notebook.tasks.run_sync_query',
+      )
+    else:
+      schedule, created = IntervalSchedule.objects.get_or_create(
+        every=15,
+        period=IntervalSchedule.SECONDS,
+      )
+
+      task, created = PeriodicTask.objects.get_or_create(
+        interval=schedule,
+        name='Scheduled query',
+        task='notebook.tasks.run_sync_query',
+      )

+ 1 - 0
desktop/core/src/desktop/settings.py

@@ -706,6 +706,7 @@ if desktop.conf.TASK_SERVER.ENABLED.get() or desktop.conf.TASK_SERVER.BEAT_ENABL
   if desktop.conf.TASK_SERVER.BEAT_ENABLED.get():
     INSTALLED_APPS.append('django_celery_beat')
     INSTALLED_APPS.append('timezone_field')
+    USE_TZ = True
 
 
 PROMETHEUS_EXPORT_MIGRATIONS = False # Needs to be there even when enable_prometheus is not enabled

+ 1 - 1
desktop/libs/notebook/src/notebook/connectors/sql_alchemy.py

@@ -146,7 +146,7 @@ class SqlAlchemyApi(Api):
           'name': col[0] if (type(col) is tuple or type(col) is dict) else col.name if hasattr(col, 'name') else col,
           'type': 'STRING_TYPE',
           'comment': ''
-        } for col in result.cursor.description]
+        } for col in result.cursor.description] if result.cursor else []
     }
     CONNECTION_CACHE[guid] = cache
 

+ 1 - 0
desktop/libs/notebook/src/notebook/tasks.py

@@ -183,6 +183,7 @@ def run_sync_query(doc_id, user):
 
   return task
 
+
 # TODO: Convert csv to excel if needed
 def download(*args, **kwargs):
   notebook = args[0]