Browse Source

[taskserver] Handle concurrent uploads using a counter and locking mechanism and add new configs (TASK_SERVER_V2) (#3733)

## What changes were proposed in this pull request?

- Handle concurrent uploads using a counter and locking mechnaism. 
  The "upload_available_space" key is set in Redis when Hue boots up. It calculates the actual available /tmp space and stores it. Whenever a new file upload is started, a new key is created in the format "upload_<unique_uuid_hash>" and stores the upload file size in Redis. Also, a second key in the format "upload_<same_uuid>_timestamp" (stores the timestamp) is created, and we reduce the "upload_available_space" value by the new upload file size. When the upload completes, the "upload_<unique_uuid_hash>" and "upload_<same_uuid>_timestamp" keys are deleted, and the reserved space is given back to "upload_available_space". 
  
  Assume we have 1GB free space:
  
  - `upload_available_space = 1GB`
  
  When file1 (100MB) upload is triggered:
  
  - `upload_<uuid1> = 100MB`
  - `upload_<uuid1>_timestamp = <timestamp>`
  - `upload_available_space = 1GB - 100MB = 900MB`
  
  When file2 (100MB) upload is triggered:
  
  - `upload_<uuid2> = 100MB`
  - `upload_<uuid2>_timestamp = <timestamp>`
  - `upload_available_space = 900MB - 100MB = 800MB`
  
  When the uploads are complete, the keys (`upload_<uuid1>`, `upload_<uuid1>_timestamp`) and (`upload_<uuid2>`, `upload_<uuid2>_timestamp`) are deleted, and `upload_available_space` is updated back:
  
  - `upload_available_space = 800MB + 100MB + 100MB = 1GB`
  
  In case of failed uploads, the same process is repeated for each retry made by the user. However, we would have leftover `upload_` keys for each retry since the keys will only be deleted on a successful upload. So we run a periodic job `cleanup_stale_uploads` to clean up these keys. This job runs every `CLEANUP_STALE_UPLOADS_IN_REDIS_PERIODIC_INTERVAL` minutes and deletes `upload_*` keys if the timestamp difference is greater than 60 minutes.
  
- Moved task server configs to TASK_SERVER_V2.
- Moved reserve, release upload space methods to filebrowser utils. New configs for periodic scheduling under task server
- Display max_file_upload_size_limit on upload modal. Parse redis broker url from configs.
- Changed task server configuration parsing from awk to grep. Set autorestart to false in redis and celery template.
- Show task_server tab in Admin Server, based on task_server_v2 configs
- Moving Uploaded chunk log message to Debug from Info
- Setting celery default log level to info
- Pull timezone from hue.ini. /tmp_cleaner job checks the timestamp of each file and deletes it based on timedelta=60mins

## How was this patch tested?
Tested on local machine and using docker builds. 

Change-Id: Id167701873d10426c7f6e5064b3feb731966b2a7

Co-authored-by: Athithyaa Selvam <aselvam@cloudera.com>
Athithyaa Selvam 1 year ago
parent
commit
52c1bf51a4
30 changed files with 1013 additions and 512 deletions
  1. 10 4
      apps/filebrowser/src/filebrowser/conf.py
  2. 101 40
      apps/filebrowser/src/filebrowser/tasks.py
  3. 3 3
      apps/filebrowser/src/filebrowser/templates/listdir.mako
  4. 39 9
      apps/filebrowser/src/filebrowser/templates/listdir_components.mako
  5. 7 2
      apps/filebrowser/src/filebrowser/urls.py
  6. 60 1
      apps/filebrowser/src/filebrowser/utils.py
  7. 101 66
      apps/filebrowser/src/filebrowser/views.py
  8. 55 0
      desktop/conf.dist/hue.ini
  9. 56 0
      desktop/conf/pseudo-distributed.ini.tmpl
  10. 115 80
      desktop/core/src/desktop/api2.py
  11. 29 19
      desktop/core/src/desktop/celery.py
  12. 113 0
      desktop/core/src/desktop/conf.py
  13. 4 1
      desktop/core/src/desktop/js/reactComponents/TaskBrowser/TaskBrowser.tsx
  14. 7 9
      desktop/core/src/desktop/lib/fs/ozone/upload.py
  15. 4 5
      desktop/core/src/desktop/lib/scheduler/api.py
  16. 14 11
      desktop/core/src/desktop/management/commands/runcelery.py
  17. 67 68
      desktop/core/src/desktop/models.py
  18. 47 21
      desktop/core/src/desktop/settings.py
  19. 3 1
      desktop/core/src/desktop/templates/about_layout.mako
  20. 16 25
      desktop/core/src/desktop/urls.py
  21. 13 23
      desktop/libs/aws/src/aws/s3/upload.py
  22. 23 25
      desktop/libs/azure/src/azure/abfs/upload.py
  23. 23 17
      desktop/libs/hadoop/src/hadoop/fs/upload.py
  24. 26 26
      desktop/libs/notebook/src/notebook/api.py
  25. 9 8
      desktop/libs/notebook/src/notebook/connectors/base.py
  26. 17 17
      desktop/libs/notebook/src/notebook/models.py
  27. 24 20
      desktop/libs/notebook/src/notebook/tasks.py
  28. 23 7
      tools/container/hue/hue.sh
  29. 2 2
      tools/container/hue/supervisor-files/hue_celery_template
  30. 2 2
      tools/container/hue/supervisor-files/hue_redis_template

+ 10 - 4
apps/filebrowser/src/filebrowser/conf.py

@@ -18,9 +18,8 @@
 import os
 import sys
 
-from desktop.conf import ENABLE_DOWNLOAD
+from desktop.conf import ENABLE_DOWNLOAD, is_oozie_enabled
 from desktop.lib.conf import Config, coerce_bool
-from desktop.conf import is_oozie_enabled
 
 if sys.version_info[0] > 2:
   from django.utils.translation import gettext_lazy as _
@@ -31,7 +30,7 @@ MAX_SNAPPY_DECOMPRESSION_SIZE = Config(
   key="max_snappy_decompression_size",
   help=_("Max snappy decompression size in bytes."),
   private=True,
-  default=1024*1024*25,
+  default=1024 * 1024 * 25,
   type=int)
 
 ARCHIVE_UPLOAD_TEMPDIR = Config(
@@ -52,10 +51,12 @@ CONCURRENT_MAX_CONNECTIONS = Config(
   type=int,
   help=_('Configure the maximum number of concurrent connections(chunks) for file uploads using the chunked file uploader.'))
 
+
 def get_desktop_enable_download():
   """Get desktop enable_download default"""
   return ENABLE_DOWNLOAD.get()
 
+
 SHOW_DOWNLOAD_BUTTON = Config(
   key="show_download_button",
   help=_("whether to show the download button in hdfs file browser."),
@@ -78,7 +79,7 @@ ENABLE_EXTRACT_UPLOADED_ARCHIVE = Config(
 
 REDIRECT_DOWNLOAD = Config(
   key="redirect_download",
-  help=_("Redirect client to WebHdfs or S3 for file download. Note: Turning this on will "\
+  help=_("Redirect client to WebHdfs or S3 for file download. Note: Turning this on will "
     "override notebook/redirect_whitelist for user selected file downloads on WebHdfs & S3."),
   type=coerce_bool,
   default=False)
@@ -96,6 +97,11 @@ MAX_FILE_SIZE_UPLOAD_LIMIT = Config(
   help=_('A limit on a file size (bytes) that can be uploaded to a filesystem. '
           'A value of -1 means there will be no limit.'))
 
+
+def max_file_size_upload_limit():
+  return MAX_FILE_SIZE_UPLOAD_LIMIT.get()
+
+
 FILE_DOWNLOAD_CACHE_CONTROL = Config(
   key="file_download_cache_control",
   type=str,

+ 101 - 40
apps/filebrowser/src/filebrowser/tasks.py

@@ -14,26 +14,29 @@
 # 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.
-import logging
-import subprocess
 import os
-
 import shutil
-import psutil
+import logging
+import subprocess
 from datetime import datetime, timedelta
 
-from celery.utils.log import get_task_logger
+import pytz
+import psutil
 from celery import states
-from django.db import transaction
+from celery.utils.log import get_task_logger
 from django.http import HttpRequest
 from django.utils import timezone
 
 from desktop.auth.backend import rewrite_user
 from desktop.celery import app
-from desktop.conf import TASK_SERVER
+from desktop.conf import TASK_SERVER_V2
 from desktop.lib import fsmanager
-from useradmin.models import User
+from filebrowser.utils import release_reserved_space_for_file_uploads, reserve_space_for_file_uploads
+
+if hasattr(TASK_SERVER_V2, 'get') and TASK_SERVER_V2.ENABLED.get():
+  from desktop.settings import TIME_ZONE, initialize_free_disk_space_in_redis, parse_broker_url
 from filebrowser.views import UPLOAD_CLASSES
+from useradmin.models import User
 
 LOG_TASK = get_task_logger(__name__)
 LOG = logging.getLogger()
@@ -52,14 +55,17 @@ STATE_MAP = {
   states.IGNORED: 'ignored'
 }
 
+
 @app.task
 def error_handler(request, exc, traceback):
   print('Task {0} raised exception: {1!r}\n{2!r}'.format(
     request.id, exc, traceback))
 
+
 @app.task()
 def upload_file_task(**kwargs):
   task_id = kwargs.get("qquuid")
+  file_size = kwargs.get("qqtotalfilesize")
   user_id = kwargs["user_id"]
   scheme = kwargs["scheme"]
   postdict = kwargs.get("postdict", None)
@@ -67,9 +73,16 @@ def upload_file_task(**kwargs):
   kwargs["username"] = request.user.username
   kwargs["task_name"] = "fileupload"
   kwargs["state"] = "STARTED"
-  now = timezone.now()
-  kwargs["task_start"] = now.strftime("%Y-%m-%dT%H:%M:%S")
+  kwargs["task_start"] = timezone.now().astimezone(pytz.timezone(TIME_ZONE)).strftime("%Y-%m-%dT%H:%M:%S")
+  kwargs["started"] = timezone.now().astimezone(pytz.timezone(TIME_ZONE)).strftime("%Y-%m-%dT%H:%M:%S")
   upload_file_task.update_state(task_id=task_id, state='STARTED', meta=kwargs)
+
+  # Reserve space for upload
+  if not reserve_space_for_file_uploads(task_id, file_size):
+    kwargs["state"] = "FAILURE"
+    upload_file_task.update_state(task_id=task_id, state='FAILURE', meta=kwargs)
+    raise Exception("Insufficient space for upload")
+
   try:
     upload_class = UPLOAD_CLASSES.get(kwargs["scheme"])
     _fs = upload_class(request, args=[], **kwargs)
@@ -82,11 +95,13 @@ def upload_file_task(**kwargs):
     raise Exception(f"Upload failed %s" % err)
 
   kwargs["state"] = "SUCCESS"
-  kwargs["started"] = now.strftime("%Y-%m-%dT%H:%M:%S")
-  kwargs["task_end"] = now.strftime("%Y-%m-%dT%H:%M:%S")
+  kwargs["task_end"] = timezone.now().astimezone(pytz.timezone(TIME_ZONE)).strftime("%Y-%m-%dT%H:%M:%S.%f")
+  kwargs["progress"] = "100%"
   upload_file_task.update_state(task_id=task_id, state='SUCCESS', meta=kwargs)
+  release_reserved_space_for_file_uploads(task_id)
   return None
 
+
 def _get_request(postdict=None, user_id=None, scheme=None):
   request = HttpRequest()
   request.POST = postdict
@@ -100,11 +115,11 @@ def _get_request(postdict=None, user_id=None, scheme=None):
 
   return request
 
+
 @app.task()
 def document_cleanup_task(**kwargs):
   keep_days = kwargs.get('keep_days')
   task_id = kwargs.get("qquuid")
-  now = timezone.now()
 
   kwargs["username"] = kwargs["username"]
   kwargs["task_name"] = "document_cleanup"
@@ -112,7 +127,7 @@ def document_cleanup_task(**kwargs):
   kwargs["parameters"] = keep_days
   kwargs["task_id"] = task_id
   kwargs["progress"] = "0%"
-  kwargs["task_start"] = now.strftime("%Y-%m-%dT%H:%M:%S")
+  kwargs["task_start"] = timezone.now().astimezone(pytz.timezone(TIME_ZONE)).strftime("%Y-%m-%dT%H:%M:%S")
 
   try:
     INSTALL_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..', '..'))
@@ -121,63 +136,109 @@ def document_cleanup_task(**kwargs):
     kwargs["state"] = "SUCCESS"
     LOG.info(f"Document_cleanup_task completed successfully.")
   except Exception as err:
+    kwargs["state"] = "FAILURE"
     document_cleanup_task.update_state(task_id=task_id, state='FAILURE', meta=kwargs)
     raise Exception(f"Upload failed %s" % err)
 
   kwargs["state"] = "SUCCESS"
   kwargs["progress"] = "100%"
-  kwargs["task_end"] = now.strftime("%Y-%m-%dT%H:%M:%S")
+  kwargs["task_end"] = timezone.now().astimezone(pytz.timezone(TIME_ZONE)).strftime("%Y-%m-%dT%H:%M:%S.%f")
   document_cleanup_task.update_state(task_id=task_id, state='SUCCESS', meta=kwargs)
 
   return None
 
+
 @app.task()
 def check_disk_usage_and_clean_task(**kwargs):
   task_id = kwargs.get("qquuid")
-  default_cleanup_threshold = 90
-  cleanup_threshold = kwargs.get('cleanup_threshold', default_cleanup_threshold)
+  cleanup_threshold = kwargs.get('cleanup_threshold', TASK_SERVER_V2.DISK_USAGE_CLEANUP_THRESHOLD)
   username = kwargs.get("username", "celery_scheduler")
-  now = timezone.now()
+  now = timezone.now().astimezone(pytz.timezone(TIME_ZONE))
   kwargs = {
     "username": username,
     "task_name": "tmp_cleanup",
     "task_id": task_id,
     "parameters": cleanup_threshold,
     "progress": "0%",
-    "task_start": now.strftime("%Y-%m-%dT%H:%M:%S")
+    "task_start": now.strftime("%Y-%m-%dT%H:%M:%S.%f"),
   }
 
+  check_disk_usage_and_clean_task.update_state(task_id=task_id, state=states.STARTED, meta=kwargs)
+
+  def delete_old_files(directory, current_time, time_delta):
+    for root, dirs, files in os.walk(directory):
+      for file in files:
+        file_path = os.path.join(root, file)
+        try:
+          file_modified_time = datetime.fromtimestamp(os.path.getmtime(file_path))
+          file_modified_time = timezone.make_aware(file_modified_time, timezone.get_default_timezone())
+          if file_modified_time < current_time - time_delta:
+            if os.path.isfile(file_path) or os.path.islink(file_path):
+              os.unlink(file_path)
+              LOG.info(f"Deleted file {file_path}")
+        except PermissionError:
+          LOG.warning(f"Permission denied: unable to delete {file_path}")
+        except Exception as err:
+          check_disk_usage_and_clean_task.update_state(task_id=task_id, state=states.FAILURE, meta=kwargs)
+          raise Exception(f"Failed to delete {file_path}. Reason: {err}")
+        finally:
+          kwargs["task_end"] = timezone.now().astimezone(pytz.timezone(TIME_ZONE)).strftime("%Y-%m-%dT%H:%M:%S.%f")
+
   disk_usage = psutil.disk_usage('/')
   if disk_usage.percent >= int(cleanup_threshold):
     LOG.info(f"Disk usage is above {cleanup_threshold}%, cleaning up /tmp directory...")
     tmp_dir = '/tmp'
-    for filename in os.listdir(tmp_dir):
-      file_path = os.path.join(tmp_dir, filename)
-      try:
-        file_modified_time = datetime.fromtimestamp(os.path.getmtime(file_path))
-        # Make file_modified_time timezone-aware
-        file_modified_time = timezone.make_aware(file_modified_time, timezone.get_default_timezone())
-        if file_modified_time < now - timedelta(minutes=15):
-          if os.path.isfile(file_path) or os.path.islink(file_path):
-            os.unlink(file_path)
-          elif os.path.isdir(file_path):
-            shutil.rmtree(file_path)
-          LOG.info(f"Deleted {file_path}")
-      except PermissionError:
-        LOG.warning(f"Permission denied: unable to delete {file_path}")
-      except Exception as err:
-        check_disk_usage_and_clean_task.update_state(task_id=task_id, state='FAILURE', meta=kwargs)
-        raise Exception(f"Failed to delete {file_path}. Reason: {err}")
-
+    delete_old_files(tmp_dir, now, timedelta(minutes=TASK_SERVER_V2.DISK_USAGE_AND_CLEAN_TASK_TIME_DELTA.get()))
     kwargs["progress"] = "100%"
-    check_disk_usage_and_clean_task.update_state(task_id=task_id, state='SUCCESS', meta=kwargs)
+    kwargs["task_end"] = timezone.now().astimezone(pytz.timezone(TIME_ZONE)).strftime("%Y-%m-%dT%H:%M:%S.%f")
+    check_disk_usage_and_clean_task.update_state(task_id=task_id, state=states.SUCCESS, meta=kwargs)
     LOG.info("/tmp directory cleaned.")
-
   else:
     kwargs["progress"] = "100%"
-    check_disk_usage_and_clean_task.update_state(task_id=task_id, state='SUCCESS', meta=kwargs)
+    kwargs["task_end"] = timezone.now().astimezone(pytz.timezone(TIME_ZONE)).strftime("%Y-%m-%dT%H:%M:%S.%f")
+    check_disk_usage_and_clean_task.update_state(task_id=task_id, state=states.SUCCESS, meta=kwargs)
     LOG.info(f"Disk usage is {disk_usage.percent}%, no need to clean up.")
 
   # Get available disk space after cleanup
   free_space = psutil.disk_usage('/tmp').free
   return {'free_space': free_space}
+
+
+@app.task()
+def cleanup_stale_uploads(**kwargs):
+  timedelta_minutes = kwargs.get("timeout_minutes", TASK_SERVER_V2.CLEANUP_STALE_UPLOADS_TASK_TIME_DELTA.get())
+  username = kwargs.get("username", "celery_scheduler")
+  task_id = kwargs.get("qquuid")
+  kwargs = {
+    "username": username,
+    "task_name": "cleanup_stale_uploads",
+    "task_id": task_id,
+    "parameters": timedelta_minutes,
+    "progress": "0%",
+    "task_start": timezone.now().astimezone(pytz.timezone(TIME_ZONE)).strftime("%Y-%m-%dT%H:%M:%S"),
+  }
+  cleanup_stale_uploads.update_state(task_id=task_id, state=states.STARTED, meta=kwargs)
+  redis_client = parse_broker_url(TASK_SERVER_V2.BROKER_URL.get())
+  try:
+    current_time = int(datetime.now().timestamp())
+
+    for key in redis_client.scan_iter('upload__*'):
+      timestamp_key = f'{key.decode()}_timestamp'
+      timestamp = redis_client.get(timestamp_key)
+
+      if timestamp:
+        timestamp = int(timestamp)
+        if current_time - timestamp > timedelta_minutes * 60:
+          file_size = int(redis_client.get(key))
+          redis_client.incrby('upload_available_space', file_size)
+          redis_client.delete(key)
+          redis_client.delete(timestamp_key)
+    initialize_free_disk_space_in_redis()
+    kwargs["progress"] = "100%"
+    kwargs["task_end"] = timezone.now().astimezone(pytz.timezone(TIME_ZONE)).strftime("%Y-%m-%dT%H:%M:%S.%f")
+    cleanup_stale_uploads.update_state(task_id=task_id, state=states.SUCCESS, meta=kwargs)
+  except Exception as e:
+    cleanup_stale_uploads.update_state(task_id=task_id, state=states.FAILURE, meta={"error": str(e)})
+    LOG.exception("Failed to cleanup stale uploads: %s", str(e))
+  finally:
+    redis_client.close()

+ 3 - 3
apps/filebrowser/src/filebrowser/templates/listdir.mako

@@ -168,7 +168,7 @@ ${ fb_components.menubar() }
             <a class="btn fileToolbarBtn" title="${_('Upload files')}" data-bind="visible: !inTrash(), css: {'disabled': isS3Root()}, click: function(){ if (!isS3Root()) { uploadFile(false) }}"><i class="fa fa-arrow-circle-o-up"></i> ${_('Upload')}</a>
             <!-- /ko -->
             <!-- ko if: isTaskServerEnabled -->
-            <a class="btn fileToolbarBtn" title="${_('Select a file to upload. The file will first be saved locally and then automatically transferred to the designated file system (e.g., S3, Azure) in the background. The upload modal closes immediately after the file is queued, allowing you to continue working. A notification, \'File upload scheduled. Please check the task server page for progress,\' will confirm the upload has started. This feature is especially useful for large files, as it eliminates the need to wait for the upload to complete.')}" data-bind="visible: !inTrash(), css: {'disabled': isS3Root()}, click: function(){ if (!isS3Root()) { uploadFile(true) }}"><i class="fa fa-arrow-circle-o-up"></i> ${_('Schedule Upload')}</a>
+            <a class="btn fileToolbarBtn" title="${_('Select a file to upload. The file will first be saved locally and then automatically transferred to the designated file system (e.g., S3, Azure) in the background. The upload modal closes immediately after the file is queued, allowing you to continue working. A notification, \'File upload scheduled. Please check the task server page for progress,\' will confirm the upload has started. This feature is especially useful for large files, as it eliminates the need to wait for the upload to complete.')}" data-bind="visible: !inTrash(), css: {'disabled': isS3Root()}, click: function(){ if (!isS3Root()) { checkAndDisplayAvailableSpace(); uploadFile(true); }}"><i class="fa fa-arrow-circle-o-up"></i> ${_('Schedule Upload')}</a>
             <!-- /ko -->
           <!-- /ko -->
           <!-- ko if: isGS -->
@@ -179,7 +179,7 @@ ${ fb_components.menubar() }
             <a class="btn fileToolbarBtn" title="${_('Upload files')}" data-bind="visible: !inTrash(), css: {'disabled': isABFSRoot()}, click: function(){ if (!isABFSRoot()) { uploadFile(false) }}"><i class="fa fa-arrow-circle-o-up"></i> ${_('Upload')}</a>
             <!-- /ko -->
             <!-- ko if: isTaskServerEnabled -->
-            <a class="btn fileToolbarBtn" title="${_('Select a file to upload. The file will first be saved locally and then automatically transferred to the designated file system (e.g., S3, Azure) in the background. The upload modal closes immediately after the file is queued, allowing you to continue working. A notification, \'File upload scheduled. Please check the task server page for progress,\' will confirm the upload has started. This feature is especially useful for large files, as it eliminates the need to wait for the upload to complete.')}" data-bind="visible: !inTrash(), css: {'disabled': isABFSRoot()}, click: function(){ if (!isABFSRoot()) { uploadFile(true) }}"><i class="fa fa-arrow-circle-o-up"></i> ${_('Schedule Upload')}</a>
+            <a class="btn fileToolbarBtn" title="${_('Select a file to upload. The file will first be saved locally and then automatically transferred to the designated file system (e.g., S3, Azure) in the background. The upload modal closes immediately after the file is queued, allowing you to continue working. A notification, \'File upload scheduled. Please check the task server page for progress,\' will confirm the upload has started. This feature is especially useful for large files, as it eliminates the need to wait for the upload to complete.')}" data-bind="visible: !inTrash(), css: {'disabled': isABFSRoot()}, click: function(){ if (!isABFSRoot()) { checkAndDisplayAvailableSpace(); uploadFile(true); }}"><i class="fa fa-arrow-circle-o-up"></i> ${_('Schedule Upload')}</a>
             <!-- /ko -->
           <!-- /ko -->
           <!-- ko ifnot: isS3() || isGS() || isABFS() -->
@@ -192,7 +192,7 @@ ${ fb_components.menubar() }
           <!-- /ko -->
           <!-- ko if: isTaskServerEnabled -->
           <div id="upload-dropdown" class="btn-group" style="vertical-align: middle">
-            <a data-hue-analytics="filebrowser:upload-btn-click" href="javascript: void(0)" class="btn upload-link dropdown-toggle" title="${_('Select a file to upload. The file will first be saved locally and then automatically transferred to the designated file system (e.g., S3, Azure) in the background. The upload modal closes immediately after the file is queued, allowing you to continue working. A notification, \'File upload scheduled. Please check the task server page for progress,\' will confirm the upload has started. This feature is especially useful for large files, as it eliminates the need to wait for the upload to complete.')}" data-bind="click: function() { uploadFile(true); }, visible: !inTrash(), css: {'disabled': (isOFS() && (isOFSRoot() || isOFSServiceID() || isOFSVol()))}">
+            <a data-hue-analytics="filebrowser:upload-btn-click" href="javascript: void(0)" class="btn upload-link dropdown-toggle" title="${_('Select a file to upload. The file will first be saved locally and then automatically transferred to the designated file system (e.g., S3, Azure) in the background. The upload modal closes immediately after the file is queued, allowing you to continue working. A notification, \'File upload scheduled. Please check the task server page for progress,\' will confirm the upload has started. This feature is especially useful for large files, as it eliminates the need to wait for the upload to complete.')}" data-bind="click: function() { checkAndDisplayAvailableSpace(); uploadFile(true);}, visible: !inTrash(), css: {'disabled': (isOFS() && (isOFSRoot() || isOFSServiceID() || isOFSVol()))}">
               <i class="fa fa-arrow-circle-o-up"></i> ${_('Schedule Upload')}
             </a>
           </div>

+ 39 - 9
apps/filebrowser/src/filebrowser/templates/listdir_components.mako

@@ -702,7 +702,8 @@ else:
     <div class="qq-uploader-selector" style="margin-left: 10px">
         <div class="qq-upload-drop-area-selector" qq-hide-dropzone><span>${_('Drop the files here to upload')}</span></div>
         <div class="qq-upload-button-selector qq-no-float">${_('Select files')}</div> &nbsp;
-        <span class="muted">${_('or drag and drop them here')}</span>
+        <span class="muted">${_('or drag and drop them here')}</span> &nbsp;
+        <span class="muted free-space-info"></span>
 
         <ul class="qq-upload-list-selector qq-upload-files unstyled qq-no-float" style="margin-right: 0;">
             <li>
@@ -2102,6 +2103,24 @@ else:
         doPoll();
       }
 
+      self.checkAndDisplayAvailableSpace = function () {
+        $.ajax({
+            url: '/filebrowser/upload/taskserver/get_available_space_for_file_uploads/',
+            success: function(response) {
+                if (typeof window.MAX_FILE_SIZE_UPLOAD_LIMIT === 'undefined') {
+                  window.MAX_FILE_SIZE_UPLOAD_LIMIT = 5 * 1024 * 1024 * 1024; // 5GB in bytes
+                }
+                var freeSpace = Math.min(response.upload_available_space, window.MAX_FILE_SIZE_UPLOAD_LIMIT);
+                $('.free-space-info').text('- Max file size upload limit: ' + formatBytes(freeSpace));
+            },
+            error: function(xhr, status, error) {
+                huePubSub.publish('hue.global.error', { message: '${ _("Error checking available space: ") }' + error});
+                $('.free-space-info').text('Error checking available space');
+            }
+        });
+      };
+ 
+
       self.uploadFile = (function () {  
           var uploader; 
           var scheduleUpload;
@@ -2112,6 +2131,7 @@ else:
             var action = "/filebrowser/upload/chunks/";
             self.taskIds = [];
             self.listItems = [];
+            self.checkAndDisplayAvailableSpace();
             uploader = new qq.FileUploader({
               element: document.getElementById("fileUploader"),
               request: {
@@ -2171,9 +2191,12 @@ else:
                       self.listItems.push(listItem);
                         if (scheduleUpload && self.pendingUploads() === 0) {
                           $('#uploadFileModal').modal('hide');
-                          $(document).trigger('info', "File upload scheduled. Please check the task server page for progress.");
+                          huePubSub.publish('hue.global.info', { message: '${ _("File upload scheduled. Please check the task server page for progress.") }'});
                         }
-                        pollForTaskProgress(response.task_id, listItem, fileName);
+                        // Add a delay of 2 seconds before calling pollForTaskProgress, to ensure the upload task is received by the task_server before checking its status. 
+                        setTimeout(function() {
+                          pollForTaskProgress(response.task_id, listItem, fileName);
+                        }, 2000);  
                       self.filesToHighlight.push(response.path);                       
                     }
                     if (self.pendingUploads() === 0) {                    
@@ -2188,13 +2211,20 @@ else:
                       
                       // Make an AJAX request to check available disk space
                       $.ajax({
-                        url: '/desktop/api2/taskserver/get_available_space/',
+                        url: '/filebrowser/upload/taskserver/get_available_space_for_file_uploads/',
                         success: function(response) {
-                          var freeSpace = response.free_space;
+                          if (typeof window.MAX_FILE_SIZE_UPLOAD_LIMIT === 'undefined' || window.MAX_FILE_SIZE_UPLOAD_LIMIT === -1) {
+                            window.MAX_FILE_SIZE_UPLOAD_LIMIT = 5 * 1024 * 1024 * 1024; // 5GB in bytes
+                          }
+                          var freeSpace = Math.min(response.upload_available_space, window.MAX_FILE_SIZE_UPLOAD_LIMIT);
                           var file = uploader.getFile(id); // Use the stored reference
-                          
-                          if ((file.size > freeSpace) && (file.size > window.MAX_FILE_SIZE_UPLOAD_LIMIT*1024*1024)) {
-                            $(document).trigger('info', "Not enough space available to upload this file.")
+                          // Update the free space display
+                          $('.free-space-info').text('- Max file size upload limit: ' + formatBytes(freeSpace));
+                          if ((file.size > freeSpace) || (file.size > window.MAX_FILE_SIZE_UPLOAD_LIMIT)) {
+                            huePubSub.publish('hue.global.error', { message: '${ _("Not enough space available to upload this file.") }'});
+                            deferred.failure(); // Reject the promise to cancel the upload
+                          } else if (file.size > window.MAX_FILE_SIZE_UPLOAD_LIMIT) {
+                            huePubSub.publish('hue.global.error', { message: '${ _("File size is bigger than MAX_FILE_SIZE_UPLOAD_LIMIT.") }'});
                             deferred.failure(); // Reject the promise to cancel the upload
                           } else {
                             var newPath = "/filebrowser/upload/chunks/file?dest=" + encodeURIComponent(self.currentPath().normalize('NFC'));
@@ -2204,7 +2234,7 @@ else:
                           }
                         },
                         error: function(xhr, status, error) {
-                          alert('Error checking available space: ' + error);
+                          huePubSub.publish('hue.global.error', { message: '${ _("Error checking available space: ") }' + error});
                           deferred.failure(); // Reject the promise to cancel the upload
                         }
                       });

+ 7 - 2
apps/filebrowser/src/filebrowser/urls.py

@@ -17,8 +17,7 @@
 
 import sys
 
-from filebrowser import views as filebrowser_views
-from filebrowser import api as filebrowser_api
+from filebrowser import api as filebrowser_api, utils as filebrowser_utils, views as filebrowser_views
 
 if sys.version_info[0] > 2:
   from django.urls import re_path
@@ -46,6 +45,12 @@ urlpatterns = [
   re_path(r'^upload/file$', filebrowser_views.upload_file, name='upload_file'),
   re_path(r'^upload/chunks', filebrowser_views.upload_chunks, name='upload_chunks'),
   re_path(r'^upload/complete', filebrowser_views.upload_complete, name='upload_complete'),
+  re_path(r'^upload/taskserver/get_available_space_for_file_uploads',
+          filebrowser_utils.get_available_space_for_file_uploads, name='get_available_space_for_file_uploads'),
+  re_path(r'^upload/taskserver/reserve_space_for_file_uploads',
+          filebrowser_utils.reserve_space_for_file_uploads, name='reserve_space_for_file_uploads'),
+  re_path(r'^upload/taskserver/release_reserved_space_for_file_uploads',
+          filebrowser_utils.release_reserved_space_for_file_uploads, name='release_reserved_space_for_file_uploads'),
   re_path(r'^extract_archive', filebrowser_views.extract_archive_using_batch_job, name='extract_archive_using_batch_job'),
   re_path(r'^compress_files', filebrowser_views.compress_files_using_batch_job, name='compress_files_using_batch_job'),
   re_path(r'^trash/restore$', filebrowser_views.trash_restore, name='trash_restore'),

+ 60 - 1
apps/filebrowser/src/filebrowser/utils.py

@@ -16,11 +16,20 @@
 import io
 import os
 import logging
-LOG = logging.getLogger()
+from datetime import datetime
 
+from desktop.conf import TASK_SERVER_V2
+from desktop.lib.django_util import JsonResponse
 from filebrowser.conf import ARCHIVE_UPLOAD_TEMPDIR
+
+if hasattr(TASK_SERVER_V2, 'get') and TASK_SERVER_V2.ENABLED.get():
+  from desktop.settings import parse_broker_url
+LOG = logging.getLogger()
+
+
 DEFAULT_WRITE_SIZE = 1024 * 1024 * 128
 
+
 def calculate_total_size(uuid, totalparts):
   total = 0
   files = [os.path.join(ARCHIVE_UPLOAD_TEMPDIR.get(), f'{uuid}_{i}') for i in range(totalparts)]
@@ -33,6 +42,7 @@ def calculate_total_size(uuid, totalparts):
       LOG.error(f"calculate_total_size: For the file '{file_path}' error occurred: {e}")
   return total
 
+
 def generate_chunks(uuid, totalparts, default_write_size=DEFAULT_WRITE_SIZE):
   fp = io.BytesIO()
   total = 0
@@ -62,3 +72,52 @@ def generate_chunks(uuid, totalparts, default_write_size=DEFAULT_WRITE_SIZE):
     fp.close()
   for file_path in files:
     os.remove(file_path)
+
+
+def get_available_space_for_file_uploads(request):
+  redis_client = parse_broker_url(TASK_SERVER_V2.BROKER_URL.get())
+  try:
+    upload_available_space = int(redis_client.get('upload_available_space'))
+    if upload_available_space is None:
+      raise ValueError("upload_available_space key not set in Redis")
+    return JsonResponse({'upload_available_space': upload_available_space})
+  except Exception as e:
+    LOG.exception("Failed to get available space: %s", str(e))
+    return JsonResponse({'error': str(e)}, status=500)
+  finally:
+    redis_client.close()
+
+
+def reserve_space_for_file_uploads(uuid, file_size):
+  redis_client = parse_broker_url(TASK_SERVER_V2.BROKER_URL.get())
+  try:
+    upload_available_space = int(redis_client.get('upload_available_space'))
+    if upload_available_space is None:
+      raise ValueError("upload_available_space key not set in Redis")
+    if upload_available_space >= file_size:
+      redis_client.decrby('upload_available_space', file_size)
+      redis_client.set(f'upload__{uuid}', file_size)
+      redis_client.set(f'upload__{uuid}_timestamp', int(datetime.now().timestamp()))
+      return True
+    else:
+      return False
+  except Exception as e:
+    LOG.exception("Failed to reserve space: %s", str(e))
+    return False
+  finally:
+    redis_client.close()
+
+
+def release_reserved_space_for_file_uploads(uuid):
+  redis_client = parse_broker_url(TASK_SERVER_V2.BROKER_URL.get())
+  try:
+    reserved_space = redis_client.get(f'upload__{uuid}')
+    if reserved_space:
+      file_size = int(redis_client.get(f'upload__{uuid}'))
+      redis_client.incrby('upload_available_space', file_size)
+      redis_client.delete(f'upload__{uuid}')
+      redis_client.delete(f'upload__{uuid}_timestamp')
+  except Exception as e:
+    LOG.exception("Failed to release reserved space: %s", str(e))
+  finally:
+    redis_client.close()

+ 101 - 66
apps/filebrowser/src/filebrowser/views.py

@@ -15,105 +15,115 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from future import standard_library
-standard_library.install_aliases()
-
-from django.views.decorators.csrf import csrf_exempt
-from filebrowser.conf import ARCHIVE_UPLOAD_TEMPDIR
-from django.core.files.uploadhandler import FileUploadHandler, StopUpload, StopFutureHandlers
-from hadoop.conf import UPLOAD_CHUNK_SIZE
-
-from builtins import object
+import os
+import re
+import sys
+import stat as stat_module
 import errno
 import logging
-import mimetypes
 import operator
-import os
+import mimetypes
 import posixpath
-import re
-import stat as stat_module
-import sys
-import urllib.request, urllib.error
-
+import urllib.error
+import urllib.request
+from builtins import object
 from bz2 import decompress
 from datetime import datetime
+from functools import partial
 
-from django.core.paginator import EmptyPage, Paginator, Page, InvalidPage
+from django.core.files.uploadhandler import FileUploadHandler, StopFutureHandlers, StopUpload
+from django.core.paginator import EmptyPage, InvalidPage, Page, Paginator
+from django.http import Http404, HttpResponse, HttpResponseForbidden, HttpResponseNotModified, HttpResponseRedirect, StreamingHttpResponse
+from django.shortcuts import redirect
+from django.template.defaultfilters import filesizeformat, stringformat
 from django.urls import reverse
-from django.template.defaultfilters import stringformat, filesizeformat
-from django.http import Http404, StreamingHttpResponse, HttpResponseNotModified,\
-    HttpResponseForbidden, HttpResponse, HttpResponseRedirect
+from django.utils.html import escape
+from django.utils.http import http_date
+from django.views.decorators.csrf import csrf_exempt
 from django.views.decorators.http import require_http_methods
 from django.views.static import was_modified_since
-from django.shortcuts import redirect
-from functools import partial
-from django.utils.http import http_date
-from django.utils.html import escape
 
 from aws.s3.s3fs import S3FileSystemException, S3ListAllBucketsException, get_s3_home_directory
+from aws.s3.upload import S3FineUploaderChunkedUpload
+from azure.abfs.upload import ABFSFineUploaderChunkedUpload
 from desktop import appmanager
 from desktop.auth.backend import is_admin
-from desktop.conf import RAZ, ENABLE_NEW_STORAGE_BROWSER, TASK_SERVER
-from desktop.lib import fsmanager
-from desktop.lib import i18n
+from desktop.conf import ENABLE_NEW_STORAGE_BROWSER, RAZ, TASK_SERVER_V2
+from desktop.lib import fsmanager, i18n
 from desktop.lib.conf import coerce_bool
-from desktop.lib.django_util import render, format_preserving_redirect
-from desktop.lib.django_util import JsonResponse
-from desktop.lib.export_csvxls import file_reader
+from desktop.lib.django_util import JsonResponse, format_preserving_redirect, render
 from desktop.lib.exceptions_renderable import PopupException
+from desktop.lib.export_csvxls import file_reader
 from desktop.lib.fs import splitpath
-from desktop.lib.fs.ozone.ofs import get_ofs_home_directory
 from desktop.lib.fs.gc.gs import get_gs_home_directory
+from desktop.lib.fs.ozone.ofs import get_ofs_home_directory
+from desktop.lib.fs.ozone.upload import OFSFineUploaderChunkedUpload
 from desktop.lib.i18n import smart_str
 from desktop.lib.paths import SAFE_CHARACTERS_URI, SAFE_CHARACTERS_URI_COMPONENTS
 from desktop.lib.tasks.compress_files.compress_utils import compress_files_in_hdfs
 from desktop.lib.tasks.extract_archive.extract_utils import extract_archive_in_hdfs
 from desktop.lib.view_util import is_ajax
 from desktop.views import serve_403_error
+from filebrowser.conf import (
+  ARCHIVE_UPLOAD_TEMPDIR,
+  ENABLE_EXTRACT_UPLOADED_ARCHIVE,
+  FILE_DOWNLOAD_CACHE_CONTROL,
+  MAX_SNAPPY_DECOMPRESSION_SIZE,
+  REDIRECT_DOWNLOAD,
+  SHOW_DOWNLOAD_BUTTON,
+  SHOW_UPLOAD_BUTTON,
+)
+from filebrowser.forms import (
+  ChmodFormSet,
+  ChownFormSet,
+  CopyFormSet,
+  EditorForm,
+  MkDirForm,
+  RenameForm,
+  RenameFormSet,
+  RestoreFormSet,
+  RmTreeFormSet,
+  SetReplicationFactorForm,
+  TouchForm,
+  TrashPurgeForm,
+  UploadArchiveForm,
+  UploadFileForm,
+)
+from filebrowser.lib import xxd
+from filebrowser.lib.archives import archive_factory
+from filebrowser.lib.rwx import filetype, rwx
+from hadoop.conf import UPLOAD_CHUNK_SIZE
 from hadoop.core_site import get_trash_interval
-from hadoop.fs.hadoopfs import Hdfs
 from hadoop.fs.exceptions import WebHdfsException
 from hadoop.fs.fsutils import do_overwrite_save
-from useradmin.models import User, Group
-
-from filebrowser.conf import ENABLE_EXTRACT_UPLOADED_ARCHIVE, MAX_SNAPPY_DECOMPRESSION_SIZE,\
-    SHOW_DOWNLOAD_BUTTON, SHOW_UPLOAD_BUTTON, REDIRECT_DOWNLOAD, FILE_DOWNLOAD_CACHE_CONTROL
-from filebrowser.lib.archives import archive_factory
-from filebrowser.lib.rwx import filetype, rwx
-from filebrowser.lib import xxd
-from filebrowser.forms import RenameForm, UploadFileForm, UploadArchiveForm, MkDirForm, EditorForm, TouchForm,\
-    RenameFormSet, RmTreeFormSet, ChmodFormSet, ChownFormSet, CopyFormSet, RestoreFormSet,\
-    TrashPurgeForm, SetReplicationFactorForm
-
+from hadoop.fs.hadoopfs import Hdfs
 from hadoop.fs.upload import HDFSFineUploaderChunkedUpload, LocalFineUploaderChunkedUpload
-from aws.s3.upload import S3FineUploaderChunkedUpload
-from azure.abfs.upload import ABFSFineUploaderChunkedUpload
-from desktop.lib.fs.ozone.upload import OFSFineUploaderChunkedUpload
+from useradmin.models import Group, User
 
 if sys.version_info[0] > 2:
   import io
-  from io import StringIO as string_io
-  from urllib.parse import quote as urllib_quote
-  from urllib.parse import unquote as urllib_unquote
-  from urllib.parse import urlparse as lib_urlparse
   from builtins import str as new_str
-  from avro import datafile, io
   from gzip import decompress as decompress_gzip
+  from io import StringIO as string_io
+  from urllib.parse import quote as urllib_quote, unquote as urllib_unquote, urlparse as lib_urlparse
+
+  from avro import datafile, io
   from django.utils.translation import gettext as _
 else:
+  from urllib import quote as urllib_quote, unquote as urllib_unquote
+
   from cStringIO import StringIO as string_io
-  from urllib import quote as urllib_quote
-  from urllib import unquote as urllib_unquote
   from urlparse import urlparse as lib_urlparse
   new_str = unicode
+  from gzip import GzipFile
+
   import parquet
   from avro import datafile, io
-  from gzip import GzipFile
   from django.utils.translation import ugettext as _
 
 
-DEFAULT_CHUNK_SIZE_BYTES = 1024 * 4 # 4KB
-MAX_CHUNK_SIZE_BYTES = 1024 * 1024 # 1MB
+DEFAULT_CHUNK_SIZE_BYTES = 1024 * 4  # 4KB
+MAX_CHUNK_SIZE_BYTES = 1024 * 1024  # 1MB
 
 # Defaults for "xxd"-style output.
 # Sentences refer to groups of bytes printed together, within a line.
@@ -124,12 +134,10 @@ BYTES_PER_SENTENCE = 2
 MAX_FILEEDITOR_SIZE = 256 * 1024
 
 INLINE_DISPLAY_MIMETYPE = re.compile(
-    'video/|image/|audio/|application/pdf|application/msword|application/excel|'
-    'application/vnd\.ms|'
-    'application/vnd\.openxmlformats'
+    r'video/|image/|audio/|application/pdf|application/msword|application/excel|application/vnd\.ms|application/vnd\.openxmlformats'
 )
 
-INLINE_DISPLAY_MIMETYPE_EXCEPTIONS = re.compile('image/svg\+xml')
+INLINE_DISPLAY_MIMETYPE_EXCEPTIONS = re.compile(r'image/svg\+xml')
 
 SCHEME_PREFIXES = {
     's3a': 's3a://',
@@ -153,6 +161,7 @@ if hasattr(ARCHIVE_UPLOAD_TEMPDIR, 'get') and not os.path.exists(ARCHIVE_UPLOAD_
 
 logger = logging.getLogger()
 
+
 class ParquetOptions(object):
   def __init__(self, col=None, format='json', no_headers=True, limit=-1):
     self.col = col
@@ -173,6 +182,7 @@ def index(request):
 
   return view(request, path)
 
+
 def _decode_slashes(path):
   # This is a fix for some installations where the path is still having the slash (/) encoded
   # as %2F while the rest of the path is actually decoded.
@@ -184,6 +194,7 @@ def _decode_slashes(path):
 
   return path
 
+
 def _normalize_path(path):
   path = _decode_slashes(path)
 
@@ -195,6 +206,7 @@ def _normalize_path(path):
 
   return path
 
+
 def get_scheme(path):
   path = _normalize_path(path)
   for scheme, prefix in SCHEME_PREFIXES.items():
@@ -202,6 +214,7 @@ def get_scheme(path):
       return scheme
   return 'hdfs'
 
+
 def download(request, path):
   """
   Downloads a file.
@@ -425,6 +438,7 @@ def edit(request, path, form=None):
     data['form'] = form
   return render("edit.mako", request, data)
 
+
 def save_file(request):
   """
   The POST endpoint to save a file in the file editor.
@@ -549,6 +563,7 @@ def _massage_page(page, paginator):
       'total_count': paginator.count
   }
 
+
 def listdir_paged(request, path):
   """
   A paginated version of listdir.
@@ -683,10 +698,12 @@ def scheme_absolute_path(root, path):
     path = splitPath._replace(scheme=splitRoot.scheme).geturl()
   return path
 
+
 def stat_absolute_path(path, stat):
   stat.path = scheme_absolute_path(path, stat.path)
   return stat
 
+
 def _massage_stats(request, stats):
   """
   Massage a stats record as returned by the filesystem implementation
@@ -1067,7 +1084,7 @@ def detect_snappy(contents):
   try:
     import snappy
     return snappy.isValidCompressed(contents)
-  except:
+  except ImportError:
     logging.exception('failed to detect snappy')
     return False
 
@@ -1087,7 +1104,7 @@ def snappy_installed():
     return True
   except ImportError:
     return False
-  except:
+  except Exception as e:
     logging.exception('failed to verify if snappy is installed')
     return False
 
@@ -1299,6 +1316,7 @@ def rename(request):
 
   return generic_op(RenameForm, request, smart_rename, ["src_path", "dest_path"], None)
 
+
 def set_replication(request):
   def smart_set_replication(src_path, replication_factor):
     result = request.fs.set_replication(urllib_unquote(src_path), replication_factor)
@@ -1307,6 +1325,7 @@ def set_replication(request):
 
   return generic_op(SetReplicationFactorForm, request, smart_set_replication, ["src_path", "replication_factor"], None)
 
+
 def mkdir(request):
   def smart_mkdir(path, name):
     # Make sure only one directory is specified at a time.
@@ -1317,6 +1336,7 @@ def mkdir(request):
 
   return generic_op(MkDirForm, request, smart_mkdir, ["path", "name"], "path")
 
+
 def touch(request):
   def smart_touch(path, name):
     # Make sure only the filename is specified.
@@ -1332,10 +1352,12 @@ def touch(request):
 
   return generic_op(TouchForm, request, smart_touch, ["path", "name"], "path")
 
+
 @require_http_methods(["POST"])
 def rmtree(request):
   recurring = []
   params = ["path"]
+
   def bulk_rmtree(*args, **kwargs):
     for arg in args:
       request.fs.do_as_user(request.user, request.fs.rmtree, arg['path'], 'skip_trash' in request.GET)
@@ -1349,6 +1371,7 @@ def rmtree(request):
 def move(request):
   recurring = ['dest_path']
   params = ['src_path']
+
   def bulk_move(*args, **kwargs):
     for arg in args:
       if arg['src_path'] == arg['dest_path']:
@@ -1367,6 +1390,7 @@ def move(request):
 def copy(request):
   recurring = ['dest_path']
   params = ['src_path']
+
   def bulk_copy(*args, **kwargs):
     ofs_skip_files = ''
     for arg in args:
@@ -1395,6 +1419,7 @@ def chmod(request):
                "group_read", "group_write", "group_execute",
                "other_read", "other_write", "other_execute"]
   params = ["path"]
+
   def bulk_chmod(*args, **kwargs):
     op = partial(request.fs.chmod, recursive=request.POST.get('recursive', False))
     for arg in args:
@@ -1420,6 +1445,7 @@ def chown(request):
 
   recurring = ["user", "group", "user_other", "group_other"]
   params = ["path"]
+
   def bulk_chown(*args, **kwargs):
     op = partial(request.fs.chown, recursive=request.POST.get('recursive', False))
     for arg in args:
@@ -1436,6 +1462,7 @@ def chown(request):
 def trash_restore(request):
   recurring = []
   params = ["path"]
+
   def bulk_restore(*args, **kwargs):
     for arg in args:
       request.fs.do_as_user(request.user, request.fs.restore, urllib_unquote(arg['path']))
@@ -1449,6 +1476,7 @@ def trash_restore(request):
 def trash_purge(request):
   return generic_op(TrashPurgeForm, request, request.fs.purge_trash, [], None)
 
+
 def _create_response(request, _fs, result="success", data="Success"):
   return {
       'path': _fs.filepath,
@@ -1461,6 +1489,7 @@ def _create_response(request, _fs, result="success", data="Success"):
       'task_id': _fs.qquuid
   }
 
+
 def perform_upload_task(request, *args, **kwargs):
   """
   Uploads a file to the specified destination.
@@ -1481,12 +1510,12 @@ def perform_upload_task(request, *args, **kwargs):
   upload_class = UPLOAD_CLASSES.get(scheme, LocalFineUploaderChunkedUpload)
   result = None
 
-  if TASK_SERVER.ENABLED.get():
+  if TASK_SERVER_V2.ENABLED.get():
     # If task server is enabled, upload the file to the task server.
     print("Uploading file to task server")
     _fs = upload_class(request, **kwargs)
     _fs.check_access()
-    from filebrowser.tasks import upload_file_task, error_handler
+    from filebrowser.tasks import error_handler, upload_file_task
     kwargs["user_id"] = request.user.id
     kwargs["scheme"] = scheme
     kwargs["filepath"] = _fs.filepath
@@ -1504,6 +1533,7 @@ def perform_upload_task(request, *args, **kwargs):
       result = "success"
   return _create_response(request, _fs, result=result, data="Success")
 
+
 def extract_upload_data(request, method):
   data = request.POST if method == "POST" else request.GET
   chunks = {
@@ -1519,6 +1549,7 @@ def extract_upload_data(request, method):
   }
   return chunks
 
+
 @require_http_methods(["POST"])
 def upload_chunks(request):
   """
@@ -1551,6 +1582,7 @@ def upload_chunks(request):
       return JsonResponse({'success': False, 'error': 'Error in upload %s' % str(e)})
   return JsonResponse({'success': False, 'error': 'Unsupported request method'})
 
+
 @require_http_methods(["POST"])
 def upload_complete(request):
   """
@@ -1567,6 +1599,7 @@ def upload_complete(request):
   except Exception as e:
     return JsonResponse({'success': False, 'error': 'Error in upload'})
 
+
 @require_http_methods(["POST"])
 def upload_file(request):
   """
@@ -1621,7 +1654,7 @@ def _upload_file(request):
       except Exception:
         pass
       if already_exists:
-        msg = _('Destination %(name)s already exists.')  % {'name': filepath}
+        msg = _('Destination %(name)s already exists.') % {'name': filepath}
       else:
         msg = _('Copy to %(name)s failed: %(error)s') % {'name': filepath, 'error': ex}
       raise PopupException(msg)
@@ -1674,7 +1707,7 @@ def compress_files_using_batch_job(request):
       except Exception as e:
         response['message'] = _('Exception occurred while compressing files: %s' % e)
     else:
-      response['message'] = _('Error: Output directory is not set.');
+      response['message'] = _('Error: Output directory is not set.')
   else:
     response['message'] = _('ERROR: Configuration parameter enable_extract_uploaded_archive ' +
                             'has to be enabled before calling this method.')
@@ -1704,9 +1737,11 @@ def truncate(toTruncate, charsToKeep=50):
   else:
     return toTruncate
 
+
 def unquote_url(url):
   url = urllib_unquote(url.encode('utf-8') if not isinstance(url, str) else url)
   return url.decode('utf-8') if isinstance(url, bytes) else url
 
+
 def _is_hdfs_superuser(request):
   return request.user.username == request.fs.superuser or request.user.groups.filter(name__exact=request.fs.supergroup).exists()

+ 55 - 0
desktop/conf.dist/hue.ini

@@ -338,6 +338,9 @@ http_500_debug_mode=false
 # Default value is true
 ## enable_help_menu=true
 
+# Enable chunked file uploader
+## enable_chunked_file_uploader=false
+
 # Administrators
 # ----------------
 [[django_admins]]
@@ -939,6 +942,58 @@ tls=no
 #   }
 # ]
 
+# Configuration options for the Task Server V2
+# ------------------------------------------------------------------------
+[[task_server_v2]]
+
+# If resource intensive or blocking can be delegated to an already running task server.
+## enabled=False
+
+# Switch on the integration with the Task Scheduler.
+## beat_enabled=False
+
+# Number of query results rows to fetch into the result storage.
+## fetch_result_limit=2000
+
+# Django file storage class to use to temporarily store query results
+## result_storage='{"backend": "django.core.files.storage.FileSystemStorage", "properties": {"location": "./logs"}}'
+
+# How the task server and tasks communicate.
+## broker_url=amqp://guest:guest@localhost//
+
+# Where to store task results. Defaults to local file system path. Celery comes with a several other backends.
+## celery_result_backend=file:///$HUE_ROOT/logs
+
+# Default options provided to the task server at startup.
+## celeryd_opts='--time-limit=300'
+
+# Django cache to use to store temporarily used data during query execution. This is in addition to result_file_storage and result_backend.
+## execution_storage='{"BACKEND": "django.core.cache.backends.locmem.LocMemCache", "LOCATION": "celery-hue"}'
+
+# Set the log level for celery
+## celery_log_level='INFO'
+
+# Switch on this peroidic cleaner which checks disk usage and makes space for file uploads
+## check_disk_usage_and_clean_task_enabled=False
+
+# Time interval in seconds to run this peroidic cleaner which checks disk usage and makes space for file uploads
+## check_disk_usage_and_clean_task_periodic_interval=1000
+
+# Clean up files in /tmp folder if the disk usage is beyond the threshold
+## disk_usage_cleanup_threshold=90
+
+# Clean up files older than timedelta. Unit of timedelta is minutes
+## disk_usage_and_clean_task_time_delta=60
+
+# Switch on this peroidic cleaner which cleans up failed upload tasks stored in redis
+## cleanup_stale_uploads_in_redis_enabled=False
+
+# Time interval in seconds to run this peroidic cleaner which cleans up failed upload tasks stored in redis
+## cleanup_stale_uploads_task_periodic_interval=900
+
+# Redis keys of format 'Upload__*' older than timedelta will be cleaned up. Unit of timedelta is minutes
+## cleanup_stale_uploads_task_time_delta=60
+
 # Settings for the Google Cloud lib
 # ------------------------------------------------------------------------
 

+ 56 - 0
desktop/conf/pseudo-distributed.ini.tmpl

@@ -343,6 +343,9 @@
   # Default value is true
   ## enable_help_menu=true
 
+  # Enable chunked file uploader
+  ## enable_chunked_file_uploader=false
+
   # Administrators
   # ----------------
   [[django_admins]]
@@ -922,6 +925,59 @@
    # Django cache to use to store temporarily used data during query execution. This is in addition to result_file_storage and result_backend.
    ## execution_storage='{"BACKEND": "django.core.cache.backends.locmem.LocMemCache", "LOCATION": "celery-hue"}'
 
+  # Configuration options for the Task Server V2
+  # ------------------------------------------------------------------------
+  [[task_server_v2]]
+
+   # If resource intensive or blocking can be delegated to an already running task server.
+   ## enabled=False
+
+   # Switch on the integration with the Task Scheduler.
+   ## beat_enabled=False
+
+   # Number of query results rows to fetch into the result storage.
+   ## fetch_result_limit=2000
+
+   # Django file storage class to use to temporarily store query results
+   ## result_storage='{"backend": "django.core.files.storage.FileSystemStorage", "properties": {"location": "./logs"}}'
+
+   # How the task server and tasks communicate.
+   ## broker_url=amqp://guest:guest@localhost//
+
+   # Where to store task results. Defaults to local file system path. Celery comes with a several other backends.
+   ## celery_result_backend=file:///$HUE_ROOT/logs
+
+   # Default options provided to the task server at startup.
+   ## celeryd_opts='--time-limit=300'
+
+   # Django cache to use to store temporarily used data during query execution. This is in addition to result_file_storage and result_backend.
+   ## execution_storage='{"BACKEND": "django.core.cache.backends.locmem.LocMemCache", "LOCATION": "celery-hue"}'
+
+   # Set the log level for celery
+   ## celery_log_level='INFO'
+
+   # Switch on this peroidic cleaner which checks disk usage and makes space for file uploads
+   ## check_disk_usage_and_clean_task_enabled=False
+
+   # Time interval in seconds to run this peroidic cleaner which checks disk usage and makes space for file uploads
+   ## check_disk_usage_and_clean_task_periodic_interval=1000
+
+   # Clean up files in /tmp folder if the disk usage is beyond the threshold
+   ## disk_usage_cleanup_threshold=90
+
+   # Clean up files older than timedelta. Unit of timedelta is minutes
+   ## disk_usage_and_clean_task_time_delta=60
+
+   # Switch on this peroidic cleaner which cleans up failed upload tasks stored in redis
+   ## cleanup_stale_uploads_in_redis_enabled=False
+
+   # Time interval in seconds to run this peroidic cleaner which cleans up failed upload tasks stored in redis
+   ## cleanup_stale_uploads_task_periodic_interval=900
+
+   # Redis keys of format 'Upload__*' older than timedelta will be cleaned up. Unit of timedelta is minutes
+   ## cleanup_stale_uploads_task_time_delta=60
+
+
   # Settings for the Google Cloud lib
   # ------------------------------------------------------------------------
 

+ 115 - 80
desktop/core/src/desktop/api2.py

@@ -15,71 +15,86 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from future import standard_library
-standard_library.install_aliases()
-from builtins import map
-import logging
 import os
 import re
-import json
 import sys
-import tempfile
+import json
+import logging
 import zipfile
+import tempfile
+from builtins import map
 
 from celery.app.control import Control
-from desktop.conf import TASK_SERVER
-if hasattr(TASK_SERVER, 'get') and TASK_SERVER.ENABLED.get():
-  from desktop.celery import app as celery_app
-import psutil
-import datetime
-import redis
-import re
-from django.http import HttpResponse, JsonResponse
 
+from desktop.conf import TASK_SERVER_V2
 
+if hasattr(TASK_SERVER_V2, 'get') and TASK_SERVER_V2.ENABLED.get():
+  from desktop.celery import app as celery_app
+  from desktop.settings import parse_broker_url
 from collections import defaultdict
 from datetime import datetime
 
 from django.core import management
 from django.db import transaction
-from django.http import HttpResponse
+from django.http import HttpResponse, JsonResponse
 from django.shortcuts import redirect
 from django.utils.html import escape
 from django.views.decorators.csrf import ensure_csrf_cookie
 from django.views.decorators.http import require_POST
 
-from metadata.conf import has_catalog
-from metadata.catalog_api import search_entities as metadata_search_entities, _highlight, \
-  search_entities_interactive as metadata_search_entities_interactive
-from notebook.connectors.base import Notebook
-from useradmin.models import User, Group
-
 from beeswax.models import Namespace
 from desktop import appmanager
 from desktop.auth.backend import is_admin
-from desktop.conf import ENABLE_CONNECTORS, ENABLE_GIST_PREVIEW, CUSTOM, get_clusters, IS_K8S_ONLY, ENABLE_SHARING
-from desktop.conf import ENABLE_NEW_STORAGE_BROWSER, ENABLE_CHUNKED_FILE_UPLOADER, TASK_SERVER
-from desktop.lib.conf import BoundContainer, GLOBAL_CONFIG, is_anonymous
+from desktop.conf import (
+  CUSTOM,
+  ENABLE_CHUNKED_FILE_UPLOADER,
+  ENABLE_CONNECTORS,
+  ENABLE_GIST_PREVIEW,
+  ENABLE_NEW_STORAGE_BROWSER,
+  ENABLE_SHARING,
+  IS_K8S_ONLY,
+  TASK_SERVER_V2,
+  get_clusters,
+)
+from desktop.lib.conf import GLOBAL_CONFIG, BoundContainer, is_anonymous
 from desktop.lib.django_util import JsonResponse, login_notrequired, render
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.lib.export_csvxls import make_response
-from desktop.lib.i18n import smart_str, force_unicode
+from desktop.lib.i18n import force_unicode, smart_str
 from desktop.lib.paths import get_desktop_root
-from desktop.models import Document2, Document, Directory, FilesystemException, uuid_default, \
-  UserPreferences, get_user_preferences, set_user_preferences, get_cluster_config, __paginate, _get_gist_document
-from desktop.views import get_banner_message, serve_403_error
 from desktop.log import DEFAULT_LOG_DIR
-from filebrowser.tasks import check_disk_usage_and_clean_task
-from filebrowser.tasks import document_cleanup_task
-
+from desktop.models import (
+  Directory,
+  Document,
+  Document2,
+  FilesystemException,
+  UserPreferences,
+  __paginate,
+  _get_gist_document,
+  get_cluster_config,
+  get_user_preferences,
+  set_user_preferences,
+  uuid_default,
+)
+from desktop.views import get_banner_message, serve_403_error
+from filebrowser.tasks import check_disk_usage_and_clean_task, document_cleanup_task
 from hadoop.cluster import is_yarn
+from metadata.catalog_api import (
+  _highlight,
+  search_entities as metadata_search_entities,
+  search_entities_interactive as metadata_search_entities_interactive,
+)
+from metadata.conf import has_catalog
+from notebook.connectors.base import Notebook
+from useradmin.models import Group, User
 
 if sys.version_info[0] > 2:
   from io import StringIO as string_io
+
   from django.utils.translation import gettext as _
 else:
-  from StringIO import StringIO as string_io
   from django.utils.translation import ugettext as _
+  from StringIO import StringIO as string_io
 
 LOG = logging.getLogger()
 
@@ -100,6 +115,7 @@ def api_error_handler(func):
 
   return decorator
 
+
 @api_error_handler
 def get_banners(request):
   banners = {
@@ -108,6 +124,7 @@ def get_banners(request):
   }
   return JsonResponse(banners)
 
+
 @api_error_handler
 def get_config(request):
   config = get_cluster_config(request.user)
@@ -115,7 +132,7 @@ def get_config(request):
   config['hue_config']['is_yarn_enabled'] = is_yarn()
   config['hue_config']['enable_new_storage_browser'] = ENABLE_NEW_STORAGE_BROWSER.get()
   config['hue_config']['enable_chunked_file_uploader'] = ENABLE_CHUNKED_FILE_UPLOADER.get()
-  config['hue_config']['enable_task_server'] = TASK_SERVER.ENABLED.get()
+  config['hue_config']['enable_task_server'] = TASK_SERVER_V2.ENABLED.get()
   config['clusters'] = list(get_clusters(request.user).values())
   config['documents'] = {
     'types': list(Document2.objects.documents(user=request.user).order_by().values_list('type', flat=True).distinct())
@@ -177,6 +194,7 @@ def get_hue_config(request):
     'apps': apps
   })
 
+
 @api_error_handler
 def get_context_namespaces(request, interface):
   '''
@@ -212,6 +230,7 @@ def get_context_namespaces(request, interface):
 
   return JsonResponse(response)
 
+
 @api_error_handler
 def get_context_computes(request, interface):
   '''
@@ -427,7 +446,7 @@ def _get_document_helper(request, uuid, with_data, with_dependencies, path):
       notebook = Notebook(document=document)
       notebook = upgrade_session_properties(request, notebook)
       data = json.loads(notebook.data)
-      if document.type == 'query-pig': # Import correctly from before Hue 4.0
+      if document.type == 'query-pig':  # Import correctly from before Hue 4.0
         properties = data['snippets'][0]['properties']
         if 'hadoopProperties' not in properties:
           properties['hadoopProperties'] = []
@@ -435,7 +454,7 @@ def _get_document_helper(request, uuid, with_data, with_dependencies, path):
           properties['parameters'] = []
         if 'resources' not in properties:
           properties['resources'] = []
-      if data.get('uuid') != document.uuid: # Old format < 3.11
+      if data.get('uuid') != document.uuid:  # Old format < 3.11
         data['uuid'] = document.uuid
 
     response['data'] = data
@@ -541,7 +560,6 @@ def update_document(request):
   })
 
 
-
 @api_error_handler
 @require_POST
 def delete_document(request):
@@ -571,6 +589,7 @@ def delete_document(request):
       'status': 0,
   })
 
+
 @api_error_handler
 @require_POST
 def copy_document(request):
@@ -596,7 +615,7 @@ def copy_document(request):
 
   # Import workspace for all oozie jobs
   if document.type == 'oozie-workflow2' or document.type == 'oozie-bundle2' or document.type == 'oozie-coordinator2':
-    from oozie.models2 import Workflow, Coordinator, Bundle, _import_workspace
+    from oozie.models2 import Bundle, Coordinator, Workflow, _import_workspace
     # Update the name field in the json 'data' field
     if document.type == 'oozie-workflow2':
       workflow = Workflow(document=document)
@@ -703,6 +722,7 @@ def share_document(request):
     'document': doc.to_dict()
   })
 
+
 @api_error_handler
 @require_POST
 def handle_submit(request):
@@ -752,40 +772,57 @@ def handle_submit(request):
     'status': 0
   })
 
+
 @api_error_handler
 def get_taskserver_tasks(request):
+  if not TASK_SERVER_V2.ENABLED.get():
+    return JsonResponse({'error': 'Task server is not enabled'}, status=400)
+
   """Retirve the tasks from the database"""
-  redis_client = redis.Redis(host='localhost', port=6379, db=0)
+  redis_client = parse_broker_url(TASK_SERVER_V2.BROKER_URL.get())
   tasks = []
+  try:
+    # Use scan_iter to efficiently iterate over keys matching the first pattern
+    for key in redis_client.scan_iter('celery-task-meta-*'):
+      task = json.loads(redis_client.get(key))
+      tasks.append(task)
 
-  # Use scan_iter to efficiently iterate over keys matching the first pattern
-  for key in redis_client.scan_iter('celery-task-meta-*'):
-    task = json.loads(redis_client.get(key))
-    tasks.append(task)
+    # Use scan_iter to efficiently iterate over keys matching the second pattern
+    for key in redis_client.scan_iter('task:*'):
+      task = json.loads(redis_client.get(key))
+      tasks.append(task)
 
-  # Use scan_iter to efficiently iterate over keys matching the second pattern
-  for key in redis_client.scan_iter('task:*'):
-    task = json.loads(redis_client.get(key))
-    tasks.append(task)
+    return JsonResponse(tasks, safe=False)
+  except Exception as e:
+    LOG.exception("Failed to retrieve tasks: %s", str(e))
+    return JsonResponse({'error': str(e)}, status=500)
+  finally:
+    redis_client.close()
 
-  return JsonResponse(tasks, safe=False)
 
 @api_error_handler
 def check_upload_status(request, task_id):
-  redis_client = redis.Redis(host='localhost', port=6379, db=0)
-  task_key = f'celery-task-meta-{task_id}'
-  task_data = redis_client.get(task_key)
+  redis_client = parse_broker_url(TASK_SERVER_V2.BROKER_URL.get())
+  try:
+    task_key = f'celery-task-meta-{task_id}'
+    task_data = redis_client.get(task_key)
 
-  if task_data is None:
-    return JsonResponse({'error': 'Task not found'}, status=404)
+    if task_data is None:
+      return JsonResponse({'error': 'Task not found'}, status=404)
 
-  task = json.loads(task_data)
-  is_finalized = task.get('status') == 'SUCCESS'
-  is_running = task.get('status') == 'RUNNING'
-  is_failure = task.get('status') == 'FAILURE'
-  is_revoked = task.get('status') == 'REVOKED'
+    task = json.loads(task_data)
+    is_finalized = task.get('status') == 'SUCCESS'
+    is_running = task.get('status') == 'RUNNING'
+    is_failure = task.get('status') == 'FAILURE'
+    is_revoked = task.get('status') == 'REVOKED'
+
+    return JsonResponse({'isFinalized': is_finalized, 'isRunning': is_running, 'isFailure': is_failure, 'isRevoked': is_revoked})
+  except Exception as e:
+    LOG.exception("Failed to check upload status: %s", str(e))
+    return JsonResponse({'error': str(e)}, status=500)
+  finally:
+    redis_client.close()
 
-  return JsonResponse({'isFinalized': is_finalized, 'isRunning': is_running, 'isFailure': is_failure, 'isRevoked': is_revoked})
 
 @api_error_handler
 def kill_task(request, task_id):
@@ -804,13 +841,9 @@ def kill_task(request, task_id):
     return JsonResponse({'status': 'error', 'message': f'Failed to terminate task {task_id}: {str(e)}'})
 
 
-def get_available_space(request):
-  free_space = psutil.disk_usage('/tmp').free
-  return JsonResponse({'free_space': free_space})
-
 def get_task_logs(request, task_id):
   log_dir = os.getenv("DESKTOP_LOG_DIR", DEFAULT_LOG_DIR)
-  log_file = "%s/rungunicornserver.log" % (log_dir)
+  log_file = "%s/celery.log" % (log_dir)
   task_log = []
   escaped_task_id = re.escape(task_id)
 
@@ -835,6 +868,7 @@ def get_task_logs(request, task_id):
 
   return HttpResponse(''.join(task_log), content_type='text/plain')
 
+
 @api_error_handler
 @require_POST
 def share_document_link(request):
@@ -926,45 +960,45 @@ def topological_sort(docs):
 
   '''There is a bug in django 1.11 (https://code.djangoproject.com/ticket/26291)
      and we are handling it via sorting the given documents in topological format.
-     
+
      Hence this function is needed only if we are using Python2 based Hue as it uses django 1.11
      and python3 based Hue don't require this method as it uses django 3.2.
 
      input => docs: -> list of documents which needs to import in Hue
-     output => serialized_doc: -> list of sorted documents 
+     output => serialized_doc: -> list of sorted documents
      (if document1 is dependent on document2 then document1 is listed after document2)'''
 
   size = len(docs)
   graph = defaultdict(list)
-  for doc in docs:     ## creating a graph, assuming a document is a node of graph
+  for doc in docs:     # creating a graph, assuming a document is a node of graph
     dep_size = len(doc['fields']['dependencies'])
     for i in range(dep_size):
       graph[(doc['fields']['dependencies'])[i][0]].append(doc['fields']['uuid'])
 
   visited = {}
   _doc = {}
-  for doc in docs:     ## making all the nodes of graph unvisited and capturing the doc in the dict with uuid as key
+  for doc in docs:     # making all the nodes of graph unvisited and capturing the doc in the dict with uuid as key
     _doc[doc['fields']['uuid']] = doc
     visited[doc['fields']['uuid']] = False
-  
+
   stack = []
-  for doc in docs:     ## calling _topological_sort function to sort the doc if node is not visited
-    if visited[doc['fields']['uuid']] == False:
+  for doc in docs:     # calling _topological_sort function to sort the doc if node is not visited
+    if not visited[doc['fields']['uuid']]:
       _topological_sort(doc['fields']['uuid'], visited, stack, graph)
-  
-  stack = stack[::-1]  ## list is in revered order so we are just reversing it
+
+  stack = stack[::-1]  # list is in revered order so we are just reversing it
 
   serialized_doc = []
   for i in range(size):
     serialized_doc.append(_doc[stack[i]])
-    
+
   return serialized_doc
 
 
 def _topological_sort(vertex, visited, stack, graph):
   visited[vertex] = True
   for i in graph[vertex]:
-    if visited[i] == False:
+    if not visited[i]:
       _topological_sort(i, visited, stack, graph)
 
   stack.append(vertex)
@@ -1024,10 +1058,10 @@ def import_documents(request):
       # Set last modified date to now
       doc['fields']['last_modified'] = datetime.now().replace(microsecond=0).isoformat()
       docs.append(doc)
-  
+
   if sys.version_info[0] < 3:
-    ## In Django 1.11 loaddata cannot deserialize fixtures with forward references hence 
-    ## calling the topological_sort function to sort the document
+    # In Django 1.11 loaddata cannot deserialize fixtures with forward references hence
+    # calling the topological_sort function to sort the document
     docs = topological_sort(docs)
 
   f = tempfile.NamedTemporaryFile(mode='w+', suffix='.json')
@@ -1036,7 +1070,7 @@ def import_documents(request):
 
   stdout = string_io()
   try:
-    with transaction.atomic(): # We wrap both commands to commit loaddata & sync
+    with transaction.atomic():  # We wrap both commands to commit loaddata & sync
       management.call_command('loaddata', f.name, verbosity=3, traceback=True, stdout=stdout)
       Document.objects.sync()
 
@@ -1064,6 +1098,7 @@ def import_documents(request):
   finally:
     stdout.close()
 
+
 def _update_imported_oozie_document(doc, uuids_map):
   for key, value in uuids_map.items():
     if value:
@@ -1342,11 +1377,11 @@ def _create_or_update_document_with_owner(doc, owner, uuids_map):
   if doc['fields']['dependencies']:
     history_deps_list = []
     for index, (uuid, version, is_history) in enumerate(doc['fields']['dependencies']):
-      if not uuid in list(uuids_map.keys()) and not is_history and \
+      if uuid not in list(uuids_map.keys()) and not is_history and \
       not Document2.objects.filter(uuid=uuid, version=version).exists():
         raise PopupException(_('Cannot import document, dependency with UUID: %s not found.') % uuid)
       elif is_history:
-        history_deps_list.insert(0, index) # Insert in decreasing order to facilitate delete
+        history_deps_list.insert(0, index)  # Insert in decreasing order to facilitate delete
         LOG.warning('History dependency with UUID: %s ignored while importing document %s' % (uuid, doc['fields']['name']))
 
     # Delete history dependencies not found in the DB

+ 29 - 19
desktop/core/src/desktop/celery.py

@@ -15,24 +15,25 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import absolute_import, unicode_literals
-from __future__ import print_function
+from __future__ import absolute_import, print_function, unicode_literals
 
-import imp
 import os
+import imp
 
 from celery import Celery
 from celery.schedules import crontab
 
-from desktop.conf import TASK_SERVER
-from desktop.settings import TIME_ZONE, INSTALLED_APPS
-if hasattr(TASK_SERVER, 'get') and TASK_SERVER.ENABLED.get():
-  from desktop.settings import CELERY_RESULT_BACKEND, CELERY_BROKER_URL
+from desktop.conf import TASK_SERVER_V2
+from desktop.settings import INSTALLED_APPS, TIME_ZONE
+
+if hasattr(TASK_SERVER_V2, 'get') and TASK_SERVER_V2.ENABLED.get():
+  from desktop.settings import CELERY_BROKER_URL, CELERY_RESULT_BACKEND
 from django.utils import timezone
 
 # Set the default Django settings module for the 'celery' program.
 os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'desktop.settings')
 
+
 class HueCelery(Celery):
   def gen_task_name(self, name, module):
     if module.endswith('.tasks'):
@@ -41,18 +42,27 @@ class HueCelery(Celery):
 
   # Method to configure the beat_schedule
   def setup_beat_schedule(self):
-    now = timezone.now()
-
-    self.conf.beat_schedule = {
-      'check_disk_usage_and_clean_task': {
-      'task': 'filebrowser.check_disk_usage_and_clean_task',
-      'schedule': 1000.0,  # Run every 1000 seconds
-      'args': (),
-      'kwargs': {'cleanup_threshold': 90},  # Provide task arguments if needed
-      },
-    }
-
-if hasattr(TASK_SERVER, 'get') and TASK_SERVER.ENABLED.get():
+
+    beat_schedule = {}
+    if TASK_SERVER_V2.CHECK_DISK_USAGE_AND_CLEAN_TASK_ENABLED.get():
+      beat_schedule['check_disk_usage_and_clean_task'] = {
+        'task': 'filebrowser.check_disk_usage_and_clean_task',
+        'schedule': TASK_SERVER_V2.CHECK_DISK_USAGE_AND_CLEAN_TASK_PERIODIC_INTERVAL.get(),
+        'args': (),
+        'kwargs': {'cleanup_threshold': TASK_SERVER_V2.DISK_USAGE_CLEANUP_THRESHOLD.get()},
+      }
+
+    if TASK_SERVER_V2.CLEANUP_STALE_UPLOADS_IN_REDIS_ENABLED.get():
+      beat_schedule['cleanup_stale_uploads'] = {
+        'task': 'filebrowser.cleanup_stale_uploads',
+        'schedule': TASK_SERVER_V2.CLEANUP_STALE_UPLOADS_IN_REDIS_PERIODIC_INTERVAL.get(),
+        'args': ()
+      }
+
+    self.conf.beat_schedule = beat_schedule
+
+
+if hasattr(TASK_SERVER_V2, 'get') and TASK_SERVER_V2.ENABLED.get():
   app = HueCelery('desktop', backend=CELERY_RESULT_BACKEND, broker=CELERY_BROKER_URL)
   app.conf.broker_transport_options = {'visibility_timeout': 3600}  # 1 hour.
   app.conf.result_key_prefix = 'desktop_'

+ 113 - 0
desktop/core/src/desktop/conf.py

@@ -2061,6 +2061,119 @@ TASK_SERVER = ConfigSection(
     ),
 ))
 
