Merge pull request #232 from mwvgroup/u/ch/store_BigQuery

Store LVK alerts in BigQuery

Author: hernandezc1

broker/cloud_functions/lvk/store_BigQuery/README.md

@@ -0,0 +1,3 @@
+# Store alerts in BigQuery
+
+This Cloud Function stores alert data in a BigQuery table.

broker/cloud_functions/lvk/store_BigQuery/deploy.sh

@@ -0,0 +1,43 @@
+#! /bin/bash
+# Deploys or deletes broker Cloud Function
+# This script will not delete Cloud Functions that are in production
+
+# "False" uses production resources
+# any other string will be appended to the names of all resources
+testid="${1:-test}"
+# "True" tearsdown/deletes resources, else setup
+teardown="${2:-False}"
+# name of the survey this broker instance will ingest
+survey="${3:-lvk}"
+# schema version
+versiontag="${4:-v1_0}"
+
+#--- GCP resources used in this script
+store_bq_trigger_topic="${survey}-alerts"
+store_bq_CF_name="${survey}-store_in_bigquery"
+
+# use test resources, if requested
+if [ "${testid}" != "False" ]; then
+    store_bq_trigger_topic="${store_bq_trigger_topic}-${testid}"
+    store_bq_CF_name="${store_bq_CF_name}-${testid}"
+fi
+
+if [ "${teardown}" = "True" ]; then
+    # ensure that we do not teardown production resources
+    if [ "${testid}" != "False" ]; then
+        gcloud functions delete "${store_bq_CF_name}"
+    fi
+
+else # Deploy the Cloud Functions
+#--- BigQuery storage cloud function
+    echo "Deploying Cloud Function: ${store_bq_CF_name}"
+    store_bq_entry_point="run"
+    memory=512MB
+
+    gcloud functions deploy "${store_bq_CF_name}" \
+        --entry-point "${store_bq_entry_point}" \
+        --runtime python312 \
+        --memory "${memory}" \
+        --trigger-topic "${store_bq_trigger_topic}" \
+        --set-env-vars TESTID="${testid}",SURVEY="${survey}",GCP_PROJECT="${GOOGLE_CLOUD_PROJECT}",VERSIONTAG="${versiontag}"
+fi

broker/cloud_functions/lvk/store_BigQuery/main.py

@@ -0,0 +1,94 @@
+#!/usr/bin/env python3
+# -*- coding: UTF-8 -*-
+
+"""This module stores alert data in BigQuery tables."""
+
+import base64
+import os
+import pittgoogle
+from google.cloud import functions_v1, pubsub_v1, logging
+
+PROJECT_ID = os.getenv("GCP_PROJECT")
+SURVEY = os.getenv("SURVEY")
+TESTID = os.getenv("TESTID")
+VERSIONTAG = os.getenv("VERSIONTAG")
+
+# connect to the cloud logger
+log_name = "store-bigquery-cloudfnc"  # same log for all broker instances
+logging_client = logging.Client()
+logger = logging_client.logger(log_name)
+
+# GCP resources used in this module
+TABLE = pittgoogle.Table.from_cloud(f"alerts_{VERSIONTAG}", survey=SURVEY, testid=TESTID)
+TOPIC = pittgoogle.Topic.from_cloud("bigquery", survey=SURVEY, testid=TESTID, projectid=PROJECT_ID)
+
+
+def run(event: dict, _context: functions_v1.context.Context) -> None:
+    """Send alert data to various BigQuery tables.
+
+    Args:
+        event: Pub/Sub message data and attributes.
+            `data` field contains the message data in a base64-encoded string.
+            `attributes` field contains the message's custom attributes in a dict.
+
+        context: Metadata describing the Cloud Function's trigging event.
+
+    'context' is an unused argument in the function that is required
+    see https://cloud.google.com/functions/1stgendocs/writing/write-event-driven-functions#background-functions
+    """
+
+    # decode the base64-encoded message data
+    decoded_data = base64.b64decode(event["data"])
+
+    # create a PubsubMessage-like object with the existing event dictionary
+    pubsub_message = pubsub_v1.types.PubsubMessage(
+        data=decoded_data, attributes=event.get("attributes", {})
+    )
+
+    # unpack the alert
+    alert = pittgoogle.Alert.from_msg(msg=pubsub_message, schema_name="default_schema")
+
+    # send the alert to BigQuery table
+    alert_table = insert_rows_alerts(alert)
+
+    # announce what's been done
+    TOPIC.publish(_create_outgoing_alert(alert, alert_table))
+
+
+def insert_rows_alerts(alert: pittgoogle.alert.Alert):
+    """Insert rows into the `alerts` table via the streaming API."""
+    # send to bigquery
+    errors = TABLE.insert_rows([alert])
+
+    # handle errors; if none, save table id for Pub/Sub message
+    if len(errors) == 0:
+        table_dict = {"alerts_table": f"{TABLE.id}"}
+    else:
+        msg = f"Error inserting to alerts table: {errors}"
+        logger.log_text(msg, severity="WARNING")
+        table_dict = {"alerts_table": None}
+
+    return table_dict
+
+
+def _create_outgoing_alert(
+    alert: pittgoogle.alert.Alert, table_dict: dict
+) -> pittgoogle.alert.Alert:
+    """Create an announcement of the table storage operation to Pub/Sub."""
+    # collect attributes
+    attrs = {
+        **alert.attributes,
+        "alerts_table": table_dict["alerts_table"],
+        "alert_type": alert.dict["alert_type"],
+        "superevent_id": alert.dict["superevent_id"],
+    }
+
+    # set empty message body; everything is in the attributes
+    msg = {}
+
+    # create outgoing alert
+    alert_out = pittgoogle.Alert.from_dict(
+        payload=msg, attributes=attrs, schema_name="default_schema"
+    )
+
+    return alert_out

