Эх сурвалжийг харах

GH-1506 [api] Stream query_store_download_bundle

sreenaths 5 жил өмнө
parent
commit
162cee02de

+ 8 - 49
apps/jobbrowser/src/jobbrowser/api2.py

@@ -18,19 +18,19 @@
 import json
 import logging
 
-from django.http import HttpResponse, FileResponse
+from urllib.request import Request, urlopen
+
+from django.http import HttpResponse
 from django.utils.translation import ugettext as _
 
-from desktop.lib.rest.http_client import RestException
 from desktop.lib.i18n import smart_unicode
 from desktop.lib.django_util import JsonResponse
-from desktop.lib.rest.http_client import HttpClient
-from desktop.lib.rest.resource import Resource
 from desktop.views import serve_403_error
-from beeswax.conf import USE_SASL
 
 from jobbrowser.apis.base_api import get_api
-from jobbrowser.conf import DISABLE_KILLING_JOBS, QUERY_STORE, USE_PROXY
+from jobbrowser.apis.query_store import query_store_proxy, stream_download_bundle
+
+from jobbrowser.conf import DISABLE_KILLING_JOBS, USE_PROXY
 
 LOG = logging.getLogger(__name__)
 
@@ -169,42 +169,13 @@ def profile(request):
     response['status'] = 0
     return JsonResponse(response)
 
-def _query_store_proxy(request, path=None):
-  response = {'status': -1}
-
-  headers = {
-    'x-do-as': request.user.username,
-    'X-Requested-By': 'das',
-    'Content-Type': 'application/json; charset=UTF-8'
-  }
-
-  client = HttpClient(QUERY_STORE.SERVER_URL.get())
-  resource = Resource(client)
-
-  if USE_SASL.get():
-    client.set_kerberos_auth()
-
-  try:
-    response = resource.invoke(request.method, path, request.GET.dict(), request.body, headers)
-  except RestException as e:
-    ex_response = e.get_parent_ex().response
-
-    if ex_response is not None:
-      response['code'] = ex_response.status_code
-      response['message'] = ex_response.reason
-      response['content'] = ex_response.text
-    else:
-      response['message'] = 'Query store not reachable!'
-      response['content'] = e.message
-
-  return response
 
 @api_error_handler
 def query_store_api(request, path=None):
   response = {'status': -1}
 
   if USE_PROXY.get():
-    response = _query_store_proxy(request, path)
+    response = query_store_proxy(request, path)
   else:
     if path == 'api/query/search':
       filters = json.loads(request.body)
@@ -216,16 +187,4 @@ def query_store_api(request, path=None):
 
 @api_error_handler
 def query_store_download_bundle(request, id=None):
-  response = {}
-
-  client = HttpClient(QUERY_STORE.SERVER_URL.get())
-  resource = Resource(client)
-  if USE_SASL.get():
-    client.set_kerberos_auth()
-
-  app = resource.get('api/data-bundle/' + id)
-
-  response = FileResponse((app, 'rb'), content_type='application/octet-stream')
-  response['Content-Disposition'] = 'attachment; filename=' + id + '.zip'
-
-  return response
+  return stream_download_bundle(request, id)

+ 79 - 0
apps/jobbrowser/src/jobbrowser/apis/query_store.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python
+# Licensed to Cloudera, Inc. under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  Cloudera, Inc. licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from desktop.lib.rest.http_client import RestException
+from desktop.lib.rest.http_client import HttpClient
+from desktop.lib.rest.resource import Resource
+
+from desktop.lib.rest.streamed_resource import stream_response
+
+from beeswax.conf import USE_SASL
+from jobbrowser.conf import QUERY_STORE
+
+def error_handler(func):
+  def decorator(*args, **kwargs):
+    response = {}
+
+    try:
+      return func(*args, **kwargs)
+    except RestException as e:
+      response = {'status': -1}
+
+      ex_response = e.get_parent_ex().response
+
+      if ex_response is not None:
+        response['code'] = ex_response.status_code
+        response['message'] = ex_response.reason
+        response['content'] = ex_response.text
+      else:
+        response['message'] = 'Query store not reachable!'
+        response['content'] = e.message
+
+    finally:
+      if response:
+        return response
+
+  return decorator
+
+@error_handler
+def query_store_proxy(request, path=None):
+  client = _create_query_store_client(request)
+
+  resource = Resource(client)
+  return resource.invoke(request.method, path, request.GET.dict(), request.body)
+
+@error_handler
+def stream_download_bundle(request, id):
+  client = _create_query_store_client(request, content_type='application/octet-stream')
+
+  url = 'api/data-bundle/' + id
+  return stream_response(client, url)
+
+def _create_query_store_client(request, content_type='application/json; charset=UTF-8'):
+  headers = {
+    'x-do-as': request.user.username,
+    'X-Requested-By': 'das',
+    'Content-Type': content_type
+  }
+
+  client = HttpClient(QUERY_STORE.SERVER_URL.get())
+  client.set_headers(headers)
+
+  if USE_SASL.get():
+    client.set_kerberos_auth()
+
+  return client

+ 4 - 1
desktop/core/src/desktop/lib/rest/http_client.py

@@ -176,7 +176,7 @@ class HttpClient(object):
     return self._session.headers.copy()
 
   def execute(self, http_method, path, params=None, data=None, headers=None, allow_redirects=False, urlencode=True,
-              files=None, clear_cookies=False, timeout=conf.REST_CONN_TIMEOUT.get()):
+              files=None, stream=False, clear_cookies=False, timeout=conf.REST_CONN_TIMEOUT.get()):
     """
     Submit an HTTP request.
     @param http_method: GET, POST, PUT, DELETE
@@ -187,6 +187,7 @@ class HttpClient(object):
     @param allow_redirects: requests should automatically resolve redirects.
     @param urlencode: percent encode paths.
     @param files: for posting Multipart-Encoded files
+    @param stream: Bool to stream the response
     @param clear_cookies: flag to force clear any cookies set in the current session
 
     @return: The result of urllib2.urlopen()
@@ -211,6 +212,8 @@ class HttpClient(object):
       request_kwargs['data'] = data
     if files:
       request_kwargs['files'] = files
+    if stream:
+      request_kwargs['stream'] = True
 
     if self._cookies and not clear_cookies:
       request_kwargs['cookies'] = self._cookies

+ 33 - 0
desktop/core/src/desktop/lib/rest/streamed_resource.py

@@ -0,0 +1,33 @@
+#!/usr/bin/env python
+# Licensed to Cloudera, Inc. under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  Cloudera, Inc. licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from django.http import StreamingHttpResponse
+
+from desktop.lib.rest.http_client import HttpClient
+
+CHUNK_SIZE = 16 * 1024
+
+def stream_response(client, url):
+  stream_resp = client.execute("GET", url, stream=True);
+  response = StreamingHttpResponse(_stream_generate(stream_resp), content_type=stream_resp.headers.get('Content-Type'))
+  response['Content-Disposition'] = stream_resp.headers.get('Content-Disposition')
+
+  return response
+
+def _stream_generate(zip_resp):
+  for chunk in zip_resp.iter_content(CHUNK_SIZE):
+    yield chunk