Skip to content

Commit f41cad5

Browse files
scotthartscotthart
committed
impl(bigtable): update query plan refresh retry logic to pass conformance tests
1 parent 19851cb commit f41cad5

File tree

9 files changed

+127
-23
lines changed

9 files changed

+127
-23
lines changed

‎google/cloud/bigtable/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -224,6 +224,7 @@ add_library(
224224
internal/rate_limiter.h
225225
internal/readrowsparser.cc
226226
internal/readrowsparser.h
227+
internal/retry_traits.cc
227228
internal/retry_traits.h
228229
internal/row_reader_impl.h
229230
internal/rpc_policy_parameters.h

‎google/cloud/bigtable/ci/run_conformance_tests_proxy_bazel.sh‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -50,8 +50,8 @@ exit_status=$?
5050
# Run all the ExecuteQuery tests that either work or we plan to skip such as
5151
# CloseClient
5252
go test -v \
53-
-run "TestExecuteQuery" \
54-
-skip "CloseClient|FailsOnEmptyMetadata|FailsOnExecuteQueryMetadata|FailsOnInvalidType|FailsOnNotEnoughData|FailsOnNotEnoughDataWithCompleteRows|FailsOnSuccesfulStreamWithNoToken|ChecksumMismatch|FailsOnTypeMismatch|FailsOnTypeMismatchWithinMap|FailsOnTypeMismatchWithinArray|FailsOnTypeMismatchWithinStruct|FailsOnStructMissingField|RetryTest_WithPlanRefresh|PlanRefresh" \
53+
-run "TestExecuteQuery|TestExecuteQuery_PlanRefresh$|TestExecuteQuery_PlanRefresh_WithMetadataChange|TestExecuteQuery_PlanRefresh_Retries|TestExecuteQuery_PlanRefresh_RecoversAfterPermanentError" \
54+
-skip "CloseClient|FailsOnEmptyMetadata|FailsOnExecuteQueryMetadata|FailsOnInvalidType|FailsOnNotEnoughData|FailsOnNotEnoughDataWithCompleteRows|FailsOnSuccesfulStreamWithNoToken|ChecksumMismatch|FailsOnTypeMismatch|FailsOnTypeMismatchWithinMap|FailsOnTypeMismatchWithinArray|FailsOnTypeMismatchWithinStruct|FailsOnStructMissingField|TestExecuteQuery_PlanRefresh_AfterResumeTokenCausesError|TestExecuteQuery_RetryTest_WithPlanRefresh|TestExecuteQuery_PlanRefresh_RespectsDeadline" \
5555
-proxy_addr=:9999
5656
exit_status=$?
5757

‎google/cloud/bigtable/google_cloud_cpp_bigtable.bzl‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -230,6 +230,7 @@ google_cloud_cpp_bigtable_srcs = [
230230
"internal/query_plan.cc",
231231
"internal/rate_limiter.cc",
232232
"internal/readrowsparser.cc",
233+
"internal/retry_traits.cc",
233234
"internal/traced_row_reader.cc",
234235
"metadata_update_policy.cc",
235236
"mutation_batcher.cc",

‎google/cloud/bigtable/internal/data_connection_impl.cc‎

Lines changed: 5 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -160,18 +160,16 @@ class DefaultPartialResultSetReader
160160
response.metadata().SerializeToString(&response_metadata_str) &&
161161
initial_metadata_str == response_metadata_str;
162162
if (response.metadata().ByteSizeLong() > 0 && !metadata_matched) {
163-
final_status_ = google::cloud::Status(
164-
google::cloud::StatusCode::kAborted,
165-
"Schema changed during ExecuteQuery operation");
163+
final_status_ = internal::AbortedError(
164+
"Schema changed during ExecuteQuery operation", GCP_ERROR_INFO());
166165
operation_context_->PostCall(*context_, final_status_);
167166
return false;
168167
}
169168
continue;
170169
}
171170

172-
final_status_ = google::cloud::Status(
173-
google::cloud::StatusCode::kInternal,
174-
"Empty ExecuteQueryResponse received from stream");
171+
final_status_ = internal::InternalError(
172+
"Empty ExecuteQueryResponse received from stream", GCP_ERROR_INFO());
175173
operation_context_->PostCall(*context_, final_status_);
176174
return false;
177175
}
@@ -964,8 +962,7 @@ bigtable::RowStream DataConnectionImpl::ExecuteQuery(
964962
}
965963
last_status = source.status();
966964

