Fix internal GetMeasurementConsumer possibly returning incorrect Certificate.

src/main/kotlin/org/wfanet/measurement/common/BUILD.bazel

@@ -6,3 +6,9 @@ kt_jvm_library(
     name = "exponential_backoff",
     srcs = ["ExponentialBackoff.kt"],
 )
+
+kt_jvm_library(
+    name = "flows",
+    srcs = ["Flows.kt"],
+    deps = ["@wfa_common_jvm//imports/kotlin/kotlinx/coroutines:core"],
+)

src/main/kotlin/org/wfanet/measurement/common/Flows.kt

@@ -0,0 +1,35 @@
+/*
+ * Copyright 2023 The Cross-Media Measurement Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.wfanet.measurement.common
+
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.take
+import kotlinx.coroutines.flow.toList
+
+/**
+ * Terminal operator that awaits for a single value to be emitted.
+ *
+ * @return the single value, or `null` if the flow is empty
+ * @throws IllegalArgumentException for a flow that contains more than one element
+ *
+ * TODO(@SanjayVas): Move this to common-jvm.
+ */
+suspend fun <T> Flow<T>.singleOrNullIfEmpty(): T? {
+  val results = take(2).toList()
+  require(results.size <= 1) { "Flow has more than one element" }
+  return results.singleOrNull()
+}

src/main/kotlin/org/wfanet/measurement/kingdom/deploy/gcloud/spanner/readers/BUILD.bazel

@@ -18,5 +18,6 @@ kt_jvm_library(
         "@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/spanner",
+        "//src/main/kotlin/org/wfanet/measurement/common:flows",
     ],
 )

src/main/kotlin/org/wfanet/measurement/kingdom/deploy/gcloud/spanner/readers/CertificateReader.kt

@@ -259,7 +259,7 @@ class CertificateReader(private val parentType: ParentType) :
       externalCertificateId = struct.getLong(parentType.externalCertificateIdColumnName)
     }
 
