Add support for querying the Riemann index

* transport: Use send and recv as method names, added recv
* client: Added send_query, create_dict, query
* client: Filter `None` values in create_dict
* command: Added query subcommand

Author: Sam Clements

riemann/client.py

@@ -31,22 +31,46 @@ def create_event(data):
         event = riemann.riemann_pb2.Event()
         event.tags.extend(data.pop('tags'))
         for name, value in data.items():
-            setattr(event, name, value)
+            if value is not None:
+                setattr(event, name, value)
         return event
 
     def send_event(self, event):
         """Wraps an event in a message and sends it to Riemann"""
         message = riemann.riemann_pb2.Msg()
         message.events.extend([event])
-        self.transport.write(message)
+        self.transport.send(message)
         return event
 
     def event(self, **data):
         """Sends an event"""
         return self.send_event(self.create_event(data))
 
-    def query(self, *args, **kwargs):
-        raise NotImplemented("Querying is not yet supported")
+    @staticmethod
+    def create_dict(event):
+        """Creates a dict from an Event"""
+        return {
+            'time': event.time,
+            'state': event.state,
+            'host': event.host,
+            'description': event.description,
+            'service': event.service,
+            'tags': list(event.tags),
+            'ttl': event.ttl,
+            'metric_f': event.metric_f,
+            'metric_d': event.metric_d,
+            'metric_sint64': event.metric_sint64
+        }
+
+    def send_query(self, query):
+        message = riemann.riemann_pb2.Msg()
+        message.query.string = query
+        self.transport.send(message)
+        return self.transport.recv()
+
+    def query(self, query):
+        response = self.send_query(query)
+        return [self.create_dict(e) for e in response.events]
 
 
 class QueuedClient(Client):
@@ -58,7 +82,7 @@ def __init__(self, *args, **kwargs):
 
     def flush(self):
         """Sends the waiting message to Riemann"""
-        self.transport.write(self.queue)
+        self.transport.send(self.queue)
         self.queue = riemann.riemann_pb2.Msg()
 
     def send_event(self, event):

riemann/command.py

@@ -3,6 +3,7 @@
 from __future__ import absolute_import, print_function
 
 import argparse
+import json
 
 import riemann.client
 import riemann.transport
@@ -60,29 +61,36 @@ def wide_formatter(*args, **kwargs):
     send.add_argument(
         *arg_names, metavar=arg_type.__name__, type=arg_type, help=arg_help)
 
-
-def filter_dict(function, dictionary):
-    return dict((k, v) for k, v in dictionary.items() if function(k, v))
+query = subparsers.add_parser('query', help='Query the Riemann index')
+query.add_argument(
+    'query', type=str,
+    help="The query to send")
+query.add_argument(
+    '-pp', '--pretty-print', action='store_true',
+    help="Pretty print output")
 
 
 def send(args, client):
-    event = client.event(**filter_dict(lambda k, v: v is not None, {
-        "time": args.time,
-        "state": args.state,
-        "host": args.event_host,
-        "description": args.description,
-        "service": args.service,
-        "tags": args.tags,
-        "ttl": args.ttl,
-        "metric_f": args.metric
-    }))
+    event = client.event(
+        time=args.time,
+        state=args.state,
+        host=args.host,
+        description=args.description,
+        service=args.service,
+        tags=args.tags,
+        ttl=args.ttl,
+        metric_f=args.metric)
 
     if args.print_message:
         print(str(event).strip())
 
 
+def query(args, client):
+    print(json.dumps(client.query(args.query), sort_keys=True, indent=2))
+
+
 def main():
     args = parser.parse_args()
     transport = TRANSPORT_CLASSES[args.transport](args.host, args.port)
     with riemann.client.Client(transport=transport) as client:
-        {'send': send}[args.subparser](args, client)
+        {'send': send, 'query': query}[args.subparser](args, client)

riemann/tests/test_riemann.py

@@ -6,16 +6,22 @@
 import py.test
 
 import riemann.client
+import riemann.riemann_pb2
 import riemann.transport
 
 
 class StringTransport(riemann.transport.Transport):
     def connect(self):
         self.string = StringIO.StringIO()
 
-    def write(self, message):
+    def send(self, message):
         self.string.write(message.SerializeToString())
 
+    def recv(self):
+        message = riemann.riemann_pb2.Msg()
+        message.ok = True
+        return message
+
     def disconnect(self):
         self.string.close()
 
@@ -33,8 +39,6 @@ def disconnect():
 
 
 class TestClient(object):
-    """Tests Client.create_event()"""
-
     def test_service(self, client):
         client.event(service='test event')
         assert 'test event' in client.transport.string.getvalue()
@@ -57,9 +61,8 @@ def test_event_cls(self, client):
         assert isinstance(client.create_event({}), riemann.riemann_pb2.Event)
         assert isinstance(client.event(), riemann.riemann_pb2.Event)
 
-    @py.test.mark.xfail
     def test_query(self, client):
-        client.query("true")
+        assert client.query("true") == []
 
 
 @py.test.fixture

riemann/transport.py

@@ -29,6 +29,10 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_value, traceback):
         self.disconnect()
 
+    @property
+    def address(self):
+        return self.host, self.port
+
     @abc.abstractmethod
     def connect(self):
         pass
@@ -38,12 +42,12 @@ def disconnect(self):
         pass
 
     @abc.abstractmethod
-    def write(self):
+    def send(self):
         pass
 
-    @property
-    def address(self):
-        return self.host, self.port
+    @abc.abstractmethod
+    def recv(self):
+        pass
 
 
 class UDPTransport(Transport):
@@ -53,29 +57,37 @@ def connect(self):
     def disconnect(self):
         self.socket.close()
 
-    def write(self, message):
+    def send(self, message):
         self.socket.sendto(message.SerializeToString(), self.address)
 
+    def recv(self):
+        raise NotImplementedError
+
 
 class TCPTransport(Transport):
     def connect(self):
         self.socket = socket.create_connection(self.address)
+        self.socket.setblocking(True)
 
     def disconnect(self):
         self.socket.close()
 
-    def write(self, message):
-        # Send the message to the server
-        data = message.SerializeToString()
-        self.socket.send(struct.pack('!I', len(data)) + data)
+    def send(self, message):
+        message = message.SerializeToString()
+        self.socket.sendall(struct.pack('!I', len(message)) + message)
+
+    def recv(self):
+        length = struct.unpack('!I', self.socket.recv(4))[0]
 
-        # Return the server's response
-        response_len = struct.unpack('!I', self.socket.recv(4))[0]
         response = riemann.riemann_pb2.Msg()
-        response.ParseFromString(self.socket.recv(response_len))
+        response.ParseFromString(self.recvall(length))
 
-        # Handle error messages
         if not response.ok:
             raise RiemannError(response.error)
-
         return response
+
+    def recvall(self, length, bufsize=4096):
+        data = ""
+        while len(data) < length:
+            data += self.socket.recv(bufsize)
+        return data

setup.py

@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name = "riemann-client",
-    version = '2.1.0',
+    version = '3.0.0-beta',
 
     author = "Sam Clements",
     author_email = "sam.clements@datasift.com",