Просмотр исходного кода

HUE-9728 [sqlAlchemy] Add views to get_tables API

ayush.goyal 4 лет назад
Родитель
Сommit
75fde49ab0

+ 7 - 1
desktop/libs/notebook/src/notebook/connectors/sql_alchemy.py

@@ -536,9 +536,15 @@ class Assist(object):
   def get_databases(self):
     return self.db.get_schema_names()
 
-  def get_tables(self, database, table_names=[]):
+  def get_table_names(self, database, table_names=[]):
     return self.db.get_table_names(database)
 
+  def get_view_names(self, database, view_names=[]):
+    return self.db.get_view_names(database)
+
+  def get_tables(self, database, table_names=[]):
+    return self.get_table_names(database) + self.get_view_names(database)
+
   def get_columns(self, database, table):
     return self.db.get_columns(table, database)
 

+ 16 - 0
desktop/libs/notebook/src/notebook/connectors/sql_alchemy_tests.py

@@ -273,6 +273,22 @@ class TestApi(object):
             [{'name': 'col1', 'type': 'STRING_TYPE', 'comment': ''}]
           )
 
+
+  def test_get_tables(self):
+    snippet = MagicMock()
+
+    with patch('notebook.connectors.sql_alchemy.create_engine') as create_engine:
+      with patch('notebook.connectors.sql_alchemy.inspect') as inspect:
+        with patch('notebook.connectors.sql_alchemy.Assist.get_table_names') as get_table_names:
+          with patch('notebook.connectors.sql_alchemy.Assist.get_view_names') as get_view_names:
+            get_table_names.return_value = ['table1']
+            get_view_names.return_value = ['view1']
+
+            response = SqlAlchemyApi(self.user, self.interpreter).autocomplete(snippet, database='database1')
+            assert_equal(response['tables_meta'][0]['name'], 'table1')
+            assert_equal(response['tables_meta'][1]['name'], 'view1')
+
+
   def test_get_sample_data_table(self):
     snippet = Mock()