Bläddra i källkod

[notebook] Escape rows properly for HTML and NULL values

Will take the opportunity to add new tests soon
Romain Rigaux 10 år sedan
förälder
incheckning
a0e0549b7d

+ 2 - 19
apps/beeswax/src/beeswax/views.py

@@ -17,7 +17,6 @@
 
 import json
 import logging
-import math
 import re
 import sys
 import time
@@ -39,9 +38,9 @@ from desktop.lib.django_util import JsonResponse
 from desktop.lib.django_util import copy_query_dict, format_preserving_redirect, render
 from desktop.lib.django_util import login_notrequired, get_desktop_uri_prefix
 from desktop.lib.exceptions_renderable import PopupException
-from desktop.lib.i18n import smart_unicode
 from desktop.models import Document
 from desktop.lib.parameterization import find_variables
+from notebook.models import escape_rows
 
 import beeswax.forms
 import beeswax.design
@@ -477,25 +476,9 @@ def view_results(request, id, first_row=0):
       downloadable = False
     else:
       results = db.fetch(handle, start_over, 100)
-      data = []
 
       # Materialize and HTML escape results
-      # TODO: use Number + list comprehension
-      for row in results.rows():
-        escaped_row = []
-        for field in row:
-          if isinstance(field, (int, long, float, complex, bool)):
-            if math.isnan(field) or math.isinf(field):
-              escaped_field = json.dumps(field)
-            else:
-              escaped_field = field
-          elif field is None:
-            escaped_field = 'NULL'
-          else:
-            field = smart_unicode(field, errors='replace') # Prevent error when getting back non utf8 like charset=iso-8859-1
-            escaped_field = escape(field).replace(' ', ' ')
-          escaped_row.append(escaped_field)
-        data.append(escaped_row)
+      data = escape_rows(results.rows())
 
       # We display the "Download" button only when we know that there are results:
       downloadable = first_row > 0 or data

+ 11 - 0
desktop/libs/notebook/src/notebook/api.py

@@ -28,6 +28,7 @@ from oozie.decorators import check_document_access_permission # Bad dependency
 
 from notebook.connectors.base import get_api, Notebook, QueryExpired
 from notebook.decorators import api_error_handler, check_document_modify_permission
+from notebook.models import escape_rows
 
 
 LOG = logging.getLogger(__name__)
@@ -81,6 +82,11 @@ def execute(request):
   snippet = json.loads(request.POST.get('snippet', '{}'))
 
   response['handle'] = get_api(request.user, snippet, request.fs, request.jt).execute(notebook, snippet)
+
+  # Materialize and HTML escape results
+  if response['handle'].get('sync') and 'data' in response['handle']['result']:
+    response['handle']['result']['data'] = escape_rows(response['handle']['result']['data'])
+
   response['status'] = 0
 
   return JsonResponse(response)
@@ -113,6 +119,11 @@ def fetch_result_data(request):
   start_over = json.loads(request.POST.get('startOver', False))
 
   response['result'] = get_api(request.user, snippet, request.fs, request.jt).fetch_result(notebook, snippet, rows, start_over)
+
+  # Materialize and HTML escape results
+  if 'data' in response['result']:
+    response['result']['data'] = escape_rows(response['result']['data'])
+
   response['status'] = 0
 
   return JsonResponse(response)

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

@@ -110,7 +110,7 @@ class HS2Api(Api):
     # No escaping...
     return {
         'has_more': results.has_more,
-        'data': list(results.rows()),
+        'data': results.rows(),
         'meta': [{
           'name': column.name,
           'type': column.type,

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

@@ -69,7 +69,7 @@ class JdbcApi(Api):
       'has_result_set': True,
       'result': {
         'has_more': False,
-        'data': list(data),
+        'data': data,
         'meta': [{
           'name': col[0],
           'type': col[1],

+ 31 - 0
desktop/libs/notebook/src/notebook/models.py

@@ -14,3 +14,34 @@
 # 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 json
+import math
+import numbers
+
+from django.utils.html import escape
+
+from desktop.lib.i18n import smart_unicode
+
+
+# Materialize and HTML escape results
+def escape_rows(rows):
+  data = []
+
+  for row in rows:
+    escaped_row = []
+    for field in row:
+      if isinstance(field, numbers.Number):
+        if math.isnan(field) or math.isinf(field):
+          escaped_field = json.dumps(field)
+        else:
+          escaped_field = field
+      elif field is None:
+        escaped_field = 'NULL'
+      else:
+        field = smart_unicode(field, errors='replace') # Prevent error when getting back non utf8 like charset=iso-8859-1
+        escaped_field = escape(field).replace(' ', ' ')
+      escaped_row.append(escaped_field)
+    data.append(escaped_row)
+
+  return data