Pārlūkot izejas kodu

HUE-9216 [sqlalchemy] Primary keys

Romain 5 gadi atpakaļ
vecāks
revīzija
756907b6e6

+ 11 - 8
desktop/libs/notebook/src/notebook/connectors/sql_alchemy.py

@@ -312,7 +312,7 @@ class SqlAlchemyApi(Api):
         }
         for col in columns
       ]
-      response['foreign_keys'] = assist.get_foreign_keys(database, table)
+      response.update(assist.get_keys(database, table))
     else:
       columns = assist.get_columns(database, table)
       response['name'] = next((col['name'] for col in columns if column == col['name']), '')
@@ -416,13 +416,16 @@ class Assist(object):
     finally:
       connection.close()
 
-  def get_foreign_keys(self, database, table):
+  def get_keys(self, database, table):
     meta = MetaData()
     metaTable = Table(table, meta, schema=database, autoload=True, autoload_with=self.engine)
 
-    return [{
-        'name': fk.parent.name,
-        'to': fk.target_fullname
-      }
-      for fk in metaTable.foreign_keys
-    ]
+    return {
+      'foreign_keys': [{
+          'name': fk.parent.name,
+          'to': fk.target_fullname
+        }
+        for fk in metaTable.foreign_keys
+      ],
+      'primary_keys': [{'name': pk.name} for pk in metaTable.primary_key.columns]
+    }

+ 17 - 31
desktop/libs/notebook/src/notebook/connectors/sql_alchemy_tests.py

@@ -269,7 +269,6 @@ class TestAutocomplete(object):
 
           assert_equal(data['databases'], ['SYSTEM', 'NULL'])
 
-
   def test_columns_with_null_type(self):
     interpreter = {
       'options': {'url': 'phoenix://'}
@@ -296,43 +295,30 @@ class TestAutocomplete(object):
           col2.__getitem__.side_effect = col2_dict
           col2.get = col2_dict
 
-          Assist.return_value=Mock(get_columns=Mock(return_value=[col1, col2]))
+          Assist.return_value=Mock(get_columns=Mock(return_value=[col1, col2]), get_keys=Mock(return_value={}))
 
           data = SqlAlchemyApi(self.user, interpreter).autocomplete(snippet, database='database', table='table')
 
           assert_equal(data['columns'], ['col1', 'col2'])
           assert_equal([col['type'] for col in data['extended_columns']], ['string', 'Null'])
 
-  def test_get_foreign_keys(self):
-
-    interpreter = {
-      'options': {'url': 'phoenix://'}
-    }
-
-    snippet = Mock()
-    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') as Assist:
-          def col1_dict(key):
-            return {
-              'name': 'col1',
-              'type': 'string'
-            }.get(key, Mock())
-          col1 = MagicMock()
-          col1.__getitem__.side_effect = col1_dict
-          col1.get = col1_dict
-          def col2_dict(key):
-            return {
-              'name': 'col2',
-              'type': NullType()
-            }.get(key, Mock())
-          col2 = MagicMock()
-          col2.__getitem__.side_effect = col2_dict
-          col2.get = col2_dict
+  def test_get_keys(self):
 
-          Assist.return_value=Mock(get_columns=Mock(return_value=[col1, col2]))
+    with patch('notebook.connectors.sql_alchemy.Table') as Table:
+      Table.return_value = Mock(
+        foreign_keys=[
+          Mock(
+            parent=Mock(name='col1'),
+            target_fullname='db2.table2.col2'
+          )
+        ],
+        primary_key=Mock(columns=[Mock(name='col2')])
+      )
 
+      db, engine, backticks = Mock(), Mock(), Mock()
+      database, table = Mock(), Mock()
 
+      keys = Assist(db, engine, backticks).get_keys(database, table)
 
-          assert_equal(data['columns'], ['col1', 'col2'])
-          assert_equal([col['type'] for col in data['extended_columns']], ['string', 'Null'])
+      assert_true(keys['primary_keys'])  # For some reason could not mock two level to get some colum names
+      assert_equal(keys['foreign_keys'][0]['to'], 'db2.table2.col2')