+TASK_SERVER_V2 = ConfigSection(
+  key="task_server_v2",
+  help=_("Task Server V2 configuration."),
+  members=dict(
+    ENABLED=Config(
+      key='enabled',
+      default=False,
+      type=coerce_bool,
+      help=_('If resource intensive or blocking can be delegated to an already running task server.')
+    ),
+    BROKER_URL=Config(
+      key='broker_url',
+      default='amqp://guest:guest@localhost//',
+      help=_('How the task server and tasks communicate.')
+    ),
+    CELERY_RESULT_BACKEND=Config(
+      key='celery_result_backend',
+      dynamic_default=task_server_default_result_directory,
+      help=_('Where to store task results. Defaults to local file system path. Celery comes with a several other backends.')
+    ),
+    RESULT_CELERYD_OPTS=Config(
+      key='celeryd_opts',
+      default='--time-limit=300',
+      help=_('Default options provided to the task server at startup.')
+    ),
+    BEAT_ENABLED=Config(
+      key='beat_enabled',
+      default=False,
+      type=coerce_bool,
+      help=_('Switch on the integration with the Task Scheduler.')
+    ),
+    BEAT_SCHEDULES_FILE=Config(
+      key='beat_schedules_file',
+      default='',
+      type=str,
+      help=_('Path to a file containing a list of beat schedules.')
+    ),
+    FETCH_RESULT_LIMIT=Config(
+      key='fetch_result_limit',
+      default=2000,
+      type=coerce_positive_integer,
+      help=_('Number of query results rows to fetch into the result storage.')
+    ),
+    RESULT_CACHE=Config(
+      key='result_cache',
+      type=str,
+      help=_('Django file cache class to use to temporarily store query results'),
+      default='{"BACKEND": "django_redis.cache.RedisCache", "LOCATION": "redis://localhost:6379/0", '
+      '"OPTIONS": {"CLIENT_CLASS": "django_redis.client.DefaultClient"},"KEY_PREFIX": "queries"}'
+    ),
+    RESULT_STORAGE=Config(
+      key='result_storage',
+      type=str,
+      help=_('Django file storage class to use to persist query results'),
+      default='{"backend": "django.core.files.storage.FileSystemStorage", "properties": {"location": "./logs"}}'
+    ),
+    EXECUTION_STORAGE=Config(
+      key='execution_storage',
+      type=str,
+      help=_('Django cache to use to store temporarily used data during query execution. '
+      'This is in addition to result_file_storage and result_backend.'),
+      default='{"BACKEND": "django.core.cache.backends.locmem.LocMemCache", "LOCATION": "celery-hue"}'
+    ),
+    CELERY_LOG_LEVEL=Config(
+      key='celery_log_level',
+      type=str,
+      help=_('Log level of celery workers.'),
+      default='INFO'
+    ),
+    CHECK_DISK_USAGE_AND_CLEAN_TASK_ENABLED=Config(
+      key='check_disk_usage_and_clean_task_enabled',
+      type=coerce_bool,
+      help=_('enable this peroidic cleaner which checks disk usage and makes space for file uploads'),
+      default=False
+    ),
+    CHECK_DISK_USAGE_AND_CLEAN_TASK_PERIODIC_INTERVAL=Config(
+      key='check_disk_usage_and_clean_task_periodic_interval',
+      type=coerce_positive_integer,
+      help=_('set the time interval in seconds to run this peroidic cleaner which checks disk usage and makes space for file uploads'),
+      default=1000
+    ),
+    DISK_USAGE_CLEANUP_THRESHOLD=Config(
+      key='disk_usage_cleanup_threshold',
+      type=coerce_positive_integer,
+      help=_('Clean up files in /tmp folder if the disk usage is beyond the threshold'),
+      default=90
+    ),
+    DISK_USAGE_AND_CLEAN_TASK_TIME_DELTA=Config(
+      key='disk_usage_and_clean_task_time_delta',
+      type=coerce_positive_integer,
+      help=_('Clean up files older than timedelta. Unit of timedelta is minutes'),
+      default=60
+    ),
+    CLEANUP_STALE_UPLOADS_IN_REDIS_ENABLED=Config(
+      key='cleanup_stale_uploads_in_redis_enabled',
+      type=coerce_bool,
+      help=_('enable this peroidic cleaner which cleans up failed upload tasks stored in redis'),
+      default=False
+    ),
+    CLEANUP_STALE_UPLOADS_IN_REDIS_PERIODIC_INTERVAL=Config(
+      key='cleanup_stale_uploads_task_periodic_interval',
+      type=coerce_positive_integer,
+      help=_('set the time interval in seconds to run this peroidic cleaner which cleans up failed upload tasks stored in redis'),
+      default=900
+    ),
+    CLEANUP_STALE_UPLOADS_TASK_TIME_DELTA=Config(
+      key='cleanup_stale_uploads_task_time_delta',
+      type=coerce_positive_integer,
+      help=_('Redis keys of format Upload__* older than timedelta will be cleaned up. Unit of timedelta is minutes'),
+      default=60
+    ),
+))
+
 
 def has_channels():
   return sys.version_info[0] > 2 and WEBSOCKETS.ENABLED.get()

