Prechádzať zdrojové kódy

[flink] Improve autocomplete for functions (#4090)

* [flink] Fix show functions when no database is specified

* [flink] Handle autocomplete function operation

---------

Co-authored-by: Ayush Goyal <36241930+agl29@users.noreply.github.com>
Grzegorz Kołakowski 7 mesiacov pred
rodič
commit
72e4f9c39a

+ 24 - 7
desktop/libs/notebook/src/notebook/connectors/flink_sql.py

@@ -180,6 +180,7 @@ class FlinkSqlApi(Api):
   def _create_session(self):
     session = self.db.create_session()
     session['id'] = session['sessionHandle']
+    session['flink_version'] = self.db.info().get('version')
 
     if self.default_database:
       self._use_database(session, self.default_catalog, self.default_database)
@@ -294,6 +295,10 @@ class FlinkSqlApi(Api):
 
     if operation == 'functions':
       response['functions'] = self._show_functions(database)
+    elif operation == 'function':
+      # When function signature is requested the following call is executed:
+      # autocomplete(db=<function_name>, table=None, col=None, nested=None, op=function)
+      response['function'] = self._show_function(database)
     elif database is None:
       response['databases'] = self._show_databases()
     elif table is None:
@@ -457,15 +462,27 @@ class FlinkSqlApi(Api):
 
   def _show_functions(self, database):
     session = self._get_session()
-    session_handle = session['id']
-
-    operation_handle = self.db.execute_statement(
-      session_handle=session_handle,
-      statement='SHOW FUNCTIONS IN `%(database)s`' % {'database': database})
-    function_list = self._check_status_and_fetch_result(session_handle, operation_handle['operationHandle'])
-
+    statement = 'SHOW FUNCTIONS IN `%(database)s`' % {'database': database} if database else 'SHOW FUNCTIONS'
+    operation_handle = self.db.execute_statement(session['id'], statement)
+    function_list = self._check_status_and_fetch_result(session['id'], operation_handle['operationHandle'])
     return [{'name': function[0]} for function in function_list]
 
+  def _show_function(self, function_name):
+    session = self._get_session()
+    if session.get('flink_version') and session['flink_version'].startswith('2.'):
+      # Describe function is available as of Flink 2.0.
+      operation_handle = self.db.execute_statement(
+        session_handle=session['id'],
+        statement='DESCRIBE FUNCTION EXTENDED %(function_name)s' % {'function_name': function_name})
+      properties = dict(self._check_status_and_fetch_result(session['id'], operation_handle['operationHandle']))
+
+      return {
+        'name': function_name,
+        'signature': properties.get('signature'),
+      }
+    else:
+      return {'name': function_name}
+
   def _use_catalog(self, session, catalog):
     self.db.configure_session(session['id'], "USE CATALOG `%s`" % catalog)
 

+ 120 - 7
desktop/libs/notebook/src/notebook/connectors/flink_sql_tests.py

@@ -24,9 +24,12 @@ from useradmin.models import User
 
 
 class TestFlinkApi(TestCase):
+  TEST_SESSION_HANDLE = '657c12d4-5509-477f-a460-ea6af927906d'
+  TEST_OPERATION_HANDLE = '1f6922e4-ec0c-4307-947b-7502757edf1a'
+
   def setup_method(self, test_method):
-    self.client = make_logged_in_client(username="hue_test", groupname="default", recreate=True, is_superuser=False)
-    self.user = User.objects.get(username="hue_test")
+    self.client = make_logged_in_client(username='hue_test', groupname='default', recreate=True, is_superuser=False)
+    self.user = User.objects.get(username='hue_test')
     self.interpreter = {
       'options': {
         'url': 'https://example.com:8081',
@@ -39,7 +42,8 @@ class TestFlinkApi(TestCase):
     # given: mock interactions
     mock_client_instance = MagicMock()
     client_mock.return_value = mock_client_instance
-    mock_client_instance.create_session.return_value = {'sessionHandle': '657c12d4-5509-477f-a460-ea6af927906d'}
+    mock_client_instance.create_session.return_value = {'sessionHandle': self.TEST_SESSION_HANDLE}
+    mock_client_instance.info.return_value = {'version': '2.0.0'}
 
     # and: FlinkSqlApi instance
     flink_api = FlinkSqlApi(self.user, interpreter=self.interpreter)
@@ -48,7 +52,7 @@ class TestFlinkApi(TestCase):
     created_session = flink_api.create_session(lang='flink', properties=None)
 
     # then
-    assert created_session == {'id': '657c12d4-5509-477f-a460-ea6af927906d', 'type': 'flink'}
+    assert created_session == {'id': self.TEST_SESSION_HANDLE, 'type': 'flink'}
     assert mock_client_instance.session_heartbeat.call_count == 1
 
   @patch('notebook.connectors.flink_sql.FlinkSqlClient')
@@ -56,7 +60,8 @@ class TestFlinkApi(TestCase):
     # given: mock interactions
     mock_client_instance = MagicMock()
     client_mock.return_value = mock_client_instance
-    mock_client_instance.create_session.return_value = {'sessionHandle': '657c12d4-5509-477f-a460-ea6af927906d'}
+    mock_client_instance.create_session.return_value = {'sessionHandle': self.TEST_SESSION_HANDLE}
+    mock_client_instance.info.return_value = {'version': '2.0.0'}
 
     # and: FlinkSqlApi instance with configuration
     self.interpreter['options']['default_catalog'] = 'default_catalog'
@@ -67,8 +72,116 @@ class TestFlinkApi(TestCase):
     created_session = flink_api.create_session(lang='flink', properties=None)
 
     # then
-    assert created_session == {'id': '657c12d4-5509-477f-a460-ea6af927906d', 'type': 'flink'}
+    assert created_session == {'id': self.TEST_SESSION_HANDLE, 'type': 'flink'}
     mock_client_instance.configure_session.assert_called_once_with(
-      '657c12d4-5509-477f-a460-ea6af927906d', "USE `default_catalog`.`default_database`"
+      self.TEST_SESSION_HANDLE, 'USE `default_catalog`.`default_database`'
     )
     assert mock_client_instance.session_heartbeat.call_count == 1
+
+  @patch('notebook.connectors.flink_sql.FlinkSqlClient')
+  def test_autocomplete_operation_functions(self, client_mock):
+    # given: mock interactions
+    mock_client_instance = MagicMock()
+    client_mock.return_value = mock_client_instance
+    mock_client_instance.create_session.return_value = {'sessionHandle': self.TEST_SESSION_HANDLE}
+    mock_client_instance.info.return_value = {'version': '2.0.0'}
+    mock_client_instance.execute_statement.return_value = {'operationHandle': self.TEST_OPERATION_HANDLE}
+    mock_client_instance.fetch_results.return_value = {
+      'resultType': 'PAYLOAD',
+      'resultKind': 'SUCCESS_WITH_CONTENT',
+      'results': {
+        'columns': [{'name': 'function name', 'logicalType': {'type': 'VARCHAR', 'nullable': True, 'length': 1000}}],
+        'rowFormat': 'JSON',
+        'data': [
+          {'kind': 'INSERT', 'fields': ['lower']},
+          {'kind': 'INSERT', 'fields': ['upper']}
+        ]},
+      'nextResultUri': f'/v3/sessions/{self.TEST_SESSION_HANDLE}/operations/{self.TEST_OPERATION_HANDLE}/result/1?rowFormat=JSON'
+    }
+
+    # and: FlinkSqlApi instance with configuration
+    flink_api = FlinkSqlApi(self.user, interpreter=self.interpreter)
+
+    # and: session is created
+    flink_api.create_session(lang='flink', properties=None)
+
+    # when
+    autocomplete_result = flink_api.autocomplete(snippet='dummy', database=None, table=None, column=None, nested=None,
+                                                 operation='functions')
+
+    # then
+    mock_client_instance.execute_statement.assert_called_once_with(self.TEST_SESSION_HANDLE, 'SHOW FUNCTIONS')
+    assert autocomplete_result == {'functions': [{'name': 'lower'}, {'name': 'upper'}]}
+
+  @patch('notebook.connectors.flink_sql.FlinkSqlClient')
+  def test_autocomplete_operation_function_flink_1_x(self, client_mock):
+    # given: mock interactions
+    mock_client_instance = MagicMock()
+    client_mock.return_value = mock_client_instance
+    mock_client_instance.create_session.return_value = {'sessionHandle': self.TEST_SESSION_HANDLE}
+    mock_client_instance.info.return_value = {'version': '1.20.0'}
+
+    # and: FlinkSqlApi instance with configuration
+    flink_api = FlinkSqlApi(self.user, interpreter=self.interpreter)
+
+    # and: session is created
+    flink_api.create_session(lang='flink', properties=None)
+
+    # when
+    autocomplete_result = flink_api.autocomplete(snippet='dummy', database='test_function', table=None, column=None,
+                                                 nested=None, operation='function')
+
+    # then
+    assert autocomplete_result == {'function': {'name': 'test_function'}}
+
+  @patch('notebook.connectors.flink_sql.FlinkSqlClient')
+  def test_autocomplete_operation_function_flink_2_x(self, client_mock):
+    # given: mock interactions
+    mock_client_instance = MagicMock()
+    client_mock.return_value = mock_client_instance
+    mock_client_instance.create_session.return_value = {'sessionHandle': self.TEST_SESSION_HANDLE}
+    mock_client_instance.info.return_value = {'version': '2.0.0'}
+    mock_client_instance.execute_statement.return_value = {'operationHandle': self.TEST_OPERATION_HANDLE}
+    mock_client_instance.fetch_results.return_value = {
+      'resultType': 'PAYLOAD',
+      'resultKind': 'SUCCESS_WITH_CONTENT',
+      'results': {
+        'columns': [
+          {'name': 'info name', 'logicalType': {'type': 'VARCHAR', 'nullable': True, 'length': 2147483647}},
+          {'name': 'info value', 'logicalType': {'type': 'VARCHAR', 'nullable': True, 'length': 2147483647}}
+        ],
+        'rowFormat': 'JSON',
+        'data': [
+          {'kind': 'INSERT', 'fields': ['is system function', 'false']},
+          {'kind': 'INSERT', 'fields': ['is temporary', 'false']},
+          {'kind': 'INSERT', 'fields': ['class name', 'com.example.flink.udf.TestFunction']},
+          {'kind': 'INSERT', 'fields': ['function language', 'JAVA']},
+          {'kind': 'INSERT', 'fields': ['resource uris', '[]']},
+          {'kind': 'INSERT', 'fields': ['kind', 'SCALAR']},
+          {'kind': 'INSERT', 'fields': ['requirements', '[]']},
+          {'kind': 'INSERT', 'fields': ['is deterministic', 'true']},
+          {'kind': 'INSERT', 'fields': ['supports constant folding', 'true']},
+          {'kind': 'INSERT', 'fields': ['signature', 'default_catalog.default_db.test_function(values <ANY>...)']},
+        ]
+      },
+      'nextResultUri': f'/v3/sessions/{self.TEST_SESSION_HANDLE}/operations/{self.TEST_OPERATION_HANDLE}/result/1?rowFormat=JSON'
+    }
+
+    # and: FlinkSqlApi instance with configuration
+    flink_api = FlinkSqlApi(self.user, interpreter=self.interpreter)
+
+    # and: session is created
+    flink_api.create_session(lang='flink', properties=None)
+
+    # when
+    autocomplete_result = flink_api.autocomplete(snippet='dummy', database='test_function', table=None, column=None,
+                                                 nested=None, operation='function')
+
+    # then
+    mock_client_instance.execute_statement.assert_called_once_with(
+      session_handle=self.TEST_SESSION_HANDLE,
+      statement='DESCRIBE FUNCTION EXTENDED test_function'
+    )
+    assert autocomplete_result == {
+      'function': {'name': 'test_function', 'signature': 'default_catalog.default_db.test_function(values <ANY>...)'}
+    }