Explorar o código

HUE-8075 [editor] Offer a download_size_limit option for result download

jdesjean %!s(int64=7) %!d(string=hai) anos
pai
achega
1d8f809cc8

+ 7 - 0
apps/beeswax/src/beeswax/conf.py

@@ -117,6 +117,13 @@ DOWNLOAD_ROW_LIMIT = Config(
   help=_t('A limit to the number of rows that can be downloaded from a query before it is truncated. '
           'A value of -1 means there will be no limit.'))
 
+DOWNLOAD_BYTES_LIMIT = Config(
+  key='download_bytes_limit',
+  default=-1,
+  type=int,
+  help=_t('A limit to the number of bytes that can be downloaded from a query before it is truncated. '
+          'A value of -1 means there will be no limit.'))
+
 APPLY_NATURAL_SORT_MAX = Config(
   key="apply_natural_sort_max",
   help=_t("The max number of records in the result set permitted to apply a natural sort to the database or tables list."),

+ 49 - 4
apps/beeswax/src/beeswax/data_export.py

@@ -17,6 +17,8 @@
 
 import json
 import logging
+import math
+import types
 
 from django.utils.translation import ugettext as _
 
@@ -42,8 +44,9 @@ def download(handle, format, db, id=None, file_name='query_result'):
     return
 
   max_rows = conf.DOWNLOAD_ROW_LIMIT.get()
+  max_bytes = conf.DOWNLOAD_BYTES_LIMIT.get()
 
-  content_generator = HS2DataAdapter(handle, db, max_rows=max_rows, start_over=True)
+  content_generator = HS2DataAdapter(handle, db, max_rows=max_rows, start_over=True, max_bytes=max_bytes)
   generator = export_csvxls.create_generator(content_generator, format)
 
   resp = export_csvxls.make_response(generator, format, file_name)
@@ -61,7 +64,7 @@ def download(handle, format, db, id=None, file_name='query_result'):
   return resp
 
 
-def upload(path, handle, user, db, fs, max_rows=-1):
+def upload(path, handle, user, db, fs, max_rows=-1, max_bytes=-1):
   """
   upload(query_model, path, user, db, fs) -> None
 
@@ -72,7 +75,7 @@ def upload(path, handle, user, db, fs, max_rows=-1):
   else:
     fs.do_as_user(user.username, fs.create, path)
 
-  content_generator = HS2DataAdapter(handle, db, max_rows=max_rows, start_over=True)
+  content_generator = HS2DataAdapter(handle, db, max_rows=max_rows, start_over=True, max_bytes=max_bytes)
   for header, data in content_generator:
     dataset = export_csvxls.dataset(None, data)
     fs.do_as_user(user.username, fs.append, path, dataset.csv)
@@ -80,24 +83,55 @@ def upload(path, handle, user, db, fs, max_rows=-1):
 
 class HS2DataAdapter:
 
-  def __init__(self, handle, db, max_rows=-1, start_over=True):
+  def __init__(self, handle, db, max_rows=-1, start_over=True, max_bytes=-1):
     self.handle = handle
     self.db = db
     self.max_rows = max_rows
+    self.max_bytes = max_bytes
     self.start_over = start_over
     self.fetch_size = FETCH_SIZE
     self.limit_rows = max_rows > -1
+    self.limit_bytes = max_bytes > -1
 
     self.first_fetched = True
     self.headers = None
     self.num_cols = None
     self.row_counter = 1
+    self.bytes_counter = 0
     self.is_truncated = False
     self.has_more = True
 
   def __iter__(self):
     return self
 
+  # Return an estimate of the size of the object using only ascii characters once serialized to string.
+  # Avoid serialization to string where possible
+  def _getsizeofascii(self, row):
+    size = 0
+    size += max(len(row) - 1, 0) # CSV commas between columns
+    size += 2 # CSV \r\n at the end of row
+    for col in row:
+      col_type = type(col)
+      if col_type == types.IntType:
+        if col == 0:
+          size += 1
+        elif col < 0:
+          size += int(math.log10(-1 * col)) + 2
+        else:
+          size += int(math.log10(col)) + 1
+      elif col_type == types.StringType:
+        size += len(col)
+      elif col_type == types.FloatType:
+        size += len(str(col))
+      elif col_type == types.BooleanType:
+        size += 4
+      elif col_type == types.NoneType:
+        size += 4
+      else:
+        size += len(str(col))
+
+    return size
+
   def next(self):
     results = self.db.fetch(self.handle, start_over=self.start_over, rows=self.fetch_size)
 
@@ -106,6 +140,10 @@ class HS2DataAdapter:
       self.start_over = False
       self.headers = results.cols()
       self.num_cols = len(self.headers)
+      if self.limit_bytes:
+        self.bytes_counter += max(self.num_cols - 1, 0)
+        for header in self.headers:
+          self.bytes_counter += len(header)
 
       # For result sets with high num of columns, fetch in smaller batches to avoid serialization cost
       if self.num_cols > 100:
@@ -118,10 +156,17 @@ class HS2DataAdapter:
 
       for row in results.rows():
         self.row_counter += 1
+        if self.limit_bytes:
+          self.bytes_counter += self._getsizeofascii(row)
+
         if self.limit_rows and self.row_counter > self.max_rows:
           LOG.warn('The query results exceeded the maximum row limit of %d and has been truncated to first %d rows.' % (self.max_rows, self.row_counter))
           self.is_truncated = True
           break
+        if self.limit_bytes and self.bytes_counter > self.max_bytes:
+          LOG.warn('The query results exceeded the maximum bytes limit of %d and has been truncated to first %d rows.' % (self.max_bytes, self.row_counter))
+          self.is_truncated = True
+          break
         data.append(row)
 
       return self.headers, data

+ 11 - 0
apps/beeswax/src/beeswax/tests.py

@@ -845,6 +845,17 @@ for x in sys.stdin:
     finally:
       finish()
 
+    finish = conf.DOWNLOAD_BYTES_LIMIT.set_for_testing(1024)
+    try:
+      hql = 'SELECT * FROM `%(db)s`.`test`' % {'db': self.db_name}
+      query = hql_query(hql)
+      handle = self.db.execute_and_wait(query)
+      resp = download(handle, 'csv', self.db)
+      content = "".join(resp.streaming_content)
+      assert_true(len(content) <= 1024)
+    finally:
+      finish()
+
 
   def test_data_upload(self):
     hql = 'SELECT * FROM `%(db)s`.`test`' % {'db': self.db_name}

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

@@ -1017,6 +1017,10 @@
   # A value of -1 means there will be no limit.
   ## download_row_limit=100000
 
+  # A limit to the number of bytes that can be downloaded from a query before it is truncated.
+  # A value of -1 means there will be no limit.
+  ## download_bytes_limit=-1
+
   # Hue will try to close the Hive query when the user leaves the editor page.
   # This will free all the query resources in HiveServer2, but also make its results inaccessible.
   ## close_queries=false

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

@@ -1019,6 +1019,10 @@
   # A value of -1 means there will be no limit.
   ## download_row_limit=100000
 
+  # A limit to the number of bytes that can be downloaded from a query before it is truncated.
+  # A value of -1 means there will be no limit.
+  ## download_bytes_limit=-1
+
   # Hue will try to close the Hive query when the user leaves the editor page.
   # This will free all the query resources in HiveServer2, but also make its results inaccessible.
   ## close_queries=false

+ 29 - 4
desktop/core/src/desktop/templates/common_notebook_ko_components.mako

@@ -27,10 +27,11 @@ from notebook.conf import ENABLE_SQL_INDEXER
 LOG = logging.getLogger(__name__)
 
 try:
-  from beeswax.conf import DOWNLOAD_ROW_LIMIT
+  from beeswax.conf import DOWNLOAD_ROW_LIMIT, DOWNLOAD_BYTES_LIMIT
 except ImportError, e:
   LOG.warn("Hive app is not enabled")
   DOWNLOAD_ROW_LIMIT = None
+  DOWNLOAD_BYTES_LIMIT = None
 %>
 
 <%def name="addSnippetMenu()">
@@ -180,12 +181,28 @@ except ImportError, e:
       </a>
       <ul class="dropdown-menu less-padding" style="z-index: 1040">
         <li>
-          <a class="download" href="javascript:void(0)" data-bind="click: downloadCsv, event: { mouseover: function(){ window.onbeforeunload = null; }, mouseout: function() { window.onbeforeunload = $(window).data('beforeunload'); } }" title="${ _('Download first %s rows as CSV') % (hasattr(DOWNLOAD_ROW_LIMIT, 'get') and DOWNLOAD_ROW_LIMIT.get()) }">
+          <a class="download" href="javascript:void(0)" data-bind="click: downloadCsv, event: { mouseover: function(){ window.onbeforeunload = null; }, mouseout: function() { window.onbeforeunload = $(window).data('beforeunload'); } }"
+          % if hasattr(DOWNLOAD_ROW_LIMIT, 'get') and DOWNLOAD_ROW_LIMIT.get() >= 0 and hasattr(DOWNLOAD_BYTES_LIMIT, 'get') and DOWNLOAD_BYTES_LIMIT.get() >= 0:
+          title="${ _('Download first %s rows or %s MB as CSV') % ( DOWNLOAD_ROW_LIMIT.get(), DOWNLOAD_BYTES_LIMIT.get() / 1024 / 1024 ) }"
+          % elif hasattr(DOWNLOAD_BYTES_LIMIT, 'get') and DOWNLOAD_BYTES_LIMIT.get():
+          title="${ _('Download first %s MB as CSV') % DOWNLOAD_BYTES_LIMIT.get() / 1024 / 1024 }"
+          % else:
+          title="${ _('Download first %s rows as CSV') % (hasattr(DOWNLOAD_ROW_LIMIT, 'get') and DOWNLOAD_ROW_LIMIT.get()) }"
+          % endif
+          >
             <i class="fa fa-fw fa-file-o"></i> ${ _('CSV') }
           </a>
         </li>
         <li>
-          <a class="download" href="javascript:void(0)" data-bind="click: downloadXls, event: { mouseover: function(){ window.onbeforeunload = null; }, mouseout: function() { window.onbeforeunload = $(window).data('beforeunload'); } }" title="${ _('Download first %s rows as XLS') % (hasattr(DOWNLOAD_ROW_LIMIT, 'get') and DOWNLOAD_ROW_LIMIT.get()) }">
+          <a class="download" href="javascript:void(0)" data-bind="click: downloadXls, event: { mouseover: function(){ window.onbeforeunload = null; }, mouseout: function() { window.onbeforeunload = $(window).data('beforeunload'); } }"
+          % if hasattr(DOWNLOAD_ROW_LIMIT, 'get') and DOWNLOAD_ROW_LIMIT.get() >= 0 and hasattr(DOWNLOAD_BYTES_LIMIT, 'get') and DOWNLOAD_BYTES_LIMIT.get() >= 0:
+          title="${ _('Download first %s rows or %s MB as XLS') % ( DOWNLOAD_ROW_LIMIT.get(), DOWNLOAD_BYTES_LIMIT.get() / 1024 / 1024 ) }"
+          % elif hasattr(DOWNLOAD_BYTES_LIMIT, 'get') and DOWNLOAD_BYTES_LIMIT.get():
+          title="${ _('Download first %s MB as XLS') % DOWNLOAD_BYTES_LIMIT.get() / 1024 / 1024 }"
+          % else:
+          title="${ _('Download first %s rows as XLS') % (hasattr(DOWNLOAD_ROW_LIMIT, 'get') and DOWNLOAD_ROW_LIMIT.get()) }"
+          % endif
+          >
             <i class="fa fa-fw fa-file-excel-o"></i> ${ _('Excel') }
           </a>
         </li>
@@ -229,8 +246,16 @@ except ImportError, e:
               <div class="controls">
                  <label class="radio">
                   <input data-bind="checked: saveTarget" type="radio" name="save-results-type" value="hdfs-file">
+                  <span style="width: 190px; overflow: hidden; text-overflow: ellipsis; display: inline-block; white-space: nowrap;">
                   &nbsp;
-                  ${ _('First %s rows') % (hasattr(DOWNLOAD_ROW_LIMIT, 'get') and DOWNLOAD_ROW_LIMIT.get()) }
+                  % if hasattr(DOWNLOAD_ROW_LIMIT, 'get') and DOWNLOAD_ROW_LIMIT.get() >= 0 and hasattr(DOWNLOAD_BYTES_LIMIT, 'get') and DOWNLOAD_BYTES_LIMIT.get() >= 0:
+                    ${ _('First %s rows or %s MB') % ( DOWNLOAD_ROW_LIMIT.get(), DOWNLOAD_BYTES_LIMIT.get() / 1024 / 1024 ) }
+                  % elif hasattr(DOWNLOAD_BYTES_LIMIT, 'get') and DOWNLOAD_BYTES_LIMIT.get() >= 0:
+                    ${ _('First %s MB') % DOWNLOAD_BYTES_LIMIT.get() / 1024 / 1024 }
+                  % else:
+                    ${ _('First %s rows') % (hasattr(DOWNLOAD_ROW_LIMIT, 'get') and DOWNLOAD_ROW_LIMIT.get()) }
+                  % endif
+                  </span>
                 </label>
                 <div data-bind="visible: saveTarget() == 'hdfs-file'" class="inline">
                   <input data-bind="value: savePath, valueUpdate: 'afterkeydown', filechooser: { value: savePath, isNestedModal: true }, filechooserOptions: { uploadFile: false, skipInitialPathIfEmpty: true, linkMarkup: true }, hdfsAutocomplete: savePath" type="text" name="target_file" placeholder="${_('Path to CSV file')}" class="pathChooser margin-left-10">

+ 3 - 2
desktop/libs/notebook/src/notebook/connectors/hiveserver2.py

@@ -46,7 +46,7 @@ LOG = logging.getLogger(__name__)
 try:
   from beeswax import conf as beeswax_conf, data_export
   from beeswax.api import _autocomplete, _get_sample_data
-  from beeswax.conf import CONFIG_WHITELIST as hive_settings, DOWNLOAD_ROW_LIMIT
+  from beeswax.conf import CONFIG_WHITELIST as hive_settings, DOWNLOAD_ROW_LIMIT, DOWNLOAD_BYTES_LIMIT
   from beeswax.data_export import upload
   from beeswax.design import hql_query, strip_trailing_semicolon, split_statements
   from beeswax.models import QUERY_TYPES, HiveServerQueryHandle, HiveServerQueryHistory, QueryHistory, Session
@@ -497,8 +497,9 @@ class HS2Api(Api):
 
     handle = self._get_handle(snippet)
     max_rows = DOWNLOAD_ROW_LIMIT.get()
+    max_bytes = DOWNLOAD_BYTES_LIMIT.get()
 
-    upload(target_file, handle, self.request.user, db, self.request.fs, max_rows=max_rows)
+    upload(target_file, handle, self.request.user, db, self.request.fs, max_rows=max_rows, max_bytes=max_bytes)
 
     return '/filebrowser/view=%s' % target_file