Merge pull request #831 from ninneko/801-download-excel

Feature: download results in Excel (XSLX) format (closes #801)

rd_ui/app/scripts/directives/query_directives.js

@@ -38,10 +38,12 @@
     }
   }
 
-  function queryResultCSVLink() {
+  function queryResultLink() {
     return {
       restrict: 'A',
-      link: function (scope, element) {
+      link: function (scope, element, attrs) {
+
+        var fileType = attrs.fileType ? attrs.fileType : "csv";
         scope.$watch('queryResult && queryResult.getData()', function(data) {
           if (!data) {
             return;
@@ -50,8 +52,8 @@
           if (scope.queryResult.getId() == null) {
             element.attr('href', '');
           } else {
-              element.attr('href', 'api/queries/' + scope.query.id + '/results/' + scope.queryResult.getId() + '.csv');
-              element.attr('download', scope.query.name.replace(" ", "_") + moment(scope.queryResult.getUpdatedAt()).format("_YYYY_MM_DD") + ".csv");
+            element.attr('href', 'api/queries/' + scope.query.id + '/results/' + scope.queryResult.getId() + '.' + fileType);
+            element.attr('download', scope.query.name.replace(" ", "_") + moment(scope.queryResult.getUpdatedAt()).format("_YYYY_MM_DD") + "." + fileType);
           }
         });
       }
@@ -286,7 +288,7 @@
   angular.module('redash.directives')
   .directive('queryLink', queryLink)
   .directive('querySourceLink', querySourceLink)
-  .directive('queryResultLink', queryResultCSVLink)
+  .directive('queryResultLink', queryResultLink)
   .directive('queryEditor', queryEditor)
   .directive('queryRefreshSelect', queryRefreshSelect)
   .directive('queryTimePicker', queryTimePicker)

rd_ui/app/views/query.html

@@ -157,12 +157,19 @@
 
             <hr>
 
+            <p>
             <p>
                 <a class="btn btn-primary btn-sm" ng-disabled="queryExecuting || !queryResult.getData()" query-result-link target="_self">
                     <span class="glyphicon glyphicon-cloud-download"></span>
-                    <span>Download Dataset</span>
+                    <span>Download CSV</span>
                 </a>
 
+                <a class="btn btn-primary btn-sm" ng-disabled="queryExecuting || !queryResult.getData()" file-type="xlsx" query-result-link target="_self">
+                    <span class="glyphicon glyphicon-cloud-download"></span>
+                    <span>Download Excel</span>
+                </a>
+            </p>
+            <p>
                 <a class="btn btn-warning btn-sm" ng-disabled="queryExecuting" data-toggle="modal" data-target="#archive-confirmation-modal"
                    ng-show="!query.is_archived && query.id != undefined && (isQueryOwner || currentUser.hasPermission('admin'))">
                     <i class="fa fa-archive" title="Archive Query"></i>
@@ -171,6 +178,7 @@
                 <button class="btn btn-default btn-sm" ng-show="query.id != undefined" ng-click="showApiKey()">
                     <i class="fa fa-key" title="Show API Key"></i>
                 </button>
+            </p>
 
                 <div class="modal fade" id="archive-confirmation-modal" tabindex="-1" role="dialog" aria-labelledby="archiveConfirmationModal" aria-hidden="true">
                   <div class="modal-dialog">

redash/handlers/query_results.py

@@ -5,7 +5,7 @@ import time
 
 from flask import make_response, request
 from flask.ext.restful import abort
-
+import xlsxwriter
 from redash import models, settings, utils
 from redash.wsgi import api
 from redash.tasks import QueryTask, record_event
@@ -105,6 +105,8 @@ class QueryResultAPI(BaseResource):
 
             if filetype == 'json':
                 response = self.make_json_response(query_result)
+            elif filetype == 'xlsx':
+                response = self.make_excel_response(query_result)
             else:
                 response = self.make_csv_response(query_result)
 
@@ -137,6 +139,28 @@ class QueryResultAPI(BaseResource):
         headers = {'Content-Type': "text/csv; charset=UTF-8"}
         return make_response(s.getvalue(), 200, headers)
 
+    @staticmethod
+    def make_excel_response(query_result):
+        s = cStringIO.StringIO()
+
+        query_data = json.loads(query_result.data)
+        book = xlsxwriter.Workbook(s)
+        sheet = book.add_worksheet("result")
+
+        column_names = []
+        for (c, col) in enumerate(query_data['columns']):
+            sheet.write(0, c, col['name'])
+            column_names.append(col['name'])
+
+        for (r, row) in enumerate(query_data['rows']):
+            for (c, name) in enumerate(column_names):
+                sheet.write(r+1, c, row[name])
+
+        book.close()
+
+        headers = {'Content-Type': "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"}
+        return make_response(s.getvalue(), 200, headers)
+
 
 api.add_org_resource(QueryResultListAPI, '/api/query_results', endpoint='query_results')
 api.add_org_resource(QueryResultAPI,

requirements.txt

@@ -37,3 +37,4 @@ funcy==1.5
 raven==5.9.2
 semver==2.2.1
 python-simple-hipchat==0.4.0
+xlsxwriter==0.8.4

tests/test_models.py

@@ -163,8 +163,9 @@ class ShouldScheduleNextTest(TestCase):
     def test_exact_time_that_needs_reschedule(self):
         now = datetime.datetime.now()
         yesterday = now - datetime.timedelta(days=1)
-        schedule = "{:02d}:00".format(now.hour - 3)
-        self.assertTrue(models.should_schedule_next(yesterday, now, schedule))
+        scheduled_datetime = now - datetime.timedelta(hours=3)
+        scheduled_time = "{:02d}:00".format(scheduled_datetime.hour)
+        self.assertTrue(models.should_schedule_next(yesterday, now, scheduled_time))
 
     def test_exact_time_that_doesnt_need_reschedule(self):
         now = date_parse("2015-10-16 20:10")