Răsfoiți Sursa

HUE-9064 [editor] Skeleton of sending back live query results via WS

Romain 6 ani în urmă
părinte
comite
e6e1692ac5

+ 2 - 1
desktop/core/src/desktop/js/apps/notebook2/snippet.js

@@ -1527,7 +1527,8 @@ export default class Snippet {
       result: {}, // TODO: Moved to executor but backend requires it
       database: this.database(),
       compute: this.compute(),
-      wasBatchExecuted: this.wasBatchExecuted()
+      wasBatchExecuted: this.wasBatchExecuted(),
+      editorWsChannel: window.WS_CHANNEL
     });
   }
 

+ 1 - 1
desktop/core/src/desktop/lib/connectors/api.py

@@ -50,7 +50,7 @@ CONNECTOR_TYPES += [
   {'nice_name': "Hive Tez", 'dialect': 'hive-tez', 'interface': 'hiveserver2', 'settings': [{'name': 'server_host', 'value': ''}, {'name': 'server_port', 'value': ''},], 'category': 'editor', 'description': '', 'properties': {'is_sql': True}},
   {'nice_name': "Hive LLAP", 'dialect': 'hive-llap', 'interface': 'hiveserver2', 'settings': [{'name': 'server_host', 'value': ''}, {'name': 'server_port', 'value': ''},], 'category': 'editor', 'description': '', 'properties': {'is_sql': True}},
   {'nice_name': "Druid", 'dialect': 'sql-druid', 'interface': 'sqlalchemy', 'settings': [{'name': 'url', 'value': 'druid://druid-host.com:8082/druid/v2/sql/'}], 'category': 'editor', 'description': '', 'properties': {'is_sql': True}},
-  {'nice_name': "Kafka SQL", 'dialect': 'kafka-sql', 'interface': 'sqlalchemy', 'settings': [], 'category': 'editor', 'description': '', 'properties': {'is_sql': True}},
+  {'nice_name': "Kafka SQL", 'dialect': 'ksql', 'interface': 'ksql', 'settings': [], 'category': 'editor', 'description': '', 'properties': {'is_sql': True}},
   {'nice_name': "SparkSQL", 'dialect': 'spark-sql', 'interface': 'sqlalchemy', 'settings': [], 'category': 'editor', 'description': '', 'properties': {'is_sql': True}},
   {'nice_name': "MySQL", 'dialect': 'mysql', 'interface': 'sqlalchemy', 'settings': [{'name': 'url', 'value': 'mysql://username:password@mysq-host:3306/hue'}], 'category': 'editor', 'description': '', 'properties': {'is_sql': True}},
   {'nice_name': "Presto", 'dialect': 'presto', 'interface': 'sqlalchemy', 'settings': [], 'category': 'editor', 'description': '', 'properties': {'is_sql': True}},

+ 0 - 14
desktop/core/src/desktop/templates/common_notebook_ko_components.mako

@@ -419,20 +419,6 @@ from notebook.conf import ENABLE_SQL_INDEXER
           $('.clipboard-content').empty();
         });
 
