feat: Write Metrics to BigQuery (#1696)

feat: Write Metrics to BigQuery
world-federation-of-advertisers · Sep 13, 2024 · 7428e7c · 7428e7c
1 parent f3b1a8a
commit 7428e7c
Show file tree

Hide file tree

Showing 21 changed files with 2,184 additions and 19 deletions.
diff --git a/MODULE.bazel b/MODULE.bazel
@@ -211,6 +211,7 @@ maven.install(
         "org.slf4j:slf4j-simple:1.7.32",
 
         # Google Cloud
+        "com.google.cloud:google-cloud-bigquerystorage",
         "com.google.cloud:google-cloud-security-private-ca",
         # TODO(googleapis/java-cloud-bom#5279): Remove when managed by BOM.
         "com.google.apis:google-api-services-storage:v1-rev20240706-2.0.0",

diff --git a/imports/java/com/google/cloud/bigquery/storage/BUILD.bazel b/imports/java/com/google/cloud/bigquery/storage/BUILD.bazel
@@ -0,0 +1,6 @@
+package(default_visibility = ["//visibility:public"])
+
+alias(
+    name = "storage",
+    actual = "@maven//:com_google_cloud_google_cloud_bigquerystorage",
+)
diff --git a/maven_install.json b/maven_install.json
@@ -1,7 +1,7 @@
 {
   "__AUTOGENERATED_FILE_DO_NOT_MODIFY_THIS_FILE_MANUALLY": "THERE_IS_NO_DATA_ONLY_ZUUL",
-  "__INPUT_ARTIFACTS_HASH": 1120915116,
-  "__RESOLVED_ARTIFACTS_HASH": -2081716671,
+  "__INPUT_ARTIFACTS_HASH": 606715834,
+  "__RESOLVED_ARTIFACTS_HASH": -1010776065,
   "artifacts": {
     "args4j:args4j": {
       "shasums": {
@@ -3426,6 +3426,44 @@
       "com.google.http-client:google-http-client",
       "com.google.http-client:google-http-client-gson"
     ],
+    "com.google.api.grpc:grpc-google-cloud-bigquerystorage-v1": [
+      "com.google.api.grpc:proto-google-common-protos",
+      "com.google.api:api-common",
+      "com.google.auto.value:auto-value-annotations",
+      "com.google.code.findbugs:jsr305",
+      "com.google.errorprone:error_prone_annotations",
+      "com.google.guava:failureaccess",
+      "com.google.guava:listenablefuture",
+      "com.google.j2objc:j2objc-annotations",
+      "io.grpc:grpc-protobuf-lite",
+      "javax.annotation:javax.annotation-api",
+      "org.checkerframework:checker-qual"
+    ],
+    "com.google.api.grpc:grpc-google-cloud-bigquerystorage-v1beta1": [
+      "com.google.api.grpc:proto-google-common-protos",
+      "com.google.api:api-common",
+      "com.google.auto.value:auto-value-annotations",
+      "com.google.code.findbugs:jsr305",
+      "com.google.errorprone:error_prone_annotations",
+      "com.google.guava:failureaccess",
+      "com.google.guava:listenablefuture",
+      "com.google.j2objc:j2objc-annotations",
+      "io.grpc:grpc-protobuf-lite",
+      "org.checkerframework:checker-qual"
+    ],
+    "com.google.api.grpc:grpc-google-cloud-bigquerystorage-v1beta2": [
+      "com.google.api.grpc:proto-google-common-protos",
+      "com.google.api:api-common",
+      "com.google.auto.value:auto-value-annotations",
+      "com.google.code.findbugs:jsr305",
+      "com.google.errorprone:error_prone_annotations",
+      "com.google.guava:failureaccess",
+      "com.google.guava:listenablefuture",
+      "com.google.j2objc:j2objc-annotations",
+      "io.grpc:grpc-protobuf-lite",
+      "javax.annotation:javax.annotation-api",
+      "org.checkerframework:checker-qual"
+    ],
     "com.google.api.grpc:grpc-google-cloud-bigtable-admin-v2": [
       "com.google.api.grpc:proto-google-cloud-bigtable-admin-v2",
       "com.google.api.grpc:proto-google-common-protos",
@@ -3548,6 +3586,36 @@
       "javax.annotation:javax.annotation-api",
       "org.checkerframework:checker-qual"
     ],
+    "com.google.api.grpc:proto-google-cloud-bigquerystorage-v1alpha": [
+      "com.google.auto.value:auto-value-annotations",
+      "com.google.code.findbugs:jsr305",
+      "com.google.errorprone:error_prone_annotations",
+      "com.google.guava:failureaccess",
+      "com.google.guava:listenablefuture",
+      "com.google.j2objc:j2objc-annotations",
+      "javax.annotation:javax.annotation-api",
+      "org.checkerframework:checker-qual"
+    ],
+    "com.google.api.grpc:proto-google-cloud-bigquerystorage-v1beta1": [
+      "com.google.auto.value:auto-value-annotations",
+      "com.google.code.findbugs:jsr305",
+      "com.google.errorprone:error_prone_annotations",
+      "com.google.guava:failureaccess",
+      "com.google.guava:listenablefuture",
+      "com.google.j2objc:j2objc-annotations",
+      "javax.annotation:javax.annotation-api",
+      "org.checkerframework:checker-qual"
+    ],
+    "com.google.api.grpc:proto-google-cloud-bigquerystorage-v1beta2": [
+      "com.google.auto.value:auto-value-annotations",
+      "com.google.code.findbugs:jsr305",
+      "com.google.errorprone:error_prone_annotations",
+      "com.google.guava:failureaccess",
+      "com.google.guava:listenablefuture",
+      "com.google.j2objc:j2objc-annotations",
+      "javax.annotation:javax.annotation-api",
+      "org.checkerframework:checker-qual"
+    ],
     "com.google.api.grpc:proto-google-cloud-bigtable-admin-v2": [
       "com.google.api.grpc:proto-google-common-protos",
       "com.google.api.grpc:proto-google-iam-v1",

diff --git a/src/main/docker/images.bzl b/src/main/docker/images.bzl
@@ -151,6 +151,11 @@ GKE_IMAGES = [
         image = "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/gcloud/postgres/tools:update_schema_image",
         repository = _PREFIX + "/duchy/gcloud-postgres-update-schema",
     ),
+    struct(
+        name = "kingdom_operational_metrics_export_image",
+        image = "//src/main/kotlin/org/wfanet/measurement/kingdom/deploy/gcloud/job:operational_metrics_export_job_image",
+        repository = _PREFIX + "/kingdom/bigquery-operational-metrics",
+    ),
 ]
 
 # List of specs for all Docker containers to push to a container registry.

diff --git a/src/main/k8s/dev/kingdom_gke.cue b/src/main/k8s/dev/kingdom_gke.cue
@@ -21,6 +21,9 @@ _systemApiAddressName: string @tag("system_api_address_name")
 // Name of K8s service account for the internal API server.
 #InternalServerServiceAccount: "internal-server"
 
+// Name of K8s service account for the operational metrics job.
+#OperationalMetricsServiceAccount: "operational-metrics"
+
 // Number of gRPC threads for the internal API server.
 #InternalServerGrpcThreads: 7
 
@@ -40,6 +43,16 @@ _systemApiAddressName: string @tag("system_api_address_name")
 	}
 }
 
+#OperationalMetricsJobResourceRequirements: ResourceRequirements=#ResourceRequirements & {
+	requests: {
+		cpu:    "10m"
+		memory: "256Mi"
+	}
+	limits: {
+		memory: ResourceRequirements.requests.memory
+	}
+}
+
 objectSets: [defaultNetworkPolicies] + [ for objectSet in kingdom {objectSet}]
 
 kingdom: #Kingdom & {
@@ -48,10 +61,30 @@ kingdom: #Kingdom & {
 
 	_verboseGrpcServerLogging: true
 
+	_imageSuffixes: [string]: string
+	_imageSuffixes: {
+		"operational-metrics": string | *"kingdom/bigquery-operational-metrics"
+	}
+	_imageConfigs: [string]: #ImageConfig
+	_imageConfigs: {
+		for name, suffix in _imageSuffixes {
+			"\(name)": {repoSuffix: suffix}
+		}
+	}
+	_images: [string]: string
+	_images: {
+		for name, config in _imageConfigs {
+			"\(name)": config.image
+		}
+	}
+
 	serviceAccounts: {
 		"\(#InternalServerServiceAccount)": #WorkloadIdentityServiceAccount & {
 			_iamServiceAccountName: "kingdom-internal"
 		}
+		"\(#OperationalMetricsServiceAccount)": #WorkloadIdentityServiceAccount & {
+			_iamServiceAccountName: "operational-metrics"
+		}
 	}
 
 	configMaps: "java": #JavaConfigMap
