def test_export_csv(): headers = ["x", "y"] data = [ ["1", "2"], ["3", "4"], ["5,6", "7"], [None, None], ["http://gethue.com", "http://gethue.com"] ] # Check CSV generator = create_generator(content_generator(headers, data), "csv") response = make_response(generator, "csv", "foo") assert_equal("application/csv", response["content-type"]) content = b''.join(response.streaming_content) assert_equal(b'x,y\r\n1,2\r\n3,4\r\n"5,6",7\r\nNULL,NULL\r\nhttp://gethue.com,http://gethue.com\r\n', content) assert_equal('attachment; filename="foo.csv"', response["content-disposition"]) # Check non-ASCII for any browser except FF or no browser info generator = create_generator(content_generator(headers, data), "csv") response = make_response(generator, "csv", u'gんtbhんjk?¥n') assert_equal("application/csv", response["content-type"]) content = b''.join(response.streaming_content) assert_equal(b'x,y\r\n1,2\r\n3,4\r\n"5,6",7\r\nNULL,NULL\r\nhttp://gethue.com,http://gethue.com\r\n', content) assert_equal('attachment; filename="g%E3%82%93tbh%E3%82%93jk%EF%BC%9F%EF%BF%A5n.csv"', response["content-disposition"]) # Check non-ASCII for FF browser generator = create_generator(content_generator(headers, data), "csv") response = make_response(generator, "csv", u'gんtbhんjk?¥n', user_agent='Mozilla / 5.0(Macintosh; Intel Mac OS X 10.12;rv:59.0) Gecko / 20100101 Firefox / 59.0)') assert_equal("application/csv", response["content-type"]) content = b''.join(response.streaming_content) assert_equal(b'x,y\r\n1,2\r\n3,4\r\n"5,6",7\r\nNULL,NULL\r\nhttp://gethue.com,http://gethue.com\r\n', content) assert_equal('attachment; filename*="g%E3%82%93tbh%E3%82%93jk%EF%BC%9F%EF%BF%A5n.csv"', response["content-disposition"])
def test_export_csv(): headers = ["x", "y"] data = [ ["1", "2"], ["3", "4"], ["5,6", "7"], [None, None], ["http://gethue.com", "http://gethue.com"] ] # Check CSV generator = create_generator(content_generator(headers, data), "csv") response = make_response(generator, "csv", "foo") assert_equal("application/csv", response["content-type"]) content = ''.join(response.streaming_content) assert_equal('x,y\r\n1,2\r\n3,4\r\n"5,6",7\r\nNULL,NULL\r\nhttp://gethue.com,http://gethue.com\r\n', content) assert_equal('attachment; filename="foo.csv"', response["content-disposition"]) # Check non-ASCII for any browser except FF or no browser info generator = create_generator(content_generator(headers, data), "csv") response = make_response(generator, "csv", u'gんtbhんjk?¥n') assert_equal("application/csv", response["content-type"]) content = ''.join(response.streaming_content) assert_equal('x,y\r\n1,2\r\n3,4\r\n"5,6",7\r\nNULL,NULL\r\nhttp://gethue.com,http://gethue.com\r\n', content) assert_equal('attachment; filename="g%E3%82%93tbh%E3%82%93jk%EF%BC%9F%EF%BF%A5n.csv"', response["content-disposition"]) # Check non-ASCII for FF browser generator = create_generator(content_generator(headers, data), "csv") response = make_response(generator, "csv", u'gんtbhんjk?¥n', user_agent='Mozilla / 5.0(Macintosh; Intel Mac OS X 10.12;rv:59.0) Gecko / 20100101 Firefox / 59.0)') assert_equal("application/csv", response["content-type"]) content = ''.join(response.streaming_content) assert_equal('x,y\r\n1,2\r\n3,4\r\n"5,6",7\r\nNULL,NULL\r\nhttp://gethue.com,http://gethue.com\r\n', content) assert_equal('attachment; filename*="g%E3%82%93tbh%E3%82%93jk%EF%BC%9F%EF%BF%A5n.csv"', response["content-disposition"])
def download_to_file(notebook, snippet, file_format='csv', postdict=None, user_id=None, max_rows=-1): from beeswax import data_export download_to_file.update_state(task_id=notebook['uuid'], state='STARTED', meta={}) request = _get_request(postdict, user_id) api = get_api(request, snippet) f, path = tempfile.mkstemp() f_log, path_log = tempfile.mkstemp() f_progress, path_progress = tempfile.mkstemp() try: os.write(f_progress, '0') meta = {'row_counter': 0, 'file_path': path, 'handle': {}, 'log_path': path_log, 'progress_path': path_progress, 'status': 'running', 'truncated': False} #TODO: Truncated result_wrapper = ResultWrapper(api, notebook, snippet, ResultWrapperCallback(notebook['uuid'], meta, f_log)) content_generator = data_export.DataAdapter(result_wrapper, max_rows=max_rows, store_data_type_in_header=True) #TODO: Move PREFETCH_RESULT_COUNT to front end response = export_csvxls.create_generator(content_generator, file_format) for chunk in response: os.write(f, chunk) meta['row_counter'] = content_generator.row_counter download_to_file.update_state(task_id=notebook['uuid'], state='AVAILABLE', meta=meta) finally: os.close(f) os.close(f_log) os.close(f_progress) return meta
def download(format, db, id=None, file_name='query_result', user_agent=None, max_rows=-1, max_bytes=-1, store_data_type_in_header=False, start_over=True): """ download(query_model, format) -> HttpResponse Retrieve the query result in the format specified. Return an HttpResponse object. """ if format not in common.DL_FORMATS: LOG.error('Unknown download format "%s"' % (format,)) return content_generator = DataAdapter(db, max_rows=max_rows, start_over=start_over, max_bytes=max_bytes, store_data_type_in_header=store_data_type_in_header) generator = export_csvxls.create_generator(content_generator, format) resp = export_csvxls.make_response(generator, format, file_name, user_agent=user_agent) if id: resp.set_cookie( 'download-%s' % id, json.dumps({ 'truncated': content_generator.is_truncated, 'row_counter': content_generator.row_counter }), max_age=DOWNLOAD_COOKIE_AGE ) return resp
def download(handle, format, db, id=None, file_name='query_result'): """ download(query_model, format) -> HttpResponse Retrieve the query result in the format specified. Return an HttpResponse object. """ if format not in common.DL_FORMATS: LOG.error('Unknown download format "%s"' % (format,)) return max_rows = conf.DOWNLOAD_ROW_LIMIT.get() content_generator = HS2DataAdapter(handle, db, max_rows=max_rows, start_over=True) generator = export_csvxls.create_generator(content_generator, format) resp = export_csvxls.make_response(generator, format, file_name) if id: resp.set_cookie( 'download-%s' % id, json.dumps({ 'truncated': content_generator.is_truncated, 'row_counter': content_generator.row_counter }), max_age=DOWNLOAD_COOKIE_AGE ) return resp
def download(handle, format, db, id=None, file_name='query_result', user_agent=None): """ download(query_model, format) -> HttpResponse Retrieve the query result in the format specified. Return an HttpResponse object. """ if format not in common.DL_FORMATS: LOG.error('Unknown download format "%s"' % (format,)) return max_rows = conf.DOWNLOAD_ROW_LIMIT.get() max_bytes = conf.DOWNLOAD_BYTES_LIMIT.get() content_generator = DataAdapter(db, handle=handle, max_rows=max_rows, max_bytes=max_bytes) generator = export_csvxls.create_generator(content_generator, format) resp = export_csvxls.make_response(generator, format, file_name, user_agent=user_agent) if id: resp.set_cookie( 'download-%s' % id, json.dumps({ 'truncated': content_generator.is_truncated, 'row_counter': content_generator.row_counter }), max_age=DOWNLOAD_COOKIE_AGE ) return resp
def upload(self, queries, source_platform='generic', workload_id=None): with NamedTemporaryFile(suffix='.csv') as f: try: content_generator = OptimizerDataAdapter(queries) queries_csv = export_csvxls.create_generator( content_generator, 'csv') # return self._root.post('/api/upload', data=data, files = {'file': ('hue-report.csv', list(queries_csv)[0])}) for row in queries_csv: f.write(row) # data = { # 'fileLocation': f.name, # 'sourcePlatform': source_platform, # # 'workloadId', # # 'version' # } args = [ '--file-location', f.name, '--file-name', os.path.basename(f.name), '--source-platform', source_platform, '--tenant', self._product_name ] if workload_id: args += ['--workload-id', workload_id] return self._exec('upload', args) ## return self._root.post('/api/upload', data=data, files={'fileLocation': ('hue-report.json', json.dumps(queries_formatted))}) except RestException, e: raise PopupException( e, title=_('Error while accessing Optimizer'))
def download(handle, format, db, id=None): """ download(query_model, format) -> HttpResponse Retrieve the query result in the format specified. Return an HttpResponse object. """ if format not in common.DL_FORMATS: LOG.error('Unknown download format "%s"' % (format, )) return max_cells = conf.DOWNLOAD_CELL_LIMIT.get() content_generator = HS2DataAdapter(handle, db, max_cells=max_cells, start_over=True) generator = export_csvxls.create_generator(content_generator, format) resp = export_csvxls.make_response(generator, format, 'query_result') if id: resp.set_cookie('download-%s' % id, json.dumps({ 'truncated': content_generator.is_truncated, 'row_counter': content_generator.row_counter }), max_age=DOWNLOAD_COOKIE_AGE) return resp
def upload(self, queries, token=None, email=None, source_platform='generic'): if token is None: token = self._authenticate() try: content_generator = OptimizerDataAdapter(queries) queries_csv = export_csvxls.create_generator( content_generator, 'csv') data = { 'email': email if email is not None else self._email, 'token': token, 'sourcePlatform': source_platform, } return self._root.post( '/api/upload', data=data, files={'file': ('hue-report.csv', list(queries_csv)[0])}) except RestException, e: raise PopupException(e, title=_('Error while accessing Optimizer'))
def download_to_file(notebook, snippet, file_format='csv', max_rows=-1, **kwargs): from beeswax import data_export download_to_file.update_state(task_id=notebook['uuid'], state='STARTED', meta={}) request = _get_request(**kwargs) api = get_api(request, snippet) meta = {'row_counter': 0, 'handle': {}, 'status': '', 'truncated': False} with storage.open(_log_key(notebook), 'wb') as f_log: result_wrapper = ExecutionWrapper( api, notebook, snippet, ExecutionWrapperCallback(notebook['uuid'], meta, f_log)) content_generator = data_export.DataAdapter( result_wrapper, max_rows=max_rows, store_data_type_in_header=True ) #TODO: Move PREFETCH_RESULT_COUNT to front end response = export_csvxls.create_generator(content_generator, file_format) with storage.open(_result_key(notebook), 'wb') as f: for chunk in response: f.write(chunk) meta['row_counter'] = content_generator.row_counter meta['truncated'] = content_generator.is_truncated download_to_file.update_state(task_id=notebook['uuid'], state='AVAILABLE', meta=meta) return meta
def download_to_file(notebook, snippet, file_format='csv', max_rows=-1, **kwargs): download_to_file.update_state(task_id=notebook['uuid'], state='STARTED', meta={}) request = _get_request(**kwargs) api = get_api(request, snippet) meta = {'row_counter': 0, 'handle': {}, 'status': '', 'truncated': False} with storage.open(_log_key(notebook), 'wb') as f_log: # TODO: use cache for editor 1000 rows and storage for result export result_wrapper = ExecutionWrapper(api, notebook, snippet, ExecutionWrapperCallback(notebook['uuid'], meta, f_log)) content_generator = data_export.DataAdapter(result_wrapper, max_rows=max_rows, store_data_type_in_header=True) # TODO: Move FETCH_RESULT_LIMIT to front end response = export_csvxls.create_generator(content_generator, file_format) with storage.open(_result_key(notebook), 'wb') as f: for chunk in response: f.write(chunk) meta['row_counter'] = content_generator.row_counter meta['truncated'] = content_generator.is_truncated download_to_file.update_state(task_id=notebook['uuid'], state='AVAILABLE', meta=meta) if TASK_SERVER.RESULT_CACHE.get(): with storage.open(_result_key(notebook)) as f: csv_reader = csv.reader(f, delimiter=','.encode('utf-8')) caches[CACHES_CELERY_QUERY_RESULT_KEY].set(_result_key(notebook), [row for row in csv_reader], 60 * 5) return meta
def upload(self, data, data_type='queries', source_platform='generic', workload_id=None): if data_type in ('table_stats', 'cols_stats'): data_suffix = '.json' if data_type == 'table_stats': extra_parameters = {'fileType': 'TABLE_STATS'} else: extra_parameters = {'fileType': 'COLUMN_STATS'} else: data_suffix = '.csv' extra_parameters = { 'fileType': 'QUERY', 'colDelim': ',', 'rowDelim': '\n', "headerFields": [ {"count": 0, "name": "SQL_ID", "coltype": "SQL_ID", "use": True, "tag": ""}, {"count": 0, "name": "ELAPSED_TIME", "coltype": "NONE", "use": True, "tag": ""}, {"count": 0, "name": "SQL_FULLTEXT", "coltype": "SQL_QUERY", "use": True, "tag": ""}, {"count": 0, "name": "DATABASE", "coltype": "NONE", "use": True, "tag": "DATABASE"} ], } f_queries_path = NamedTemporaryFile(suffix=data_suffix) f_queries_path.close() # Reopened as real file below to work well with the command try: f_queries = open(f_queries_path.name, 'w+') try: # Queries if data_suffix == '.csv': content_generator = OptimizerQueryDataAdapter(data) queries_csv = export_csvxls.create_generator(content_generator, 'csv') for row in queries_csv: f_queries.write(row) LOG.debug(row[:1000]) else: # Table, column stats f_queries.write(json.dumps(data)) LOG.debug(json.dumps(data[:10])) finally: f_queries.close() parameters = { 'tenant' : self._tenant_id, 'fileLocation': f_queries.name, 'sourcePlatform': source_platform, } parameters.update(extra_parameters) response = self._api.call_api('upload', parameters) status = json.loads(response) status['count'] = len(data) return status except RestException as e: raise PopupException(e, title=_('Error while accessing Optimizer')) finally: os.remove(f_queries_path.name)
def download(api, session, cell, format): if format not in DL_FORMATS: LOG.error('Unknown download format "%s"' % format) return content_generator = SparkDataAdapter(api, session, cell) generator = export_csvxls.create_generator(content_generator, format) return export_csvxls.make_response(generator, format, "script_result")
def download(api, session, cell, format, user_agent=None): if format not in DL_FORMATS: LOG.error('Unknown download format "%s"' % format) return content_generator = SparkDataAdapter(api, session, cell) generator = export_csvxls.create_generator(content_generator, format) return export_csvxls.make_response(generator, format, 'script_result', user_agent=user_agent)
def upload(self, data, data_type='queries', source_platform='generic', workload_id=None): if data_type in ('table_stats', 'cols_stats'): data_suffix = '.json' if data_type == 'table_stats': extra_parameters = {'fileType': 'TABLE_STATS'} else: extra_parameters = {'fileType': 'COLUMN_STATS'} else: data_suffix = '.csv' extra_parameters = { 'colDelim': ',', 'rowDelim': '\n', } f_queries_path = NamedTemporaryFile(suffix=data_suffix) f_queries_path.close( ) # Reopened as real file below to work well with the command try: f_queries = open(f_queries_path.name, 'w+') try: # Queries if data_suffix == '.csv': content_generator = OptimizerQueryDataAdapter(data) queries_csv = export_csvxls.create_generator( content_generator, 'csv') for row in queries_csv: f_queries.write(row) LOG.debug(row) else: # Table, column stats f_queries.write(json.dumps(data)) LOG.debug(json.dumps(data)) finally: f_queries.close() parameters = { 'tenant': self._product_name, 'fileLocation': f_queries.name, 'sourcePlatform': source_platform, } parameters.update(extra_parameters) response = self._api.call_api('upload', parameters) status = json.loads(response) status['count'] = len(data) return status except RestException, e: raise PopupException(e, title=_('Error while accessing Optimizer'))
def test_export_csv(): header = ["x", "y"] data = [ ["1", "2"], ["3", "4"], ["5,6", "7"], [None, None] ] # Check CSV generator = create_generator(content_generator(header, data), "csv") response = make_response(generator, "csv", "foo") assert_equal("application/csv", response["content-type"]) assert_equal('x,y\r\n1,2\r\n3,4\r\n"5,6",7\r\nNULL,NULL\r\n', response.content) assert_equal("attachment; filename=foo.csv", response["content-disposition"])
def test_export_csv(): header = ["x", "y"] data = [ ["1", "2"], ["3", "4"], ["5,6", "7"], [None, None] ] # Check CSV generator = create_generator(content_generator(header, data), "csv") response = make_response(generator, "csv", "foo") assert_equal("application/csv", response["content-type"]) content = ''.join(response.streaming_content) assert_equal('x,y\r\n1,2\r\n3,4\r\n"5,6",7\r\nNULL,NULL\r\n', content) assert_equal("attachment; filename=foo.csv", response["content-disposition"])
def download(self, notebook, snippet, file_format='csv'): from beeswax import data_export #TODO: Move to notebook? from beeswax import conf result_wrapper = ExecutionWrapper(self, notebook, snippet) max_rows = conf.DOWNLOAD_ROW_LIMIT.get() max_bytes = conf.DOWNLOAD_BYTES_LIMIT.get() content_generator = data_export.DataAdapter(result_wrapper, max_rows=max_rows, max_bytes=max_bytes) return export_csvxls.create_generator(content_generator, file_format)
def test_export_xls_truncate_rows(): headers = ["a"] data = [["1"]] * (MAX_XLS_ROWS + 1) sheet = [headers] + data # Check XLS generator = create_generator(content_generator(headers, data), "xls") response = make_response(generator, "xls", "foo") assert_equal("application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", response["content-type"]) sheet_data = _read_xls_sheet_data(response) assert_equal(len(sheet[:MAX_XLS_ROWS + 1]), len(sheet_data)) assert_equal("attachment; filename=foo.xlsx", response["content-disposition"])
def download(results, format, collection): """ download(results, format) -> HttpResponse Transform the search result set to the specified format and download. """ if format not in DL_FORMATS: LOG.error('Unknown download format "%s"' % format) return content_generator = SearchDataAdapter(results, format, collection) generator = export_csvxls.create_generator(content_generator, format) return export_csvxls.make_response(generator, format, 'query_result')
def download(results, format, collection): """ download(results, format) -> HttpResponse Transform the search result set to the specified format and dwonload. """ if format not in DL_FORMATS: LOG.error('Unknown download format "%s"' % format) return content_generator = SearchDataAdapter(results, format, collection) generator = export_csvxls.create_generator(content_generator, format) return export_csvxls.make_response(generator, format, 'query_result')
def download(handle, format, db): """ download(query_model, format) -> HttpResponse Retrieve the query result in the format specified. Return an HttpResponse object. """ if format not in common.DL_FORMATS: LOG.error('Unknown download format "%s"' % (format,)) return content_generator = HS2DataAdapter(handle, db, conf.DOWNLOAD_ROW_LIMIT.get()) generator = export_csvxls.create_generator(content_generator, format) return export_csvxls.make_response(generator, format, 'query_result')
def test_export_xls(): header = ["x", "y"] data = [ ["1", "2"], ["3", "4"], ["5,6", "7"], [None, None] ] dataset = tablib.Dataset(headers=header) for row in data: dataset.append([cell is not None and cell or "NULL" for cell in row]) # Check XLS generator = create_generator(content_generator(header, data), "xls") response = make_response(generator, "xls", "foo") assert_equal("application/xls", response["content-type"]) assert_equal(dataset.xls, response.content) assert_equal("attachment; filename=foo.xls", response["content-disposition"])
def test_export_xls_truncate_cols(): headers = ["a"] * (MAX_XLS_COLS + 1) data = [["1"] * (MAX_XLS_COLS + 1)] dataset = tablib.Dataset(headers=headers[:MAX_XLS_COLS]) dataset.extend([data[0][:MAX_XLS_COLS]]) # Check XLS generator = create_generator(content_generator(headers, data), "xls") response = make_response(generator, "xls", "foo") assert_equal("application/xls", response["content-type"]) content = "".join(response.streaming_content) assert_equal(dataset.xls, content) assert_equal("attachment; filename=foo.xls", response["content-disposition"])
def test_export_xls(): headers = ["x", "y"] data = [ ["1", "2"], ["3", "4"], ["5,6", "7"], [None, None], ["http://gethue.com", "http://gethue.com"] ] sheet = [headers] + data # Check XLS generator = create_generator(content_generator(headers, data), "xls") response = make_response(generator, "xls", "foo") assert_equal("application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", response["content-type"]) expected_data = [[cell is not None and cell.replace("http://gethue.com", '=HYPERLINK("http://gethue.com")') or "NULL" for cell in row] for row in sheet] sheet_data = _read_xls_sheet_data(response) assert_equal(expected_data, sheet_data) assert_equal('attachment; filename="foo.xlsx"', response["content-disposition"])
def test_export_xls(): headers = ["x", "y"] data = [["1", "2"], ["3", "4"], ["5,6", "7"], [None, None]] sheet = [headers] + data # Check XLS generator = create_generator(content_generator(headers, data), "xls") response = make_response(generator, "xls", "foo") assert_equal("application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", response["content-type"]) expected_data = [[cell is not None and cell or "NULL" for cell in row] for row in sheet] sheet_data = _read_xls_sheet_data(response) assert_equal(expected_data, sheet_data) assert_equal("attachment; filename=foo.xlsx", response["content-disposition"])
def test_export_xls_truncate_cols(): headers = ["a"] * (MAX_XLS_COLS + 1) data = [["1"] * (MAX_XLS_COLS + 1)] dataset = tablib.Dataset(headers=headers[:MAX_XLS_COLS]) dataset.extend([data[0][:MAX_XLS_COLS]]) # Check XLS generator = create_generator(content_generator(headers, data), "xls") response = make_response(generator, "xls", "foo") assert_equal("application/xls", response["content-type"]) content = ''.join(response.streaming_content) assert_equal(dataset.xls, content) assert_equal("attachment; filename=foo.xls", response["content-disposition"])
def test_export_xls(): header = ["x", "y"] data = [ ["1", "2"], ["3", "4"], ["5,6", "7"], [None, None] ] dataset = tablib.Dataset(headers=header) for row in data: dataset.append([cell is not None and cell or "NULL" for cell in row]) # Check XLS generator = create_generator(content_generator(header, data), "xls") response = make_response(generator, "xls", "foo") assert_equal("application/xls", response["content-type"]) content = ''.join(response.streaming_content) assert_equal(dataset.xls, content) assert_equal("attachment; filename=foo.xls", response["content-disposition"])
def test_export_xls_truncate_rows(): headers = ["a"] data = [["1"]] * (MAX_XLS_ROWS + 1) sheet = [headers] + data # Check XLS generator = create_generator(content_generator(headers, data), "xls") response = make_response(generator, "xls", "foo") assert_equal( "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", response["content-type"]) sheet_data = _read_xls_sheet_data(response) assert_equal(len(sheet[:MAX_XLS_ROWS + 1]), len(sheet_data)) assert_equal("attachment; filename=foo.xlsx", response["content-disposition"])
def upload(self, queries, token=None, email=None, source_platform='generic'): if token is None: token = self._authenticate() try: content_generator = OptimizerDataAdapter(queries) queries_csv = export_csvxls.create_generator(content_generator, 'csv') data = { 'email': email if email is not None else self._email, 'token': token, 'sourcePlatform': source_platform, } return self._root.post('/api/upload', data=data, files = {'file': ('hue-report.csv', list(queries_csv)[0])}) except RestException, e: raise PopupException(e, title=_('Error while accessing Optimizer'))
def download_to_file(notebook, snippet, file_format='csv', max_rows=-1, **kwargs): from beeswax import data_export download_to_file.update_state(task_id=notebook['uuid'], state='STARTED', meta={}) request = _get_request(**kwargs) api = get_api(request, snippet) f, path = tempfile.mkstemp() f_log, path_log = tempfile.mkstemp() try: #TODO: We need to move this metadata somewhere else, it gets erased on exception and we can no longer cleanup the files. meta = { 'row_counter': 0, 'file_path': path, 'handle': {}, 'log_path': path_log, 'status': 'running', 'truncated': False } result_wrapper = ResultWrapper( api, notebook, snippet, ResultWrapperCallback(notebook['uuid'], meta, f_log)) content_generator = data_export.DataAdapter( result_wrapper, max_rows=max_rows, store_data_type_in_header=True ) #TODO: Move PREFETCH_RESULT_COUNT to front end response = export_csvxls.create_generator(content_generator, file_format) for chunk in response: os.write(f, chunk) meta['row_counter'] = content_generator.row_counter meta['truncated'] = content_generator.is_truncated download_to_file.update_state(task_id=notebook['uuid'], state='AVAILABLE', meta=meta) finally: os.close(f) os.close(f_log) return meta
def upload(self, data, data_type='queries', source_platform='generic', workload_id=None): data_headers = OptimizerApi.UPLOAD[data_type]['file_headers'] if data_type in ('table_stats', 'cols_stats'): data_suffix = '.log' else: data_suffix = '.csv' f_queries_path = NamedTemporaryFile(suffix=data_suffix) f_format_path = NamedTemporaryFile(suffix='.json') f_queries_path.close() f_format_path.close() # Reopened as real file below to work well with the command try: f_queries = open(f_queries_path.name, 'w+') f_format = open(f_format_path.name, 'w+') try: content_generator = OptimizerDataAdapter(data, data_type=data_type) queries_csv = export_csvxls.create_generator(content_generator, 'csv') for row in queries_csv: f_queries.write(row) f_format.write(data_headers % { 'source_platform': source_platform, 'tenant': self._product_name, 'query_file': f_queries.name, 'query_file_name': os.path.basename(f_queries.name) }) finally: f_queries.close() f_format.close() args = [ '--cli-input-json', 'file://%s' % f_format.name ] if workload_id: args += ['--workload-id', workload_id] return self._exec('upload', args) except RestException, e: raise PopupException(e, title=_('Error while accessing Optimizer'))
def download_to_file(notebook, snippet, file_format='csv', max_rows=-1, **kwargs): task_id = notebook['uuid'] result_key = _result_key(task_id) download_to_file.update_state(task_id=task_id, state='STARTED', meta={}) request = _get_request(**kwargs) api = get_api(request, snippet) meta = {'row_counter': 0, 'handle': {}, 'status': '', 'truncated': False} with storage.open(_log_key(notebook, snippet), 'wb') as f_log: result_wrapper = ExecutionWrapper( api, notebook, snippet, ExecutionWrapperCallback(notebook['uuid'], meta, f_log)) content_generator = DataAdapter(result_wrapper, max_rows=max_rows, store_data_type_in_header=True) response = export_csvxls.create_generator(content_generator, file_format) with storage.open(result_key, 'wb') as f: for chunk in response: f.write(chunk.encode('utf-8')) if TASK_SERVER.RESULT_CACHE.get(): with storage.open(result_key, 'rb') as store: 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) LOG.info('Caching results %s.' % result_key) meta['row_counter'] = content_generator.row_counter meta['truncated'] = content_generator.is_truncated download_to_file.update_state(task_id=task_id, state='AVAILABLE', meta=meta) return meta
def test_export_xls_truncate_cols(): headers = [u"a"] * (MAX_XLS_COLS + 1) data = [[u"1"] * (MAX_XLS_COLS + 1)] sheet = [headers] + data # Check XLS generator = create_generator(content_generator(headers, data), "xls") response = make_response(generator, "xls", "foo") assert_equal("application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", response["content-type"]) expected_data = [[cell is not None and cell or "NULL" for cell in row][:MAX_XLS_COLS] for row in sheet] sheet_data = _read_xls_sheet_data(response) sheet_data[0] = sheet_data[0][:MAX_XLS_COLS] sheet_data[1] = sheet_data[1][:MAX_XLS_COLS] assert_equal(len(expected_data), len(sheet_data)) assert_equal(len(expected_data[0]), len(sheet_data[0])) assert_equal(expected_data, sheet_data) assert_equal("attachment; filename=foo.xlsx", response["content-disposition"])
def test_export_xls_truncate_cols(): headers = [u"a"] * (MAX_XLS_COLS + 1) data = [[u"1"] * (MAX_XLS_COLS + 1)] sheet = [headers] + data # Check XLS generator = create_generator(content_generator(headers, data), "xls") response = make_response(generator, "xls", "foo") assert_equal( "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", response["content-type"]) expected_data = [[cell is not None and cell or "NULL" for cell in row][:MAX_XLS_COLS] for row in sheet] sheet_data = _read_xls_sheet_data(response) sheet_data[0] = sheet_data[0][:MAX_XLS_COLS] sheet_data[1] = sheet_data[1][:MAX_XLS_COLS] assert_equal(len(expected_data), len(sheet_data)) assert_equal(len(expected_data[0]), len(sheet_data[0])) assert_equal(expected_data, sheet_data) assert_equal("attachment; filename=foo.xlsx", response["content-disposition"])