From cafde1536ffd474d072cc432d2660e50e30a0032 Mon Sep 17 00:00:00 2001
From: JamshedRahman <31935303+JamshedRahman@users.noreply.github.com>
Date: Thu, 12 Jul 2018 10:19:50 -0400
Subject: [PATCH] Adding Druid Time Granularities (#5379)

* Adding Druid Time Granularities

* fixed a linter error
---
 superset/assets/src/explore/controls.jsx |  2 +
 superset/connectors/druid/models.py      |  4 +-
 tests/druid_tests.py                     | 56 ++++++++++++++++++++++++
 3 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/superset/assets/src/explore/controls.jsx b/superset/assets/src/explore/controls.jsx
index 310a0b605..353b24e45 100644
--- a/superset/assets/src/explore/controls.jsx
+++ b/superset/assets/src/explore/controls.jsx
@@ -765,6 +765,8 @@ export const controls = {
       ['week_starting_sunday', 'week starting Sunday'],
       ['week_ending_saturday', 'week ending Saturday'],
       ['P1M', 'month'],
+      ['P3M', 'quarter'],
+      ['P1Y', 'year'],
     ],
     description: t('The time granularity for the visualization. Note that you ' +
     'can type and use simple natural language as in `10 seconds`, ' +
diff --git a/superset/connectors/druid/models.py b/superset/connectors/druid/models.py
index 6fb44b308..de38ee9c9 100644
--- a/superset/connectors/druid/models.py
+++ b/superset/connectors/druid/models.py
@@ -536,7 +536,7 @@ class DruidDatasource(Model, BaseDatasource):
                 'all', '5 seconds', '30 seconds', '1 minute', '5 minutes'
                 '30 minutes', '1 hour', '6 hour', '1 day', '7 days',
                 'week', 'week_starting_sunday', 'week_ending_saturday',
-                'month',
+                'month', 'quarter', 'year',
             ],
             'time_grains': ['now'],
         }
@@ -744,6 +744,8 @@ class DruidDatasource(Model, BaseDatasource):
             'week_starting_sunday': 'P1W',
             'week_ending_saturday': 'P1W',
             'month': 'P1M',
+            'quarter': 'P3M',
+            'year': 'P1Y',
         }
 
         granularity = {'type': 'period'}
diff --git a/tests/druid_tests.py b/tests/druid_tests.py
index bb932f317..426fa9a86 100644
--- a/tests/druid_tests.py
+++ b/tests/druid_tests.py
@@ -460,6 +460,62 @@ class DruidTests(SupersetTestCase):
             cluster.get_base_broker_url(),
             'http://localhost:7980/druid/v2')
 
+    @patch('superset.connectors.druid.models.PyDruid')
+    def test_druid_time_granularities(self, PyDruid):
+        self.login(username='admin')
+        cluster = self.get_cluster(PyDruid)
+        cluster.refresh_datasources()
+        cluster.refresh_datasources(merge_flag=True)
+        datasource_id = cluster.datasources[0].id
+        db.session.commit()
+
+        nres = [
+            list(v['event'].items()) + [('timestamp', v['timestamp'])]
+            for v in GB_RESULT_SET]
+        nres = [dict(v) for v in nres]
+        import pandas as pd
+        df = pd.DataFrame(nres)
+        instance = PyDruid.return_value
+        instance.export_pandas.return_value = df
+        instance.query_dict = {}
+        instance.query_builder.last_query.query_dict = {}
+
+        form_data = {
+            'viz_type': 'table',
+            'since': '7+days+ago',
+            'until': 'now',
+            'metrics': ['count'],
+            'groupby': [],
+            'include_time': 'true',
+        }
+
+        granularity_map = {
+            '5 seconds': 'PT5S',
+            '30 seconds': 'PT30S',
+            '1 minute': 'PT1M',
+            '5 minutes': 'PT5M',
+            '1 hour': 'PT1H',
+            '6 hour': 'PT6H',
+            'one day': 'P1D',
+            '1 day': 'P1D',
+            '7 days': 'P7D',
+            'week': 'P1W',
+            'week_starting_sunday': 'P1W',
+            'week_ending_saturday': 'P1W',
+            'month': 'P1M',
+            'quarter': 'P3M',
+            'year': 'P1Y',
+        }
+        url = ('/superset/explore_json/druid/{}/'.format(datasource_id))
+
+        for granularity_mapping in granularity_map:
+            form_data['granularity'] = granularity_mapping
+            self.get_json_resp(url, {'form_data': json.dumps(form_data)})
+            self.assertEqual(
+                granularity_map[granularity_mapping],
+                instance.timeseries.call_args[1]['granularity']['period'],
+            )
+
 
 if __name__ == '__main__':
     unittest.main()