@@ -73,6 +106,45 @@ kingdom: #Kingdom & {
 		}
 	}
 
+	cronJobs: {
+		"operational-metrics": {
+			_container: {
+				_javaOptions: maxHeapSize: "48M"
+				image:     _images["operational-metrics"]
+				resources: #OperationalMetricsJobResourceRequirements
+				args: [
+					"--bigquery-project=\(#GCloudProject)",
+					"--bigquery-dataset=operational_metrics",
+					"--measurements-table=measurements",
+					"--requisitions-table=requisitions",
+					"--latest-measurement-read-table=latest_measurement_read",
+					"--tls-cert-file=/var/run/secrets/files/kingdom_tls.pem",
+					"--tls-key-file=/var/run/secrets/files/kingdom_tls.key",
+					"--cert-collection-file=/var/run/secrets/files/kingdom_root.pem",
+					"--internal-api-target=" + (#Target & {name: "gcp-kingdom-data-server"}).target,
+					"--internal-api-cert-host=localhost",
+				]
+			}
+			spec: {
+				concurrencyPolicy: "Forbid"
+				schedule:          "30 * * * *" // Hourly, 30 minutes past the hour
+				jobTemplate: spec: template: spec: #ServiceAccountPodSpec & {
+					serviceAccountName: #OperationalMetricsServiceAccount
+				}
+			}
+		}
+	}
+
+	networkPolicies: {
+		"operational-metrics": {
+			_app_label: "operational-metrics-app"
+			_egresses: {
+				// Need to send external traffic to BigQuery.
+				any: {}
+			}
+		}
+	}
+
 	services: {
 		"system-api-server": _ipAddressName:         _systemApiAddressName
 		"v2alpha-public-api-server": _ipAddressName: _publicApiAddressName

diff --git a/src/main/k8s/kingdom.cue b/src/main/k8s/kingdom.cue
@@ -258,6 +258,7 @@ import ("strings")
 				"completed-measurements-deletion-app",
 				"pending-measurements-cancellation-app",
 				"exchanges-deletion-app",
+				"operational-metrics-app",
 			]
 			_egresses: {
 				// Need to send external traffic to Spanner.

diff --git a/src/main/kotlin/org/wfanet/measurement/kingdom/deploy/common/server/BUILD.bazel b/src/main/kotlin/org/wfanet/measurement/kingdom/deploy/common/server/BUILD.bazel
@@ -23,7 +23,10 @@ kt_jvm_library(
 kt_jvm_library(
     name = "utils",
     srcs = ["Utils.kt"],
-    visibility = ["//src/main/kotlin/org/wfanet/measurement/kingdom/deploy/common/job:__subpackages__"],
+    visibility = [
+        "//src/main/kotlin/org/wfanet/measurement/kingdom/deploy/common/job:__subpackages__",
+        "//src/main/kotlin/org/wfanet/measurement/kingdom/deploy/gcloud/job:__subpackages__",
+    ],
     deps = [
         "//src/main/kotlin/org/wfanet/measurement/common/identity",
         "//src/main/kotlin/org/wfanet/measurement/kingdom/deploy/common:flags",

diff --git a/src/main/kotlin/org/wfanet/measurement/kingdom/deploy/gcloud/job/BUILD.bazel b/src/main/kotlin/org/wfanet/measurement/kingdom/deploy/gcloud/job/BUILD.bazel
@@ -0,0 +1,63 @@
+load("@rules_java//java:defs.bzl", "java_binary")
+load("@wfa_rules_kotlin_jvm//kotlin:defs.bzl", "kt_jvm_library")
+load("//src/main/docker:macros.bzl", "java_image")
+
+kt_jvm_library(
+    name = "operational_metrics_export_job",
+    srcs = ["OperationalMetricsExportJob.kt"],
+    runtime_deps = ["@wfa_common_jvm//imports/java/io/grpc/netty"],
+    deps = [
+        ":operational_metrics_export",
+        "//imports/java/com/google/cloud/bigquery/storage",
+        "//src/main/kotlin/org/wfanet/measurement/kingdom/deploy/common/server:utils",
+        "//src/main/proto/wfa/measurement/internal/kingdom:measurements_service_kt_jvm_grpc_proto",
+        "@wfa_common_jvm//imports/java/com/google/cloud/bigquery",
+        "@wfa_common_jvm//imports/java/picocli",
+        "@wfa_common_jvm//imports/kotlin/kotlinx/coroutines:core",
+        "@wfa_common_jvm//src/main/kotlin/org/wfanet/measurement/common",
+        "@wfa_common_jvm//src/main/kotlin/org/wfanet/measurement/common/crypto:signing_certs",
+        "@wfa_common_jvm//src/main/kotlin/org/wfanet/measurement/common/grpc",
+    ],
+)
+
+java_binary(
+    name = "OperationalMetricsExportJob",
+    main_class = "org.wfanet.measurement.kingdom.deploy.gcloud.job.OperationalMetricsExportJobKt",
+    runtime_deps = [":operational_metrics_export_job"],
+)
+
+java_image(
+    name = "operational_metrics_export_job_image",
+    binary = ":OperationalMetricsExportJob",
+    main_class = "org.wfanet.measurement.kingdom.deploy.gcloud.job.OperationalMetricsExportJobKt",
+    visibility = ["//src:docker_image_deployment"],
+)
+
+kt_jvm_library(
+    name = "operational_metrics_export",
+    srcs = ["OperationalMetricsExport.kt"],
+    visibility = ["//src/test/kotlin/org/wfanet/measurement/kingdom/deploy/gcloud/job:__pkg__"],
+    deps = [
+        ":stream_writer_factory",
+        "//imports/java/com/google/cloud/bigquery/storage",
+        "//src/main/kotlin/org/wfanet/measurement/api:public_api_version",
+        "//src/main/kotlin/org/wfanet/measurement/api/v2alpha:packed_messages",
+        "//src/main/proto/wfa/measurement/api/v2alpha:measurement_spec_kt_jvm_proto",
+        "//src/main/proto/wfa/measurement/internal/kingdom:measurements_service_kt_jvm_grpc_proto",
+        "//src/main/proto/wfa/measurement/internal/kingdom/bigquerytables:operational_metrics_dataset_kt_jvm_proto",
+        "@wfa_common_jvm//imports/java/com/google/cloud/bigquery",
+        "@wfa_common_jvm//imports/kotlin/kotlinx/coroutines:core",
+        "@wfa_common_jvm//src/main/kotlin/org/wfanet/measurement/common",
+        "@wfa_common_jvm//src/main/kotlin/org/wfanet/measurement/common/identity",
+        "@wfa_common_jvm//src/main/kotlin/org/wfanet/measurement/gcloud/common",
+    ],
+)
+
+kt_jvm_library(
+    name = "stream_writer_factory",
+    srcs = ["StreamWriterFactory.kt"],
+    visibility = ["//src/test/kotlin/org/wfanet/measurement/kingdom/deploy/gcloud/job:__pkg__"],
+    deps = [
+        "//imports/java/com/google/cloud/bigquery/storage",
+    ],
+)