浏览代码

GH-1682 [sql_alchemy] adding get_logs() function and unit test

ayush.goyal 4 年之前
父节点
当前提交
ad73a67a22

+ 6 - 2
desktop/libs/notebook/src/notebook/connectors/sql_alchemy.py

@@ -253,7 +253,9 @@ class SqlAlchemyApi(Api):
 
     result = connection.execute(statement)
 
+    logs = [message for message in result.cursor.fetch_logs()] if result.cursor and hasattr(result.cursor, 'fetch_logs') else []
     cache = {
+      'logs': logs,
       'connection': connection,
       'result': result,
       'meta': [
@@ -394,8 +396,10 @@ class SqlAlchemyApi(Api):
 
   @query_error_handler
   def get_log(self, notebook, snippet, startFrom=None, size=None):
-    return ''
-
+    guid = snippet['result']['handle']['guid']
+    cache = CONNECTIONS.get(guid)
+    if cache:
+      return '\n'.join(cache['logs'])
 
   @query_error_handler
   def close_statement(self, notebook, snippet):

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

@@ -339,6 +339,24 @@ class TestApi(object):
           execute.assert_called_with('SELECT 1;')
 
 
+  def test_get_log(self):
+    notebook = Mock()
+    snippet = MagicMock()
+
+    with patch('notebook.connectors.sql_alchemy.CONNECTIONS') as CONNECTIONS:
+      log = [
+        'INFO  : Compiling command(queryId=hive_20210217124246_d3ac774b-cdca-48d2-bfca-f23951ee2097): SELECT 1',
+        'INFO  : Semantic Analysis Completed (retrial = false)',
+        'INFO  : Completed compiling command(queryId=hive_20210217124246_d3ac774b-cdca-48d2-bfca-f23951ee2097); Time taken: 0.092 seconds',
+        'INFO  : Completed executing command(queryId=hive_20210217124246_d3ac774b-cdca-48d2-bfca-f23951ee2097); Time taken: 0.006 seconds',
+        'INFO  : OK'
+      ]
+      CONNECTIONS.get.return_value = {'logs': log}
+
+      data = SqlAlchemyApi(self.user, self.interpreter).get_log(notebook, snippet)
+      assert_equal(data, '\n'.join(log))
+
+
 class TestDialects(object):
 
   def setUp(self):