Browse Source

[notebook] Add autocomplete for JDBC interface

Romain Rigaux 10 years ago
parent
commit
72ff279947

+ 21 - 2
desktop/libs/librdbms/src/librdbms/jdbc.py

@@ -30,6 +30,22 @@ except ImportError, e:
   LOG.exception('Failed to import py4j')
   LOG.exception('Failed to import py4j')
 
 
 
 
+def query_and_fetch(db, statement, n=None):
+  try:
+    db.connect()
+
+    curs = db.cursor()
+
+    try:
+      curs.execute(statement)
+
+      return curs.fetchmany(n), curs.description
+    finally:
+      curs.close()
+  finally:
+    db.close()  
+
+
 class Jdbc():
 class Jdbc():
 
 
   def __init__(self, driver_name, url, username, password):
   def __init__(self, driver_name, url, username, password):
@@ -72,10 +88,10 @@ class Cursor():
     self.rs = self.stmt.executeQuery(statement)
     self.rs = self.stmt.executeQuery(statement)
     self._meta = self.rs.getMetaData()
     self._meta = self.rs.getMetaData()
 
 
-  def fetchmany(self, n):
+  def fetchmany(self, n=None):
     res = []
     res = []
 
 
-    while self.rs.next() and n > 0:
+    while self.rs.next() and (n is not None and n > 0):
       row = []
       row = []
       for c in xrange(self._meta.getColumnCount()):
       for c in xrange(self._meta.getColumnCount()):
         cell = self.rs.getObject(c + 1)
         cell = self.rs.getObject(c + 1)
@@ -89,6 +105,9 @@ class Cursor():
 
 
     return res
     return res
 
 
+  def fetchall(self, n=None):
+    return self.fetchmany()
+
   @property
   @property
   def description(self):
   def description(self):
     return [[
     return [[

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

@@ -185,6 +185,7 @@ class HS2Api(Api):
   def get_jobs(self, log):
   def get_jobs(self, log):
     return _parse_out_hadoop_jobs(log)
     return _parse_out_hadoop_jobs(log)
 
 
+  @query_error_handler
   def autocomplete(self, snippet, database=None, table=None, column=None, nested=None):
   def autocomplete(self, snippet, database=None, table=None, column=None, nested=None):
     db = self._get_db(snippet)
     db = self._get_db(snippet)
     return _autocomplete(db, database, table, column, nested)
     return _autocomplete(db, database, table, column, nested)

+ 51 - 18
desktop/libs/notebook/src/notebook/connectors/jdbc.py

@@ -21,7 +21,7 @@ from django.utils.translation import ugettext as _
 
 
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.lib.i18n import force_unicode
 from desktop.lib.i18n import force_unicode
-from librdbms.jdbc import Jdbc
+from librdbms.jdbc import Jdbc, query_and_fetch
 
 
 from notebook.connectors.base import Api, QueryError
 from notebook.connectors.base import Api, QueryError
 
 
@@ -44,25 +44,14 @@ def query_error_handler(func):
 
 
 class JdbcApi(Api):
 class JdbcApi(Api):
 
 
-  @query_error_handler
-  def execute(self, notebook, snippet):
-
-    db = Jdbc(self.options['driver'], self.options['url'], self.options['user'], self.options['password'])
-
-    try:
-      db.connect()
-
-      curs = db.cursor()
+  def __init__(self, user, fs=None, jt=None, options=None):
+    Api.__init__(self, user, fs=fs, jt=jt, options=options)
 
 
-      try:
-        curs.execute(snippet['statement'])
+    self.db = Jdbc(self.options['driver'], self.options['url'], self.options['user'], self.options['password'])
 
 
-        data = curs.fetchmany(100)
-        description = curs.description
-      finally:
-        curs.close()
-    finally:
-      db.close()
+  @query_error_handler
+  def execute(self, notebook, snippet):
+    data, description = query_and_fetch(self.db, snippet['statement'], 100)
 
 
     return {
     return {
       'sync': True,
       'sync': True,
@@ -106,3 +95,47 @@ class JdbcApi(Api):
   @query_error_handler
   @query_error_handler
   def close_statement(self, snippet):
   def close_statement(self, snippet):
     return {'status': -1}
     return {'status': -1}
+
+  @query_error_handler
+  def autocomplete(self, snippet, database=None, table=None, column=None, nested=None):
+    assist = Assist(self.db)
+    response = {'error': 0}
+
+    try:
+      if database is None:
+        response['databases'] = assist.get_databases()
+      elif table is None:
+        response['tables'] = assist.get_tables(database)
+      else:
+        columns = assist.get_columns(database, table)
+        response['columns'] = [col[0] for col in columns]
+        response['extended_columns'] = [{
+            'name': col[0],
+            'type': col[1],
+            'comment': col[5]
+          } for col in columns
+        ]
+    except Exception, e:
+      LOG.warn('Autocomplete data fetching error: %s' % e)
+      response['code'] = -1
+      response['error'] = e.message
+
+    return response
+
+
+class Assist():
+
+  def __init__(self, db):
+    self.db = db
+
+  def get_databases(self):
+    databases, description = query_and_fetch(self.db, 'SHOW DATABASES')
+    return databases
+
+  def get_tables(self, database, table_names=[]):
+    tables, description = query_and_fetch(self.db, 'SHOW TABLES')
+    return tables
+
+  def get_columns(self, database, table):
+    columns, description = query_and_fetch(self.db, 'SHOW COLUMNS FROM %s.%s' % (database, table))
+    return columns