967-
if (SafeGrpcRetryAllowingQueryPlanRefresh::IsQueryPlanExpired(
968-
source.status())) {
965+
if (QueryPlanRefreshRetry::IsQueryPlanExpired(source.status())) {
969966
query_plan->Invalidate(source.status(),
970967
query_plan_data->prepared_query());
971968
}

‎google/cloud/bigtable/internal/retry_traits.cc‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
// Copyright 2025 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// https://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
#include "google/cloud/bigtable/internal/retry_traits.h"
16+
#include "google/cloud/internal/status_payload_keys.h"
17+
#include "google/cloud/status.h"
18+
#include "absl/strings/match.h"
19+
#include <google/rpc/error_details.pb.h>
20+
#include <google/rpc/status.pb.h>
21+
22+
namespace google {
23+
namespace cloud {
24+
namespace bigtable_internal {
25+
GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
26+
27+
bool QueryPlanRefreshRetry::IsQueryPlanExpired(Status const& s) {
28+
if (s.code() == StatusCode::kFailedPrecondition) {
29+
if (absl::StrContains(s.message(), "PREPARED_QUERY_EXPIRED")) {
30+
return true;
31+
}
32+
auto payload = google::cloud::internal::GetPayload(
33+
s, google::cloud::internal::StatusPayloadGrpcProto());
34+
google::rpc::Status proto;
35+
if (payload && proto.ParseFromString(*payload)) {
36+
google::rpc::PreconditionFailure failure;
37+
for (google::protobuf::Any const& any : proto.details()) {
38+
if (any.UnpackTo(&failure)) {
39+
for (auto const& v : failure.violations()) {
40+
if (absl::StrContains(v.type(), "PREPARED_QUERY_EXPIRED")) {
41+
return true;
42+
}
43+
}
44+
}
45+
}
46+
}
47+
}
48+
return false;
49+
}
50+
51+
GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
52+
} // namespace bigtable_internal
53+
} // namespace cloud
54+
} // namespace google

‎google/cloud/bigtable/internal/retry_traits.h‎

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@
1919
#include "google/cloud/grpc_error_delegate.h"
2020
#include "google/cloud/internal/retry_policy_impl.h"
2121
#include "google/cloud/status.h"
22-
#include "absl/strings/match.h"
2322
#include <grpcpp/grpcpp.h>
2423

2524
namespace google {
@@ -49,18 +48,13 @@ struct SafeGrpcRetry {
4948
}
5049
};
5150

52-
struct SafeGrpcRetryAllowingQueryPlanRefresh {
53-
static bool IsQueryPlanExpired(Status const& s) {
54-
return (s.code() == StatusCode::kFailedPrecondition &&
55-
absl::StrContains(s.message(), "PREPARED_QUERY_EXPIRED"));
56-
}
57-
51+
struct QueryPlanRefreshRetry {
52+
static bool IsQueryPlanExpired(Status const& s);
5853
static bool IsOk(Status const& status) { return status.ok(); }
5954
static bool IsTransientFailure(Status const& status) {
6055
auto const code = status.code();
6156
return code == StatusCode::kAborted || code == StatusCode::kUnavailable ||
62-
IsQueryPlanExpired(status) ||
63-
google::cloud::internal::IsTransientInternalError(status);
57+
code == StatusCode::kInternal || IsQueryPlanExpired(status);
6458
}
6559
static bool IsPermanentFailure(Status const& status) {
6660
return !IsOk(status) && !IsTransientFailure(status);

‎google/cloud/bigtable/internal/retry_traits_test.cc‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,10 @@
1313
// limitations under the License.
1414

1515
#include "google/cloud/bigtable/internal/retry_traits.h"
16+
#include "google/cloud/internal/make_status.h"
17+
#include "google/cloud/internal/status_payload_keys.h"
18+
#include <google/rpc/error_details.pb.h>
19+
#include <google/rpc/status.pb.h>
1620
#include <gmock/gmock.h>
1721

1822
namespace google {
@@ -29,6 +33,61 @@ TEST(SafeGrpcRetry, RstStreamRetried) {
2933
Status(StatusCode::kInternal, "RST_STREAM")));
3034
}
3135

36+
TEST(QueryPlanRefreshRetry, IsQueryPlanExpiredNoStatusPayload) {
37+
auto non_query_plan_failed_precondition =
38+
internal::FailedPreconditionError("not the query plan");
39+
EXPECT_FALSE(QueryPlanRefreshRetry::IsQueryPlanExpired(
40+
non_query_plan_failed_precondition));
41+
42+
auto query_plan_expired =
43+
internal::FailedPreconditionError("PREPARED_QUERY_EXPIRED");
44+
EXPECT_TRUE(QueryPlanRefreshRetry::IsQueryPlanExpired(query_plan_expired));
45+
}
46+
47+
TEST(QueryPlanRefreshRetry, QueryPlanExpiredStatusPayload) {
48+
auto query_plan_expired_violation =
49+
internal::FailedPreconditionError("failed precondition");
50+
google::rpc::PreconditionFailure_Violation violation;
51+
violation.set_type("PREPARED_QUERY_EXPIRED");
52+
violation.set_description(
53+
"The prepared query has expired. Please re-issue the ExecuteQuery with a "
54+
"valid prepared query.");
55+
google::rpc::PreconditionFailure precondition;
56+
*precondition.add_violations() = violation;
57+
google::rpc::Status status;
58+
status.set_code(9);
59+
status.set_message("failed precondition");
60+
google::protobuf::Any any;
61+
ASSERT_TRUE(any.PackFrom(precondition));
62+
*status.add_details() = any;
63+
internal::SetPayload(query_plan_expired_violation,
64+
google::cloud::internal::StatusPayloadGrpcProto(),
65+
status.SerializeAsString());
66+
EXPECT_TRUE(
67+
QueryPlanRefreshRetry::IsQueryPlanExpired(query_plan_expired_violation));
68+
}
69+
70+
TEST(QueryPlanRefreshRetry, QueryPlanNotExpiredStatusPayload) {
71+
auto query_plan_not_expired_violation =
72+
internal::FailedPreconditionError("failed precondition");
73+
google::rpc::PreconditionFailure_Violation violation;
74+
violation.set_type("something else");
75+
violation.set_description("This is not the violation you are looking for");
76+
google::rpc::PreconditionFailure precondition;
77+
*precondition.add_violations() = violation;
78+
google::rpc::Status status;
79+
status.set_code(9);
80+
status.set_message("failed precondition");
81+
google::protobuf::Any any;
82+
ASSERT_TRUE(any.PackFrom(precondition));
83+
*status.add_details() = any;
84+
internal::SetPayload(query_plan_not_expired_violation,
85+
google::cloud::internal::StatusPayloadGrpcProto(),
86+
status.SerializeAsString());
87+
EXPECT_FALSE(QueryPlanRefreshRetry::IsQueryPlanExpired(
88+
query_plan_not_expired_violation));
89+
}
90+
3291
} // namespace
3392
GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
3493
} // namespace bigtable_internal

‎google/cloud/bigtable/retry_policy.h‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -207,7 +207,7 @@ class QueryPlanRefreshLimitedErrorCountRetryPolicy : public DataRetryPolicy {
207207

208208
private:
209209
google::cloud::internal::LimitedErrorCountRetryPolicy<
210-
bigtable_internal::SafeGrpcRetryAllowingQueryPlanRefresh>
210+
bigtable_internal::QueryPlanRefreshRetry>
211211
impl_;
212212
};
213213

@@ -280,7 +280,7 @@ class QueryPlanRefreshLimitedTimeRetryPolicy : public DataRetryPolicy {
280280

281281
private:
282282
google::cloud::internal::LimitedTimeRetryPolicy<
283-
bigtable_internal::SafeGrpcRetryAllowingQueryPlanRefresh>
283+
bigtable_internal::QueryPlanRefreshRetry>
284284
impl_;
285285
};
286286

‎google/cloud/bigtable/test_proxy/cbt_test_proxy.cc‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -373,7 +373,6 @@ grpc::Status CbtTestProxy::ExecuteQuery(
373373

374374
// Call prepare query
375375
auto instance = MakeInstanceResource(request_proto.instance_name());
376-
// NOLINTNEXTLINE(deprecated-declarations)
377376
bigtable::SqlStatement sql_statement{request_proto.query()};
378377
auto prepared_query =
379378
client.PrepareQuery(*std::move(instance), sql_statement);
@@ -390,7 +389,6 @@ grpc::Status CbtTestProxy::ExecuteQuery(
390389
params.insert(std::make_pair(param.first, std::move(value)));
391390
}
392391
auto bound_query = prepared_query->BindParameters(params);
393-
auto bound_query_metadata = bound_query.response()->metadata();
394392

395393
RowStream result = client.ExecuteQuery(std::move(bound_query), {});
396394

0 commit comments

Comments
 (0)