Merge pull request influxdata#292 from amancevice/master

aviau · aviau · commit 1e8e3ddf79f9 · 2016-02-26T11:57:12.000-05:00
Allow multiple statements in single query
diff --git a/influxdb/_dataframe_client.py b/influxdb/_dataframe_client.py
@@ -101,6 +101,8 @@ def get_list_series(self, database=None):
 
     def _to_dataframe(self, rs):
         result = {}
+        if isinstance(rs, list):
+            return map(self._to_dataframe, rs)
         for key, data in rs.items():
             name, tags = key
             if tags is None:
diff --git a/influxdb/tests/dataframe_client_test.py b/influxdb/tests/dataframe_client_test.py
@@ -244,6 +244,56 @@ def test_query_into_dataframe(self):
             for k in expected:
                 assert_frame_equal(expected[k], result[k])
 
+    def test_multiquery_into_dataframe(self):
+        data = {
+            "results": [
+                {
+                    "series": [
+                        {
+                            "name": "cpu_load_short",
+                            "columns": ["time", "value"],
+                            "values": [
+                                ["2015-01-29T21:55:43.702900257Z", 0.55],
+                                ["2015-01-29T21:55:43.702900257Z", 23422],
+                                ["2015-06-11T20:46:02Z", 0.64]
+                            ]
+                        }
+                    ]
+                }, {
+                    "series": [
+                        {
+                            "name": "cpu_load_short",
+                            "columns": ["time", "count"],
+                            "values": [
+                                ["1970-01-01T00:00:00Z", 3]
+                            ]
+                        }
+                    ]
+                }
+            ]
+        }
+
+        pd1 = pd.DataFrame(
+            [[0.55], [23422.0], [0.64]], columns=['value'],
+            index=pd.to_datetime([
+                "2015-01-29 21:55:43.702900257+0000",
+                "2015-01-29 21:55:43.702900257+0000",
+                "2015-06-11 20:46:02+0000"])).tz_localize('UTC')
+        pd2 = pd.DataFrame(
+            [[3]], columns=['count'],
+            index=pd.to_datetime(["1970-01-01 00:00:00+00:00"]))\
+            .tz_localize('UTC')
+        expected = [{'cpu_load_short': pd1}, {'cpu_load_short': pd2}]
+
+        cli = DataFrameClient('host', 8086, 'username', 'password', 'db')
+        iql = "SELECT value FROM cpu_load_short WHERE region='us-west';"\
+            "SELECT count(value) FROM cpu_load_short WHERE region='us-west'"
+        with _mocked_session(cli, 'GET', 200, data):
+            result = cli.query(iql)
+            for r, e in zip(result, expected):
+                for k in e:
+                    assert_frame_equal(e[k], r[k])
+
     def test_query_with_empty_result(self):
         cli = DataFrameClient('host', 8086, 'username', 'password', 'db')
         with _mocked_session(cli, 'GET', 200, {"results": [{}]}):
