Ver Fonte

HUE-8758 [sqlalchemy] Add support for connector semicolon statement trim

Romain há 5 anos atrás
pai
commit
0a33d80099

+ 21 - 0
desktop/core/src/desktop/lib/connectors/types.py

@@ -52,6 +52,7 @@ CONNECTOR_TYPES = [
       'has_auto_limit': False,
       'has_reference_language': True,
       'has_reference_functions': True,
+      'trim_statement_semicolon': False,
     }
   },
   {
@@ -78,6 +79,7 @@ CONNECTOR_TYPES = [
       'has_auto_limit': False,
       'has_reference_language': True,
       'has_reference_functions': True,
+      'trim_statement_semicolon': False,
     }
   },
   {
@@ -104,6 +106,7 @@ CONNECTOR_TYPES = [
       'has_auto_limit': False,
       'has_reference_language': False,
       'has_reference_functions': False,
+      'trim_statement_semicolon': False,
     }
   },
   {
@@ -130,6 +133,7 @@ CONNECTOR_TYPES = [
       'has_auto_limit': False,
       'has_reference_language': False,
       'has_reference_functions': False,
+      'trim_statement_semicolon': False,
     }
   },
   {
@@ -156,6 +160,7 @@ CONNECTOR_TYPES = [
       'has_auto_limit': False,
       'has_reference_language': False,
       'has_reference_functions': False,
+      'trim_statement_semicolon': False,
     }
   },
   {
@@ -182,6 +187,7 @@ CONNECTOR_TYPES = [
       'has_auto_limit': False,
       'has_reference_language': False,
       'has_reference_functions': False,
+      'trim_statement_semicolon': False,
     }
   },
   {
@@ -208,6 +214,7 @@ CONNECTOR_TYPES = [
       'has_auto_limit': False,
       'has_reference_language': False,
       'has_reference_functions': False,
+      'trim_statement_semicolon': False,
     }
   },
   {
@@ -234,6 +241,7 @@ CONNECTOR_TYPES = [
       'has_auto_limit': False,
       'has_reference_language': False,
       'has_reference_functions': False,
+      'trim_statement_semicolon': False,
     }
   },
   {
@@ -260,6 +268,7 @@ CONNECTOR_TYPES = [
       'has_auto_limit': False,
       'has_reference_language': False,
       'has_reference_functions': False,
+      'trim_statement_semicolon': False,
     }
   },
   {
@@ -286,6 +295,7 @@ CONNECTOR_TYPES = [
       'has_auto_limit': False,
       'has_reference_language': False,
       'has_reference_functions': False,
+      'trim_statement_semicolon': False,
     }
   },
   {
@@ -312,6 +322,7 @@ CONNECTOR_TYPES = [
       'has_auto_limit': False,
       'has_reference_language': False,
       'has_reference_functions': False,
+      'trim_statement_semicolon': True,
     }
   },
   {
@@ -338,6 +349,7 @@ CONNECTOR_TYPES = [
       'has_auto_limit': False,
       'has_reference_language': False,
       'has_reference_functions': False,
+      'trim_statement_semicolon': False,
     }
   },
   {
@@ -365,6 +377,7 @@ CONNECTOR_TYPES = [
       'has_auto_limit': False,
       'has_reference_language': False,
       'has_reference_functions': False,
+      'trim_statement_semicolon': False,
     }
   },
   {
@@ -389,6 +402,7 @@ CONNECTOR_TYPES = [
       'has_auto_limit': False,
       'has_reference_language': False,
       'has_reference_functions': False,
+      'trim_statement_semicolon': False,
     }
   },
   {
@@ -413,6 +427,7 @@ CONNECTOR_TYPES = [
       'has_auto_limit': False,
       'has_reference_language': False,
       'has_reference_functions': False,
+      'trim_statement_semicolon': False,
     }
   },
   {
@@ -437,6 +452,7 @@ CONNECTOR_TYPES = [
       'has_auto_limit': False,
       'has_reference_language': False,
       'has_reference_functions': False,
+      'trim_statement_semicolon': False,
     }
   },
   {
@@ -461,6 +477,7 @@ CONNECTOR_TYPES = [
       'has_auto_limit': False,
       'has_reference_language': False,
       'has_reference_functions': False,
+      'trim_statement_semicolon': False,
     }
   },
   {
@@ -487,6 +504,7 @@ CONNECTOR_TYPES = [
       'has_auto_limit': False,
       'has_reference_language': False,
       'has_reference_functions': False,
+      'trim_statement_semicolon': False,
     }
   },
   {
@@ -513,6 +531,7 @@ CONNECTOR_TYPES = [
       'has_auto_limit': False,
       'has_reference_language': False,
       'has_reference_functions': False,
+      'trim_statement_semicolon': False,
     }
   },
   {
@@ -539,6 +558,7 @@ CONNECTOR_TYPES = [
       'has_auto_limit': False,
       'has_reference_language': False,
       'has_reference_functions': False,
+      'trim_statement_semicolon': False,
     }
   },
   {
@@ -568,6 +588,7 @@ CONNECTOR_TYPES = [
       'has_auto_limit': False,
       'has_reference_language': False,
       'has_reference_functions': False,
+      'trim_statement_semicolon': False,
     }
   },
 

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

@@ -202,8 +202,13 @@ class SqlAlchemyApi(Api):
 
     engine = self._get_engine()
     connection = engine.connect()
+    statement = snippet['statement']
 
-    result = connection.execute(snippet['statement'])
+    if self.options['url'].startswith('presto://') or \
+        self.interpreter.get('dialect_properties') and interpreter['dialect_properties']['trim_statement_semicolon']:
+      statement = statement.strip().rstrip(';')
+
+    result = connection.execute(statement)
 
     cache = {
       'connection': connection,

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

@@ -246,6 +246,41 @@ class TestApi(object):
           assert_equal(response['rows'], [[1], [2]])
 
 
+  def test_dialect_trim_statement_semicolon(self):
+    interpreter = {
+      'name': 'presto',
+      'options': {
+        'url': 'presto://hue:8080/hue',
+        'session': {},
+      }
+    }
+
+    with patch('notebook.connectors.sql_alchemy.SqlAlchemyApi._create_engine') as _create_engine:
+      with patch('notebook.connectors.sql_alchemy.SqlAlchemyApi._get_session') as _get_session:
+        execute = Mock(return_value=Mock(cursor=None))
+        _create_engine.return_value = Mock(
+          connect=Mock(
+            return_value=Mock(
+              execute=execute
+            )
+          )
+        )
+        notebook = {}
+        snippet = {'statement': 'SELECT 1;'}
+
+        # Trim
+        engine = SqlAlchemyApi(self.user, interpreter).execute(notebook, snippet)
+
+        execute.assert_called_with('SELECT 1')
+
+        # No Trim
+        interpreter['options']['url'] = 'mysql://hue:3306/hue'
+
+        engine = SqlAlchemyApi(self.user, interpreter).execute(notebook, snippet)
+
+        execute.assert_called_with('SELECT 1;')
+
+
 class TestDialects(object):
 
   def setUp(self):