Browse Source

HUE-3797 [scheduler] Skeleton of a schedulable SQL task

Romain 6 years ago
parent
commit
3ff0918ce4

+ 3 - 2
apps/useradmin/src/useradmin/management/commands/useradmin_sync_with_unix.py

@@ -14,12 +14,13 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from useradmin.views import sync_unix_users_and_groups
 
 from django.core.management.base import BaseCommand
-
 from django.utils.translation import ugettext_lazy as _
 
+from useradmin.views import sync_unix_users_and_groups
+
+
 class Command(BaseCommand):
   """
   Handler for syncing the Hue database with Unix users and groups

+ 25 - 1
desktop/core/src/desktop/celery.py

@@ -1,3 +1,20 @@
+#!/usr/bin/env python
+# Licensed to Cloudera, Inc. under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  Cloudera, Inc. licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from __future__ import absolute_import, unicode_literals
 
 import os
@@ -35,9 +52,16 @@ if 'django_celery_beat' in INSTALLED_APPS:
 
   if True:
     app.conf.beat_schedule.update({
-      'add-every-monday-morning2': {
+      'add-every-monday-morning': {
         'task': 'desktop.celery.debug_task',
         'schedule': crontab(minute='*'),
         #'args': (16, 16),
       },
     })
+    app.conf.beat_schedule.update({
+      'customer_count_query': {
+        'task': 'notebook.tasks.run_sync_query',
+        'schedule': crontab(minute='*'),
+        'args': (None, None),
+      },
+    })

+ 16 - 1
desktop/core/src/desktop/lib/scheduler/lib/base.py

@@ -1,4 +1,19 @@
-
+#!/usr/bin/env python
+# Licensed to Cloudera, Inc. under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  Cloudera, Inc. licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 
 from desktop.lib.scheduler.lib.beat import CeleryBeatApi
 from desktop.lib.scheduler.lib.oozie import OozieApi

+ 16 - 0
desktop/core/src/desktop/lib/scheduler/tasks.py

@@ -1,3 +1,19 @@
+#!/usr/bin/env python
+# Licensed to Cloudera, Inc. under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  Cloudera, Inc. licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 
 from desktop.celery import app
 

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

@@ -23,7 +23,7 @@ import uuid
 
 from django.utils.translation import ugettext as _
 
-from desktop.conf import has_multi_cluster
+from desktop.conf import has_multi_cluster, TASK_SERVER
 from desktop.lib import export_csvxls
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.lib.i18n import smart_unicode
@@ -275,7 +275,7 @@ class Notebook(object):
 def get_api(request, snippet):
   from notebook.connectors.oozie_batch import OozieApi
 
-  if snippet.get('wasBatchExecuted'):
+  if snippet.get('wasBatchExecuted') and not TASK_SERVER.ENABLED.get():
     return OozieApi(user=request.user, request=request)
 
   if snippet['type'] == 'report':

+ 1 - 1
desktop/libs/notebook/src/notebook/models.py

@@ -72,7 +72,7 @@ def make_notebook(name='Browse', description='', editor_type='hive', statement='
                   namespace=None, compute=None):
   '''
   skip_historify: do not add the task to the query history. e.g. SQL Dashboard
-  isManaged: true when being a managed by Hue operation (include_managed=True in document), e.g. exporting query result, dropping some tables
+  is_task / isManaged: true when being a managed by Hue operation (include_managed=True in document), e.g. exporting query result, dropping some tables
   '''
   from notebook.connectors.hiveserver2 import HS2Api
 

+ 43 - 9
desktop/libs/notebook/src/notebook/tasks.py

@@ -17,10 +17,12 @@
 from __future__ import absolute_import, unicode_literals
 
 import csv
+import datetime
 import json
 import logging
 import StringIO
 import sys
+import time
 
 from celery.utils.log import get_task_logger
 from celery import states
@@ -87,8 +89,8 @@ class ExecutionWrapperCallback(object):
     self.meta['status'] = status
     download_to_file.update_state(task_id=self.uuid, state='PROGRESS', meta=self.meta)
 
-#TODO: Add periodic cleanup task
-#TODO: UI should be able to close a query that is available, but not expired
+# TODO: Add periodic cleanup task
+# TODO: UI should be able to close a query that is available, but not expired
 @app.task()
 def download_to_file(notebook, snippet, file_format='csv', max_rows=-1, **kwargs):
   download_to_file.update_state(task_id=notebook['uuid'], state='STARTED', meta={})
@@ -123,7 +125,7 @@ def close_statement_async(notebook, snippet, **kwargs):
 
 
 @app.task(ignore_result=True)
-def batch_execute_query(doc_id, user):
+def run_sync_query(doc_id, user):
   '''Independently run a query as a user and insert the result into another table.'''
   # get SQL
   # Add INSERT INTO table
@@ -131,9 +133,37 @@ def batch_execute_query(doc_id, user):
   # execute query
   # return when done. send email notification. get taskid.
   # see in Flower API for listing runs?
+  from django.contrib.auth.models import User
+  from notebook.models import make_notebook, MockedDjangoRequest
 
+  from desktop.auth.backend import rewrite_user
 
-#TODO: Convert csv to excel if needed
+  editor_type = 'impala'
+  sql = 'INSERT into customer_scheduled SELECT * FROM default.customers LIMIT 100;'
+  request = MockedDjangoRequest(user=rewrite_user(User.objects.get(username='romain')))
+
+  notebook = make_notebook(
+      name='Scheduler query N',
+      editor_type=editor_type,
+      statement=sql,
+      status='ready',
+      #on_success_url=on_success_url,
+      last_executed=time.mktime(datetime.datetime.now().timetuple()) * 1000,
+      is_task=True
+  )
+
+  task = notebook.execute(request, batch=True)
+
+  task['uuid'] = task['history_uuid']
+  status = check_status(task)
+
+  while status['status'] in ('waiting', 'running'):
+    status = check_status(task)
+    time.sleep(3)
+
+  return task
+
+# TODO: Convert csv to excel if needed
 def download(*args, **kwargs):
   notebook = args[0]
   result = download_to_file.AsyncResult(args[0]['uuid'])
@@ -159,12 +189,15 @@ def execute(*args, **kwargs):
   snippet = args[1]
   kwargs['max_rows'] = TASK_SERVER.FETCH_RESULT_LIMIT.get()
   _patch_status(notebook)
-  download_to_file.apply_async(args=args, kwargs=kwargs, task_id=notebook['uuid'])
+
+  task = download_to_file.apply_async(args=args, kwargs=kwargs, task_id=notebook['uuid'])
 
   should_close, resp = get_current_statement(snippet) # This redoes some of the work in api.execute. Other option is to pass statement, but then we'd have to modify notebook.api.
-  #if should_close: #front end already calls close_statement for multi statement execution no need to do here. In addition, we'd have to figure out what was the previous guid.
+  # if should_close: #front end already calls close_statement for multi statement execution no need to do here.
+  # In addition, we'd have to figure out what was the previous guid.
 
-  resp.update({'sync': False,
+  resp.update({
+      'sync': False,
       'has_result_set': True,
       'modified_row_count': 0,
       'guid': '',
@@ -173,7 +206,8 @@ def execute(*args, **kwargs):
         'data': [],
         'meta': [],
         'type': 'table'
-      }})
+      }}
+    )
   return resp
 
 def check_status(*args, **kwargs):
@@ -211,7 +245,7 @@ def get_log(notebook, snippet, startFrom=None, size=None, postdict=None, user_id
         output.write(line)
     return output.getvalue()
 
-def get_jobs(notebook, snippet, logs, **kwargs): #Re implement to fetch updated guid in download_to_file from DB
+def get_jobs(notebook, snippet, logs, **kwargs): # Re implementation to fetch updated guid in download_to_file from DB
   result = download_to_file.AsyncResult(notebook['uuid'])
   state = result.state
   if state == states.PENDING:

+ 3 - 5
desktop/libs/notebook/src/notebook/views.py

@@ -18,14 +18,13 @@
 import json
 import logging
 
-from beeswax.data_export import DOWNLOAD_COOKIE_AGE
-
 from django.urls import reverse
 from django.db.models import Q
 from django.shortcuts import redirect
 from django.utils.translation import ugettext as _
 from django.views.decorators.clickjacking import xframe_options_exempt
 
+from beeswax.data_export import DOWNLOAD_COOKIE_AGE
 from desktop.conf import ENABLE_DOWNLOAD, USE_NEW_EDITOR, TASK_SERVER
 from desktop.lib import export_csvxls
 from desktop.lib.django_util import render, JsonResponse
@@ -33,9 +32,9 @@ from desktop.lib.exceptions_renderable import PopupException
 from desktop.lib.json_utils import JSONEncoderForHTML
 from desktop.models import Document2, Document, FilesystemException
 from desktop.views import serve_403_error
-
 from metadata.conf import has_optimizer, has_catalog, has_workload_analytics
 
+from notebook import tasks as ntasks
 from notebook.conf import get_ordered_interpreters, SHOW_NOTEBOOKS
 from notebook.connectors.base import Notebook, get_api as _get_api, _get_snippet_name
 from notebook.connectors.spark_shell import SparkApi
@@ -43,10 +42,9 @@ from notebook.decorators import check_editor_access_permission, check_document_a
 from notebook.management.commands.notebook_setup import Command
 from notebook.models import make_notebook
 
+
 LOG = logging.getLogger(__name__)
 
-if TASK_SERVER.ENABLED.get():
-  import notebook.tasks as ntasks
 
 class ApiWrapper(object):
   def __init__(self, request, snippet):