Bläddra i källkod

HUE-4856 [core] Allow beeswax_install_examples to accept optional DB name

Jenny Kim 9 år sedan
förälder
incheckning
927aa37

+ 11 - 7
apps/beeswax/src/beeswax/management/commands/beeswax_install_examples.py

@@ -46,15 +46,17 @@ class InstallException(Exception):
 
 
 class Command(BaseCommand):
-  args = '<beeswax|impala>'
+  args = '<beeswax|impala> <db_name>'
   help = 'Install examples but do not overwrite them.'
 
   def handle(self, *args, **options):
     if args:
       app_name = args[0]
+      db_name = args[1] if len(args) > 1 else 'default'
       user = User.objects.get(username=pwd.getpwuid(os.getuid()).pw_name)
     else:
       app_name = options['app_name']
+      db_name = options.get('db_name', 'default')
       user = options['user']
 
     tables = options['tables'] if 'tables' in options else 'tables.json'
@@ -65,7 +67,7 @@ class Command(BaseCommand):
     try:
       sample_user = install_sample_user()
       self._install_queries(sample_user, app_name)
-      self._install_tables(user, app_name, tables)
+      self._install_tables(user, app_name, db_name, tables)
     except Exception, ex:
       exception = ex
 
@@ -84,14 +86,14 @@ class Command(BaseCommand):
       else: 
         raise exception
 
-  def _install_tables(self, django_user, app_name, tables):
+  def _install_tables(self, django_user, app_name, db_name, tables):
     data_dir = beeswax.conf.LOCAL_EXAMPLES_DATA_DIR.get()
     table_file = file(os.path.join(data_dir, tables))
     table_list = json.load(table_file)
     table_file.close()
 
     for table_dict in table_list:
-      table = SampleTable(table_dict, app_name)
+      table = SampleTable(table_dict, app_name, db_name)
       try:
         table.install(django_user)
       except Exception, ex:
@@ -118,7 +120,7 @@ class SampleTable(object):
   """
   Represents a table loaded from the tables.json file
   """
-  def __init__(self, data_dict, app_name):
+  def __init__(self, data_dict, app_name, db_name='default'):
     self.name = data_dict['table_name']
     if 'partition_files' in data_dict:
       self.partition_files = data_dict['partition_files']
@@ -128,6 +130,7 @@ class SampleTable(object):
     self.hql = data_dict['create_hql']
     self.query_server = get_query_server_config(app_name)
     self.app_name = app_name
+    self.db_name = db_name
 
     # Sanity check
     self._data_dir = beeswax.conf.LOCAL_EXAMPLES_DATA_DIR.get()
@@ -159,14 +162,15 @@ class SampleTable(object):
     try:
       # Already exists?
       if self.app_name == 'impala':
-        db.invalidate(database='default', flush_all=False)
-      db.get_table('default', self.name)
+        db.invalidate(database=self.db_name, flush_all=False)
+      db.get_table(self.db_name, self.name)
       msg = _('Table "%(table)s" already exists.') % {'table': self.name}
       LOG.error(msg)
       return False
     except Exception:
       query = hql_query(self.hql)
       try:
+        db.use(self.db_name)
         results = db.execute_and_wait(query)
         if not results:
           msg = _('Error creating table %(table)s: Operation timeout.') % {'table': self.name}

+ 2 - 2
apps/beeswax/src/beeswax/tests.py

@@ -1217,7 +1217,7 @@ for x in sys.stdin:
     resp = self.client.get('/beeswax/install_examples')
     assert_true('POST request is required.' in json.loads(resp.content)['message'])
 
-    self.client.post('/beeswax/install_examples')
+    self.client.post('/beeswax/install_examples', {'db_name': self.db_name})
 
     # New tables exists
     resp = self.client.get('/metastore/tables/%s?format=json' % self.db_name)
@@ -1259,7 +1259,7 @@ for x in sys.stdin:
       assert_true('Sample: Customers' in resp.content)
 
       # Now install it a second time, and no error
-      resp = self.client.post('/beeswax/install_examples')
+      resp = self.client.post('/beeswax/install_examples', {'db_name': self.db_name})
       assert_equal(0, json.loads(resp.content)['status'])
       assert_equal('', json.loads(resp.content)['message'])
 

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

@@ -585,7 +585,8 @@ def install_examples(request):
   if request.method == 'POST':
     try:
       app_name = get_app_name(request)
-      beeswax.management.commands.beeswax_install_examples.Command().handle(app_name=app_name, user=request.user)
+      db_name = request.POST.get('db_name', 'default')
+      beeswax.management.commands.beeswax_install_examples.Command().handle(app_name=app_name, db_name=db_name, user=request.user)
       response['status'] = 0
     except Exception, err:
       LOG.exception(err)