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

[importer] adding excel support for remote file importer

ayush.goyal 4 жил өмнө
parent
commit
4982a291ab

+ 1 - 0
desktop/core/requirements.txt

@@ -44,6 +44,7 @@ Markdown==3.1
 mysqlclient==1.4.6
 nose==1.3.7
 openpyxl==2.6.2
+pandas==1.1.5
 phoenixdb==1.1.0
 prompt-toolkit==2.0.10
 protobuf==3.17.0

+ 38 - 1
desktop/libs/indexer/src/indexer/api3.py

@@ -71,6 +71,7 @@ if sys.version_info[0] > 2:
   from io import StringIO as string_io
   from urllib.parse import urlparse, unquote as urllib_unquote
   from django.utils.translation import gettext as _
+  import pandas as pd
 else:
   from StringIO import StringIO as string_io
   from urllib import unquote as urllib_unquote
@@ -127,7 +128,7 @@ def guess_format(request):
     path = urllib_unquote(file_format["path"])
     if 'xlsx' in path:
       format_ = {
-        "type": "xlsx",
+        "type": "excel",
         "hasHeader": True
       }
     else:
@@ -141,6 +142,20 @@ def guess_format(request):
 
   elif file_format['inputFormat'] == 'file':
     path = urllib_unquote(file_format["path"])
+    if path[-3:] == 'xls' or path[-4:] == 'xlsx':
+      if sys.version_info[0] > 2:
+        file_obj = request.fs.open(path)
+        if path[-3:] == 'xls':
+          df = pd.read_excel(file_obj.read(1024 * 1024 * 1024), engine='xlrd')
+        else:
+          df = pd.read_excel(file_obj.read(1024 * 1024 * 1024), engine='openpyxl')
+        _csv_data = df.to_csv(index=False)
+
+        path = excel_to_csv_file_name_change(path)
+        request.fs.create(path, overwrite=True, data=_csv_data)
+      else:
+        return JsonResponse({'status': -1, 'message': 'Python2 based Hue does not support Excel file importer'})
+
     indexer = MorphlineIndexer(request.user, request.fs)
     if not request.fs.isfile(path):
       raise PopupException(_('Path %(path)s is not a file') % file_format)
@@ -153,6 +168,16 @@ def guess_format(request):
       }
     })
     _convert_format(format_)
+
+    if file_format["path"][-3:] == 'xls' or file_format["path"][-4:] == 'xlsx': 
+      format_ = {
+          "quoteChar": "\"",
+          "recordSeparator": '\\n',
+          "type": "excel",
+          "hasHeader": True,
+          "fieldSeparator": ","
+        }
+
   elif file_format['inputFormat'] == 'table':
     db = dbms.get(request.user)
     try:
@@ -271,6 +296,8 @@ def guess_field_types(request):
   elif file_format['inputFormat'] == 'file':
     indexer = MorphlineIndexer(request.user, request.fs)
     path = urllib_unquote(file_format["path"])
+    if path[-3:] == 'xls' or path[-4:] == 'xlsx':
+      path = excel_to_csv_file_name_change(path)
     stream = request.fs.open(path)
     encoding = check_encoding(stream.read(10000))
     LOG.debug('File %s encoding is %s' % (path, encoding))
@@ -432,6 +459,8 @@ def importer_submit(request):
   if source['inputFormat'] == 'file':
     if source['path']:
       path = urllib_unquote(source['path'])
+      if path[-3:] == 'xls' or path[-4:] == 'xlsx':
+        path = excel_to_csv_file_name_change(path)
       source['path'] = request.fs.netnormpath(path)
       stream = request.fs.open(path)
       file_encoding = check_encoding(stream.read(10000))
@@ -721,6 +750,14 @@ def save_pipeline(request):
   return JsonResponse(response)
 
 
+def excel_to_csv_file_name_change(path):
+  if path[-4:] == 'xlsx':
+    path = path[:-4] + 'csv'
+  elif path[-3:] == 'xls':
+    path = path[:-3] + 'csv'
+  return path
+
+
 def upload_local_file_drag_and_drop(request):
   response = {'status': -1, 'data': ''}
   form = UploadLocalFileForm(request.POST, request.FILES)

+ 23 - 1
desktop/libs/indexer/src/indexer/api3_tests.py

@@ -17,12 +17,13 @@
 
 import json
 import sys
+from nose.plugins.skip import SkipTest
 from nose.tools import assert_equal, assert_true
 from django.utils.datastructures import MultiValueDict
 from django.core.files.uploadhandler import InMemoryUploadedFile
 
 from desktop.settings import BASE_DIR
-from indexer.api3 import upload_local_file, guess_field_types
+from indexer.api3 import upload_local_file, guess_field_types, guess_format
 
 if sys.version_info[0] > 2:
   from urllib.parse import unquote as urllib_unquote
@@ -104,3 +105,24 @@ def test_col_names():
   assert_true('date_1_' in columns_name)
   assert_true('hour_1' in columns_name)
   assert_true('minute' in columns_name)
+
+
+def test_guess_format_excel_remote_file():
+  if sys.version_info[0] > 2:
+    with patch('indexer.api3.pd') as pd:
+      with patch('indexer.api3.MorphlineIndexer') as MorphlineIndexer:
+        file_format = {
+          'inputFormat': 'file',
+          'path': 's3a://gethue/example1.xlsx'
+        }
+        file_format = json.dumps(file_format)
+        request = Mock(
+          POST={'fileFormat': file_format}
+        )
+
+        response = guess_format(request)
+        response = json.loads(response.content)
+
+        assert_equal(response['type'], "excel")
+  else:
+    raise SkipTest

+ 9 - 9
desktop/libs/indexer/src/indexer/file_format.py

@@ -304,15 +304,6 @@ class ParquetFormat(FileFormat):
   _description = _("Parquet Table")
 
 
-class XLSXFormat(GrokkedFormat):
-  _name = "xlsx"
-  _description = _("XLSX File")
-  _args = [
-    CheckboxArgument("hasHeader", "Has Header")
-  ]
-  _extensions = ["xlsx"]
-
-
 class CSVFormat(FileFormat):
   _name = "csv"
   _description = _("CSV File")
@@ -612,6 +603,15 @@ class CSVFormat(FileFormat):
     return fields
 
 
+class XLSXFormat(CSVFormat):
+  _name = "excel"
+  _description = _("Excel File")
+  _args = [
+    CheckboxArgument("hasHeader", "Has Header")
+  ]
+  _extensions = ["xlsx", "xls"]
+
+
 class JsonFormat(CSVFormat):
   _name = "json"
   _description = _("Json")