From 0ce16bcfd2dd666ff5294469cc47263acafca1d7 Mon Sep 17 00:00:00 2001
From: Shrivardhan Rao <Shrivardhan_Rao@intuit.com>
Date: Tue, 29 Aug 2023 17:05:11 -0700
Subject: [PATCH] explicit intervals and datasource

---
 .pre-commit-config.yaml | 10 ++++++++--
 pydruid/query.py        |  4 +++-
 tests/test_client.py    | 14 ++++++++++----
 tests/test_query.py     | 12 +++++++-----
 4 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index fee7b2a5..5e35b785 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1,5 +1,5 @@
 repos:
--   repo: https://github.com/ambv/black
+-   repo: https://github.com/psf/black
     rev: 19.10b0
     hooks:
     -   id: black
@@ -21,7 +21,13 @@ repos:
     -   id: check-yaml
     -   id: debug-statements
 
--   repo: https://gitlab.com/pycqa/flake8
+-   repo: https://github.com/pycqa/flake8
     rev: 3.8.2
     hooks:
     -   id: flake8
+        exclude:
+            docs
+            env
+            tests
+            .eggs
+            build
diff --git a/pydruid/query.py b/pydruid/query.py
index 09ec3260..a33e7d50 100644
--- a/pydruid/query.py
+++ b/pydruid/query.py
@@ -247,7 +247,7 @@ def parse_datasource(datasource, query_type):
                 "dict or list of strings"
             )
         if isinstance(datasource, str):
-            return datasource
+            return {"type": "table", "name": datasource}
         else:
             return {"type": "union", "dataSources": datasource}
 
@@ -310,6 +310,8 @@ def build_query(self, query_type, args):
                 query_dict[key] = build_dimension(val)
             elif key == "dimensions":
                 query_dict[key] = [build_dimension(v) for v in val]
+            elif key == "intervals":
+                query_dict[key] = {"type": "intervals", key: val}
             else:
                 query_dict[key] = val
 
diff --git a/tests/test_client.py b/tests/test_client.py
index 440a9204..e6fff7be 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -89,8 +89,8 @@ def test_druid_returns_html_error(self, mock_urlopen):
             str(e.value)
             == textwrap.dedent(
                 """
-            HTTP Error 500: Internal Server Error 
-             Druid Error: javax.servlet.ServletException: java.lang.OutOfMemoryError: GC overhead limit exceeded 
+            HTTP Error 500: Internal Server Error
+             Druid Error: javax.servlet.ServletException: java.lang.OutOfMemoryError: GC overhead limit exceeded
              Query is: {
                 "aggregations": [
                     {
@@ -102,7 +102,10 @@ def test_druid_returns_html_error(self, mock_urlopen):
                 "context": {
                     "timeout": 1000
                 },
-                "dataSource": "testdatasource",
+                "dataSource": {
+                    "name": "testdatasource",
+                    "type": "table"
+                },
                 "dimension": "user_name",
                 "filter": {
                     "dimension": "user_lang",
@@ -110,7 +113,10 @@ def test_druid_returns_html_error(self, mock_urlopen):
                     "value": "en"
                 },
                 "granularity": "all",
-                "intervals": "2015-12-29/pt1h",
+                "intervals": {
+                    "intervals": "2015-12-29/pt1h",
+                    "type": "intervals"
+                },
                 "metric": "count",
                 "queryType": "topN",
                 "threshold": 1
diff --git a/tests/test_query.py b/tests/test_query.py
index dbeef448..1c899012 100644
--- a/tests/test_query.py
+++ b/tests/test_query.py
@@ -16,7 +16,6 @@
 #
 
 import csv
-import os
 
 import pandas
 import pytest
@@ -62,7 +61,7 @@ def test_build_query(self):
         # given
         expected_query_dict = {
             "queryType": None,
-            "dataSource": "things",
+            "dataSource": {"name": "things", "type": "table"},
             "aggregations": [{"fieldName": "thing", "name": "count", "type": "count"}],
             "postAggregations": [
                 {
@@ -106,7 +105,7 @@ def test_build_query_none_type(self):
         # given
         expected_query_dict = {
             "queryType": None,
-            "dataSource": "things",
+            "dataSource": {"name": "things", "type": "table"},
             "aggregations": [{"fieldName": "thing", "name": "count", "type": "count"}],
             "filter": {"dimension": "one", "type": "selector", "value": 1},
             "having": {"aggregation": "sum", "type": "greaterThan", "value": 1},
@@ -153,7 +152,10 @@ def test_validate_query(self):
 
     def test_union_datasource(self):
         # Given
-        expected_query_dict = {"queryType": None, "dataSource": "things"}
+        expected_query_dict = {
+            "queryType": None,
+            "dataSource": {"name": "things", "type": "table"},
+        }
         builder = QueryBuilder()
         # when
         builder_dict = {"datasource": "things"}
@@ -187,7 +189,7 @@ def test_build_subquery(self):
         expected_query_dict = {
             "query": {
                 "queryType": "groupBy",
-                "dataSource": "things",
+                "dataSource": {"name": "things", "type": "table"},
                 "aggregations": [
                     {"fieldName": "thing", "name": "count", "type": "count"}
                 ],