Add in process postgres duchy correctness tests (#1152)

This adds in process postgres duchy correctness tests

src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/postgres/BUILD.bazel

@@ -4,7 +4,7 @@ kt_jvm_library(
     name = "services",
     srcs = glob(["*Service.kt"]),
     visibility = [
-        "//src/main/kotlin/org/wfanet/measurement/duchy/common/deploy/postgres:__pkg__",
+        "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/service:__pkg__",
         "//src/test/kotlin/org/wfanet/measurement/duchy/deploy/common/postgres:__pkg__",
     ],
     deps = [

src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/postgres/testing/BUILD.bazel

@@ -4,6 +4,7 @@ package(
     default_testonly = True,
     default_visibility = [
         "//src/test/kotlin/org/wfanet/measurement/duchy/deploy/common/postgres:__subpackages__",
+        "//src/test/kotlin/org/wfanet/measurement/integration/common/duchy:__pkg__",
     ],
 )
 

src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/postgres/tools/BUILD.bazel

@@ -0,0 +1,20 @@
+load("@rules_java//java:defs.bzl", "java_binary")
+load("//src/main/docker:macros.bzl", "java_image")
+
+java_binary(
+    name = "UpdateSchema",
+    args = ["--changelog=duchy/postgres/changelog.yaml"],
+    main_class = "org.wfanet.measurement.common.db.postgres.tools.UpdateSchema",
+    resources = ["//src/main/resources/duchy/postgres"],
+    runtime_deps = [
+        "@wfa_common_jvm//src/main/kotlin/org/wfanet/measurement/common/db/postgres/tools:update_schema",
+    ],
+)
+
+java_image(
+    name = "update_schema_image",
+    args = ["--changelog=duchy/postgres/changelog.yaml"],
+    binary = ":UpdateSchema",
+    main_class = "org.wfanet.measurement.common.db.postgres.tools.UpdateSchema",
+    visibility = ["//src:docker_image_deployment"],
+)

src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/server/BUILD.bazel

@@ -131,3 +131,23 @@ java_image(
     main_class = "org.wfanet.measurement.duchy.deploy.common.server.ForwardedStorageRequisitionFulfillmentServerKt",
     visibility = ["//src:docker_image_deployment"],
 )
+
+kt_jvm_library(
+    name = "duchy_data_server",
+    srcs = ["DuchyDataServer.kt"],
+    visibility = [
+        "//src/main/kotlin/org/wfanet/measurement/duchy/deploy:__subpackages__",
+        "//src/main/kotlin/org/wfanet/measurement/integration/common:__subpackages__",
+    ],
+    deps = [
+        "//src/main/kotlin/org/wfanet/measurement/common/identity",
+        "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/common:flags",
+        "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/service:duchy_data_services",
+        "//src/main/proto/wfa/measurement/system/v1alpha:computation_log_entries_service_kt_jvm_grpc_proto",
+        "@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/grpc",
+        "@wfa_common_jvm//src/main/kotlin/org/wfanet/measurement/common/identity",
+    ],
+)

src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/server/DuchyDataServer.kt

@@ -0,0 +1,81 @@
+// 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.duchy.deploy.common.server
+
+import io.grpc.ManagedChannel
+import java.time.Duration
+import kotlinx.coroutines.runInterruptible
+import org.wfanet.measurement.common.crypto.SigningCerts
+import org.wfanet.measurement.common.grpc.CommonServer
+import org.wfanet.measurement.common.grpc.buildMutualTlsChannel
+import org.wfanet.measurement.common.grpc.withShutdownTimeout
+import org.wfanet.measurement.common.identity.DuchyInfoFlags
+import org.wfanet.measurement.common.identity.withDuchyId
+import org.wfanet.measurement.duchy.deploy.common.CommonDuchyFlags
+import org.wfanet.measurement.duchy.deploy.common.SystemApiFlags
+import org.wfanet.measurement.duchy.deploy.common.service.DuchyDataServices
+import org.wfanet.measurement.duchy.deploy.common.service.toList
+import org.wfanet.measurement.system.v1alpha.ComputationLogEntriesGrpcKt.ComputationLogEntriesCoroutineStub
+import picocli.CommandLine
+
+abstract class DuchyDataServer : Runnable {
+  @CommandLine.Mixin
+  lateinit var serverFlags: CommonServer.Flags
+    private set
+
+  @CommandLine.Mixin
+  lateinit var duchyFlags: CommonDuchyFlags
+    private set
+
+  @CommandLine.Mixin
+  lateinit var duchyInfoFlags: DuchyInfoFlags
+    private set
+
+  @CommandLine.Option(
+    names = ["--channel-shutdown-timeout"],
+    defaultValue = "3s",
+    description = ["How long to allow for the gRPC channel to shutdown."],
+    required = true
+  )
+  lateinit var channelShutdownTimeout: Duration
+    private set
+
+  @CommandLine.Mixin
+  lateinit var systemApiFlags: SystemApiFlags
+    private set
+
+  val clientCerts =
+    SigningCerts.fromPemFiles(
+      certificateFile = serverFlags.tlsFlags.certFile,
+      privateKeyFile = serverFlags.tlsFlags.privateKeyFile,
+      trustedCertCollectionFile = serverFlags.tlsFlags.certCollectionFile
+    )
+  val channel: ManagedChannel =
+    buildMutualTlsChannel(systemApiFlags.target, clientCerts, systemApiFlags.certHost)
+      .withShutdownTimeout(channelShutdownTimeout)
+
+  val computationLogEntriesClient =
+    ComputationLogEntriesCoroutineStub(channel).withDuchyId(duchyFlags.duchyName)
+
+  protected suspend fun run(services: DuchyDataServices) {
+    val server = CommonServer.fromFlags(serverFlags, this::class.simpleName!!, services.toList())
+
+    runInterruptible { server.start().blockUntilShutdown() }
+  }
+
+  companion object {
+    const val SERVICE_NAME = "Computations"
+  }
+}

src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/service/BUILD.bazel

@@ -0,0 +1,32 @@
+load("@io_bazel_rules_kotlin//kotlin:jvm.bzl", "kt_jvm_library")
+
+kt_jvm_library(
+    name = "duchy_data_services",
+    srcs = ["DuchyDataServices.kt"],
+    visibility = [
+        "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/server:__pkg__",
+        "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/gcloud/service:__pkg__",
+        "//src/test/kotlin/org/wfanet/measurement/integration/common/duchy:__pkg__",
+    ],
+    deps = [
+        "//src/main/proto/wfa/measurement/internal/duchy:computation_stats_service_kt_jvm_grpc_proto",
+        "//src/main/proto/wfa/measurement/internal/duchy:computations_service_kt_jvm_grpc_proto",
+        "//src/main/proto/wfa/measurement/internal/duchy:continuation_tokens_service_kt_jvm_grpc_proto",
+    ],
+)
+
+kt_jvm_library(
+    name = "postgres_duchy_data_services",
+    srcs = ["PostgresDuchyDataServices.kt"],
+    visibility = [
+        "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/gcloud/server:__pkg__",
+        "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/postgres/server:__pkg__",
+        "//src/test/kotlin/org/wfanet/measurement/integration/common/duchy:__pkg__",
+    ],
+    deps = [
+        ":duchy_data_services",
+        "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/postgres:services",
+        "@wfa_common_jvm//src/main/kotlin/org/wfanet/measurement/common/db/r2dbc",
+        "@wfa_common_jvm//src/main/kotlin/org/wfanet/measurement/common/identity",
+    ],
+)

src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/service/DuchyDataServices.kt

@@ -0,0 +1,31 @@
+// 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.duchy.deploy.common.service
+
+import io.grpc.BindableService
+import kotlin.reflect.full.declaredMemberProperties
+import org.wfanet.measurement.internal.duchy.ComputationStatsGrpcKt.ComputationStatsCoroutineImplBase
+import org.wfanet.measurement.internal.duchy.ComputationsGrpcKt.ComputationsCoroutineImplBase
+import org.wfanet.measurement.internal.duchy.ContinuationTokensGrpcKt.ContinuationTokensCoroutineImplBase
+
+data class DuchyDataServices(
+  val computationsService: ComputationsCoroutineImplBase,
+  val computationStatsService: ComputationStatsCoroutineImplBase,
+  val continuationTokensService: ContinuationTokensCoroutineImplBase
+)
+
+fun DuchyDataServices.toList(): List<BindableService> {
+  return DuchyDataServices::class.declaredMemberProperties.map { it.get(this) as BindableService }
+}

src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/service/PostgresDuchyDataServices.kt

@@ -0,0 +1,74 @@
+// 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.duchy.deploy.common.service
+
+import java.time.Clock
+import org.wfanet.measurement.common.db.r2dbc.DatabaseClient
+import org.wfanet.measurement.common.identity.IdGenerator
+import org.wfanet.measurement.duchy.db.computation.ComputationProtocolStageDetails
+import org.wfanet.measurement.duchy.db.computation.ComputationProtocolStageDetailsHelper
+import org.wfanet.measurement.duchy.db.computation.ComputationProtocolStages
+import org.wfanet.measurement.duchy.db.computation.ComputationProtocolStagesEnumHelper
+import org.wfanet.measurement.duchy.db.computation.ComputationTypeEnumHelper
+import org.wfanet.measurement.duchy.db.computation.ComputationTypes
+import org.wfanet.measurement.duchy.deploy.common.postgres.PostgresComputationStatsService
+import org.wfanet.measurement.duchy.deploy.common.postgres.PostgresComputationsService
+import org.wfanet.measurement.duchy.deploy.common.postgres.PostgresContinuationTokensService
+import org.wfanet.measurement.duchy.storage.ComputationStore
+import org.wfanet.measurement.duchy.storage.RequisitionStore
+import org.wfanet.measurement.internal.duchy.ComputationDetails
+import org.wfanet.measurement.internal.duchy.ComputationStage
+import org.wfanet.measurement.internal.duchy.ComputationStageDetails
+import org.wfanet.measurement.internal.duchy.ComputationTypeEnum.ComputationType
+import org.wfanet.measurement.storage.StorageClient
+import org.wfanet.measurement.system.v1alpha.ComputationLogEntriesGrpcKt.ComputationLogEntriesCoroutineStub
+
+object PostgresDuchyDataServices {
+  @JvmStatic
+  fun create(
+    storageClient: StorageClient,
+    computationLogEntriesClient: ComputationLogEntriesCoroutineStub,
+    duchyName: String,
+    idGenerator: IdGenerator,
+    client: DatabaseClient
+  ): DuchyDataServices {
+    val computationTypeEnumHelper: ComputationTypeEnumHelper<ComputationType> = ComputationTypes
+    val protocolStagesEnumHelper:
+      ComputationProtocolStagesEnumHelper<ComputationType, ComputationStage> =
+      ComputationProtocolStages
+    val computationProtocolStageDetailsHelper:
+      ComputationProtocolStageDetailsHelper<
+        ComputationType, ComputationStage, ComputationStageDetails, ComputationDetails
+      > =
+      ComputationProtocolStageDetails
+
+    return DuchyDataServices(
+      PostgresComputationsService(
+        computationTypeEnumHelper = computationTypeEnumHelper,
+        protocolStagesEnumHelper = protocolStagesEnumHelper,
+        computationProtocolStageDetailsHelper = computationProtocolStageDetailsHelper,
+        computationStore = ComputationStore(storageClient),
+        requisitionStore = RequisitionStore(storageClient),
+        computationLogEntriesClient = computationLogEntriesClient,
+        duchyName = duchyName,
+        clock = Clock.systemUTC(),
+        client = client,
+        idGenerator = idGenerator,
+      ),
+      PostgresComputationStatsService(client, idGenerator),
+      PostgresContinuationTokensService(client, idGenerator)
+    )
+  }
+}

src/main/kotlin/org/wfanet/measurement/duchy/deploy/gcloud/service/BUILD.bazel

@@ -0,0 +1,18 @@
+load("@io_bazel_rules_kotlin//kotlin:jvm.bzl", "kt_jvm_library")
+
+kt_jvm_library(
+    name = "spanner_duchy_data_services",
+    srcs = ["SpannerDuchyDataServices.kt"],
+    visibility = [
+        "//src/test/kotlin/org/wfanet/measurement/integration/common/duchy:__pkg__",
+    ],
+    deps = [
+        "//src/main/kotlin/org/wfanet/measurement/duchy/db/computation",
+        "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/service:duchy_data_services",
+        "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/gcloud/spanner/computation",
+        "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/gcloud/spanner/continuationtoken:service",
+        "//src/main/kotlin/org/wfanet/measurement/duchy/service/internal/computations",
+        "//src/main/kotlin/org/wfanet/measurement/duchy/service/internal/computationstats:service",
+        "@wfa_common_jvm//src/main/kotlin/org/wfanet/measurement/common/db/r2dbc",
+    ],
+)