Эх сурвалжийг харах

HUE-3228 [editor] Add test for column name backtick escapating

Romain 6 жил өмнө
parent
commit
e573871a68

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

@@ -48,6 +48,7 @@ import datetime
 import json
 import logging
 import uuid
+import re
 import sys
 import textwrap
 
@@ -96,7 +97,7 @@ class SqlAlchemyApi(Api):
   def __init__(self, user, interpreter):
     self.user = user
     self.options = interpreter['options']
-    self.backticks = '"' if self.options['url'].startswith('postgresql://') or self.options['url'].startswith('awsathena') else '`'
+    self.backticks = '"' if re.match('^(postgresql://|awsathena)', self.options.get('url', '')) else '`'
 
   def _create_engine(self):
     if '${' in self.options['url']: # URL parameters substitution

+ 25 - 4
desktop/libs/notebook/src/notebook/connectors/sql_alchemy_tests.py

@@ -42,9 +42,27 @@ class TestApi():
     grant_access("test", "default", "notebook")
 
 
+  def test_column_backticks_escaping(self):
+    interpreter = {
+      'options': {
+        'url': 'mysql://'
+      }
+    }
+    assert_equal(SqlAlchemyApi(self.user, interpreter).backticks, '`')
+
+    interpreter = {
+      'options': {
+        'url': 'postgresql://'
+      }
+    }
+    assert_equal(SqlAlchemyApi(self.user, interpreter).backticks, '"')
+
+
   def test_create_athena_engine(self):
     interpreter = {
-      'options': {"url": "awsathena+rest://XXXXXXXXXXXXXXXXXXXX:XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX@athena.us-west-2.amazonaws.com:443/default?s3_staging_dir=s3://gethue-athena/scratch"}
+      'options': {
+        'url': 'awsathena+rest://XXXXXXXXXXXXXXXXXXXX:XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX@athena.us-west-2.amazonaws.com:443/default?s3_staging_dir=s3://gethue-athena/scratch'
+      }
     }
 
     with patch('notebook.connectors.sql_alchemy.create_engine') as create_engine:
@@ -53,8 +71,9 @@ class TestApi():
 
   def test_fetch_result_empty(self):
     interpreter = {
-      'options': {},
-      'url': 'mysql://hue:localhost@hue:3306/hue'
+      'options': {
+        'url': 'mysql://hue:localhost@hue:3306/hue'
+      },
     }
 
     notebook = Mock()
@@ -87,7 +106,9 @@ class TestApi():
 
   def test_fetch_result_rows(self):
     interpreter = {
-      'options': {}
+      'options': {
+        'url': 'mysql://hue:localhost@hue:3306/hue'
+      }
     }
 
     notebook = Mock()