Bläddra i källkod

HUE-3797 [scheduler] Improve sync task to properly work with beat

Romain 6 år sedan
förälder
incheckning
b481d1d5c1

+ 2 - 1
desktop/core/src/desktop/lib/scheduler/lib/beat.py

@@ -33,7 +33,8 @@ class CeleryBeatApi(Api):
     mapping
       {u'oozie.use.system.libpath': u'True', 'dryrun': False, u'start_date': u'2019-08-10T17:02', u'end_date': u'2019-08-17T17:02'}
     '''
-    is_cron = True # IntervalSchedule is buggy https://github.com/celery/django-celery-beat/issues/279
+    # IntervalSchedule is buggy https://github.com/celery/django-celery-beat/issues/279
+    is_cron = True
 
     # Assumes SQL queries currently
     document = Document2.objects.get(uuid=coordinator.get_data_for_json()['properties']['document'])

+ 8 - 2
desktop/libs/notebook/src/notebook/connectors/base.py

@@ -106,8 +106,14 @@ class Notebook(object):
     return _data
 
   def get_str(self, from_oozie_action=False):
-    return '\n\n\n'.join(['USE %s;\n\n%s' % (snippet['database'], snippet['statement_raw'] if from_oozie_action else Notebook.statement_with_variables(snippet))
-                          for snippet in self.get_data()['snippets']])
+    return '\n\n\n'.join([
+        'USE %s;\n\n%s' % (
+          snippet['database'],
+          snippet['statement_raw'] if from_oozie_action else Notebook.statement_with_variables(snippet)
+        )
+        for snippet in self.get_data()['snippets']
+      ]
+    )
 
   @staticmethod
   def statement_with_variables(snippet):

+ 4 - 3
desktop/libs/notebook/src/notebook/tasks.py

@@ -45,7 +45,7 @@ from useradmin.models import User
 
 from notebook.api import _get_statement
 from notebook.connectors.base import get_api, QueryExpired, ExecutionWrapper
-from notebook.models import make_notebook, MockedDjangoRequest
+from notebook.models import make_notebook, MockedDjangoRequest, Notebook
 from notebook.sql_utils import get_current_statement
 
 if sys.version_info[0] > 2:
@@ -152,9 +152,10 @@ def run_sync_query(doc_id, user):
   # Add variable substitution
   # Send notifications: done/on failure
   if type(user) is str:
-    user = User.objects.get(username=user)
+    lookup = {'email' if ENABLE_ORGANIZATIONS.get() else 'username': user}
+    user = User.objects.get(**lookup)
 
-  query_document = Document2.objects.document(user=user, doc_id=doc_id)
+  query_document = Document2.objects.get_by_uuid(user=user, uuid=doc_id)
   notebook = Notebook(document=query_document).get_data()
   snippet = notebook['snippets'][0]