Browse Source

Implement download for jdbc notebook connector

Rentao Wu 6 years ago
parent
commit
ca2ea6c5e3
1 changed files with 28 additions and 3 deletions
  1. 28 3
      desktop/libs/notebook/src/notebook/connectors/jdbc.py

+ 28 - 3
desktop/libs/notebook/src/notebook/connectors/jdbc.py

@@ -20,11 +20,11 @@ import sys
 
 from django.utils.translation import ugettext as _
 
-from desktop.lib.exceptions_renderable import PopupException
+from beeswax import data_export
 from desktop.lib.i18n import force_unicode, smart_str
 from librdbms.jdbc import Jdbc, query_and_fetch
 
-from notebook.connectors.base import Api, QueryError, AuthenticationRequired
+from notebook.connectors.base import Api, QueryError, AuthenticationRequired, _get_snippet_name
 
 
 LOG = logging.getLogger(__name__)
@@ -125,7 +125,10 @@ class JdbcApi(Api):
     return {'status': 0}
 
   def download(self, notebook, snippet, format, user_agent=None, max_rows=None, store_data_type_in_header=False):
-    raise PopupException('Downloading is not supported yet')
+    file_name = _get_snippet_name(notebook)
+    data, description = query_and_fetch(self.db, snippet['statement'])
+    db = FixedResult(data, description)
+    return data_export.download(None, format, db, id=snippet['id'], file_name=file_name)
 
   def progress(self, snippet, logs):
     return 50
@@ -219,3 +222,25 @@ class Assist():
     #response['rows'] = data
     #response['columns'] = []
     return query_and_fetch(self.db, 'SELECT %s FROM %s.%s limit 100' % (column, database, table))
+
+class FixedResultSet():
+
+  def __init__(self, data, metadata):
+    self.data = data
+    self.metadata = metadata
+    self.has_more = False
+
+  def cols(self):
+    return [str(col[0]) for col in self.metadata]
+
+  def rows(self):
+    return self.data if self.data is not None else []
+
+class FixedResult():
+
+  def __init__(self, data, metadata):
+    self.data = data
+    self.metadata = metadata
+
+  def fetch(self, handle=None, start_over=None, rows=None):
+    return FixedResultSet(self.data, self.metadata)