Quellcode durchsuchen

HUE-8738 [task] Switching to code get reader to load back data

Romain vor 5 Jahren
Ursprung
Commit
dcb2703e8b

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

@@ -77,7 +77,8 @@ class TestApi(object):
     interpreter = {
       'name': 'hive',
       'options': {
-        'url': 'awsathena+rest://XXXXXXXXXXXXXXX:XXXXXXXXXXXXXXXXXXX@athena.us-west-2.amazonaws.com:443/default?s3_staging_dir=s3://gethue-athena/scratch'
+        'url': 'awsathena+rest://XXXXXXXXXXXXXXX:XXXXXXXXXXXXXXXXXXX@athena.us-west-2.amazonaws.com:443/default?'
+            's3_staging_dir=s3://gethue-athena/scratch'
       }
     }
 
@@ -332,7 +333,7 @@ class TestAutocomplete(object):
     with patch('notebook.connectors.sql_alchemy.create_engine') as create_engine:
       with patch('notebook.connectors.sql_alchemy.inspect') as inspect:
         with patch('notebook.connectors.sql_alchemy.Assist') as Assist:
-          Assist.return_value=Mock(get_databases=Mock(return_value=['SYSTEM', None]))
+          Assist.return_value = Mock(get_databases=Mock(return_value=['SYSTEM', None]))
 
           data = SqlAlchemyApi(self.user, interpreter).autocomplete(snippet)
 
@@ -365,7 +366,7 @@ class TestAutocomplete(object):
           col2.__getitem__.side_effect = col2_dict
           col2.get = col2_dict
 
-          Assist.return_value=Mock(get_columns=Mock(return_value=[col1, col2]), get_keys=Mock(return_value={}))
+          Assist.return_value = Mock(get_columns=Mock(return_value=[col1, col2]), get_keys=Mock(return_value={}))
 
           data = SqlAlchemyApi(self.user, interpreter).autocomplete(snippet, database='database', table='table')
 

+ 2 - 2
desktop/libs/notebook/src/notebook/tasks.py

@@ -21,9 +21,9 @@ standard_library.install_aliases()
 
 from builtins import next, object
 
+import codecs
 import csv
 import datetime
-import io
 import json
 import logging
 import sys
@@ -137,7 +137,7 @@ def download_to_file(notebook, snippet, file_format='csv', max_rows=-1, **kwargs
 
     if TASK_SERVER.RESULT_CACHE.get():
       with storage.open(result_key, 'rb') as store:
-        with io.TextIOWrapper(store, encoding='utf-8') as text_file:
+        with codecs.getreader('utf-8')(store) as text_file:
           delimiter = ',' if sys.version_info[0] > 2 else ','.encode('utf-8')
           csv_reader = csv.reader(text_file, delimiter=delimiter)
           caches[CACHES_CELERY_QUERY_RESULT_KEY].set(result_key, [row for row in csv_reader], 60 * 5)