11#!/usr/bin/env python3
22
3- import firebase_admin
4- from firebase_admin import credentials
5- from firebase_admin import firestore
6- from datetime import datetime , timezone
73import json
84import argparse
95import re
106import os
7+ from datetime import datetime , timezone
8+
9+ # --- FIRESTORE (remove when deprecating) ---
10+ import firebase_admin
11+ from firebase_admin import credentials
12+ from firebase_admin import firestore
13+ # --- END FIRESTORE ---
14+
15+ # --- PUBSUB ---
16+ from google .cloud import pubsub_v1
17+ from google .oauth2 import service_account
18+ # --- END PUBSUB ---
1119
1220# make sure the working dir is flow/
1321os .chdir (os .path .join (os .path .dirname (os .path .abspath (__file__ )), ".." ))
2735parser .add_argument ("--cred" , type = str , help = "Service account credentials file" )
2836parser .add_argument ("--variant" , type = str , default = "base" )
2937
38+ # --- PUBSUB args ---
39+ parser .add_argument ("--pubsubProjectID" , type = str , help = "GCP project ID for Pub/Sub" )
40+ parser .add_argument (
41+ "--pubsubTopicID" ,
42+ type = str ,
43+ default = "ci-metrics-reports-topics" ,
44+ help = "Pub/Sub topic ID" ,
45+ )
46+ parser .add_argument (
47+ "--pubsubCred" , type = str , help = "Service account credentials file for Pub/Sub"
48+ )
49+ # --- END PUBSUB args ---
50+
3051# Parse the arguments
3152args = parser .parse_args ()
3253
3354
55+ # --- FIRESTORE (remove when deprecating) ---
3456def upload_data (db , dataFile , platform , design , variant , args , rules ):
3557 # Set the document data
3658 key = args .commitSHA + "-" + platform + "-" + design + "-" + variant
@@ -166,6 +188,39 @@ def upload_data(db, dataFile, platform, design, variant, args, rules):
166188 raise Exception (f"Failed to upload data for { platform } { design } { variant } ." )
167189
168190
191+ # --- END FIRESTORE ---
192+
193+
194+ # --- PUBSUB ---
195+ def publish_to_pubsub (publisher , topic_path , dataFile , platform , design , variant , args , rules ):
196+ """Publish a single design's metrics to Pub/Sub as a JSON message."""
197+ with open (dataFile ) as f :
198+ data = json .load (f )
199+
200+ # Build the payload: CLI args + metrics with ':' replaced by '__'
201+ payload = {
202+ "build_id" : args .buildID ,
203+ "branch_name" : args .branchName ,
204+ "pipeline_id" : args .pipelineID ,
205+ "change_branch" : args .changeBranch ,
206+ "commit_sha" : args .commitSHA ,
207+ "jenkins_url" : args .jenkinsURL ,
208+ "rules" : rules ,
209+ }
210+
211+ for k , v in data .items ():
212+ new_key = re .sub (":" , "__" , k )
213+ payload [new_key ] = v
214+
215+ message_data = json .dumps (payload ).encode ("utf-8" )
216+ future = publisher .publish (topic_path , data = message_data )
217+ message_id = future .result ()
218+ print (f"[INFO] Published to Pub/Sub (message ID: { message_id } ) for { platform } { design } { variant } ." )
219+
220+
221+ # --- END PUBSUB ---
222+
223+
169224def get_rules (dataFile ):
170225 data = {}
171226 if os .path .exists (dataFile ):
@@ -175,10 +230,29 @@ def get_rules(dataFile):
175230 return data
176231
177232
178- # Initialize Firebase Admin SDK with service account credentials
179- firebase_admin .initialize_app (credentials .Certificate (args .cred ))
180- # Initialize Firestore client
181- db = firestore .client ()
233+ # --- FIRESTORE init (remove when deprecating) ---
234+ db = None
235+ if args .cred :
236+ firebase_admin .initialize_app (credentials .Certificate (args .cred ))
237+ db = firestore .client ()
238+ # --- END FIRESTORE init ---
239+
240+ # --- PUBSUB init ---
241+ publisher = None
242+ topic_path = None
243+ if args .pubsubCred and args .pubsubProjectID :
244+ pubsub_credentials = service_account .Credentials .from_service_account_file (
245+ args .pubsubCred
246+ )
247+ publisher = pubsub_v1 .PublisherClient (credentials = pubsub_credentials )
248+ topic_path = publisher .topic_path (args .pubsubProjectID , args .pubsubTopicID )
249+ print (f"[INFO] Pub/Sub publisher initialized for topic: { topic_path } " )
250+ elif args .pubsubProjectID :
251+ # No credentials file — use default credentials (e.g., emulator or ADC)
252+ publisher = pubsub_v1 .PublisherClient ()
253+ topic_path = publisher .topic_path (args .pubsubProjectID , args .pubsubTopicID )
254+ print (f"[INFO] Pub/Sub publisher initialized (default creds) for topic: { topic_path } " )
255+ # --- END PUBSUB init ---
182256
183257RUN_FILENAME = "metadata.json"
184258
@@ -199,6 +273,20 @@ def get_rules(dataFile):
199273 print (f"[WARN] Skiping upload { platform } { design } { variant } ." )
200274 continue
201275 print (f"[INFO] Get rules for { platform } { design } { variant } ." )
202- rules = get_rules (os .path .join ("designs" , platform , design , RUN_FILENAME ))
203- print (f"[INFO] Upload data for { platform } { design } { variant } ." )
204- upload_data (db , dataFile , platform , design , variant , args , rules )
276+ rules = get_rules (os .path .join ("designs" , platform , design , f"rules-{ variant } .json" ))
277+
278+ # --- FIRESTORE (remove when deprecating) ---
279+ if db :
280+ print (f"[INFO] Upload data for { platform } { design } { variant } ." )
281+ upload_data (db , dataFile , platform , design , variant , args , rules )
282+ # --- END FIRESTORE ---
283+
284+ # --- PUBSUB ---
285+ if publisher :
286+ try :
287+ publish_to_pubsub (
288+ publisher , topic_path , dataFile , platform , design , variant , args , rules
289+ )
290+ except Exception as e :
291+ print (f"[WARN] Pub/Sub publish failed for { platform } { design } { variant } : { e } " )
292+ # --- END PUBSUB ---
0 commit comments