-        % if conf.WEBSOCKETS.ENABLED.get():
-        var chatSocket = new WebSocket('ws://' + window.location.host + '/ws/editor/results/' + self.snippet.id() + '/');
-
-        chatSocket.onmessage = function(e) {
-          var data = JSON.parse(e.data);
-          var message = data['message'];
-          console.log(message);
-        };
-
-        chatSocket.onclose = function(e) {
-          console.error('Chat socket closed unexpectedly');
-        };
-        % endif
-
         self.trySaveResults = function () {
           if (self.isValidDestination()) {
             self.saveResults();

+ 30 - 5
desktop/libs/kafka/src/kafka/ksql_client.py

@@ -25,9 +25,13 @@ from django.utils.translation import ugettext as _
 
 from desktop.lib.i18n import smart_unicode
 from desktop.lib.rest.http_client import RestException
+from desktop.conf import has_channels
 
 from kafka.conf import KAFKA
 
+if has_channels():
+  from notebook.consumer import _send_to_channel
+
 
 LOG = logging.getLogger(__name__)
 
@@ -48,6 +52,10 @@ class KSqlApi(object):
   https://pypi.org/project/ksql/
 
   pip install ksql
+
+  https://github.com/bryanyang0528/ksql-python/pull/60 fixes:
+  - STREAMS requires a LIMIT currently or will hang or run forever
+  - https://github.com/bryanyang0528/ksql-python/issues/57
   """
 
   def __init__(self, user=None, security_enabled=False, ssl_cert_ca_verify=False):
@@ -76,15 +84,25 @@ class KSqlApi(object):
     return response[0]
 
 
-  def query(self, statement):
+  def query(self, statement, channel_name=None):
     data = []
     metadata = []
 
     is_select = statement.strip().lower().startswith('select')
     if is_select or statement.strip().lower().startswith('print'):
-      # STREAMS requires a LIMIT currently or will hang without https://github.com/bryanyang0528/ksql-python/pull/60
+
       result = self.client.query(statement)
-      for line in ''.join(list(result)).split('\n'): # Until https://github.com/bryanyang0528/ksql-python/issues/57
+
+      metadata = [['Row', 'STRING']]
+
+      if has_channels() and channel_name:
+        _send_to_channel(
+            channel_name,
+            message_type='task.progress',
+            message_data={'status': 'running', 'query_id': 1111}
+        )
+
+      for line in result:
         # columns = line.keys()
         # data.append([line[col] for col in columns])
         if is_select and line: # Empty first 2 lines?
@@ -93,8 +111,15 @@ class KSqlApi(object):
             data.append(data_line['row']['columns'])
         else:
           data.append([line])
-        # TODO: WS to plug-in
-      metadata = [['Row', 'STRING']]
+
+        if has_channels() and channel_name:
+          _send_to_channel(
+              channel_name,
+              message_type='task.result',
+              message_data={'data': data, 'metadata': metadata, 'query_id': 1111}
+          )
+          # TODO: special message when end of stream
+          data = []
     else:
       data, metadata = self._decode_result(
         self.ksql(statement)

+ 5 - 1
desktop/libs/notebook/src/notebook/connectors/ksql.py

@@ -51,7 +51,11 @@ class KSqlApi(Api):
 
   @query_error_handler
   def execute(self, notebook, snippet):
-    data, description = self.db.query(snippet['statement'])
+
+    data, description = self.db.query(
+        snippet['statement'],
+        channel_name=snippet.get('editorWsChannel')
+    )
     has_result_set = data is not None
 
     return {

+ 34 - 11
desktop/libs/notebook/src/notebook/consumer.py

@@ -16,11 +16,18 @@
 # limitations under the License.
 
 import json
+import logging
 
 from desktop.conf import has_channels
 
+
+LOG = logging.getLogger(__name__)
+
+
 if has_channels():
+  from asgiref.sync import async_to_sync
   from channels.generic.websocket import AsyncWebsocketConsumer
+  from channels.layers import get_channel_layer
 
 
   class EditorConsumer(AsyncWebsocketConsumer):
@@ -28,23 +35,39 @@ if has_channels():
     async def connect(self):
       await self.accept()
 
-      await self.send(text_data=json.dumps({
-        'type': 'channel_name',
-        'data': self.channel_name,
-        'accept': True
-      }))
+      LOG.info('User %(user)s connected to WS Editor.' % self.scope)
+
+      await self.send(
+        text_data=json.dumps({
+          'type': 'channel_name',
+          'data': self.channel_name,
+          'accept': True
+        })
+      )
 
 
     async def task_progress(self, event):
-      await self.send(text_data=json.dumps({
-        'type': 'task_progress',
-        'data': event["data"]
-      }))
+      await self.send(
+        text_data=json.dumps({
+          'type': 'query_progress',
+          'data': event["data"]
+        })
+      )
+
+    async def task_result(self, event):
+      await self.send(
+        text_data=json.dumps({
+          'type': 'query_result',
+          'data': event["data"]
+        })
+      )
 
 
   def _send_to_channel(channel_name, message_type, message_data):
     channel_layer = get_channel_layer()
-    async_to_sync(channel_layer.send)(channel_name, {
+    async_to_sync(channel_layer.send)(
+      channel_name, {
         "type": message_type,
         "data": message_data,
-    })
+      }
+    )

+ 20 - 0
desktop/libs/notebook/src/notebook/templates/editor_components2.mako

@@ -1694,6 +1694,26 @@
       }
     });
 
+    % if conf.WEBSOCKETS.ENABLED.get():
+        var editorWs = new WebSocket('ws://' + window.location.host + '/ws/editor/results/' + 'userA' + '/');
+
+        editorWs.onopen = function(e) {
+          console.info('Notification socket open.');
+        };
+
+        editorWs.onmessage = function(e) {
+          var data = JSON.parse(e.data);
+          if (data['type'] == 'channel_name') {
+            window.WS_CHANNEL = data['data'];
+          }
+          console.log(data);
+        };
+
+        editorWs.onclose = function(e) {
+          console.error('Chat socket closed unexpectedly');
+        };
+    % endif
+
     window.EDITOR_ENABLE_QUERY_SCHEDULING = '${ ENABLE_QUERY_SCHEDULING.get() }' === 'True';
 
     window.EDITOR_ID = ${ editor_id or 'null' };