broker/cloud_functions/lvk/store_BigQuery/requirements.txt

@@ -0,0 +1,8 @@
+# As explained here
+# https://cloud.google.com/functions/docs/writing/specifying-dependencies-python
+# dependencies for a Cloud Function must be specified in a `requirements.txt`
+# file (or packaged with the function) in the same directory as `main.py`
+
+google-cloud-functions
+google-cloud-logging
+pittgoogle-client>=0.3.11

broker/consumer/lvk/ps-connector.properties

@@ -27,7 +27,7 @@ name=CPSSinkConnector
 # Tha Java class for the Pub/Sub sink connector.
 connector.class=com.google.pubsub.kafka.sink.CloudPubSubSinkConnector
 # The maximum number of tasks that should be created for this connector.
-tasks.max=10
+tasks.max=1
 # Set the key converter for the Pub/Sub sink connector.
 key.converter=org.apache.kafka.connect.converters.ByteArrayConverter
 # Set the value converter for the Pub/Sub sink connector.

broker/consumer/lvk/psconnect-worker-authenticated.properties

@@ -6,6 +6,7 @@
 bootstrap.servers=kafka.gcn.nasa.gov:9092
 plugin.path=/usr/local/share/kafka/plugins
 offset.storage.file.filename=/tmp/connect.offsets
+connections.max.idle.ms=5400000
 
 # ByteArrayConverter provides a “pass-through” option that does no conversion.
 key.converter=org.apache.kafka.connect.converters.ByteArrayConverter
@@ -22,7 +23,7 @@ sasl.jaas.config=org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginMo
      clientSecret="CLIENT_SECRET";
 
 # settings with `consumer.` prefixes are passed through to the Kafka consumer
-consumer.group.id=pitt-google-broker-test
+consumer.group.id=GROUP_ID
 consumer.auto.offset.reset=earliest
 consumer.sasl.mechanism=OAUTHBEARER
 consumer.sasl.kerberos.service.name=kafka

broker/consumer/lvk/vm_startup.sh

@@ -23,7 +23,7 @@ fi
 
 #--- GCP resources used in this script
 broker_bucket="${PROJECT_ID}-${survey}-broker_files"
-PS_TOPIC_DEFAULT="${survey}-alerts_raw"
+PS_TOPIC_DEFAULT="${survey}-alerts"
 # use test resources, if requested
 if [ "$testid" != "False" ]; then
     broker_bucket="${broker_bucket}-${testid}"
