impl(generator): format variables for lro await (#14330)

googleapis · Jun 13, 2024 · 54274c0 · 54274c0
1 parent 1adeea2
commit 54274c0
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 0 deletions.
diff --git a/generator/internal/longrunning.cc b/generator/internal/longrunning.cc
@@ -95,6 +95,9 @@ bool IsLongrunningMetadataTypeUsedAsResponse(MethodDescriptor const& method) {
 void SetLongrunningOperationMethodVars(
     google::protobuf::MethodDescriptor const& method,
     VarsDictionary& method_vars) {
+  method_vars["longrunning_operation_type"] =
+      ProtoNameToCppName(method.output_type()->full_name());
+
   if (method.output_type()->full_name() == "google.longrunning.Operation") {
     auto operation_info =
         method.options().GetExtension(google::longrunning::operation_info);
@@ -193,6 +196,10 @@ void SetLongrunningOperationServiceVars(
         service_vars["longrunning_set_operation_fields"] = R"""(
       r.set_project(request.project());
       r.set_operation(op);
+)""";
+        service_vars["longrunning_await_set_operation_fields"] = R"""(
+      r.set_project(info.project);
+      r.set_operation(info.operation);
 )""";
         auto global_lro_path = absl::StrFormat(
             R"""(absl::StrCat("/compute/",
@@ -216,6 +223,9 @@ void SetLongrunningOperationServiceVars(
             "DeleteOperationRequest";
         service_vars["longrunning_set_operation_fields"] = R"""(
       r.set_operation(op);
+)""";
+        service_vars["longrunning_await_set_operation_fields"] = R"""(
+      r.set_operation(info.operation);
 )""";
         auto global_org_lro_path = absl::StrFormat(
             R"""(absl::StrCat("/compute/",
@@ -239,6 +249,11 @@ void SetLongrunningOperationServiceVars(
       r.set_project(request.project());
       r.set_region(request.region());
       r.set_operation(op);
+)""";
+        service_vars["longrunning_await_set_operation_fields"] = R"""(
+      r.set_project(info.project);
+      r.set_region(info.region);
+      r.set_operation(info.operation);
 )""";
         auto region_lro_path = absl::StrFormat(
             R"""(absl::StrCat("/compute/",
@@ -263,6 +278,11 @@ void SetLongrunningOperationServiceVars(
       r.set_project(request.project());
       r.set_zone(request.zone());
       r.set_operation(op);
+)""";
+        service_vars["longrunning_await_set_operation_fields"] = R"""(
+      r.set_project(info.project);
+      r.set_zone(info.zone);
+      r.set_operation(info.operation);
 )""";
         auto zone_lro_path = absl::StrFormat(
             R"""(absl::StrCat("/compute/",

diff --git a/generator/internal/longrunning_test.cc b/generator/internal/longrunning_test.cc
@@ -535,6 +535,8 @@ TEST_F(LongrunningVarsTest,
       service_file_descriptor->service(0)->method(1);
   VarsDictionary vars;
   SetLongrunningOperationMethodVars(*method, vars);
+  EXPECT_THAT(vars, Contains(Pair("longrunning_operation_type",
+                                  "google::longrunning::Operation")));
   EXPECT_THAT(vars, Contains(Pair("longrunning_metadata_type",
                                   "google::protobuf::Struct")));
   EXPECT_THAT(vars, Contains(Pair("longrunning_response_type",
@@ -580,6 +582,8 @@ TEST_F(LongrunningVarsTest, SetLongrunningOperationMethodVarsBespokeLRO) {
   EXPECT_FALSE(IsGRPCLongrunningOperation(*method));
   EXPECT_TRUE(IsHttpLongrunningOperation(*method));
   SetLongrunningOperationMethodVars(*method, vars);
+  EXPECT_THAT(vars, Contains(Pair("longrunning_operation_type",
+                                  "my::service::v1::Operation")));
   EXPECT_THAT(vars, Contains(Pair("longrunning_response_type",
                                   "my::service::v1::Operation")));
   EXPECT_THAT(vars, Contains(Pair("longrunning_deduced_response_message_type",
@@ -637,6 +641,11 @@ TEST_F(LongrunningVarsTest, SetLongrunningOperationServiceVarsNonGRPCGlobal) {
   EXPECT_THAT(vars, Contains(Pair("longrunning_set_operation_fields", R"""(
       r.set_project(request.project());
       r.set_operation(op);
+)""")));
+  EXPECT_THAT(vars,
+              Contains(Pair("longrunning_await_set_operation_fields", R"""(
+      r.set_project(info.project);
+      r.set_operation(info.operation);
 )""")));
   EXPECT_THAT(vars, Contains(Pair("longrunning_get_operation_path_rest",
                                   R"""(absl::StrCat("/compute/",
@@ -673,6 +682,10 @@ TEST_F(LongrunningVarsTest,
           "DeleteOperationRequest")));
   EXPECT_THAT(vars, Contains(Pair("longrunning_set_operation_fields", R"""(
       r.set_operation(op);
+)""")));
+  EXPECT_THAT(vars,
+              Contains(Pair("longrunning_await_set_operation_fields", R"""(
+      r.set_operation(info.operation);
 )""")));
   EXPECT_THAT(vars, Contains(Pair("longrunning_get_operation_path_rest",
                                   R"""(absl::StrCat("/compute/",
@@ -706,6 +719,12 @@ TEST_F(LongrunningVarsTest, SetLongrunningOperationServiceVarsNonGRPCRegion) {
       r.set_project(request.project());
       r.set_region(request.region());
       r.set_operation(op);
+)""")));
+  EXPECT_THAT(vars,
+              Contains(Pair("longrunning_await_set_operation_fields", R"""(
+      r.set_project(info.project);
+      r.set_region(info.region);
+      r.set_operation(info.operation);
 )""")));
   EXPECT_THAT(vars, Contains(Pair("longrunning_get_operation_path_rest",
                                   R"""(absl::StrCat("/compute/",
@@ -743,6 +762,12 @@ TEST_F(LongrunningVarsTest, SetLongrunningOperationServiceVarsNonGRPCZone) {
       r.set_project(request.project());
       r.set_zone(request.zone());
       r.set_operation(op);
+)""")));
+  EXPECT_THAT(vars,
+              Contains(Pair("longrunning_await_set_operation_fields", R"""(
+      r.set_project(info.project);
+      r.set_zone(info.zone);
+      r.set_operation(info.operation);
 )""")));
   EXPECT_THAT(vars, Contains(Pair("longrunning_get_operation_path_rest",
                                   R"""(absl::StrCat("/compute/",