-    private fun buildMeasurementConsumerCertificate(struct: Struct) = certificate {
+    fun buildMeasurementConsumerCertificate(struct: Struct) = certificate {
       fillCommon(struct)
 
       val parentType = ParentType.MEASUREMENT_CONSUMER

src/main/kotlin/org/wfanet/measurement/kingdom/deploy/gcloud/spanner/readers/DataProviderReader.kt

@@ -16,9 +16,9 @@ package org.wfanet.measurement.kingdom.deploy.gcloud.spanner.readers
 
 import com.google.cloud.spanner.Key
 import com.google.cloud.spanner.Struct
-import kotlinx.coroutines.flow.singleOrNull
 import org.wfanet.measurement.common.identity.ExternalId
 import org.wfanet.measurement.common.identity.InternalId
+import org.wfanet.measurement.common.singleOrNullIfEmpty
 import org.wfanet.measurement.gcloud.spanner.AsyncDatabaseClient
 import org.wfanet.measurement.gcloud.spanner.appendClause
 import org.wfanet.measurement.gcloud.spanner.getInternalId
@@ -68,10 +68,9 @@ class DataProviderReader : SpannerReader<DataProviderReader.Result>() {
     return fillStatementBuilder {
         appendClause("WHERE ExternalDataProviderId = @externalDataProviderId")
         bind("externalDataProviderId").to(externalDataProviderId.value)
-        appendClause("LIMIT 1")
       }
       .execute(readContext)
-      .singleOrNull()
+      .singleOrNullIfEmpty()
   }
 
   private fun buildDataProvider(struct: Struct): DataProvider =

src/main/kotlin/org/wfanet/measurement/kingdom/deploy/gcloud/spanner/readers/MeasurementConsumerReader.kt

@@ -16,16 +16,13 @@ package org.wfanet.measurement.kingdom.deploy.gcloud.spanner.readers
 
 import com.google.cloud.spanner.Key
 import com.google.cloud.spanner.Struct
-import kotlinx.coroutines.flow.singleOrNull
 import org.wfanet.measurement.common.identity.ExternalId
 import org.wfanet.measurement.common.identity.InternalId
+import org.wfanet.measurement.common.singleOrNullIfEmpty
 import org.wfanet.measurement.gcloud.spanner.AsyncDatabaseClient
 import org.wfanet.measurement.gcloud.spanner.appendClause
-import org.wfanet.measurement.gcloud.spanner.getBytesAsByteString
 import org.wfanet.measurement.gcloud.spanner.getInternalId
-import org.wfanet.measurement.gcloud.spanner.getProtoEnum
 import org.wfanet.measurement.gcloud.spanner.getProtoMessage
-import org.wfanet.measurement.internal.kingdom.Certificate
 import org.wfanet.measurement.internal.kingdom.MeasurementConsumer
 
 class MeasurementConsumerReader : SpannerReader<MeasurementConsumerReader.Result>() {
@@ -39,14 +36,16 @@ class MeasurementConsumerReader : SpannerReader<MeasurementConsumerReader.Result
       MeasurementConsumers.MeasurementConsumerDetails,
       MeasurementConsumers.MeasurementConsumerDetailsJson,
       MeasurementConsumerCertificates.ExternalMeasurementConsumerCertificateId,
-      Certificates.CertificateId,
       Certificates.SubjectKeyIdentifier,
       Certificates.NotValidBefore,
       Certificates.NotValidAfter,
       Certificates.RevocationState,
       Certificates.CertificateDetails
     FROM MeasurementConsumers
-    JOIN MeasurementConsumerCertificates USING (MeasurementConsumerId)
+    JOIN MeasurementConsumerCertificates ON (
+      MeasurementConsumerCertificates.MeasurementConsumerId = MeasurementConsumers.MeasurementConsumerId
+      AND MeasurementConsumerCertificates.CertificateId = MeasurementConsumers.PublicKeyCertificateId
+    )
     JOIN Certificates USING (CertificateId)
     """
       .trimIndent()
@@ -60,7 +59,7 @@ class MeasurementConsumerReader : SpannerReader<MeasurementConsumerReader.Result
         externalMeasurementConsumerId = struct.getLong("ExternalMeasurementConsumerId")
         details =
           struct.getProtoMessage("MeasurementConsumerDetails", MeasurementConsumer.Details.parser())
-        certificate = buildCertificate(struct)
+        certificate = CertificateReader.buildMeasurementConsumerCertificate(struct)
       }
       .build()
 
@@ -71,27 +70,11 @@ class MeasurementConsumerReader : SpannerReader<MeasurementConsumerReader.Result
     return fillStatementBuilder {
         appendClause("WHERE ExternalMeasurementConsumerId = @externalMeasurementConsumerId")
         bind("externalMeasurementConsumerId").to(externalMeasurementConsumerId.value)
-        appendClause("LIMIT 1")
       }
       .execute(readContext)
-      .singleOrNull()
+      .singleOrNullIfEmpty()
   }
 
-  // TODO(uakyol) : Move this function to CertificateReader when it is implemented.
-  private fun buildCertificate(struct: Struct): Certificate =
-    Certificate.newBuilder()
-      .apply {
-        externalMeasurementConsumerId = struct.getLong("ExternalMeasurementConsumerId")
-        externalCertificateId = struct.getLong("ExternalMeasurementConsumerCertificateId")
-        subjectKeyIdentifier = struct.getBytesAsByteString("SubjectKeyIdentifier")
-        notValidBefore = struct.getTimestamp("NotValidBefore").toProto()
-        notValidAfter = struct.getTimestamp("NotValidAfter").toProto()
-        revocationState =
-          struct.getProtoEnum("RevocationState", Certificate.RevocationState::forNumber)
-        details = struct.getProtoMessage("CertificateDetails", Certificate.Details.parser())
-      }
-      .build()
-
   companion object {
     /** Reads the [InternalId] for a MeasurementConsumer given its [ExternalId]. */
     suspend fun readMeasurementConsumerId(

src/main/kotlin/org/wfanet/measurement/kingdom/service/internal/testing/PublicKeysServiceTest.kt

@@ -37,7 +37,6 @@ import org.wfanet.measurement.internal.kingdom.getMeasurementConsumerRequest
 import org.wfanet.measurement.internal.kingdom.updatePublicKeyRequest
 import org.wfanet.measurement.kingdom.deploy.common.testing.DuchyIdSetter
 
-private const val RANDOM_SEED = 1
 private const val API_VERSION = "v2alpha"
 
 private val PUBLIC_KEY = ByteString.copyFromUtf8("new public key")
@@ -88,10 +87,20 @@ abstract class PublicKeysServiceTest<T : PublicKeysCoroutineImplBase> {
 
   @Test
   fun `updatePublicKey updates the public key info for a data provider`() = runBlocking {
-    val dataProvider = population.createDataProvider(dataProvidersService)
-    val certificate = population.createDataProviderCertificate(certificatesService, dataProvider)
+    val now = clock.instant()
+    val dataProvider = population.createDataProvider(dataProvidersService, notValidBefore = now)
+    val certificate =
+      population.createDataProviderCertificate(
+        certificatesService,
+        dataProvider,
+        notValidBefore = now
+      )
     // Insert another certificate to make sure it's not just using the most recent one.
-    population.createDataProviderCertificate(certificatesService, dataProvider)
+    population.createDataProviderCertificate(
+      certificatesService,
+      dataProvider,
+      notValidBefore = now
+    )
 
     publicKeysService.updatePublicKey(
       updatePublicKeyRequest {
@@ -114,10 +123,25 @@ abstract class PublicKeysServiceTest<T : PublicKeysCoroutineImplBase> {
 
   @Test
   fun `updatePublicKey updates the public key info for a measurement consumer`() = runBlocking {
+    val now = clock.instant()
     val measurementConsumer =
-      population.createMeasurementConsumer(measurementConsumersService, accountsService)
+      population.createMeasurementConsumer(
+        measurementConsumersService,
+        accountsService,
+        notValidBefore = now
+      )
     val certificate =
-      population.createMeasurementConsumerCertificate(certificatesService, measurementConsumer)
+      population.createMeasurementConsumerCertificate(
+        certificatesService,
+        measurementConsumer,
+        notValidBefore = now
+      )
+    // Insert another certificate to make sure it's not just using the most recent one.
+    population.createMeasurementConsumerCertificate(
+      certificatesService,
+      measurementConsumer,
+      notValidBefore = now
+    )
 
     publicKeysService.updatePublicKey(
       updatePublicKeyRequest {