Ver código fonte

HUE-7 [importer] Accept to read gzip files in source

Romain Rigaux 8 anos atrás
pai
commit
50f1ffc
1 arquivos alterados com 78 adições e 35 exclusões
  1. 78 35
      desktop/libs/indexer/src/indexer/file_format.py

+ 78 - 35
desktop/libs/indexer/src/indexer/file_format.py

@@ -14,12 +14,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.import logging
 import csv
+import gzip
 import operator
 import itertools
 import logging
 
 from django.utils.translation import ugettext as _
 
+from desktop.lib import i18n
+
 from indexer.argument import CheckboxArgument, TextDelimiterArgument
 from indexer.conf import ENABLE_NEW_INDEXER
 from indexer.fields import Field, guess_field_type_from_samples
@@ -29,6 +32,10 @@ from indexer.operations import get_operator
 LOG = logging.getLogger(__name__)
 
 
+IMPORT_PEEK_SIZE = 1024 * 1024 * 5
+IMPORT_PEEK_NLINES = 20
+
+
 def get_format_types():
   formats = [CSVFormat]
 
@@ -331,48 +338,56 @@ class CSVFormat(FileFormat):
 
   @classmethod
   def _guess_from_file_stream(cls, file_stream):
-    sample = cls._get_sample(file_stream)
-
-    try:
-      dialect, has_header = cls._guess_dialect(sample)
-      delimiter = dialect.delimiter
-      line_terminator = dialect.lineterminator
-      quote_char = dialect.quotechar
-    except Exception:
-      # guess dialect failed, fall back to defaults:
-      return cls()
-
-    return cls(**{
-      "delimiter": delimiter,
-      "line_terminator": line_terminator,
-      "quote_char": quote_char,
-      "has_header": has_header,
-      "sample": sample
-    })
+    for sample in  cls._get_sample(file_stream):
+      try:
+        dialect, has_header = cls._guess_dialect(sample)
+        delimiter = dialect.delimiter
+        line_terminator = dialect.lineterminator
+        quote_char = dialect.quotechar
+
+        return cls(**{
+          "delimiter": delimiter,
+          "line_terminator": line_terminator,
+          "quote_char": quote_char,
+          "has_header": has_header,
+          "sample": sample
+        })
+      except Exception:
+        LOG.exception('Warning, cannot read the file format.')
+
+    # Guess dialect failed, fall back to defaults:
+    return cls()
 
   @classmethod
   def _get_sample(cls, file_stream):
-    file_stream.seek(0)
-    sample = '\n'.join(file_stream.read(1024 * 1024 * 5).splitlines()[:20])
-    file_stream.seek(0)
+    encoding = i18n.get_site_encoding()
 
-    return sample
+    for reader in [TextFileReader, GzipFileReader]:
+      file_stream.seek(0)
+      lines = reader.readlines(file_stream, encoding)
+      file_stream.seek(0)
+
+      if lines is not None:
+        yield '\n'.join(lines)
 
   @classmethod
   def _from_format(cls, file_stream, format_):
-    sample = cls._get_sample(file_stream)
-
-    delimiter = format_["fieldSeparator"].encode('utf-8')
-    line_terminator = format_["recordSeparator"].encode('utf-8')
-    quote_char = format_["quoteChar"].encode('utf-8')
-    has_header = format_["hasHeader"]
-    return cls(**{
-      "delimiter": delimiter,
-      "line_terminator": line_terminator,
-      "quote_char": quote_char,
-      "has_header": has_header,
-      "sample": sample
-    })
+    for sample in cls._get_sample(file_stream):
+      try:
+        delimiter = format_["fieldSeparator"].encode('utf-8')
+        line_terminator = format_["recordSeparator"].encode('utf-8')
+        quote_char = format_["quoteChar"].encode('utf-8')
+        has_header = format_["hasHeader"]
+
+        return cls(**{
+          "delimiter": delimiter,
+          "line_terminator": line_terminator,
+          "quote_char": quote_char,
+          "has_header": has_header,
+          "sample": sample
+        })
+      except Exception:
+        LOG.exception('Warning, cannot read the file format.')
 
   @classmethod
   def get_instance(cls, file_stream, format_):
@@ -483,6 +498,34 @@ class CSVFormat(FileFormat):
     return fields
 
 
+class GzipFileReader(object):
+  TYPE = 'gzip'
+
+  @staticmethod
+  def readlines(fileobj, encoding):
+    gz = gzip.GzipFile(fileobj=fileobj, mode='rb')
+    try:
+      data = gz.read(IMPORT_PEEK_SIZE)
+    except IOError:
+      return None
+    try:
+      return data.splitlines()[:IMPORT_PEEK_NLINES]
+    except UnicodeError:
+      return None
+
+
+class TextFileReader(object):
+  TYPE = 'text'
+
+  @staticmethod
+  def readlines(fileobj, encoding):
+    try:
+      data = fileobj.read(IMPORT_PEEK_SIZE)
+      return data.splitlines()[:IMPORT_PEEK_NLINES]
+    except UnicodeError:
+      return None
+
+
 class HiveFormat(CSVFormat):
   FIELD_TYPE_TRANSLATE = {
     "BOOLEAN_TYPE": "string",