Added a timeout to the cli and TCPTransport (closes #10)

Author: borntyping

riemann_client/command.py

@@ -14,11 +14,14 @@
 
 
 def udp_transport_factory(args):
+    if args.timeout is not None:
+        parser.error('--timeout cannot be used with the UDP transport')
     return riemann_client.transport.UDPTransport(args.host, args.port)
 
 
 def tcp_transport_factory(args):
-    return riemann_client.transport.TCPTransport(args.host, args.port)
+    return riemann_client.transport.TCPTransport(
+        args.host, args.port, args.timeout)
 
 
 def tls_transport_factory(args):
@@ -68,6 +71,10 @@ def tls_transport_factory(args):
     '-t', '--transport', choices=TRANSPORT_FACTORIES.keys(), default='tcp',
     help="The transport to use (default: %(default)s)")
 
+parser.add_argument(
+    '-T', '--timeout', type=float,
+    help="Timeout for TCP connections (default: %(default)s)")
+
 subparsers = parser.add_subparsers(dest='subparser')
 
 
@@ -79,7 +86,7 @@ def send_function(args, client):
         'host': args.event_host,
         'description': args.description,
         'service': args.service,
-        'tags': map(str.strip, args.tags.split(',')),
+        'tags': args.tags and map(str.strip, args.tags.split(',')),
         'ttl': args.ttl,
         'metric_f': args.metric})
 

riemann_client/transport.py

@@ -10,6 +10,12 @@
 import riemann_client.riemann_pb2
 
 
+# Default arguments
+HOST = 'localhost'
+PORT = 5555
+TIMEOUT = None
+
+
 def socket_recvall(socket, length, bufsize=4096):
     """Recives bytes from a socket until the buffer is the requested length"""
     data = ""
@@ -51,7 +57,7 @@ def send(self):
 class SocketTransport(Transport):
     """Provides common functionality for Transports using sockets"""
 
-    def __init__(self, host='localhost', port=5555):
+    def __init__(self, host=HOST, port=PORT):
         self.host = host
         self.port = port
 
@@ -74,6 +80,7 @@ def socket(self, value):
 class UDPTransport(SocketTransport):
     def connect(self):
         self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        self.socket.settimeout(self.timeout)
 
     def disconnect(self):
         self.socket.close()
@@ -84,9 +91,12 @@ def send(self, message):
 
 
 class TCPTransport(SocketTransport):
+    def __init__(self, host=HOST, port=PORT, timeout=TIMEOUT):
+        super(TCPTransport, self).__init__(host, port)
+        self.timeout = timeout
+
     def connect(self):
-        self.socket = socket.create_connection(self.address)
-        self.socket.setblocking(True)
+        self.socket = socket.create_connection(self.address, self.timeout)
 
     def disconnect(self):
         self.socket.close()
@@ -106,8 +116,8 @@ def send(self, message):
 
 
 class TLSTransport(TCPTransport):
-    def __init__(self, host='localhost', port=5554, ca_certs=None):
-        super(TLSTransport, self).__init__(host, port)
+    def __init__(self, host=HOST, port=PORT, timeout=TIMEOUT, ca_certs=None):
+        super(TLSTransport, self).__init__(host, port, timeout)
         self.ca_certs = ca_certs
 
     def connect(self):