Make some minor improvements to error handling

Author: JeremyMcCormick

cloud_functions/stage_chunk/main.py

@@ -70,8 +70,11 @@ def trigger_stage_chunk(event, context):
         message = base64.b64decode(event["data"]).decode("utf-8")
         data = json.loads(message)
 
-        bucket = data.get("bucket")
-        name = data.get("name")
+        try:
+            bucket = data["bucket"]
+            name = data["name"]
+        except KeyError as e:
+            raise ValueError(f"Missing expected key in Pub/Sub message: {e}")
 
         if not bucket or not name:
             raise ValueError("Missing 'bucket' or 'name' in Pub/Sub message")

cloud_functions/stage_chunk/stage_chunk_beam_job.py

@@ -26,7 +26,6 @@
 from apache_beam.io.filesystems import FileSystems
 import argparse
 import logging
-import sys
 
 
 class BeamSuppressUpdateDestinationSchemaWarning(logging.Filter):
@@ -159,6 +158,8 @@ def run(argv=None):
     input_path = custom_options.input_path
     triggering_frequency = custom_options.triggering_frequency
     priority = custom_options.priority
+    if priority not in ["BATCH", "INTERACTIVE"]:
+        raise ValueError(f"Invalid BigQuery priority: {priority}")
 
     # Validate available input files
     table_names = ["DiaObject", "DiaSource", "DiaForcedSource"]