Better zgrep.

utils/zgrep.py

@@ -1,66 +1,95 @@
 #!/usr/bin/env python
 
-from optparse import OptionParser
 import os
+import time
 import select
 import subprocess
 import sys
-from vendor import yaml
+from optparse import OptionParser
+from requests.exceptions import ConnectionError
 
+sys.path.insert(0, '/srv/newsblur')
+os.environ['DJANGO_SETTINGS_MODULE'] = 'settings'
+import fabfile
 
+NEWSBLUR_USERNAME = 'sclay'
 IGNORE_HOSTS = [
     'push',
 ]
 
-def main(role="app", role2="dev", command=None, path=None):
+def main(role="app", role2="work", command=None, path=None):
+    delay = 1
+
+    while True:
+        try:
+            streams = create_streams_for_roles(role, role2, command=command, path=path)
+            print " --- Loading %s App Log Tails ---" % len(streams)
+            read_streams(streams)
+        except UnicodeDecodeError: # unexpected end of data
+            print " --- Lost connections - Retrying... ---"
+            time.sleep(1)
+            continue
+        except ConnectionError:
+            print " --- Retrying in %s seconds... ---" % delay
+            time.sleep(delay)
+            delay += 1
+            continue
+        except KeyboardInterrupt:
+            print " --- End of Logging ---"
+            break
+
+def create_streams_for_roles(role, role2, command=None, path=None):
     streams = list()
+    hosts = fabfile.do(split=True)
+    found = set()
+
     if not path:
         path = "/srv/newsblur/logs/newsblur.log"
     if not command:
         command = "tail -f"
-    hosts_path = os.path.expanduser(os.path.join('../secrets-newsblur/configs/hosts.yml'))
-    hosts = yaml.load(open(hosts_path))
-    
-    for r in [role, role2]:
-        if isinstance(hosts[r], dict):
-            hosts[r] = ["%s:%s" % (hosts[r][k][-1], k) for k in hosts[r].keys()]
-    
-    for hostname in set(hosts[role] + hosts[role2]):
-        if any(h in hostname for h in IGNORE_HOSTS): continue
-        if ':' in hostname:
+    for hostname in (hosts[role] + hosts[role2]):
+        if isinstance(hostname, dict):
+            address = hostname['address']
+            hostname = hostname['name']
+        elif ':' in hostname:
             hostname, address = hostname.split(':', 1)
+        elif isinstance(hostname, tuple):
+            hostname, address = hostname[0], hostname[1]
         else:
             address = hostname
+        if any(h in hostname for h in IGNORE_HOSTS): continue
+        if hostname in found: continue
         if 'ec2' in hostname:
-            s = subprocess.Popen(["ssh", "-i", os.path.expanduser("~/.ec2/sclay.pem"), 
+            s = subprocess.Popen(["ssh", 
+                                  "-i", os.path.expanduser(os.path.join(fabfile.env.SECRETS_PATH,
+                                                                        "keys/ec2.pem")),
                                   address, "%s %s" % (command, path)], stdout=subprocess.PIPE)
         else:
-            s = subprocess.Popen(["ssh", address, "%s %s" % (command, path)], stdout=subprocess.PIPE)
+            s = subprocess.Popen(["ssh", "-l", NEWSBLUR_USERNAME, 
+                                  "-i", os.path.expanduser(os.path.join(fabfile.env.SECRETS_PATH,
+                                                                        "keys/newsblur.key")),
+                                  address, "%s %s" % (command, path)], stdout=subprocess.PIPE)
         s.name = hostname
         streams.append(s)
+        found.add(hostname)
+        
+    return streams
 
-    try:
-        i = 0
-        while True:
-            i += 1
-            r, _, _ = select.select(
-                [stream.stdout.fileno() for stream in streams], [], [])
-            for fileno in r:
-                for stream in streams:
-                    if stream.stdout.fileno() != fileno:
-                        continue
-                    data = os.read(fileno, 4096)
-                    if not data:
-                        streams.remove(stream)
-                        break
-                    combination_message = "%s" % (data)
-                    sys.stdout.write(combination_message)
+def read_streams(streams):
+    while True:
+        r, _, _ = select.select(
+            [stream.stdout.fileno() for stream in streams], [], [])
+        for fileno in r:
+            for stream in streams:
+                if stream.stdout.fileno() != fileno:
+                    continue
+                data = os.read(fileno, 4096)
+                if not data:
+                    streams.remove(stream)
                     break
-            if i > 1000:
+                combination_message = "[%-6s] %s" % (stream.name[:6], data)
+                sys.stdout.write(combination_message)
                 break
-    except KeyboardInterrupt:
-        print " --- End of Logging ---"
-
 
 if __name__ == "__main__":
     parser = OptionParser()
@@ -71,4 +100,4 @@ if __name__ == "__main__":
     path = options.path
     find = options.find
     command = "zgrep \"%s\"" % find
-    main(role="app", role2="dev", command=command, path=path)
+    main(role="app", role2="dev", command=command, path=path)