impl(bigtable): update query plan refresh retry logic to pass conformance tests

google/cloud/bigtable/CMakeLists.txt

@@ -224,6 +224,7 @@ add_library(
     internal/rate_limiter.h
     internal/readrowsparser.cc
     internal/readrowsparser.h
+    internal/retry_traits.cc
     internal/retry_traits.h
     internal/row_reader_impl.h
     internal/rpc_policy_parameters.h

google/cloud/bigtable/ci/run_conformance_tests_proxy_bazel.sh

@@ -50,8 +50,8 @@ exit_status=$?
 # Run all the ExecuteQuery tests that either work or we plan to skip such as
 # CloseClient
 go test -v \
-  -run "TestExecuteQuery" \
-  -skip "CloseClient|FailsOnEmptyMetadata|FailsOnExecuteQueryMetadata|FailsOnInvalidType|FailsOnNotEnoughData|FailsOnNotEnoughDataWithCompleteRows|FailsOnSuccesfulStreamWithNoToken|ChecksumMismatch|FailsOnTypeMismatch|FailsOnTypeMismatchWithinMap|FailsOnTypeMismatchWithinArray|FailsOnTypeMismatchWithinStruct|FailsOnStructMissingField|RetryTest_WithPlanRefresh|PlanRefresh" \
+  -run "TestExecuteQuery|TestExecuteQuery_PlanRefresh$|TestExecuteQuery_PlanRefresh_WithMetadataChange|TestExecuteQuery_PlanRefresh_Retries|TestExecuteQuery_PlanRefresh_RecoversAfterPermanentError" \
+  -skip "CloseClient|FailsOnEmptyMetadata|FailsOnExecuteQueryMetadata|FailsOnInvalidType|FailsOnNotEnoughData|FailsOnNotEnoughDataWithCompleteRows|FailsOnSuccesfulStreamWithNoToken|ChecksumMismatch|FailsOnTypeMismatch|FailsOnTypeMismatchWithinMap|FailsOnTypeMismatchWithinArray|FailsOnTypeMismatchWithinStruct|FailsOnStructMissingField|TestExecuteQuery_PlanRefresh_AfterResumeTokenCausesError|TestExecuteQuery_RetryTest_WithPlanRefresh|TestExecuteQuery_PlanRefresh_RespectsDeadline" \
   -proxy_addr=:9999
 exit_status=$?
 

google/cloud/bigtable/google_cloud_cpp_bigtable.bzl

@@ -230,6 +230,7 @@ google_cloud_cpp_bigtable_srcs = [
     "internal/query_plan.cc",
     "internal/rate_limiter.cc",
     "internal/readrowsparser.cc",
+    "internal/retry_traits.cc",
     "internal/traced_row_reader.cc",
     "metadata_update_policy.cc",
     "mutation_batcher.cc",

google/cloud/bigtable/internal/data_connection_impl.cc

@@ -160,18 +160,16 @@ class DefaultPartialResultSetReader
             response.metadata().SerializeToString(&response_metadata_str) &&
             initial_metadata_str == response_metadata_str;
         if (response.metadata().ByteSizeLong() > 0 && !metadata_matched) {
-          final_status_ = google::cloud::Status(
-              google::cloud::StatusCode::kAborted,
-              "Schema changed during ExecuteQuery operation");
+          final_status_ = internal::AbortedError(
+              "Schema changed during ExecuteQuery operation", GCP_ERROR_INFO());
           operation_context_->PostCall(*context_, final_status_);
           return false;
         }
         continue;
       }
 
-      final_status_ = google::cloud::Status(
-          google::cloud::StatusCode::kInternal,
-          "Empty ExecuteQueryResponse received from stream");
+      final_status_ = internal::InternalError(
+          "Empty ExecuteQueryResponse received from stream", GCP_ERROR_INFO());
       operation_context_->PostCall(*context_, final_status_);
       return false;
     }
@@ -964,8 +962,7 @@ bigtable::RowStream DataConnectionImpl::ExecuteQuery(
       }
       last_status = source.status();
 