+ 4 - 1
desktop/core/src/desktop/js/reactComponents/TaskBrowser/TaskBrowser.tsx

@@ -262,6 +262,9 @@ export const TaskBrowserTable: React.FC<TaskBrowserTableProps> = ({
           {record.result?.task_name === 'tmp_cleanup' && (
             <span>{`{cleanup threshold: ${record.result?.parameters}}`}</span>
           )}
+          {record.result?.task_name === 'cleanup_stale_uploads' && (
+            <span>{`{cleanup timedelta: ${record.result?.parameters}}`}</span>
+          )}
         </div>
       )
     },
@@ -280,7 +283,7 @@ export const TaskBrowserTable: React.FC<TaskBrowserTableProps> = ({
     {
       title: t('Duration'),
       key: 'duration',
-      render: (_, record) => calculateDuration(record.result?.task_start, record.date_done)
+      render: (_, record) => calculateDuration(record.result?.task_start, record.result?.task_end)
     }
   ];
 

+ 7 - 9
desktop/core/src/desktop/lib/fs/ozone/upload.py

@@ -15,8 +15,8 @@
 # limitations under the License.
 
 import io
-import logging
 import sys
+import logging
 import unicodedata
 
 from django.core.files.uploadedfile import SimpleUploadedFile
@@ -31,12 +31,13 @@ if sys.version_info[0] > 2:
 else:
   from django.utils.translation import ugettext as _
 
+from desktop.conf import TASK_SERVER_V2
 from desktop.lib.exceptions_renderable import PopupException
-from filebrowser.utils import generate_chunks, calculate_total_size
-from desktop.conf import TASK_SERVER
+from filebrowser.utils import calculate_total_size, generate_chunks
 
 LOG = logging.getLogger()
 
+
 class OFSFineUploaderChunkedUpload(object):
   def __init__(self, request, *args, **kwargs):
     self.qquuid = kwargs.get('qquuid')
@@ -44,9 +45,9 @@ class OFSFineUploaderChunkedUpload(object):
     self.totalfilesize = kwargs.get('qqtotalfilesize')
     self.file_name = kwargs.get('qqfilename')
     if self.file_name:
-      self.file_name = unicodedata.normalize('NFC', self.file_name) # Normalize unicode
+      self.file_name = unicodedata.normalize('NFC', self.file_name)  # Normalize unicode
     self.chunk_size = UPLOAD_CHUNK_SIZE.get()
-    if kwargs.get('chunk_size', None) != None:
+    if kwargs.get('chunk_size', None):
       self.chunk_size = kwargs.get('chunk_size')
     self.destination = kwargs.get('dest', None)  # GET param avoids infinite looping
     self.target_path = None
@@ -76,7 +77,7 @@ class OFSFineUploaderChunkedUpload(object):
   def upload_chunks(self):
     LOG.debug("OFSFineUploaderChunkedUpload: upload_chunks")
 
-    if TASK_SERVER.ENABLED.get():
+    if TASK_SERVER_V2.ENABLED.get():
       if self._is_ofs_upload():
         self._fs = self._get_ofs(self._request)
 
@@ -198,7 +199,6 @@ class OFSFileUploadHandler(FileUploadHandler):
 
     LOG.debug("Chunk size = %d" % UPLOAD_CHUNK_SIZE.get())
 
-
   def new_file(self, field_name, file_name, *args, **kwargs):
     if self._is_ofs_upload():
       super(OFSFileUploadHandler, self).new_file(field_name, file_name, *args, **kwargs)
@@ -217,7 +217,6 @@ class OFSFileUploadHandler(FileUploadHandler):
         self.request.META['upload_failed'] = e
         raise StopUpload()
 
-
   def receive_data_chunk(self, raw_data, start):
     if self._is_ofs_upload():
       LOG.debug("OFSfileUploadHandler receive_data_chunk")
@@ -231,7 +230,6 @@ class OFSFileUploadHandler(FileUploadHandler):
     else:
       return raw_data
 
-
   def file_complete(self, file_size):
     if self._is_ofs_upload():
       # Finish the upload

+ 4 - 5
desktop/core/src/desktop/lib/scheduler/api.py

@@ -15,25 +15,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import logging
 import sys
+import logging
 
 from django.forms.formsets import formset_factory
 from django.urls import reverse
 
 from desktop.auth.backend import is_admin
-from desktop.conf import TASK_SERVER
-from desktop.models import Document2
 from desktop.lib.django_util import JsonResponse, render
 from desktop.lib.i18n import force_unicode
 from desktop.lib.scheduler.lib.api import get_api
+from desktop.models import Document2
 
 LOG = logging.getLogger()
 
 try:
   from oozie.decorators import check_document_access_permission
   from oozie.forms import ParameterForm
-  from oozie.views.editor2 import edit_coordinator, new_coordinator, Coordinator
+  from oozie.views.editor2 import Coordinator, edit_coordinator, new_coordinator
 except Exception as e:
   LOG.warning('Oozie application is not enabled: %s' % e)
 
@@ -51,7 +50,7 @@ def get_schedule(request):
 
 
 # To move to lib in case oozie is blacklisted
-#@check_document_access_permission()
+# @check_document_access_permission()
 def submit_schedule(request, doc_id):
   interface = request.GET.get('interface', request.POST.get('interface', 'hive'))
 

+ 14 - 11
desktop/core/src/desktop/management/commands/runcelery.py

@@ -15,17 +15,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import logging
 import os
 import sys
+import logging
+
+from celery.bin.celery import CeleryCommand, main as celery_main
+from django.core.management.base import BaseCommand
+from django.utils import autoreload
 
 from desktop import conf
+from desktop.conf import TASK_SERVER_V2
 from desktop.lib.daemon_utils import drop_privileges_if_necessary
 from desktop.log import DEFAULT_LOG_DIR
 
-from django.core.management.base import BaseCommand
-from django.utils import autoreload
-
 if sys.version_info[0] > 2:
   from django.utils.translation import gettext as _
 else:
@@ -35,17 +37,16 @@ SERVER_HELP = r"""
   Run celery worker.
 """
 
-from celery.bin.celery import CeleryCommand
-from celery.bin.celery import main as celery_main
-
 LOG = logging.getLogger()
 CELERY_OPTIONS = {
   'server_user': conf.SERVER_USER.get(),
   'server_group': conf.SERVER_GROUP.get(),
 }
 
+
 class Command(BaseCommand):
   help = SERVER_HELP
+
   def add_arguments(self, parser):
     parser.add_argument('worker')
     parser.add_argument(
@@ -61,7 +62,7 @@ class Command(BaseCommand):
     parser.add_argument(
         '--loglevel',
         type=str,
-        default='DEBUG'
+        default='INFO'
     )
     parser.add_argument('--beat')
     parser.add_argument(
@@ -78,17 +79,19 @@ class Command(BaseCommand):
   def usage(self, subcommand):
     return SERVER_HELP
 
+
 def runcelery(*args, **options):
   # Native does not load Hue's config
   log_dir = os.getenv("DESKTOP_LOG_DIR", DEFAULT_LOG_DIR)
-  log_file = "%s/rungunicornserver.log" % (log_dir)
-  concurrency = int(conf.GUNICORN_NUMBER_OF_WORKERS.get()/4) or options['concurrency']
+  log_file = "%s/celery.log" % (log_dir)
+  concurrency = max(int(conf.GUNICORN_NUMBER_OF_WORKERS.get() / 4), 1) or options['concurrency']
   schedule_file = options['schedule_file']
+  celery_log_level = TASK_SERVER_V2.CELERY_LOG_LEVEL.get()
   opts = [
     'celery',
     '--app=' + options['app'],
     'worker',
-    '--loglevel=' + options['loglevel'],
+    '--loglevel=' + str(celery_log_level),
     '--concurrency=' + str(concurrency),
     '--beat',
     '-s', schedule_file,

+ 67 - 68
desktop/core/src/desktop/models.py

@@ -16,59 +16,70 @@
 # limitations under the License.
 
 from __future__ import absolute_import
-from future import standard_library
-standard_library.install_aliases()
-from builtins import next
-from builtins import object
-import calendar
-import json
-import logging
+
 import os
 import sys
+import json
 import uuid
-
+import logging
+import calendar
+from builtins import next, object
 from collections import OrderedDict
 from itertools import chain
 
-from django.db import connection, models, transaction
-from django.db.models import Q
-from django.db.models.query import QuerySet
 from django.contrib.auth.validators import UnicodeUsernameValidator
-from django.contrib.contenttypes.fields import GenericRelation, GenericForeignKey
+from django.contrib.contenttypes.fields import GenericForeignKey, GenericRelation
 from django.contrib.contenttypes.models import ContentType
 from django.contrib.staticfiles.storage import staticfiles_storage
-from django.urls import reverse, NoReverseMatch
-
-from dashboard.conf import get_engines, HAS_REPORT_ENABLED, IS_ENABLED as DASHBOARD_ENABLED
-from hadoop.core_site import get_raz_api_url, get_raz_s3_default_bucket
-from kafka.conf import has_kafka
-from indexer.conf import ENABLE_DIRECT_UPLOAD
-from metadata.conf import get_optimizer_mode
-from notebook.conf import DEFAULT_LIMIT, SHOW_NOTEBOOKS, get_ordered_interpreters, DEFAULT_INTERPRETER
-from useradmin.models import User, Group, get_organization
-from useradmin.organization import _fitered_queryset
+from django.db import connection, models, transaction
+from django.db.models import Q
+from django.db.models.query import QuerySet
+from django.urls import NoReverseMatch, reverse
 
+from dashboard.conf import HAS_REPORT_ENABLED, IS_ENABLED as DASHBOARD_ENABLED, get_engines
 from desktop import appmanager
 from desktop.auth.backend import is_admin
-from desktop.conf import APP_BLACKLIST, COLLECT_USAGE, DISABLE_SOURCE_AUTOCOMPLETE, ENABLE_CONNECTORS, \
-    ENABLE_ORGANIZATIONS, ENABLE_PROMETHEUS, ENABLE_SHARING, ENABLE_UNIFIED_ANALYTICS, get_clusters, has_connectors, \
-    HUE_HOST_NAME, HUE_IMAGE_VERSION, IS_MULTICLUSTER_ONLY, RAZ, TASK_SERVER
+from desktop.conf import (
+  APP_BLACKLIST,
+  COLLECT_USAGE,
+  DISABLE_SOURCE_AUTOCOMPLETE,
+  ENABLE_CONNECTORS,
+  ENABLE_ORGANIZATIONS,
+  ENABLE_PROMETHEUS,
+  ENABLE_SHARING,
+  ENABLE_UNIFIED_ANALYTICS,
+  HUE_HOST_NAME,
+  HUE_IMAGE_VERSION,
+  IS_MULTICLUSTER_ONLY,
+  RAZ,
+  TASK_SERVER_V2,
+  get_clusters,
+  has_connectors,
+)
 from desktop.lib import fsmanager
 from desktop.lib.connectors.api import _get_installed_connectors
 from desktop.lib.connectors.models import Connector
-from desktop.lib.i18n import force_unicode
 from desktop.lib.exceptions_renderable import PopupException
-from desktop.lib.paths import get_run_root, SAFE_CHARACTERS_URI_COMPONENTS
+from desktop.lib.i18n import force_unicode
+from desktop.lib.paths import SAFE_CHARACTERS_URI_COMPONENTS, get_run_root
 from desktop.redaction import global_redaction_engine
 from desktop.settings import DOCUMENT2_SEARCH_MAX_LENGTH, HUE_DESKTOP_VERSION
-
 from filebrowser.conf import REMOTE_STORAGE_HOME
+from hadoop.core_site import get_raz_api_url, get_raz_s3_default_bucket
+from indexer.conf import ENABLE_DIRECT_UPLOAD
+from kafka.conf import has_kafka
+from metadata.conf import get_optimizer_mode
+from notebook.conf import DEFAULT_INTERPRETER, DEFAULT_LIMIT, SHOW_NOTEBOOKS, get_ordered_interpreters
+from useradmin.models import Group, User, get_organization
+from useradmin.organization import _fitered_queryset
 
 if sys.version_info[0] > 2:
   from urllib.parse import quote as urllib_quote
+
   from django.utils.translation import gettext as _, gettext_lazy as _t
 else:
   from urllib import quote as urllib_quote
+
   from django.utils.translation import ugettext as _, ugettext_lazy as _t
 
 
@@ -86,6 +97,7 @@ IMAGE_VERSION = None
 def uuid_default():
   return str(uuid.uuid4())
 
+
 def hue_version():
   global HUE_VERSION
 
@@ -94,6 +106,7 @@ def hue_version():
 
   return HUE_VERSION
 
+
 def hue_image_version():
   global IMAGE_VERSION
 
@@ -109,12 +122,15 @@ def hue_image_version():
 
   return IMAGE_VERSION
 
+
 def name_of_hue_host():
   return HUE_HOST_NAME.get()
 
+
 def _version_from_properties(f):
   return dict(line.strip().split('=') for line in f.readlines() if len(line.strip().split('=')) == 2).get('cloudera.cdh.release')
 
+
 def get_sample_user_install(user):
   if ENABLE_ORGANIZATIONS.get():
     organization = get_organization(email=user.email if user else None)
@@ -207,7 +223,6 @@ class DefaultConfiguration(models.Model):
   class Meta(object):
     ordering = ["app", "-is_default", "user"]
 
-
   @property
   def properties_list(self):
     """
@@ -327,11 +342,11 @@ class DocumentTag(models.Model):
   owner = models.ForeignKey(User, on_delete=models.CASCADE, db_index=True)
   tag = models.SlugField()
 
-  DEFAULT = 'default' # Always there
-  TRASH = 'trash' # There when the document is trashed
-  HISTORY = 'history' # There when the document is a submission history
-  EXAMPLE = 'example' # Hue examples
-  IMPORTED2 = 'imported2' # Was imported to document2
+  DEFAULT = 'default'  # Always there
+  TRASH = 'trash'  # There when the document is trashed
+  HISTORY = 'history'  # There when the document is a submission history
+  EXAMPLE = 'example'  # Hue examples
+  IMPORTED2 = 'imported2'  # Was imported to document2
 
   RESERVED = (DEFAULT, TRASH, HISTORY, EXAMPLE, IMPORTED2)
 
@@ -340,7 +355,6 @@ class DocumentTag(models.Model):
   class Meta(object):
     unique_together = ('owner', 'tag')
 
-
   def __unicode__(self):
     return force_unicode('%s') % (self.tag,)
 
@@ -456,7 +470,7 @@ class DocumentManager(models.Manager):
     table_names = connection.introspection.table_names()
 
     try:
-      from oozie.models import Workflow, Coordinator, Bundle
+      from oozie.models import Bundle, Coordinator, Workflow
 
       if \
           Workflow._meta.db_table in table_names or \
@@ -512,7 +526,7 @@ class DocumentManager(models.Manager):
           for dashboard in Collection.objects.all():
             if 'collection' in dashboard.properties_dict:
               col_dict = dashboard.properties_dict['collection']
-              if not 'uuid' in col_dict:
+              if 'uuid' not in col_dict:
                 _uuid = str(uuid.uuid4())
                 col_dict['uuid'] = _uuid
                 dashboard.update_properties({'collection': col_dict})
@@ -549,7 +563,6 @@ class DocumentManager(models.Manager):
     except Exception as e:
       LOG.exception('error syncing Document2')
 
-
     if not doc2_only and Document._meta.db_table in table_names:
       # Make sure doc have at least a tag
       try:
@@ -837,7 +850,7 @@ class Document(models.Model):
       'owner': self.owner.username,
       'name': self.name,
       'description': self.description,
-      'uuid': None, # no uuid == v1
+      'uuid': None,  # no uuid == v1
       'id': self.id,
       'doc1_id': self.id,
       'object_id': self.object_id,
@@ -856,9 +869,8 @@ class DocumentPermissionManager(models.Manager):
       perms_string = ' and '.join(', '.join(perms).rsplit(', ', 1))
       raise PopupException(_('Only %s permissions are supported, not %s.') % (perms_string, name))
 
-
   def share_to_default(self, document, name='read'):
-    from useradmin.models import get_default_user_group # Remove build dependency
+    from useradmin.models import get_default_user_group  # Remove build dependency
 
     self._check_perm(name)
 
@@ -920,7 +932,7 @@ class DocumentPermission(models.Model):
 
   users = models.ManyToManyField(User, db_index=True, db_table='documentpermission_users')
   groups = models.ManyToManyField(Group, db_index=True, db_table='documentpermission_groups')
-  perms = models.CharField(default=READ_PERM, max_length=10, choices=( # one perm
+  perms = models.CharField(default=READ_PERM, max_length=10, choices=(  # one perm
     (READ_PERM, 'read'),
     (WRITE_PERM, 'write'),
   ))
@@ -980,7 +992,6 @@ class Document2QueryMixin(object):
 
     return docs.defer('description', 'data', 'extra', 'search').order_by('-last_modified')
 
-
   def search_documents(self, types=None, search_text=None, order_by=None):
     """
     Search for documents based on type filters, search_text or order_by and return a queryset of document objects
@@ -1215,7 +1226,7 @@ class Document2(models.Model):
 
   parent_directory = models.ForeignKey('self', blank=True, null=True, related_name='children', on_delete=models.CASCADE)
 
-  doc = GenericRelation(Document, related_query_name='doc_doc') # Compatibility with Hue 3
+  doc = GenericRelation(Document, related_query_name='doc_doc')  # Compatibility with Hue 3
 
   objects = Document2Manager()
 
@@ -1255,7 +1266,7 @@ class Document2(models.Model):
 
   @property
   def is_home_directory(self):
-    return self.is_directory and self.parent_directory == None and self.name == self.HOME_DIR
+    return self.is_directory and self.parent_directory is None and self.name == self.HOME_DIR
 
   @property
   def is_trash_directory(self):
@@ -1332,7 +1343,7 @@ class Document2(models.Model):
     return self.dependencies.filter(is_history=True).order_by('-last_modified')
 
   def add_to_history(self, user, data_dict):
-    doc_id = self.id # Need to copy as the clone messes it
+    doc_id = self.id  # Need to copy as the clone messes it
 
     history_doc = self.copy(name=self.name, owner=user)
     history_doc.update_data({'history': data_dict})
@@ -1374,14 +1385,12 @@ class Document2(models.Model):
       raise FilesystemException(_('Cannot save document %s under parent directory %s due to circular dependency') %
                                 (self.name, self.parent_directory.uuid))
 
-
   def inherit_permissions(self):
     if self.parent_directory is not None:
       parent_perms = Document2Permission.objects.filter(doc=self.parent_directory)
       for perm in parent_perms:
         self.share(self.owner, name=perm.perms, users=perm.users.all(), groups=perm.groups.all())
 
-
   def move(self, directory, user):
     if not directory.is_directory:
       raise FilesystemException(_('Target with UUID %s is not a directory') % directory.uuid)
@@ -1681,7 +1690,6 @@ class Directory(Document2):
 
     return documents.defer('description', 'data', 'extra', 'search').distinct().order_by('-last_modified')
 
-
   def save(self, *args, **kwargs):
     self.type = 'directory'
     super(Directory, self).save(*args, **kwargs)
@@ -1704,10 +1712,10 @@ class Document2Permission(models.Model):
   users = models.ManyToManyField(User, db_index=True, db_table='documentpermission2_users')
   groups = models.ManyToManyField(Group, db_index=True, db_table='documentpermission2_groups')
 
-  perms = models.CharField(default=READ_PERM, max_length=10, db_index=True, choices=( # One perm
+  perms = models.CharField(default=READ_PERM, max_length=10, db_index=True, choices=(  # One perm
     (READ_PERM, 'read'),
     (WRITE_PERM, 'write'),
-    (COMMENT_PERM, 'comment'), # Unused
+    (COMMENT_PERM, 'comment'),  # Unused
     (LINK_READ_PERM, 'link read'),
     (LINK_WRITE_PERM, 'link write'),
   ))
@@ -1740,6 +1748,7 @@ class Document2Permission(models.Model):
 def get_cluster_config(user):
   return Cluster(user).get_app_config().get_config()
 
+
 def get_remote_home_storage(user=None):
   remote_home_storage = REMOTE_STORAGE_HOME.get() if hasattr(REMOTE_STORAGE_HOME, 'get') and REMOTE_STORAGE_HOME.get() else None
 
@@ -1773,7 +1782,6 @@ class ClusterConfig(object):
     self.apps = appmanager.get_apps_dict(self.user) if apps is None else apps
     self.cluster_type = cluster_type
 
-
   def get_config(self):
     app_config = self.get_apps()
     editors = app_config.get('editor')
@@ -1800,7 +1808,7 @@ class ClusterConfig(object):
       ],
       'default_sql_interpreter': default_sql_interpreter,
       'cluster_type': self.cluster_type,
-      'has_computes': self.cluster_type in ('cdw', 'altus', 'snowball'), # or any grouped engine connectors
+      'has_computes': self.cluster_type in ('cdw', 'altus', 'snowball'),  # or any grouped engine connectors
       'hue_config': {
         'enable_sharing': ENABLE_SHARING.get(),
         'collect_usage': COLLECT_USAGE.get()
@@ -1810,7 +1818,6 @@ class ClusterConfig(object):
       'hue_version': version_of_hue
     }
 
-
   def get_apps(self):
     apps = OrderedDict([
       app for app in [
@@ -1826,7 +1833,6 @@ class ClusterConfig(object):
 
     return apps
 
-
   def get_main_quick_action(self, apps):
     if not apps:
       raise PopupException(_('No permission to any app.'))
@@ -1863,7 +1869,6 @@ class ClusterConfig(object):
     else:
       return default_app
 
-
   def _get_home(self):
     return {
       'name': 'home',
@@ -1873,7 +1878,6 @@ class ClusterConfig(object):
       'page': '/home'
     }
 
-
   def displayName(self, dialect, name):
     if ENABLE_UNIFIED_ANALYTICS.get() and dialect == 'hive':
       return 'Unified Analytics'
@@ -1882,7 +1886,6 @@ class ClusterConfig(object):
     else:
       return name
 
-
   def _get_editor(self):
     interpreters = []
 
@@ -2010,7 +2013,6 @@ class ClusterConfig(object):
     elif 'filebrowser' in self.apps and fsmanager.is_enabled_and_has_access('hdfs', self.user):
       hdfs_connectors.append(_('Files'))
 
-
     remote_home_storage = get_remote_home_storage(self.user)
 
     for hdfs_connector in hdfs_connectors:
@@ -2173,7 +2175,6 @@ class ClusterConfig(object):
     else:
       return None
 
-
   def _get_scheduler(self):
     interpreters = []
 
@@ -2199,7 +2200,7 @@ class ClusterConfig(object):
         }
       ])
 
-    if TASK_SERVER.BEAT_ENABLED.get():
+    if TASK_SERVER_V2.BEAT_ENABLED.get():
       interpreters.append({
           'type': 'celery-beat',
           'displayName': _('Scheduled Tasks'),
@@ -2221,7 +2222,6 @@ class ClusterConfig(object):
     else:
       return None
 
-
   def _get_sdk_apps(self):
     current_app, other_apps, apps_list = _get_apps(self.user)
 
@@ -2258,9 +2258,9 @@ class Cluster(object):
     self.clusters = get_clusters(user)
 
     if IS_MULTICLUSTER_ONLY.get():
-      self.data = self.clusters['Altus'] # Backward compatibility
+      self.data = self.clusters['Altus']  # Backward compatibility
     else:
-      self.data = list(self.clusters.values())[0] # Next: CLUSTER_ID.get() or user persisted
+      self.data = list(self.clusters.values())[0]  # Next: CLUSTER_ID.get() or user persisted
 
   def get_type(self):
     return self.data['type']
@@ -2309,7 +2309,6 @@ def get_user_preferences(user, key=None):
     return dict((x.key, x.value) for x in UserPreferences.objects.filter(user=user))
 
 
-
 def set_user_preferences(user, key, value):
   try:
     x = UserPreferences.objects.get(user=user, key=key)
@@ -2332,17 +2331,17 @@ def get_data_link(meta):
     elif 'fam' in meta:
       link += '[%(fam)s]' % meta
   elif meta['type'] == 'hdfs':
-    link = '/filebrowser/view=%(path)s' % meta # Could add a byte #
+    link = '/filebrowser/view=%(path)s' % meta  # Could add a byte #
   elif meta['type'] == 'link':
     link = meta['link']
   elif meta['type'] == 'hive':
-    link = '/metastore/table/%(database)s/%(table)s' % meta # Could also add col=val
+    link = '/metastore/table/%(database)s/%(table)s' % meta  # Could also add col=val
 
   return link
 
 
 def _get_gist_document(uuid):
-  return Document2.objects.get(uuid=uuid, type='gist') # Workaround until there is a share to all permission
+  return Document2.objects.get(uuid=uuid, type='gist')  # Workaround until there is a share to all permission
 
 
 def __paginate(page, limit, queryset):

+ 47 - 21
desktop/core/src/desktop/settings.py

@@ -20,27 +20,29 @@
 # Local customizations are done by symlinking a file
 # as local_settings.py.
 
-from builtins import map, zip
-import datetime
 import gc
-import json
-import logging
 import os
-import pkg_resources
 import sys
+import json
 import uuid
+import logging
+import datetime
+from builtins import map, zip
+from urllib.parse import urlparse
+
+import redis
+import psutil
+import pkg_resources
 
 import desktop.redaction
-from desktop.lib import conf
+from aws.conf import is_enabled as is_s3_enabled
+from azure.conf import is_abfs_enabled
 from desktop import appmanager
-
+from desktop.conf import TASK_SERVER_V2, is_chunked_fileuploader_enabled, is_gs_enabled, is_ofs_enabled
+from desktop.lib import conf
 from desktop.lib.paths import get_desktop_root, get_run_root
 from desktop.lib.python_util import force_dict_to_strings
 
-from aws.conf import is_enabled as is_s3_enabled
-from azure.conf import is_abfs_enabled
-from desktop.conf import is_ofs_enabled, is_gs_enabled, is_chunked_fileuploader_enabled
-
 if sys.version_info[0] > 2:
   from django.utils.translation import gettext_lazy as _
 else:
@@ -452,10 +454,10 @@ CACHES[CACHES_HIVE_DISCOVERY_KEY] = {
 
 CACHES_CELERY_KEY = 'celery'
 CACHES_CELERY_QUERY_RESULT_KEY = 'celery_query_results'
-if desktop.conf.TASK_SERVER.ENABLED.get():
-  CACHES[CACHES_CELERY_KEY] = json.loads(desktop.conf.TASK_SERVER.EXECUTION_STORAGE.get())
-  if desktop.conf.TASK_SERVER.RESULT_CACHE.get():
-    CACHES[CACHES_CELERY_QUERY_RESULT_KEY] = json.loads(desktop.conf.TASK_SERVER.RESULT_CACHE.get())
+if desktop.conf.TASK_SERVER_V2.ENABLED.get():
+  CACHES[CACHES_CELERY_KEY] = json.loads(desktop.conf.TASK_SERVER_V2.EXECUTION_STORAGE.get())
+  if desktop.conf.TASK_SERVER_V2.RESULT_CACHE.get():
+    CACHES[CACHES_CELERY_QUERY_RESULT_KEY] = json.loads(desktop.conf.TASK_SERVER_V2.RESULT_CACHE.get())
 
 # Configure sessions
 SESSION_COOKIE_NAME = desktop.conf.SESSION.COOKIE_NAME.get()
@@ -730,16 +732,16 @@ DOCUMENT2_MAX_ENTRIES = 100000
 # Celery settings
 ################################################################
 
-if desktop.conf.TASK_SERVER.ENABLED.get() or desktop.conf.TASK_SERVER.BEAT_ENABLED.get():
-  CELERY_BROKER_URL = desktop.conf.TASK_SERVER.BROKER_URL.get()
+if desktop.conf.TASK_SERVER_V2.ENABLED.get() or desktop.conf.TASK_SERVER_V2.BEAT_ENABLED.get():
+  CELERY_BROKER_URL = desktop.conf.TASK_SERVER_V2.BROKER_URL.get()
 
   CELERY_ACCEPT_CONTENT = ['json']
-  CELERY_RESULT_BACKEND = desktop.conf.TASK_SERVER.CELERY_RESULT_BACKEND.get()
+  CELERY_RESULT_BACKEND = desktop.conf.TASK_SERVER_V2.CELERY_RESULT_BACKEND.get()
   CELERY_TASK_SERIALIZER = 'json'
   CELERY_ENABLE_UTC = True
-  CELERY_TIMEZONE = "America/Los_Angeles"
+  CELERY_TIMEZONE = TIME_ZONE
 
-  CELERYD_OPTS = desktop.conf.TASK_SERVER.RESULT_CELERYD_OPTS.get()
+  CELERYD_OPTS = desktop.conf.TASK_SERVER_V2.RESULT_CELERYD_OPTS.get()
   CELERY_TASK_DEFAULT_QUEUE = 'default'
 
   CELERY_TASK_QUEUES = {
@@ -757,6 +759,30 @@ if desktop.conf.TASK_SERVER.ENABLED.get() or desktop.conf.TASK_SERVER.BEAT_ENABL
     },
   }
 
+  if desktop.conf.TASK_SERVER_V2.ENABLED.get():
+    def initialize_free_disk_space_in_redis():
+      redis_client = parse_broker_url(desktop.conf.TASK_SERVER_V2.BROKER_URL.get())
+      free_space = psutil.disk_usage('/tmp').free
+      available_space = redis_client.get('upload_available_space')
+      if available_space is None:
+        available_space = free_space
+      else:
+        available_space = int(available_space)
+      upload_keys_exist = any(redis_client.scan_iter('upload__*'))
+      redis_client.delete('upload_available_space')
+      if not upload_keys_exist:
+        redis_client.setnx('upload_available_space', free_space)
+      else:
+        redis_client.setnx('upload_available_space', min(free_space, available_space))
+
+    def parse_broker_url(broker_url):
+      parsed_url = urlparse(broker_url)
+      host = parsed_url.hostname
+      port = parsed_url.port
+      db = int(parsed_url.path.lstrip('/'))
+      return redis.Redis(host=host, port=port, db=db)
+
+    initialize_free_disk_space_in_redis()
 
 # %n will be replaced with the first part of the nodename.
 # CELERYD_LOG_FILE="/var/log/celery/%n%I.log"
@@ -765,7 +791,7 @@ if desktop.conf.TASK_SERVER.ENABLED.get() or desktop.conf.TASK_SERVER.BEAT_ENABL
 # CELERYD_USER = desktop.conf.SERVER_USER.get()
 # CELERYD_GROUP = desktop.conf.SERVER_GROUP.get()
 
-  if desktop.conf.TASK_SERVER.BEAT_ENABLED.get():
+  if desktop.conf.TASK_SERVER_V2.BEAT_ENABLED.get():
     INSTALLED_APPS.append('django_celery_beat')
     INSTALLED_APPS.append('timezone_field')
     USE_TZ = True

+ 3 - 1
desktop/core/src/desktop/templates/about_layout.mako

@@ -22,7 +22,7 @@ else:
   from django.utils.translation import ugettext as _
 
 from desktop.auth.backend import is_admin, is_hue_admin
-from desktop.conf import METRICS, has_connectors, ANALYTICS
+from desktop.conf import METRICS, has_connectors, ANALYTICS, TASK_SERVER_V2
 
 def is_selected(section, matcher):
   if section == matcher:
@@ -75,9 +75,11 @@ def is_selected(section, matcher):
                   <a href="${ url('desktop.lib.metrics.views.index') }">${_('Metrics')}</a>
                 </li>
                 % endif
+                % if hasattr(TASK_SERVER_V2, 'get') and TASK_SERVER_V2.ENABLED.get():
                 <li class="${is_selected(section, 'task_server')}">
                   <a href="${ url('desktop.views.task_server_view') }">${_('Task Server')}</a>
                 </li>
+                % endif
               % endif
             </ul>
           </div>

+ 16 - 25
desktop/core/src/desktop/urls.py

@@ -17,42 +17,34 @@
 
 from __future__ import absolute_import
 
-import logging
 import re
 import sys
+import logging
 
+from django.conf import settings
+from django.views.static import serve
+from rest_framework_simplejwt.views import (
+  TokenObtainPairView,
+  TokenRefreshView,
+  TokenVerifyView,
+)
 
 # FIXME: This could be replaced with hooking into the `AppConfig.ready()`
 # signal in Django 1.7:
 #
 # https://docs.djangoproject.com/en/1.7/ref/applications/#django.apps.AppConfig.ready
 #
-
 import desktop.lib.metrics.file_reporter
-desktop.lib.metrics.file_reporter.start_file_reporter()
-
-from django.conf import settings
-from django.views.static import serve
-from rest_framework_simplejwt.views import (
-    TokenObtainPairView,
-    TokenRefreshView,
-    TokenVerifyView,
-)
-
-from notebook import views as notebook_views
-from useradmin import views as useradmin_views
-
-from desktop import appmanager
-
-from desktop import views as desktop_views
-from desktop import api as desktop_api
-from desktop import api2 as desktop_api2
-from desktop import api_public_urls_v1
+from desktop import api as desktop_api, api2 as desktop_api2, api_public_urls_v1, appmanager, views as desktop_views
 from desktop.auth import views as desktop_auth_views
-from desktop.conf import METRICS, USE_NEW_EDITOR, ANALYTICS, has_connectors, ENABLE_PROMETHEUS, SLACK
+from desktop.conf import ANALYTICS, ENABLE_PROMETHEUS, METRICS, SLACK, USE_NEW_EDITOR, has_connectors
 from desktop.configuration import api as desktop_configuration_api
 from desktop.lib.vcs import api as desktop_lib_vcs_api
 from desktop.settings import is_oidc_configured
+from notebook import views as notebook_views
+from useradmin import views as useradmin_views
+
+desktop.lib.metrics.file_reporter.start_file_reporter()
 
 if sys.version_info[0] > 2:
   from django.urls import include, re_path
@@ -71,7 +63,7 @@ handler500 = 'desktop.views.serve_500_error'
 dynamic_patterns = [
   re_path(r'^hue/accounts/login', desktop_auth_views.dt_login, name='desktop_auth_views_dt_login'),
   re_path(r'^hue/accounts/logout/?$', desktop_auth_views.dt_logout, {'next_page': '/'}),
-  re_path(r'^accounts/login/?$', desktop_auth_views.dt_login), # Deprecated
+  re_path(r'^accounts/login/?$', desktop_auth_views.dt_login),  # Deprecated
   re_path(r'^accounts/logout/?$', desktop_auth_views.dt_logout, {'next_page': '/'}),
   re_path(r'^profile$', desktop_auth_views.profile),
   re_path(r'^login/oauth/?$', desktop_auth_views.oauth_login),
@@ -101,7 +93,7 @@ dynamic_patterns += [
   re_path(r'^desktop/get_debug_level', desktop_views.get_debug_level),
   re_path(r'^desktop/set_all_debug', desktop_views.set_all_debug),
   re_path(r'^desktop/reset_all_debug', desktop_views.reset_all_debug),
-  re_path(r'^bootstrap.js$', desktop_views.bootstrap), # unused
+  re_path(r'^bootstrap.js$', desktop_views.bootstrap),  # unused
 
   re_path(r'^desktop/status_bar/?$', desktop_views.status_bar),
   re_path(r'^desktop/debug/is_alive$', desktop_views.is_alive),
@@ -174,7 +166,6 @@ dynamic_patterns += [
   re_path(r'^desktop/api2/taskserver/check_upload_status/(?P<task_id>[^/]+)/?$', desktop_api2.check_upload_status,
           name="desktop_api2.check_upload_status"),
   re_path(r'^desktop/api2/taskserver/kill_task/(?P<task_id>[^/]+)/?$', desktop_api2.kill_task, name="desktop_api2.kill_task"),
-  re_path(r'^desktop/api2/taskserver/get_available_space/?$', desktop_api2.get_available_space, name="desktop_api2.get_available_space"),
 
   re_path(r'^desktop/api2/get_config/?$', desktop_api2.get_config),
   re_path(r'^desktop/api2/get_hue_config/?$', desktop_api2.get_hue_config),

+ 13 - 23
desktop/libs/aws/src/aws/s3/upload.py

@@ -22,24 +22,24 @@ See http://docs.djangoproject.com/en/1.9/topics/http/file-uploads/
 """
 
 import io
-from future import standard_library
-standard_library.install_aliases()
-import logging
 import sys
+import logging
 import unicodedata
 
-if sys.version_info[0] > 2:
-  from io import BytesIO as stream_io
-else:
-  from cStringIO import StringIO as stream_io
-
 from django.core.files.uploadedfile import SimpleUploadedFile
 from django.core.files.uploadhandler import FileUploadHandler, SkipFile, StopFutureHandlers, StopUpload, UploadFileException
 
-from desktop.conf import TASK_SERVER
-from desktop.lib.fsmanager import get_client
 from aws.s3 import parse_uri
 from aws.s3.s3fs import S3FileSystemException
+from desktop.conf import TASK_SERVER_V2
+from desktop.lib.exceptions_renderable import PopupException
+from desktop.lib.fsmanager import get_client
+from filebrowser.utils import calculate_total_size, generate_chunks
+
+if sys.version_info[0] > 2:
+  from io import BytesIO as stream_io
+else:
+  from cStringIO import StringIO as stream_io
 
 if sys.version_info[0] > 2:
   from django.utils.translation import gettext as _
@@ -50,8 +50,6 @@ DEFAULT_WRITE_SIZE = 1024 * 1024 * 128  # TODO: set in configuration (currently
 
 LOG = logging.getLogger()
 
-from desktop.lib.exceptions_renderable import PopupException
-from filebrowser.utils import generate_chunks, calculate_total_size
 
 class S3FineUploaderChunkedUpload(object):
   def __init__(self, request, *args, **kwargs):
@@ -62,7 +60,7 @@ class S3FineUploaderChunkedUpload(object):
     self.totalfilesize = kwargs.get('qqtotalfilesize')
     self.file_name = kwargs.get('qqfilename')
     if self.file_name:
-      self.file_name = unicodedata.normalize('NFC', self.file_name) # Normalize unicode
+      self.file_name = unicodedata.normalize('NFC', self.file_name)  # Normalize unicode
     self.destination = kwargs.get('dest', None)  # GET param avoids infinite looping
     self._fs = get_client(fs='s3a', user=self._request.user.username)
     self.bucket_name, self.key_name = parse_uri(self.destination)[:2]
@@ -70,7 +68,7 @@ class S3FineUploaderChunkedUpload(object):
     self._fs._stats(self.destination)
     self._bucket = self._fs._get_bucket(self.bucket_name)
     self.filepath = self._fs.join(self.key_name, self.file_name)
-    if kwargs.get('chunk_size', None) != None:
+    if kwargs.get('chunk_size', None):
       self.chunk_size = kwargs.get('chunk_size')
 
   def check_access(self):
@@ -94,7 +92,7 @@ class S3FineUploaderChunkedUpload(object):
                             {'name': self.file_name, 'qquuid': self.qquuid, 'size': self.totalfilesize})
 
   def upload_chunks(self):
-    if TASK_SERVER.ENABLED.get():
+    if TASK_SERVER_V2.ENABLED.get():
       try:
         self._mp = self._bucket.initiate_multipart_upload(self.filepath)
       except (S3FileUploadError, S3FileSystemException) as e:
@@ -165,7 +163,6 @@ class S3FileUploadHandler(FileUploadHandler):
       self._fs._stats(self.destination)
       self._bucket = self._fs._get_bucket(self.bucket_name)
 
-
   def new_file(self, field_name, file_name, *args, **kwargs):
     if self._is_s3_upload():
       super(S3FileUploadHandler, self).new_file(field_name, file_name, *args, **kwargs)
@@ -186,7 +183,6 @@ class S3FileUploadHandler(FileUploadHandler):
         self.request.META['upload_failed'] = e
         raise StopUpload()
 
-
   def receive_data_chunk(self, raw_data, start):
     if self._is_s3_upload():
       try:
@@ -202,7 +198,6 @@ class S3FileUploadHandler(FileUploadHandler):
     else:
       return raw_data
 
-
   def file_complete(self, file_size):
     if self._is_s3_upload():
       # Finish the upload
@@ -213,7 +208,6 @@ class S3FileUploadHandler(FileUploadHandler):
     else:
       return None
 
-
   def _get_s3fs(self, request):
     # Pre 6.0 request.fs did not exist, now it does. The logic for assigning request.fs is not correct for FileUploadHandler.
     fs = get_client(user=request.user.username)
@@ -223,16 +217,13 @@ class S3FileUploadHandler(FileUploadHandler):
 
     return fs
 
-
   def _is_s3_upload(self):
     return self._get_scheme() and self._get_scheme().startswith('S3')
 
-
   def _check_access(self):
     if not self._fs.check_access(self.destination, permission='WRITE'):
       raise S3FileSystemException('Insufficient permissions to write to S3 path "%s".' % self.destination)
 
-
   def _get_scheme(self):
     if self.destination:
       dst_parts = self.destination.split('://')
@@ -243,7 +234,6 @@ class S3FileUploadHandler(FileUploadHandler):
     else:
       return None
 
-
   def _get_file_part(self, raw_data):
     fp = stream_io()
     fp.write(raw_data)

+ 23 - 25
desktop/libs/azure/src/azure/abfs/upload.py

@@ -13,36 +13,35 @@
 # 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 standard_library
-standard_library.install_aliases()
-import logging
+
 import sys
+import logging
 import unicodedata
 
-if sys.version_info[0] > 2:
-  from io import StringIO as string_io
-else:
-  from cStringIO import StringIO as string_io
-
 from django.core.files.uploadedfile import SimpleUploadedFile
 from django.core.files.uploadhandler import FileUploadHandler, SkipFile, StopFutureHandlers, StopUpload, UploadFileException
-from desktop.conf import TASK_SERVER
 
-from desktop.lib.exceptions_renderable import PopupException
-from desktop.lib.fsmanager import get_client
 from azure.abfs.__init__ import parse_uri
 from azure.abfs.abfs import ABFSFileSystemException
+from desktop.conf import TASK_SERVER_V2
+from desktop.lib.exceptions_renderable import PopupException
+from desktop.lib.fsmanager import get_client
+from filebrowser.utils import calculate_total_size, generate_chunks
+
+if sys.version_info[0] > 2:
+  from io import StringIO as string_io
+else:
+  from cStringIO import StringIO as string_io
 
 if sys.version_info[0] > 2:
   from django.utils.translation import gettext as _
 else:
   from django.utils.translation import ugettext as _
 
-DEFAULT_WRITE_SIZE = 100 * 1024 * 1024 # As per Azure doc, maximum blob size is 100MB
+DEFAULT_WRITE_SIZE = 100 * 1024 * 1024  # As per Azure doc, maximum blob size is 100MB
 
 LOG = logging.getLogger()
 
-from filebrowser.utils import generate_chunks, calculate_total_size
 
 class ABFSFineUploaderChunkedUpload(object):
   def __init__(self, request, *args, **kwargs):
@@ -53,17 +52,17 @@ class ABFSFineUploaderChunkedUpload(object):
     self.totalfilesize = kwargs.get('qqtotalfilesize')
     self.file_name = kwargs.get('qqfilename')
     if self.file_name:
-      self.file_name = unicodedata.normalize('NFC', self.file_name) # Normalize unicode
+      self.file_name = unicodedata.normalize('NFC', self.file_name)  # Normalize unicode
     self.destination = kwargs.get('dest', None)  # GET param avoids infinite looping
     self.target_path = None
 
-    if kwargs.get('chunk_size', None) != None:
+    if kwargs.get('chunk_size', None):
       self.chunk_size = kwargs.get('chunk_size')
 
     if self._is_abfs_upload():
       self._fs = self._get_abfs(request)
       self.filesystem, self.directory = parse_uri(self.destination)[:2]
-       # Verify that the path exists
+      # Verify that the path exists
       self._fs.stats(self.destination)
 
   def check_access(self):
@@ -75,9 +74,9 @@ class ABFSFineUploaderChunkedUpload(object):
 
     try:
       # Check access permissions before attempting upload
-      #self._check_access() #implement later
+      # self._check_access() #implement later
       LOG.debug("ABFSFineUploaderChunkedUpload: Initiating ABFS upload to target path: %s" % self.target_path)
-      if not TASK_SERVER.ENABLED.get():
+      if not TASK_SERVER_V2.ENABLED.get():
         self._fs.create(self.target_path)
     except (ABFSFileUploadError, ABFSFileSystemException) as e:
       LOG.error("ABFSFineUploaderChunkedUpload: Encountered error in ABFSUploadHandler check_access: %s" % e)
@@ -89,7 +88,7 @@ class ABFSFineUploaderChunkedUpload(object):
                             {'name': self.file_name, 'qquuid': self.qquuid, 'size': self.totalfilesize})
 
   def upload_chunks(self):
-    if TASK_SERVER.ENABLED.get():
+    if TASK_SERVER_V2.ENABLED.get():
       self.target_path = self._fs.join(self.destination, self.file_name)
       try:
         LOG.debug("ABFSFineUploaderChunkedUpload: Initiating ABFS upload to target path: %s" % self.target_path)
@@ -109,7 +108,7 @@ class ABFSFineUploaderChunkedUpload(object):
       LOG.exception('ABFSFineUploaderChunkedUpload: Failed to upload file to ABFS at %s: %s' % (self.target_path, e))
       raise PopupException("ABFSFineUploaderChunkedUpload: S3FileUploadHandler uploading file %s part: %d failed" % (self.filepath, i))
     finally:
-      #finish the upload
+      # finish the upload
       self._fs.flush(self.target_path, {'position': self.totalfilesize})
       LOG.info("ABFSFineUploaderChunkedUpload: has completed file upload to ABFS, total file size is: %d." % self.totalfilesize)
       LOG.debug("%s" % self._fs.stats(self.target_path))
@@ -140,6 +139,7 @@ class ABFSFineUploaderChunkedUpload(object):
     else:
       return None
 
+
 class ABFSFileUploadError(UploadFileException):
   pass
 
@@ -162,12 +162,11 @@ class ABFSFileUploadHandler(FileUploadHandler):
     if self._is_abfs_upload():
       self._fs = self._get_abfs(request)
       self.filesystem, self.directory = parse_uri(self.destination)[:2]
-       # Verify that the path exists
+      # Verify that the path exists
       self._fs.stats(self.destination)
 
     LOG.debug("Chunk size = %d" % DEFAULT_WRITE_SIZE)
 
-
   def new_file(self, field_name, file_name, *args, **kwargs):
     if self._is_abfs_upload():
       super(ABFSFileUploadHandler, self).new_file(field_name, file_name, *args, **kwargs)
@@ -177,7 +176,7 @@ class ABFSFileUploadHandler(FileUploadHandler):
 
       try:
         # Check access permissions before attempting upload
-        #self._check_access() #implement later
+        # self._check_access() #implement later
         LOG.debug("Initiating ABFS upload to target path: %s" % self.target_path)
         self._fs.create(self.target_path)
         self.file = SimpleUploadedFile(name=file_name, content='')
@@ -187,7 +186,6 @@ class ABFSFileUploadHandler(FileUploadHandler):
         self.request.META['upload_failed'] = e
         raise StopUpload()
 
-
   def receive_data_chunk(self, raw_data, start):
     if self._is_abfs_upload():
       try:
@@ -203,7 +201,7 @@ class ABFSFileUploadHandler(FileUploadHandler):
 
   def file_complete(self, file_size):
     if self._is_abfs_upload():
-      #finish the upload
+      # finish the upload
       self._fs.flush(self.target_path, {'position': int(file_size)})
       LOG.info("ABFSFileUploadHandler has completed file upload to ABFS, total file size is: %d." % file_size)
       self.file.size = file_size

+ 23 - 17
desktop/libs/hadoop/src/hadoop/fs/upload.py

@@ -25,25 +25,24 @@ which is triggered by a magic prefix ("HDFS") in the field name.
 See http://docs.djangoproject.com/en/1.2/topics/http/file-uploads/
 """
 
-from builtins import object
+import os
+import sys
+import time
 import errno
 import logging
-import os
 import posixpath
-import sys
 import unicodedata
-import time
+from builtins import object
 
-from django.core.files.uploadhandler import FileUploadHandler, StopFutureHandlers, StopUpload, UploadFileException, SkipFile
-
-from desktop.lib import fsmanager
+from django.core.files.uploadhandler import FileUploadHandler, SkipFile, StopFutureHandlers, StopUpload, UploadFileException
 
 import hadoop.cluster
-from hadoop.conf import UPLOAD_CHUNK_SIZE
-from hadoop.fs.exceptions import WebHdfsException
+from desktop.lib import fsmanager
 from desktop.lib.exceptions_renderable import PopupException
 from filebrowser.conf import ARCHIVE_UPLOAD_TEMPDIR
-from filebrowser.utils import generate_chunks, calculate_total_size
+from filebrowser.utils import calculate_total_size, generate_chunks
+from hadoop.conf import UPLOAD_CHUNK_SIZE
+from hadoop.fs.exceptions import WebHdfsException
 
 if sys.version_info[0] > 2:
   from django.utils.translation import gettext as _
@@ -54,6 +53,7 @@ LOG = logging.getLogger()
 
 UPLOAD_SUBDIR = 'hue-uploads'
 
+
 class LocalFineUploaderChunkedUpload(object):
   def __init__(self, request, *args, **kwargs):
     self._request = request
@@ -62,7 +62,7 @@ class LocalFineUploaderChunkedUpload(object):
     self.totalfilesize = kwargs.get('qqtotalfilesize')
     self.file_name = kwargs.get('qqfilename')
     if self.file_name:
-      self.file_name = unicodedata.normalize('NFC', self.file_name) # Normalize unicode
+      self.file_name = unicodedata.normalize('NFC', self.file_name)  # Normalize unicode
     local = "local:/"
     if local in kwargs.get('dest', ""):
       self.dest = kwargs.get('dest')[len(local):]
@@ -72,14 +72,18 @@ class LocalFineUploaderChunkedUpload(object):
     self.filepath = request.fs.join(self.dest, self.file_name)
     self._file = None
     self.chunk_size = 0
+
   def check_access(self):
     pass
+
   def upload_chunks(self):
     pass
+
   def upload(self):
     self.check_access()
     self.upload_chunks()
 
+
 class HDFSFineUploaderChunkedUpload(object):
   def __init__(self, request, *args, **kwargs):
     self._request = request
@@ -88,16 +92,16 @@ class HDFSFineUploaderChunkedUpload(object):
     self.totalfilesize = kwargs.get('qqtotalfilesize')
     self.file_name = kwargs.get('qqfilename')
     if self.file_name:
-      self.file_name = unicodedata.normalize('NFC', self.file_name) # Normalize unicode
+      self.file_name = unicodedata.normalize('NFC', self.file_name)  # Normalize unicode
     self.dest = kwargs.get('dest')
     self.file_name = kwargs.get('qqfilename')
-    if kwargs.get('filepath', None) != None:
+    if kwargs.get('filepath', None):
       self.filepath = kwargs.get('filepath')
     else:
       self.filepath = request.fs.join(self.dest, self.file_name)
       kwargs['filepath'] = self.filepath
     self._file = None
-    if kwargs.get('chunk_size', None) != None:
+    if kwargs.get('chunk_size', None):
       self.chunk_size = kwargs.get('chunk_size')
 
   def check_access(self):
@@ -141,7 +145,7 @@ class HDFSFineUploaderChunkedUpload(object):
       except Exception:
         pass
       if already_exists:
-        msg = _('Destination %(name)s already exists.')  % {'name': self.filepath}
+        msg = _('Destination %(name)s already exists.') % {'name': self.filepath}
       else:
         msg = _('Copy to %(name)s failed: %(error)s') % {'name': self.filepath, 'error': ex}
       raise PopupException(msg)
@@ -224,6 +228,7 @@ class HDFStemporaryUploadedFile(object):
   def close(self):
     self._file.close()
 
+
 class FineUploaderChunkedUploadHandler(FileUploadHandler):
   """
   A custom file upload handler for handling chunked uploads using FineUploader.
@@ -268,7 +273,8 @@ class FineUploaderChunkedUploadHandler(FileUploadHandler):
     - file_size (int): The total size of the uploaded file.
     """
     elapsed = time.time() - self._starttime
-    LOG.info('Uploaded %s bytes %s to in %s seconds' % (file_size, self.chunk_file_path, elapsed))
+    LOG.debug('Uploaded %s bytes %s to in %s seconds' % (file_size, self.chunk_file_path, elapsed))
+
 
 class HDFSfileUploadHandler(FileUploadHandler):
   """
@@ -289,7 +295,7 @@ class HDFSfileUploadHandler(FileUploadHandler):
     self._file = None
     self._starttime = 0
     self._activated = False
-    self._destination = request.GET.get('dest', None) # GET param avoids infinite looping
+    self._destination = request.GET.get('dest', None)  # GET param avoids infinite looping
     self.request = request
     fs = fsmanager.get_filesystem('default')
     if not fs:

+ 26 - 26
desktop/libs/notebook/src/notebook/api.py

@@ -15,42 +15,40 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from future import standard_library
-standard_library.install_aliases()
+import sys
 import json
 import logging
 
 import sqlparse
-import sys
-
-from django.urls import reverse
+import opentracing.tracer
 from django.db.models import Q
+from django.urls import reverse
 from django.views.decorators.http import require_GET, require_POST
-import opentracing.tracer
 
 from azure.abfs.__init__ import abfspath
-from desktop.conf import TASK_SERVER, ENABLE_CONNECTORS
-from desktop.lib.i18n import smart_str
+from desktop.conf import ENABLE_CONNECTORS, TASK_SERVER_V2
 from desktop.lib.django_util import JsonResponse
 from desktop.lib.exceptions_renderable import PopupException
-from desktop.models import Document2, Document, __paginate, _get_gist_document, FilesystemException
-from indexer.file_format import HiveFormat
+from desktop.lib.i18n import smart_str
+from desktop.models import Document, Document2, FilesystemException, __paginate, _get_gist_document
 from indexer.fields import Field
+from indexer.file_format import HiveFormat
 from metadata.conf import OPTIMIZER
-
 from notebook.conf import EXAMPLES
-from notebook.connectors.base import Notebook, QueryExpired, SessionExpired, QueryError, _get_snippet_name, patch_snippet_for_connector
+from notebook.connectors.base import Notebook, QueryError, QueryExpired, SessionExpired, _get_snippet_name, patch_snippet_for_connector
 from notebook.connectors.hiveserver2 import HS2Api
 from notebook.decorators import api_error_handler, check_document_access_permission, check_document_modify_permission
-from notebook.models import escape_rows, make_notebook, upgrade_session_properties, get_api, _get_dialect_example
+from notebook.models import _get_dialect_example, escape_rows, get_api, make_notebook, upgrade_session_properties
 
 if sys.version_info[0] > 2:
   from urllib.parse import unquote as urllib_unquote
+
   from django.utils.translation import gettext as _
 else:
-  from django.utils.translation import ugettext as _
   from urllib import unquote as urllib_unquote
 
+  from django.utils.translation import ugettext as _
+
 
 LOG = logging.getLogger()
 
@@ -142,8 +140,8 @@ def _execute_notebook(request, notebook, snippet):
 
   try:
     try:
-      sessions = notebook.get('sessions') and notebook['sessions'] # Session reference for snippet execution without persisting it
-      active_executable = json.loads(request.POST.get('executable', '{}')) # Editor v2
+      sessions = notebook.get('sessions') and notebook['sessions']  # Session reference for snippet execution without persisting it
+      active_executable = json.loads(request.POST.get('executable', '{}'))  # Editor v2
       # TODO: Use statement, database etc. from active_executable
 
       if historify:
@@ -168,7 +166,7 @@ def _execute_notebook(request, notebook, snippet):
       if historify:
         _snippet = [s for s in notebook['snippets'] if s['id'] == snippet['id']][0]
 
-        if 'id' in active_executable: # Editor v2
+        if 'id' in active_executable:  # Editor v2
           # notebook_executable is the 1-to-1 match of active_executable in the notebook structure
           notebook_executable = [e for e in _snippet['executor']['executables'] if e['id'] == active_executable['id']][0]
           if 'handle' in response:
@@ -180,28 +178,28 @@ def _execute_notebook(request, notebook, snippet):
             }
             notebook_executable['operationId'] = history.uuid
 
-        if response.get('handle'): # No failure
-          if 'result' not in _snippet: # Editor v2
+        if response.get('handle'):  # No failure
+          if 'result' not in _snippet:  # Editor v2
             _snippet['result'] = {}
           _snippet['result']['handle'] = response['handle']
           _snippet['result']['statements_count'] = response['handle'].get('statements_count', 1)
           _snippet['result']['statement_id'] = response['handle'].get('statement_id', 0)
           _snippet['result']['handle']['statement'] = response['handle'].get(
               'statement', snippet['statement']
-          ).strip() # For non HS2, as non multi query yet
+          ).strip()  # For non HS2, as non multi query yet
         else:
           _snippet['status'] = 'failed'
 
-        if history: # If _historify failed, history will be None.
+        if history:  # If _historify failed, history will be None.
           # If we get Atomic block exception, something underneath interpreter.execute() crashed and is not handled.
           history.update_data(notebook)
           history.save()
 
           response['history_id'] = history.id
           response['history_uuid'] = history.uuid
-          if notebook['isSaved']: # Keep track of history of saved queries
+          if notebook['isSaved']:  # Keep track of history of saved queries
             response['history_parent_uuid'] = history.dependencies.filter(type__startswith='query-').latest('last_modified').uuid
-  except QueryError as ex: # We inject the history information from _historify() to the failed queries
+  except QueryError as ex:  # We inject the history information from _historify() to the failed queries
     if response.get('history_id'):
       ex.extra['history_id'] = response['history_id']
     if response.get('history_uuid'):
@@ -473,6 +471,7 @@ def get_logs(request):
 
   return JsonResponse(response)
 
+
 def _save_notebook(notebook, user):
   if notebook.get('type') != 'notebook' and notebook['snippets'][0].get('connector') and \
           notebook['snippets'][0]['connector'].get('dialect'):  # TODO Connector unification
@@ -542,7 +541,7 @@ def _clear_sessions(notebook):
 def _historify(notebook, user):
   query_type = 'query-%(dialect)s' % notebook if ENABLE_CONNECTORS.get() else notebook['type']
   name = notebook['name'] if (notebook['name'] and notebook['name'].strip() != '') else DEFAULT_HISTORY_NAME
-  is_managed = notebook.get('isManaged') == True  # Prevents None
+  is_managed = bool(notebook.get('isManaged'))  # Prevents None
 
   if is_managed and Document2.objects.filter(uuid=notebook['uuid']).exists():
     history_doc = Document2.objects.get(uuid=notebook['uuid'])
@@ -597,6 +596,7 @@ def _get_statement(notebook):
       LOG.warning('Could not get statement from query history: %s' % e)
   return ''
 
+
 @require_GET
 @api_error_handler
 @check_document_access_permission
@@ -809,7 +809,7 @@ def format(request):
   response = {'status': 0}
 
   statements = request.POST.get('statements', '')
-  response['formatted_statements'] = sqlparse.format(statements, reindent=True, keyword_case='upper') # SQL only currently
+  response['formatted_statements'] = sqlparse.format(statements, reindent=True, keyword_case='upper')  # SQL only currently
 
   return JsonResponse(response)
 
@@ -831,7 +831,7 @@ def export_result(request):
 
   api = get_api(request, snippet)
 
-  if data_format == 'hdfs-file': # Blocking operation, like downloading
+  if data_format == 'hdfs-file':  # Blocking operation, like downloading
     if request.fs.isdir(destination):
       if notebook.get('name'):
         destination += '/%(name)s.csv' % notebook

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

@@ -15,24 +15,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from builtins import object
-import json
-import logging
 import re
 import sys
+import json
 import time
 import uuid
+import logging
+from builtins import object
 
 from django.utils.encoding import smart_str
 
 from beeswax.common import find_compute, is_compute
 from desktop.auth.backend import is_admin
-from desktop.conf import TASK_SERVER, has_connectors
+from desktop.conf import TASK_SERVER_V2, has_connectors
 from desktop.lib import export_csvxls
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.lib.i18n import smart_unicode
 from metadata.optimizer.base import get_api as get_optimizer_api
-
 from notebook.conf import get_ordered_interpreters
 from notebook.sql_utils import get_current_statement
 
@@ -426,7 +425,7 @@ def patch_snippet_for_connector(snippet, user=None):
 def get_api(request, snippet):
   from notebook.connectors.oozie_batch import OozieApi
 
-  if snippet.get('wasBatchExecuted') and not TASK_SERVER.ENABLED.get():
+  if snippet.get('wasBatchExecuted') and not TASK_SERVER_V2.ENABLED.get():
     return OozieApi(user=request.user, request=request)
 
   if snippet.get('type') == 'report':
@@ -576,8 +575,10 @@ class Api(object):
     raise OperationNotSupported()
 
   def download(self, notebook, snippet, file_format='csv'):
-    from beeswax import data_export  # TODO: Move to notebook?
-    from beeswax import conf
+    from beeswax import (
+      conf,
+      data_export,  # TODO: Move to notebook?
+    )
 
     result_wrapper = ExecutionWrapper(self, notebook, snippet)
 

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

@@ -15,17 +15,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from future import standard_library
-standard_library.install_aliases()
-from builtins import str, object
-import datetime
+import sys
 import json
-import logging
 import math
-import numbers
-import sys
 import uuid
-
+import logging
+import numbers
+import datetime
+from builtins import object, str
 from datetime import timedelta
 
 from django.contrib.sessions.models import Session
@@ -33,23 +30,24 @@ from django.db.models import Count
 from django.db.models.functions import Trunc
 from django.utils.html import escape
 
-from desktop.conf import has_connectors, TASK_SERVER
+from desktop.conf import TASK_SERVER_V2, has_connectors
 from desktop.lib.connectors.models import _get_installed_connectors
 from desktop.lib.i18n import smart_unicode
 from desktop.lib.paths import SAFE_CHARACTERS_URI
 from desktop.models import Directory, Document2
-from useradmin.models import User, install_sample_user
-
 from notebook.conf import EXAMPLES, get_ordered_interpreters
 from notebook.connectors.base import Notebook, get_api as _get_api, get_interpreter
+from useradmin.models import User, install_sample_user
 
 if sys.version_info[0] > 2:
   from urllib.parse import quote as urllib_quote
+
   from django.utils.translation import gettext as _
 else:
-  from django.utils.translation import ugettext as _
   from urllib import quote as urllib_quote
 
+  from django.utils.translation import ugettext as _
+
 
 LOG = logging.getLogger()
 
@@ -121,7 +119,7 @@ def make_notebook(
     if settings is not None:
       _update_property_value(sessions_properties, 'files', files)
   elif editor_type == 'java':
-    sessions_properties = [] # Java options
+    sessions_properties = []  # Java options
   else:
     sessions_properties = []
 
@@ -136,7 +134,7 @@ def make_notebook(
          'id': None
       }
     ],
-    'selectedSnippet': editor_connector, # TODO: might need update in notebook.ko.js
+    'selectedSnippet': editor_connector,  # TODO: might need update in notebook.ko.js
     'type': 'notebook' if is_notebook else 'query-%s' % editor_type,
     'showHistory': True,
     'isSaved': is_saved,
@@ -213,7 +211,7 @@ def make_notebook2(name='Browse', description='', is_saved=False, snippets=None)
         'type': _snippet['type'],
         'properties': HS2Api.get_properties(snippet['type']),
         'id': None
-      } for _snippet in _snippets # Non unique types currently
+      } for _snippet in _snippets  # Non unique types currently
     ],
     'selectedSnippet': _snippets[0]['type'],
     'showHistory': False,
@@ -523,7 +521,7 @@ def _convert_type(btype, bdata):
   elif btype == RDBMS:
     data = json.loads(bdata)
     return data['query']['server']
-  elif btype == SPARK: # We should not import
+  elif btype == SPARK:  # We should not import
     return 'spark'
   else:
     return 'hive'
@@ -552,6 +550,7 @@ def _get_example_directory(user):
   )
   return examples_dir
 
+
 def _get_dialect_example(dialect):
   sample_user = install_sample_user()
   examples_dir = _get_example_directory(sample_user)
@@ -732,10 +731,11 @@ class MockRequest():
     self.POST = {}
     self.GET = {}
 
+
 def install_custom_examples():
   if EXAMPLES.AUTO_LOAD.get():
-    from desktop.auth.backend import rewrite_user
     from beeswax.management.commands import beeswax_install_examples
+    from desktop.auth.backend import rewrite_user
     from useradmin.models import install_sample_user
 
     user = rewrite_user(

+ 24 - 20
desktop/libs/notebook/src/notebook/tasks.py

@@ -16,21 +16,18 @@
 # limitations under the License.
 
 from __future__ import absolute_import, unicode_literals
-from future import standard_library
-standard_library.install_aliases()
 
-from builtins import next, object
-
-import codecs
 import csv
-import datetime
-import json
-import logging
 import sys
+import json
 import time
+import codecs
+import logging
+import datetime
+from builtins import next, object
 
-from celery.utils.log import get_task_logger
 from celery import states
+from celery.utils.log import get_task_logger
 from django.core.cache import caches
 from django.core.files.storage import get_storage_class
 from django.db import transaction
@@ -39,16 +36,15 @@ from django.http import FileResponse, HttpRequest
 from beeswax.data_export import DataAdapter
 from desktop.auth.backend import rewrite_user
 from desktop.celery import app
-from desktop.conf import ENABLE_HUE_5, TASK_SERVER
+from desktop.conf import ENABLE_HUE_5, TASK_SERVER_V2
 from desktop.lib import export_csvxls, fsmanager
 from desktop.models import Document2
 from desktop.settings import CACHES_CELERY_KEY, CACHES_CELERY_QUERY_RESULT_KEY
-from useradmin.models import User
-
 from notebook.api import _get_statement
-from notebook.connectors.base import get_api, QueryExpired, ExecutionWrapper, QueryError
-from notebook.models import make_notebook, MockedDjangoRequest, Notebook
+from notebook.connectors.base import ExecutionWrapper, QueryError, QueryExpired, get_api
+from notebook.models import MockedDjangoRequest, Notebook, make_notebook
 from notebook.sql_utils import get_current_statement
+from useradmin.models import User
 
 if sys.version_info[0] > 2:
   from io import StringIO as string_io
@@ -72,7 +68,7 @@ STATE_MAP = {
   states.REJECTED: 'rejected',
   states.IGNORED: 'ignored'
 }
-storage_info = json.loads(TASK_SERVER.RESULT_STORAGE.get())
+storage_info = json.loads(TASK_SERVER_V2.RESULT_STORAGE.get())
 storage = get_storage_class(storage_info.get('backend'))(**storage_info.get('properties', {}))
 
 
@@ -134,7 +130,7 @@ def download_to_file(notebook, snippet, file_format='csv', max_rows=-1, **kwargs
       for chunk in response:
         f.write(chunk.encode('utf-8'))
 
-    if TASK_SERVER.RESULT_CACHE.get():
+    if TASK_SERVER_V2.RESULT_CACHE.get():
       with storage.open(result_key, 'rb') as store:
         with codecs.getreader('utf-8')(store) as text_file:
           delimiter = ',' if sys.version_info[0] > 2 else ','.encode('utf-8')
@@ -238,7 +234,7 @@ def _patch_status(notebook):
 def execute(*args, **kwargs):
   notebook = args[0]
   snippet = args[1]
-  kwargs['max_rows'] = TASK_SERVER.FETCH_RESULT_LIMIT.get()
+  kwargs['max_rows'] = TASK_SERVER_V2.FETCH_RESULT_LIMIT.get()
   _patch_status(notebook)
 
   task = download_to_file.apply_async(args=args, kwargs=kwargs, task_id=notebook['uuid'])
@@ -287,7 +283,7 @@ def get_log(notebook, snippet, startFrom=None, size=None, postdict=None, user_id
   elif state in states.EXCEPTION_STATES:
     return ''
 
-  if TASK_SERVER.RESULT_CACHE.get():
+  if TASK_SERVER_V2.RESULT_CACHE.get():
     return ''
   else:
     if not startFrom:
@@ -305,7 +301,7 @@ def get_log(notebook, snippet, startFrom=None, size=None, postdict=None, user_id
       return output.getvalue()
 
 
-def get_jobs(notebook, snippet, logs, **kwargs): # Re implementation 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
 
@@ -399,7 +395,7 @@ def fetch_result(notebook, snippet, rows, start_over, **kwargs):
 def _get_data(task_id):
   result_key = _result_key(task_id)
 
-  if TASK_SERVER.RESULT_CACHE.get():
+  if TASK_SERVER_V2.RESULT_CACHE.get():
     csv_reader = caches[CACHES_CELERY_QUERY_RESULT_KEY].get(result_key)  # TODO check if expired
     if csv_reader is None:
       raise QueryError('Cached results %s not found.' % result_key)
@@ -430,6 +426,7 @@ def fetch_result_size(*args, **kwargs):
   info = result.info
   return {'rows': info.get('row_counter', 0)}
 
+
 def cancel(*args, **kwargs):
   notebook = args[0]
   snippet = args[1]
@@ -490,6 +487,7 @@ def _cleanup(notebook, snippet):
   storage.delete(_log_key(notebook, snippet))
   caches[CACHES_CELERY_KEY].delete(_fetch_progress_key(notebook, snippet))
 
+
 def _get_query_key(notebook, snippet):
   if ENABLE_HUE_5.get():
     if snippet.get('executable'):
@@ -506,21 +504,27 @@ def _get_query_key(notebook, snippet):
   else:
     return query_key
 
+
 def _log_key(notebook, snippet):
   return _get_query_key(notebook, snippet) + '_log'
 
+
 def _result_key(task_id):
   return task_id + '_result'
 
+
 def _fetch_progress_key(notebook, snippet):
   return _get_query_key(notebook, snippet) + '_fetch_progress'
 
+
 def _cancel_statement_async_id(notebook, snippet):
   return _get_query_key(notebook, snippet) + '_cancel'
 
+
 def _close_statement_async_id(notebook, snippet):
   return _get_query_key(notebook, snippet) + '_close'
 
+
 def _get_request(postdict=None, user_id=None):
   request = HttpRequest()
   request.POST = postdict

+ 23 - 7
tools/container/hue/hue.sh

@@ -111,12 +111,28 @@ function set_samlcert() {
   export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib64
 }
 
+function task_server_enabled() {
+  for file in "$HUE_CONF_DIR/zhue_safety_valve.ini" "$HUE_CONF_DIR/zhue.ini" "$HUE_CONF_DIR/hue.ini"; do
+    echo "Checking task_server configs in: $file"
+    if grep -A 1 '^\[\[task_server_v2\]\]' "$file" | grep -qi 'enabled=True'; then
+      echo "Task server configs enabled in $file"
+      return 0
+    elif grep -A 1 '^\[\[task_server_v2\]\]' "$file" | grep -qi 'enabled=False'; then
+      echo "Task server configs disabled in $file"
+      return 1
+    else
+      echo "Task server configs not found in $file"
+    fi
+  done
+  return 2
+}
+
 function start_celery() {
   echo "Starting Celery worker..."
   # The schedule of periodic tasks performed by celery-beat is stored in the celerybeat-schedule file.
   touch $HUE_HOME/celerybeat-schedule
   chmod 644 $HUE_HOME/celerybeat-schedule
-  $HUE_BIN/hue runcelery worker --app desktop.celery --loglevel=DEBUG --schedule_file $HUE_HOME/celerybeat-schedule
+  $HUE_BIN/hue runcelery worker --app desktop.celery --loglevel=INFO --schedule_file $HUE_HOME/celerybeat-schedule
 }
 
 function start_redis() {
@@ -148,16 +164,16 @@ elif [[ $1 == rungunicornserver ]]; then
   fix_column_issue "axes_accessattempt" "trusted"
   $HUE_BIN/hue rungunicornserver
 elif [[ $1 == start_celery ]]; then
-  if awk '/^\[\[task_server\]\]/{flag=1;next}/^\[/{flag=0}flag && /enabled=True/{print "Task server enabled"; exit}' $HUE_CONF_DIR/zhue_safety_valve.ini || \
-     awk '/^\[\[task_server\]\]/{flag=1;next}/^\[/{flag=0}flag && /enabled=True/{print "Task server enabled"; exit}' $HUE_CONF_DIR/zhue.ini || \
-     awk '/^\[\[task_server\]\]/{flag=1;next}/^\[/{flag=0}flag && /enabled=True/{print "Task server enabled"; exit}' $HUE_CONF_DIR/hue.ini; then
+  if task_server_enabled; then
     start_celery
+  else
+    exit 1
   fi
 elif [[ $1 == start_redis ]]; then
-  if awk '/^\[\[task_server\]\]/{flag=1;next}/^\[/{flag=0}flag && /enabled=True/{print "Task server enabled"; exit}' $HUE_CONF_DIR/zhue_safety_valve.ini || \
-     awk '/^\[\[task_server\]\]/{flag=1;next}/^\[/{flag=0}flag && /enabled=True/{print "Task server enabled"; exit}' $HUE_CONF_DIR/zhue.ini || \
-     awk '/^\[\[task_server\]\]/{flag=1;next}/^\[/{flag=0}flag && /enabled=True/{print "Task server enabled"; exit}' $HUE_CONF_DIR/hue.ini; then
+  if task_server_enabled; then
     start_redis
+  else
+    exit 1
   fi
 fi
 

+ 2 - 2
tools/container/hue/supervisor-files/hue_celery_template

@@ -11,8 +11,8 @@ environment=DESKTOP_LOG_DIR="/var/log/${HUEUSER}",PYTHON_EGG_CACHE="${HUE_CONF_D
 user=${HUEUSER}
 group=${HUEUSER}
 exitcodes=0,2
-autorestart=true
-startsecs=10
+autorestart=false
+startsecs=20
 stopwaitsecs=30
 killasgroup=true
 # rlimits

+ 2 - 2
tools/container/hue/supervisor-files/hue_redis_template

@@ -11,8 +11,8 @@ environment=DESKTOP_LOG_DIR="/var/log/${HUEUSER}",PYTHON_EGG_CACHE="${HUE_CONF_D
 user=${HUEUSER}
 group=${HUEUSER}
 exitcodes=0,2
-autorestart=true
-startsecs=10
+autorestart=false
+startsecs=20
 stopwaitsecs=30
 killasgroup=true
 # rlimits