influxdata · aviau · Apr 9, 2015 · Apr 1, 2015 · Apr 1, 2015 · Apr 1, 2015
@@ -8,6 +8,7 @@
 import requests
 import requests.exceptions
 
+from influxdb.resultset import ResultSet
 
 try:
     xrange
@@ -109,7 +110,7 @@ def __init__(self,
     # if one doesn't care in that, then it can simply change its client
     # instance 'keep_json_response_order' attribute value (to a falsy one).
     # This will then eventually help for performance considerations.
-    _keep_json_response_order = True
+    _keep_json_response_order = False
     # NB: For "group by" query type :
     # This setting is actually necessary in order to have a consistent and
     # reproducible rsp format if you "group by" on more than 1 tag.
@@ -122,33 +123,6 @@ def keep_json_response_order(self):
     def keep_json_response_order(self, new_value):
         self._keep_json_response_order = new_value
 
-    @staticmethod
-    def format_query_response(response):
-        """Returns a list of items from a query response"""
-        series = {}
-        if 'results' in response:
-            for result in response['results']:
-                if 'series' in result:
-                    for row in result['series']:
-                        items = []
-                        if 'name' in row:
-                            name = row['name']
-                            tags = row.get('tags', None)
-                            if tags:
-                                name = (row['name'], tuple(tags.items()))
-                            assert name not in series
-                            series[name] = items
-                        else:
-                            series = items  # Special case for system queries.
-                        if 'columns' in row and 'values' in row:
-                            columns = row['columns']
-                            for value in row['values']:
-                                item = {}
-                                for cur_col, field in enumerate(value):
-                                    item[columns[cur_col]] = field
-                                items.append(item)
-        return series
-
     def switch_database(self, database):
         """
         switch_database()
@@ -234,15 +208,16 @@ def query(self,
               query,
               params={},
               expected_response_code=200,
-              database=None,
-              raw=False):
+              database=None):
         """
         Query data
 
         :param params: Additional parameters to be passed to requests.
         :param database: Database to query, default to None.
         :param expected_response_code: Expected response code. Defaults to 200.
-        :param raw: Wether or not to return the raw influxdb response.
+
+        :rtype : ResultSet
+
         """
 
         params['q'] = query
@@ -261,8 +236,7 @@ def query(self,
             json_kw.update(object_pairs_hook=OrderedDict)
         data = response.json(**json_kw)
 
-        return (data if raw
-                else self.format_query_response(data))
+        return ResultSet(data)
 
     def write_points(self,
                      points,
@@ -327,8 +301,7 @@ def get_list_database(self):
         """
         Get the list of databases
         """
-        rsp = self.query("SHOW DATABASES")
-        return [db['name'] for db in rsp['databases']]
+        return list(self.query("SHOW DATABASES")['databases'])
 
     def create_database(self, dbname):
         """
@@ -368,21 +341,31 @@ def get_list_retention_policies(self, database=None):
         """
         Get the list of retention policies
         """
-        return self.query(
+        rsp = self.query(
             "SHOW RETENTION POLICIES %s" % (database or self._database)
         )
+        return list(rsp['results'])
 
     def get_list_series(self, database=None):
         """
         Get the list of series
         """
-        return self.query("SHOW SERIES", database=database)
+        rsp = self.query("SHOW SERIES", database=database)
+        series = []
+        for serie in rsp.items():
+            series.append(
+                {
+                    "name": serie[0][0],
+                    "tags": list(serie[1])
+                }
+            )
+        return series
 
     def get_list_users(self):
         """
         Get the list of users
         """
-        return self.query("SHOW USERS")
+        return list(self.query("SHOW USERS"))
 
     def delete_series(self, name, database=None):
         database = database or self._database

@@ -0,0 +1,118 @@
+# -*- coding: utf-8 -*-
+
+_sentinel = object()
+
+
+class ResultSet(object):
+    """A wrapper around series results """
+
+    def __init__(self, series):
+        self.raw = series
+
+    def __getitem__(self, key):
+        """
+        :param key: Either a serie name or a 2-tuple(serie_name, tags_dict)
+                    If the given serie name is None then any serie (matching
+                    the eventual given tags) will be given its points one
+                    after the other.
+        :return: A generator yielding `Point`s matching the given key.
+        NB:
+        The order in which the points are yielded is actually undefined but
+        it might change..
+        """
+        if isinstance(key, tuple):
+            if 2 != len(key):
+                raise TypeError('only 2-tuples allowed')
+            name = key[0]
+            tags = key[1]
+            if not isinstance(tags, dict) and tags is not None:
+                raise TypeError('tags should be a dict')
+        else:
+            name = key
+            tags = None
+        # TODO(aviau): Fix for python 3.2
+        # if not isinstance(name, (str, bytes, type(None))) \
+        #        and not isinstance(name, type("".decode("utf-8"))):
+        #    raise TypeError('serie_name must be an str or None')
+
+        for serie in self._get_series():
+            serie_name = serie.get('name', 'results')
+            if serie_name is None:
+                # this is a "system" query or a query which
+                # doesn't return a name attribute.
+                # like 'show retention policies' ..
+                if key is None:
+                    for point in serie['values']:
+                        yield self.point_from_cols_vals(
+                            serie['columns'],
+                            point
+                        )
+
+            elif name in (None, serie_name):
+                # by default if no tags was provided then
+                # we will matches every returned serie
+                serie_tags = serie.get('tags', {})
+                if tags is None or self._tag_matches(serie_tags, tags):
+                    for point in serie.get('values', []):
+                        yield self.point_from_cols_vals(
+                            serie['columns'],
+                            point
+                        )
+
+    def __repr__(self):
+        return str(self.raw)
+
+    def __iter__(self):
+        """ Iterating a ResultSet will yield one dict instance per serie result.
+        """
+        for key in self.keys():
+            yield list(self.__getitem__(key))
+
+    def _tag_matches(self, tags, filter):
+        """Checks if all key/values in filter match in tags"""
+        for tag_name, tag_value in filter.items():
+            # using _sentinel as I'm not sure that "None"
+            # could be used, because it could be a valid
+            # serie_tags value : when a serie has no such tag
+            # then I think it's set to /null/None/.. TBC..
+            serie_tag_value = tags.get(tag_name, _sentinel)
+            if serie_tag_value != tag_value:
+                return False
+        return True
+
+    def _get_series(self):
+        """Returns all series"""
+        series = []
+        try:
+            for result in self.raw['results']:
+                series.extend(result['series'])
+        except KeyError:
+            pass
+        return series
+
+    def __len__(self):
+        return len(self.keys())
+
+    def keys(self):
+        keys = []
+        for serie in self._get_series():
+            keys.append(
+                (serie.get('name', 'results'), serie.get('tags', None))
+            )
+        return keys
+
+    def items(self):
+        items = []
+        for serie in self._get_series():
+            serie_key = (serie.get('name', 'results'), serie.get('tags', None))
+            items.append(
+                (serie_key, self[serie_key])
+            )
+        return items
+
+    @staticmethod
+    def point_from_cols_vals(cols, vals):
+        point = {}
+        for col_index, col_name in enumerate(cols):
+            point[col_name] = vals[col_index]
+        return point
@@ -265,12 +265,11 @@ def test_query(self):
                 "http://localhost:8086/query",
                 text=example_response
             )
-            self.assertDictEqual(
-                self.cli.query('select * from foo'),
-                {'cpu_load_short':
-                    [{'value': 0.64, 'time': '2009-11-10T23:00:00Z'}],
-                 'sdfsdfsdf':
-                    [{'value': 0.64, 'time': '2009-11-10T23:00:00Z'}]}
+            rs = self.cli.query('select * from foo')
+
+            self.assertListEqual(
+                list(rs['cpu_load_short']),
+                [{'value': 0.64, 'time': '2009-11-10T23:00:00Z'}]
             )
 
     @unittest.skip('Not implemented for 0.9')
@@ -349,14 +348,19 @@ def test_drop_database_fails(self):
             cli.drop_database('old_db')
 
     def test_get_list_database(self):
-        data = {'results': [{'series': [
-            {'name': 'databases', 'columns': ['name'],
-             'values': [['mydb'], ['myotherdb']]}]}]}
+        data = {'results': [
+            {'series': [
+                {'name': 'databases',
+                 'values': [
+                     ['new_db_1'],
+                     ['new_db_2']],
+                 'columns': ['name']}]}
+        ]}
 
         with _mocked_session(self.cli, 'get', 200, json.dumps(data)):
             self.assertListEqual(
                 self.cli.get_list_database(),
-                ['mydb', 'myotherdb']
+                [{'name': 'new_db_1'}, {'name': 'new_db_2'}]
             )
 
     @raises(Exception)
@@ -367,19 +371,23 @@ def test_get_list_database_fails(self):
 
     def test_get_list_series(self):
         example_response = \
-            '{"results": [{"series": [{"values": [["fsfdsdf", "24h0m0s", 2]],'\
-            ' "columns": ["name", "duration", "replicaN"]}]}]}'
+            '{"results": [{"series": [{"name": "cpu_load_short", "columns": ' \
+            '["_id", "host", "region"], "values": ' \
+            '[[1, "server01", "us-west"]]}]}]}'
 
         with requests_mock.Mocker() as m:
             m.register_uri(
                 requests_mock.GET,
                 "http://localhost:8086/query",
                 text=example_response
             )
+
             self.assertListEqual(
                 self.cli.get_list_series(),
-                [{'duration': '24h0m0s',
-                  'name': 'fsfdsdf', 'replicaN': 2}]
+                [{'name': 'cpu_load_short',
+                  'tags': [
+                      {'host': 'server01', '_id': 1, 'region': 'us-west'}
+                  ]}]
             )
 
     def test_create_retention_policy_default(self):