@@ -60,21 +60,22 @@ client_id="${survey}-${PROJECT_ID}-client-id"
 client_secret="${survey}-${PROJECT_ID}-client-secret"
 CLIENT_ID=$(gcloud secrets versions access latest --secret="${client_id}")
 CLIENT_SECRET=$(gcloud secrets versions access latest --secret="${client_secret}")
+group_id="pittgooglebroker"
+# use test resources, if requested
+if [ "$testid" != "False" ]; then
+    group_id="${group_id}-${testid}"
+fi
 
 cd "${workingdir}" || exit
 
 fconfig=admin.properties
-sed -i "s/CLIENT_ID/${CLIENT_ID}/g" ${fconfig}
-sed -i "s/CLIENT_SECRET/${CLIENT_SECRET}/g" ${fconfig}
+sed -i "s/CLIENT_ID/${CLIENT_ID}/g" ${fconfig} && sed -i "s/CLIENT_SECRET/${CLIENT_SECRET}/g" ${fconfig}
 
 fconfig=psconnect-worker-authenticated.properties
-sed -i "s/CLIENT_ID/${CLIENT_ID}/g" ${fconfig}
-sed -i "s/CLIENT_SECRET/${CLIENT_SECRET}/g" ${fconfig}
+sed -i "s/CLIENT_ID/${CLIENT_ID}/g" ${fconfig} && sed -i "s/CLIENT_SECRET/${CLIENT_SECRET}/g" ${fconfig} && sed -i "s/GROUP_ID/${group_id}/g" ${fconfig}
 
 fconfig=ps-connector.properties
-sed -i "s/PROJECT_ID/${PROJECT_ID}/g" ${fconfig}
-sed -i "s/PS_TOPIC/${PS_TOPIC}/g" ${fconfig}
-sed -i "s/KAFKA_TOPIC/${KAFKA_TOPIC}/g" ${fconfig}
+sed -i "s/PROJECT_ID/${PROJECT_ID}/g" ${fconfig} && sed -i "s/PS_TOPIC/${PS_TOPIC}/g" ${fconfig} && sed -i "s/KAFKA_TOPIC/${KAFKA_TOPIC}/g" ${fconfig}
 
 #--- Check until alerts start streaming into the topic
 alerts_flowing=false

broker/setup_broker/lvk/README.md

@@ -99,9 +99,10 @@ Initialize parameters and call the deployment script:
 testid="mytest"
 teardown="False"
 survey="lvk"
+schema_version="1.0"
 region="us-central1"
 
-./setup_broker.sh "${testid}" "${teardown}" "${survey}" "${region}"
+./setup_broker.sh "${testid}" "${teardown}" "${survey}" "${schema_version}" "${region}"
 ```
 
 This will create all of the necessary GCP resources. Allow the consumer VM to finish its installation process. Once
@@ -141,7 +142,8 @@ Similar to [deploy broker instance](#deploy-broker-instance). Initialize paramet
 testid="mytest"
 teardown="True"
 survey="lvk"
+schema_version="1.0"
 region="us-central1"
 
-./setup_broker.sh "${testid}" "${teardown}" "${survey}" "${region}"
+./setup_broker.sh "${testid}" "${teardown}" "${survey}" "${schema_version}" "${region}"
 ```

broker/setup_broker/lvk/setup_broker.sh

@@ -1,14 +1,16 @@
 #! /bin/bash
 # Create and configure GCP resources needed to run the nightly broker.
 
-testid="${1:-test}"
 # "False" uses production resources
 # any other string will be appended to the names of all resources
-teardown="${2:-False}"
+testid="${1:-test}"
 # "True" tearsdown/deletes resources, else setup
-survey="${3:-lvk}"
+teardown="${2:-False}"
 # name of the survey this broker instance will ingest
-region="${4:-us-central1}"
+survey="${3:-lvk}"
+schema_version="${4:-1.0}"
+versiontag=v$(echo "${schema_version}" | tr . _) # 1.0 -> v1_0
+region="${5:-us-central1}"
 zone="${region}-a"  # just use zone "a" instead of adding another script arg
 
 PROJECT_ID=$GOOGLE_CLOUD_PROJECT # get the environment variable
