api3.py 6.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170
  1. #!/usr/bin/env python
  2. # Licensed to Cloudera, Inc. under one
  3. # or more contributor license agreements. See the NOTICE file
  4. # distributed with this work for additional information
  5. # regarding copyright ownership. Cloudera, Inc. licenses this file
  6. # to you under the Apache License, Version 2.0 (the
  7. # "License"); you may not use this file except in compliance
  8. # with the License. You may obtain a copy of the License at
  9. #
  10. # http://www.apache.org/licenses/LICENSE-2.0
  11. #
  12. # Unless required by applicable law or agreed to in writing, software
  13. # distributed under the License is distributed on an "AS IS" BASIS,
  14. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  15. # See the License for the specific language governing permissions and
  16. # limitations under the License.
  17. import json
  18. import logging
  19. from django.utils.translation import ugettext as _
  20. from beeswax.server import dbms
  21. from desktop.lib.django_util import JsonResponse
  22. from desktop.lib.exceptions_renderable import PopupException
  23. from desktop.models import Document2
  24. from notebook.connectors.base import get_api, Notebook
  25. from indexer.controller import CollectionManagerController
  26. from indexer.file_format import HiveFormat
  27. from indexer.fields import Field
  28. from indexer.smart_indexer import Indexer
  29. LOG = logging.getLogger(__name__)
  30. def _escape_white_space_characters(s, inverse = False):
  31. MAPPINGS = {
  32. "\n": "\\n",
  33. "\t": "\\t",
  34. "\r": "\\r",
  35. " ": "\\s"
  36. }
  37. to = 1 if inverse else 0
  38. from_ = 0 if inverse else 1
  39. for pair in MAPPINGS.iteritems():
  40. s = s.replace(pair[to], pair[from_]).encode('utf-8')
  41. return s
  42. def _convert_format(format_dict, inverse=False):
  43. for field in format_dict:
  44. if isinstance(format_dict[field], basestring):
  45. format_dict[field] = _escape_white_space_characters(format_dict[field], inverse)
  46. def guess_format(request):
  47. file_format = json.loads(request.POST.get('fileFormat', '{}'))
  48. if file_format['inputFormat'] == 'file':
  49. indexer = Indexer(request.user, request.fs)
  50. stream = request.fs.open(file_format["path"])
  51. format_ = indexer.guess_format({
  52. "file":{
  53. "stream": stream,
  54. "name": file_format['path']
  55. }
  56. })
  57. _convert_format(format_)
  58. elif file_format['inputFormat'] == 'table':
  59. db = dbms.get(request.user)
  60. table_metadata = db.get_table(database=file_format['databaseName'], table_name=file_format['tableName'])
  61. storage = dict([(delim['data_type'], delim['comment']) for delim in table_metadata.storage_details])
  62. if table_metadata.details['properties']['format'] == 'text':
  63. format_ = {"quoteChar": "\"", "recordSeparator": '\\n', "type": "csv", "hasHeader": False, "fieldSeparator": storage['serialization.format']}
  64. elif table_metadata.details['properties']['format'] == 'parquet':
  65. format_ = {"type": "parquet", "hasHeader": False,}
  66. else:
  67. raise PopupException('Hive table format %s is not supported.' % table_metadata.details['properties']['format'])
  68. elif file_format['inputFormat'] == 'query':
  69. format_ = {"quoteChar": "\"", "recordSeparator": "\\n", "type": "csv", "hasHeader": False, "fieldSeparator": "\u0001"}
  70. return JsonResponse(format_)
  71. def guess_field_types(request):
  72. file_format = json.loads(request.POST.get('fileFormat', '{}'))
  73. if file_format['inputFormat'] == 'file':
  74. indexer = Indexer(request.user, request.fs)
  75. stream = request.fs.open(file_format["path"])
  76. _convert_format(file_format["format"], inverse=True)
  77. format_ = indexer.guess_field_types({
  78. "file": {
  79. "stream": stream,
  80. "name": file_format['path']
  81. },
  82. "format": file_format['format']
  83. })
  84. elif file_format['inputFormat'] == 'table':
  85. sample = get_api(request, {'type': 'hive'}).get_sample_data({'type': 'hive'}, database=file_format['databaseName'], table=file_format['tableName'])
  86. db = dbms.get(request.user)
  87. table_metadata = db.get_table(database=file_format['databaseName'], table_name=file_format['tableName'])
  88. format_ = {
  89. "sample": sample['rows'][:4],
  90. "columns": [
  91. Field(col.name, HiveFormat.FIELD_TYPE_TRANSLATE.get(col.type, 'string')).to_dict()
  92. for col in table_metadata.cols
  93. ]
  94. }
  95. elif file_format['inputFormat'] == 'query': # Only support open query history
  96. # TODO get schema from explain query, which is not possible
  97. notebook = Notebook(document=Document2.objects.get(id=file_format['query'])).get_data()
  98. snippet = notebook['snippets'][0]
  99. sample = get_api(request, snippet).fetch_result(notebook, snippet, 4, start_over=True)
  100. format_ = {
  101. "sample": sample['rows'][:4],
  102. "sample_cols": sample.meta,
  103. "columns": [
  104. Field(col['name'], HiveFormat.FIELD_TYPE_TRANSLATE.get(col['type'], 'string')).to_dict()
  105. for col in sample.meta
  106. ]
  107. }
  108. return JsonResponse(format_)
  109. def index_file(request):
  110. file_format = json.loads(request.POST.get('fileFormat', '{}'))
  111. _convert_format(file_format["format"], inverse=True)
  112. collection_name = file_format["name"]
  113. job_handle = _index(request, file_format, collection_name)
  114. return JsonResponse(job_handle)
  115. def _index(request, file_format, collection_name, query=None):
  116. indexer = Indexer(request.user, request.fs)
  117. unique_field = indexer.get_unique_field(file_format)
  118. is_unique_generated = indexer.is_unique_generated(file_format)
  119. schema_fields = indexer.get_kept_field_list(file_format['columns'])
  120. if is_unique_generated:
  121. schema_fields += [{"name": unique_field, "type": "string"}]
  122. morphline = indexer.generate_morphline_config(collection_name, file_format, unique_field)
  123. collection_manager = CollectionManagerController(request.user)
  124. if not collection_manager.collection_exists(collection_name):
  125. collection_manager.create_collection(collection_name, schema_fields, unique_key_field=unique_field)
  126. if file_format['inputFormat'] == 'table':
  127. db = dbms.get(request.user)
  128. table_metadata = db.get_table(database=file_format['databaseName'], table_name=file_format['tableName'])
  129. input_path = table_metadata.path_location
  130. elif file_format['inputFormat'] == 'file':
  131. input_path = '${nameNode}%s' % file_format["path"]
  132. else:
  133. input_path = None
  134. return indexer.run_morphline(request, collection_name, morphline, input_path, query)