gcb: add wait_for_build. (#449)

infra/gcb/build.py

@@ -7,10 +7,11 @@ Usage: build.py <project_dir>
 
 import datetime
 import os
-import pprint
 import sys
 import yaml
 
+from google.cloud import logging
+from google.cloud import pubsub
 from oauth2client.client import GoogleCredentials
 from googleapiclient.discovery import build
 
@@ -42,10 +43,41 @@ def load_project_yaml(project_dir):
     return project_yaml
 
 
+def create_log_topic(build_id):
+  pubsub_client = pubsub.Client()
+  log_topic = pubsub_client.topic('build-logs-' + build_id)
+  if log_topic.exists():
+    log_topic.delete()
+  log_topic.create()
+
+  policy = log_topic.get_iam_policy()
+  policy.owners.add(policy.group('cloud-logs@google.com'))
+  log_topic.set_iam_policy(policy)
+
+  return log_topic
+
+
+def create_sink(log_topic, build_id):
+  log_destination = 'pubsub.googleapis.com/' + log_topic.full_name
+  log_filter = (
+      'resource.type="build"\n'
+      'resource.labels.build_id="{0}"\n'
+  ).format(build_id)
+
+  logging_client = logging.Client()
+  sink = logging_client.sink(
+      'sink-build-logs-' + build_id, filter_=log_filter,
+      destination=log_destination)
+  if sink.exists():
+    sink.delete()
+  sink.create()
+
+  return sink
+
+
 def get_build_steps(project_yaml):
   name = project_yaml['name']
   image = project_yaml['image']
-  print "Building " + image
 
   ts = datetime.datetime.now().strftime('%Y%m%d%H%M')
 
@@ -106,8 +138,6 @@ def main():
   options = {}
   if "GCB_OPTIONS" in os.environ:
     options = yaml.safe_load(os.environ["GCB_OPTIONS"])
-    print "Using options", options
-
 
   build_body = {
       'source': {
@@ -125,8 +155,14 @@ def main():
 
   credentials = GoogleCredentials.get_application_default()
   cloudbuild = build('cloudbuild', 'v1', credentials=credentials)
-  pp = pprint.PrettyPrinter(indent=4)
-  pp.pprint(cloudbuild.projects().builds().create(projectId='clusterfuzz-external', body=build_body).execute())
+  build_info = cloudbuild.projects().builds().create(projectId='clusterfuzz-external', body=build_body).execute()
+  build_id =  build_info['metadata']['build']['id']
+
+  print build_id
+
+  # Create pub/sub topic for build logs.
+  log_topic = create_log_topic(build_id)
+  create_sink(log_topic, build_id)
 
 
 if __name__ == "__main__":

infra/gcb/requirements.txt

@@ -1,3 +1,5 @@
 google-api-python-client
 PyYAML
+google-cloud-pubsub
+google-cloud-logging
 

infra/gcb/wait_for_build.py

@@ -0,0 +1,87 @@
+#!/usr/bin/python2
+
+"""Waits for project build on Google Cloud Builder.
+
+Usage: wait_for_build.py <build_id>
+"""
+
+import json
+import sys
+import time
+import thread
+import threading
+
+from google.cloud import logging
+from google.cloud import pubsub
+from googleapiclient.discovery import build
+from oauth2client.client import GoogleCredentials
+
+POLL_INTERVAL = 15
+status = None
+
+
+def usage():
+  sys.stderr.write(
+    "Usage: " + sys.argv[0] + " <build_id>\n")
+  exit(1)
+
+
+def create_log_subscription(log_topic, build_id):
+  log_sub = log_topic.subscription('build-sub-' + build_id)
+  if log_sub.exists():
+    log_sub.delete()
+
+  log_sub.create()
+  return log_sub
+
+
+def poll_build_status_thread(build_id):
+  global status
+  credentials = GoogleCredentials.get_application_default()
+  cloudbuild = build('cloudbuild', 'v1', credentials=credentials)
+
+  while True:
+    build_info = cloudbuild.projects().builds().get(
+        projectId='clusterfuzz-external', id=build_id).execute()
+    status = build_info['status']
+    if status == 'SUCCESS' or status == 'FAILURE':
+      thread.interrupt_main()
+      return
+
+    time.sleep(POLL_INTERVAL)
+
+
+def main():
+  if len(sys.argv) != 2:
+    usage()
+
+  build_id = sys.argv[1]
+
+  pubsub_client = pubsub.Client()
+  log_topic = pubsub_client.topic('build-logs-' + build_id)
+  assert log_topic.exists()
+
+  status_thread = threading.Thread(target=poll_build_status_thread,
+                                   args=(build_id,))
+  status_thread.daemon = True
+  status_thread.start()
+
+  try:
+    log_sub = create_log_subscription(log_topic, build_id)
+    while True:
+      pulled = log_sub.pull(max_messages=32)
+      for ack_id, message in pulled:
+        print json.loads(message.data)['textPayload']
+        log_sub.acknowledge([ack_id])
+  except KeyboardInterrupt:
+    if status:
+      print status
+
+    if status == 'SUCCESS':
+      sys.exit(0)
+
+    sys.exit(1)
+
+
+if __name__ == "__main__":
+  main()