@@ -20,6 +22,7 @@ echo
 echo "GOOGLE_CLOUD_PROJECT = ${PROJECT_ID}"
 echo "survey = ${survey}"
 echo "testid = ${testid}"
+echo "schema_version = ${schema_version}"
 echo "teardown = ${teardown}"
 echo
 echo "Continue?  [y/(n)]: "
@@ -34,15 +37,32 @@ fi
 
 #--- GCP resources used directly in this script
 broker_bucket="${PROJECT_ID}-${survey}-broker_files"
+bq_dataset="${survey}"
 topic_alerts="${survey}-alerts"
+topic_storebigquery="${survey}-bigquery"
+
 # use test resources, if requested
 if [ "$testid" != "False" ]; then
     broker_bucket="${broker_bucket}-${testid}"
+    bq_dataset="${bq_dataset}_${testid}"
     topic_alerts="${topic_alerts}-${testid}"
+    topic_storebigquery="${topic_storebigquery}-${testid}"
 fi
 
-#--- Create (or delete) GCS, Pub/Sub resources
+alerts_table="alerts_${versiontag}"
+
+#--- Create (or delete) BigQuery, GCS, Pub/Sub resources
+echo
+echo "Configuring BigQuery, GCS, Pub/Sub resources..."
 if [ "${teardown}" != "True" ]; then
+    # create bigquery dataset and table
+    bq --location="${region}" mk --dataset "${bq_dataset}"
+
+    cd templates || exit 5
+    bq mk --table "${PROJECT_ID}:${bq_dataset}.${alerts_table}" "bq_${survey}_${alerts_table}_schema.json" || exit 5
+    bq update --description "Alert data from LIGO/Virgo/KAGRA. This table is an archive of the lvk-alerts Pub/Sub stream. It has the same schema as the original alert bytes, including nested and repeated fields." "${PROJECT_ID}:${bq_dataset}.${alerts_table}"
+    cd .. || exit 5
+
     # create broker bucket and upload files
     echo "Creating broker_bucket and uploading files..."
     gsutil mb -b on -l "${region}" "gs://${broker_bucket}"
@@ -51,22 +71,37 @@ if [ "${teardown}" != "True" ]; then
     # create pubsub
     echo "Configuring Pub/Sub resources..."
     gcloud pubsub topics create "${topic_alerts}"
+    gcloud pubsub topics create "${topic_storebigquery}"
 
     # Set IAM policies on resources
     user="allUsers"
     roleid="projects/${GOOGLE_CLOUD_PROJECT}/roles/userPublic"
     gcloud pubsub topics add-iam-policy-binding "${topic_alerts}" --member="${user}" --role="${roleid}"
+    gcloud pubsub topics add-iam-policy-binding "${topic_storebigquery}" --member="${user}" --role="${roleid}"
 
 else
     # ensure that we do not teardown production resources
     if [ "${testid}" != "False" ]; then
         o="GSUtil:parallel_process_count=1" # disable multiprocessing for Macs
         gsutil -m -o "${o}" rm -r "gs://${broker_bucket}"
+        bq rm -r -f "${PROJECT_ID}:${bq_dataset}"
         gcloud pubsub topics delete "${topic_alerts}"
+        gcloud pubsub topics delete "${topic_storebigquery}"
     fi
 fi
 
 #--- Create VM instances
 echo
 echo "Configuring VMs..."
 ./create_vms.sh "${broker_bucket}" "${testid}" "${teardown}" "${survey}" "${zone}"
+
+#--- Deploy Cloud Functions
+echo
+echo "Configuring Cloud Functions..."
+cd .. && cd .. && cd cloud_functions && cd lvk || exit 5
+
+#--- BigQuery storage cloud function
+cd store_BigQuery && ./deploy.sh "$testid" "$teardown" "$survey" "$versiontag" || exit 5
+
+#--- return to setup_broker/lvk directory
+cd .. && cd .. && cd .. && cd setup_broker && cd lvk || exit 5