-      if (SafeGrpcRetryAllowingQueryPlanRefresh::IsQueryPlanExpired(
-              source.status())) {
+      if (QueryPlanRefreshRetry::IsQueryPlanExpired(source.status())) {
         query_plan->Invalidate(source.status(),
                                query_plan_data->prepared_query());
       }

google/cloud/bigtable/internal/retry_traits.cc

@@ -0,0 +1,54 @@
+// Copyright 2025 Google LLC
+//
+// 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
+//
+//     https://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.
+
+#include "google/cloud/bigtable/internal/retry_traits.h"
+#include "google/cloud/internal/status_payload_keys.h"
+#include "google/cloud/status.h"
+#include "absl/strings/match.h"
+#include <google/rpc/error_details.pb.h>
+#include <google/rpc/status.pb.h>
+
+namespace google {
+namespace cloud {
+namespace bigtable_internal {
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
+
+bool QueryPlanRefreshRetry::IsQueryPlanExpired(Status const& s) {
+  if (s.code() == StatusCode::kFailedPrecondition) {
+    if (absl::StrContains(s.message(), "PREPARED_QUERY_EXPIRED")) {
+      return true;
+    }
+    auto payload = google::cloud::internal::GetPayload(
+        s, google::cloud::internal::StatusPayloadGrpcProto());
+    google::rpc::Status proto;
+    if (payload && proto.ParseFromString(*payload)) {
+      google::rpc::PreconditionFailure failure;
+      for (google::protobuf::Any const& any : proto.details()) {
+        if (any.UnpackTo(&failure)) {
+          for (auto const& v : failure.violations()) {
+            if (absl::StrContains(v.type(), "PREPARED_QUERY_EXPIRED")) {
+              return true;
+            }
+          }
+        }
+      }
+    }
+  }
+  return false;
+}
+
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
+}  // namespace bigtable_internal
+}  // namespace cloud
+}  // namespace google

google/cloud/bigtable/internal/retry_traits.h

@@ -19,7 +19,6 @@
 #include "google/cloud/grpc_error_delegate.h"
 #include "google/cloud/internal/retry_policy_impl.h"
 #include "google/cloud/status.h"
-#include "absl/strings/match.h"
 #include <grpcpp/grpcpp.h>
 
 namespace google {
@@ -49,18 +48,13 @@ struct SafeGrpcRetry {
   }
 };
 
-struct SafeGrpcRetryAllowingQueryPlanRefresh {
-  static bool IsQueryPlanExpired(Status const& s) {
-    return (s.code() == StatusCode::kFailedPrecondition &&
-            absl::StrContains(s.message(), "PREPARED_QUERY_EXPIRED"));
-  }
-
+struct QueryPlanRefreshRetry {
+  static bool IsQueryPlanExpired(Status const& s);
   static bool IsOk(Status const& status) { return status.ok(); }
   static bool IsTransientFailure(Status const& status) {
     auto const code = status.code();
     return code == StatusCode::kAborted || code == StatusCode::kUnavailable ||
-           IsQueryPlanExpired(status) ||
-           google::cloud::internal::IsTransientInternalError(status);
+           code == StatusCode::kInternal || IsQueryPlanExpired(status);
   }
   static bool IsPermanentFailure(Status const& status) {
     return !IsOk(status) && !IsTransientFailure(status);

google/cloud/bigtable/internal/retry_traits_test.cc

@@ -13,7 +13,11 @@
 // limitations under the License.
 
 #include "google/cloud/bigtable/internal/retry_traits.h"
+#include "google/cloud/internal/make_status.h"
+#include "google/cloud/internal/status_payload_keys.h"
 #include <gmock/gmock.h>
+#include <google/rpc/error_details.pb.h>
+#include <google/rpc/status.pb.h>
 
 namespace google {
 namespace cloud {
@@ -29,6 +33,61 @@ TEST(SafeGrpcRetry, RstStreamRetried) {
       Status(StatusCode::kInternal, "RST_STREAM")));
 }
 
+TEST(QueryPlanRefreshRetry, IsQueryPlanExpiredNoStatusPayload) {
+  auto non_query_plan_failed_precondition =
+      internal::FailedPreconditionError("not the query plan");
+  EXPECT_FALSE(QueryPlanRefreshRetry::IsQueryPlanExpired(
+      non_query_plan_failed_precondition));
+
+  auto query_plan_expired =
+      internal::FailedPreconditionError("PREPARED_QUERY_EXPIRED");
+  EXPECT_TRUE(QueryPlanRefreshRetry::IsQueryPlanExpired(query_plan_expired));
+}
+
+TEST(QueryPlanRefreshRetry, QueryPlanExpiredStatusPayload) {
+  auto query_plan_expired_violation =
+      internal::FailedPreconditionError("failed precondition");
+  google::rpc::PreconditionFailure_Violation violation;
+  violation.set_type("PREPARED_QUERY_EXPIRED");
+  violation.set_description(
+      "The prepared query has expired. Please re-issue the ExecuteQuery with a "
+      "valid prepared query.");
+  google::rpc::PreconditionFailure precondition;
+  *precondition.add_violations() = violation;
+  google::rpc::Status status;
+  status.set_code(9);
+  status.set_message("failed precondition");
+  google::protobuf::Any any;
+  ASSERT_TRUE(any.PackFrom(precondition));
+  *status.add_details() = any;
+  internal::SetPayload(query_plan_expired_violation,
+                       google::cloud::internal::StatusPayloadGrpcProto(),
+                       status.SerializeAsString());
+  EXPECT_TRUE(
+      QueryPlanRefreshRetry::IsQueryPlanExpired(query_plan_expired_violation));
+}
+
+TEST(QueryPlanRefreshRetry, QueryPlanNotExpiredStatusPayload) {
+  auto query_plan_not_expired_violation =
+      internal::FailedPreconditionError("failed precondition");
+  google::rpc::PreconditionFailure_Violation violation;
+  violation.set_type("something else");
+  violation.set_description("This is not the violation you are looking for");
+  google::rpc::PreconditionFailure precondition;
+  *precondition.add_violations() = violation;
+  google::rpc::Status status;
+  status.set_code(9);
+  status.set_message("failed precondition");
+  google::protobuf::Any any;
+  ASSERT_TRUE(any.PackFrom(precondition));
+  *status.add_details() = any;
+  internal::SetPayload(query_plan_not_expired_violation,
+                       google::cloud::internal::StatusPayloadGrpcProto(),
+                       status.SerializeAsString());
+  EXPECT_FALSE(QueryPlanRefreshRetry::IsQueryPlanExpired(
+      query_plan_not_expired_violation));
+}
+
 }  // namespace
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace bigtable_internal

google/cloud/bigtable/retry_policy.h

@@ -207,7 +207,7 @@ class QueryPlanRefreshLimitedErrorCountRetryPolicy : public DataRetryPolicy {
 
  private:
   google::cloud::internal::LimitedErrorCountRetryPolicy<
-      bigtable_internal::SafeGrpcRetryAllowingQueryPlanRefresh>
+      bigtable_internal::QueryPlanRefreshRetry>
       impl_;
 };
 
@@ -280,7 +280,7 @@ class QueryPlanRefreshLimitedTimeRetryPolicy : public DataRetryPolicy {
 
  private:
   google::cloud::internal::LimitedTimeRetryPolicy<
-      bigtable_internal::SafeGrpcRetryAllowingQueryPlanRefresh>
+      bigtable_internal::QueryPlanRefreshRetry>
       impl_;
 };
 

google/cloud/bigtable/test_proxy/cbt_test_proxy.cc

@@ -373,7 +373,6 @@ grpc::Status CbtTestProxy::ExecuteQuery(
 
   // Call prepare query
   auto instance = MakeInstanceResource(request_proto.instance_name());
-  // NOLINTNEXTLINE(deprecated-declarations)
   bigtable::SqlStatement sql_statement{request_proto.query()};
   auto prepared_query =
       client.PrepareQuery(*std::move(instance), sql_statement);
@@ -390,7 +389,6 @@ grpc::Status CbtTestProxy::ExecuteQuery(
     params.insert(std::make_pair(param.first, std::move(value)));
   }
   auto bound_query = prepared_query->BindParameters(params);
-  auto bound_query_metadata = bound_query.response()->metadata();
 
   RowStream result = client.ExecuteQuery(std::move(bound_query), {});