Skip to content

Commit 73903cd

Browse files
scotthartscotthart
authored
Merge branch 'main' into bigtable_query_plan_refresh_1
2 parents 0e6cb38 + cfa8052 commit 73903cd

File tree

5 files changed

+243
-14
lines changed

5 files changed

+243
-14
lines changed

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

Lines changed: 32 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -622,20 +622,28 @@ StatusOr<bigtable::PreparedQuery> DataConnectionImpl::PrepareQuery(
622622
bigtable::PrepareQueryParams const& params) {
623623
auto current = google::cloud::internal::SaveCurrentOptions();
624624
google::bigtable::v2::PrepareQueryRequest request;
625-
request.set_instance_name(params.instance.FullName());
625+
auto instance_full_name = params.instance.FullName();
626+
request.set_instance_name(instance_full_name);
626627
request.set_app_profile_id(app_profile_id(*current));
627628
request.set_query(params.sql_statement.sql());
628629
for (auto const& p : params.sql_statement.params()) {
629630
(*request.mutable_param_types())[p.first] = p.second.type();
630631
}
632+
auto operation_context = operation_context_factory_->PrepareQuery(
633+
instance_full_name, app_profile_id(*current));
631634
auto response = google::cloud::internal::RetryLoop(
632635
retry_policy(*current), backoff_policy(*current),
633636
Idempotency::kIdempotent,
634-
[this](grpc::ClientContext& context, Options const& options,
635-
google::bigtable::v2::PrepareQueryRequest const& request) {
636-
return stub_->PrepareQuery(context, options, request);
637+
[this, operation_context](
638+
grpc::ClientContext& context, Options const& options,
639+
google::bigtable::v2::PrepareQueryRequest const& request) {
640+
operation_context->PreCall(context);
641+
auto const& result = stub_->PrepareQuery(context, options, request);
642+
operation_context->PostCall(context, result.status());
643+
return result;
637644
},
638645
*current, request, __func__);
646+
operation_context->OnDone(response.status());
639647
if (!response) {
640648
return std::move(response).status();
641649
}
@@ -665,30 +673,42 @@ future<StatusOr<bigtable::PreparedQuery>> DataConnectionImpl::AsyncPrepareQuery(
665673
bigtable::PrepareQueryParams const& params) {
666674
auto current = google::cloud::internal::SaveCurrentOptions();
667675
google::bigtable::v2::PrepareQueryRequest request;
668-
request.set_instance_name(params.instance.FullName());
676+
auto instance_full_name = params.instance.FullName();
677+
request.set_instance_name(instance_full_name);
669678
request.set_app_profile_id(app_profile_id(*current));
670679
request.set_query(params.sql_statement.sql());
671680
for (auto const& p : params.sql_statement.params()) {
672681
(*request.mutable_param_types())[p.first] = p.second.type();
673682
}
674683
auto retry = retry_policy(*current);
675684
auto backoff = backoff_policy(*current);
676-
auto const* func = __func__;
685+
auto operation_context = operation_context_factory_->PrepareQuery(
686+
instance_full_name, app_profile_id(*current));
687+
auto const* func = __func__;
677688
return google::cloud::internal::AsyncRetryLoop(
678689
std::move(retry), std::move(backoff), Idempotency::kIdempotent,
679690
background_->cq(),
680-
[this](CompletionQueue& cq,
681-
std::shared_ptr<grpc::ClientContext> context,
682-
google::cloud::internal::ImmutableOptions options,
683-
google::bigtable::v2::PrepareQueryRequest const& request) {
684-
return stub_->AsyncPrepareQuery(cq, std::move(context),
685-
std::move(options), request);
691+
[this, operation_context](
692+
CompletionQueue& cq,
693+
std::shared_ptr<grpc::ClientContext> context,
694+
google::cloud::internal::ImmutableOptions options,
695+
google::bigtable::v2::PrepareQueryRequest const& request) {
696+
operation_context->PreCall(*context);
697+
auto f = stub_->AsyncPrepareQuery(cq, context,
698+
std::move(options), request);
699+
return f.then(
700+
[operation_context, context = std::move(context)](auto f) {
701+
auto s = f.get();
702+
operation_context->PostCall(*context, s.status());
703+
return s;
704+
});
686705
},
687706
current, request, func)
688707
.then([this, request, current, params = std::move(params),
689708
func](future<StatusOr<google::bigtable::v2::PrepareQueryResponse>>
690709
future) -> StatusOr<bigtable::PreparedQuery> {
691710
auto response = future.get();
711+
operation_context->OnDone(response.status());
692712
if (!response) {
693713
return std::move(response).status();
694714
}

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

Lines changed: 68 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -371,6 +371,16 @@ class FakeOperationContextFactory : public OperationContextFactory {
371371
return Helper(name, app_profile);
372372
}
373373

374+
std::shared_ptr<OperationContext> PrepareQuery(
375+
std::string const&, std::string const& app_profile) override {
376+
return Helper("", app_profile);
377+
}
378+
379+
std::shared_ptr<OperationContext> ExecuteQuery(
380+
std::string const&, std::string const& app_profile) override {
381+
return Helper("", app_profile);
382+
}
383+
374384
private:
375385
std::shared_ptr<OperationContext> Helper(std::string const& name,
376386
std::string const& app_profile) {
@@ -2780,6 +2790,20 @@ TEST_F(DataConnectionTest, ExecuteQuery) {
27802790
}
27812791

27822792
TEST_F(DataConnectionTest, PrepareQuerySuccess) {
2793+
#ifdef GOOGLE_CLOUD_CPP_BIGTABLE_WITH_OTEL_METRICS
2794+
auto mock_metric = std::make_unique<MockMetric>();
2795+
EXPECT_CALL(*mock_metric, PreCall).Times(1);
2796+
EXPECT_CALL(*mock_metric, PostCall).Times(1);
2797+
EXPECT_CALL(*mock_metric, OnDone).Times(1);
2798+
EXPECT_CALL(*mock_metric, ElementRequest).Times(0);
2799+
EXPECT_CALL(*mock_metric, ElementDelivery).Times(0);
2800+
auto fake_metric = std::make_shared<CloningMetric>(std::move(mock_metric));
2801+
auto clock = std::make_shared<testing_util::FakeSteadyClock>();
2802+
auto factory = std::make_unique<FakeOperationContextFactory>(
2803+
ResourceLabels{}, DataLabels{}, fake_metric, clock);
2804+
#else
2805+
auto factory = std::make_unique<SimpleOperationContextFactory>();
2806+
#endif
27832807
auto mock = std::make_shared<MockBigtableStub>();
27842808
EXPECT_CALL(*mock, PrepareQuery)
27852809
.WillOnce([](grpc::ClientContext&, Options const&,
@@ -2816,13 +2840,27 @@ TEST_F(DataConnectionTest, PrepareQuerySuccess) {
28162840
}
28172841

28182842
TEST_F(DataConnectionTest, PrepareQueryPermanentError) {
2843+
#ifdef GOOGLE_CLOUD_CPP_BIGTABLE_WITH_OTEL_METRICS
2844+
auto mock_metric = std::make_unique<MockMetric>();
2845+
EXPECT_CALL(*mock_metric, PreCall).Times(1);
2846+
EXPECT_CALL(*mock_metric, PostCall).Times(1);
2847+
EXPECT_CALL(*mock_metric, OnDone).Times(1);
2848+
EXPECT_CALL(*mock_metric, ElementRequest).Times(0);
2849+
EXPECT_CALL(*mock_metric, ElementDelivery).Times(0);
2850+
auto fake_metric = std::make_shared<CloningMetric>(std::move(mock_metric));
2851+
auto clock = std::make_shared<testing_util::FakeSteadyClock>();
2852+
auto factory = std::make_unique<FakeOperationContextFactory>(
2853+
ResourceLabels{}, DataLabels{}, fake_metric, clock);
2854+
#else
2855+
auto factory = std::make_unique<SimpleOperationContextFactory>();
2856+
#endif
28192857
auto mock = std::make_shared<MockBigtableStub>();
28202858
EXPECT_CALL(*mock, PrepareQuery)
28212859
.WillOnce(
28222860
[](grpc::ClientContext&, Options const&,
28232861
v2::PrepareQueryRequest const&) { return PermanentError(); });
28242862

2825-
auto conn = TestConnection(std::move(mock));
2863+
auto conn = TestConnection(std::move(mock), std::move(factory));
28262864
internal::OptionsSpan span(CallOptions());
28272865
auto result = conn->PrepareQuery(bigtable::PrepareQueryParams{
28282866
bigtable::InstanceResource(google::cloud::Project("the-project"),
@@ -2832,6 +2870,20 @@ TEST_F(DataConnectionTest, PrepareQueryPermanentError) {
28322870
}
28332871

28342872
TEST_F(DataConnectionTest, AsyncPrepareQuerySuccess) {
2873+
#ifdef GOOGLE_CLOUD_CPP_BIGTABLE_WITH_OTEL_METRICS
2874+
auto mock_metric = std::make_unique<MockMetric>();
2875+
EXPECT_CALL(*mock_metric, PreCall).Times(1);
2876+
EXPECT_CALL(*mock_metric, PostCall).Times(1);
2877+
EXPECT_CALL(*mock_metric, OnDone).Times(1);
2878+
EXPECT_CALL(*mock_metric, ElementRequest).Times(0);
2879+
EXPECT_CALL(*mock_metric, ElementDelivery).Times(0);
2880+
auto fake_metric = std::make_shared<CloningMetric>(std::move(mock_metric));
2881+
auto clock = std::make_shared<testing_util::FakeSteadyClock>();
2882+
auto factory = std::make_unique<FakeOperationContextFactory>(
2883+
ResourceLabels{}, DataLabels{}, fake_metric, clock);
2884+
#else
2885+
auto factory = std::make_unique<SimpleOperationContextFactory>();
2886+
#endif
28352887
auto mock = std::make_shared<MockBigtableStub>();
28362888
EXPECT_CALL(*mock, AsyncPrepareQuery)
28372889
.WillOnce([](CompletionQueue const&, auto, auto,
@@ -2864,6 +2916,20 @@ TEST_F(DataConnectionTest, AsyncPrepareQuerySuccess) {
28642916
}
28652917

28662918
TEST_F(DataConnectionTest, AsyncPrepareQueryPermanentError) {
2919+
#ifdef GOOGLE_CLOUD_CPP_BIGTABLE_WITH_OTEL_METRICS
2920+
auto mock_metric = std::make_unique<MockMetric>();
2921+
EXPECT_CALL(*mock_metric, PreCall).Times(1);
2922+
EXPECT_CALL(*mock_metric, PostCall).Times(1);
2923+
EXPECT_CALL(*mock_metric, OnDone).Times(1);
2924+
EXPECT_CALL(*mock_metric, ElementRequest).Times(0);
2925+
EXPECT_CALL(*mock_metric, ElementDelivery).Times(0);
2926+
auto fake_metric = std::make_shared<CloningMetric>(std::move(mock_metric));
2927+
auto clock = std::make_shared<testing_util::FakeSteadyClock>();
2928+
auto factory = std::make_unique<FakeOperationContextFactory>(
2929+
ResourceLabels{}, DataLabels{}, fake_metric, clock);
2930+
#else
2931+
auto factory = std::make_unique<SimpleOperationContextFactory>();
2932+
#endif
28672933
auto mock = std::make_shared<MockBigtableStub>();
28682934
EXPECT_CALL(*mock, AsyncPrepareQuery)
28692935
.WillOnce(
@@ -2872,7 +2938,7 @@ TEST_F(DataConnectionTest, AsyncPrepareQueryPermanentError) {
28722938
PermanentError());
28732939
});
28742940

2875-
auto conn = TestConnection(std::move(mock));
2941+
auto conn = TestConnection(std::move(mock), std::move(factory));
28762942
internal::OptionsSpan span(CallOptions());
28772943
auto params = bigtable::PrepareQueryParams{
28782944
bigtable::InstanceResource(google::cloud::Project("the-project"),

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

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,19 @@ ResourceLabels ResourceLabelsFromTableName(std::string const& table_name) {
4848
return resource_labels;
4949
}
5050

51+
ResourceLabels ResourceLabelsFromInstanceName(
52+
std::string const& instance_name) {
53+
// split instance_name into component pieces
54+
// projects/<project>/instances/<instance>
55+
std::vector<absl::string_view> name_parts =
56+
absl::StrSplit(instance_name, '/');
57+
if (name_parts.size() < 4) return {};
58+
ResourceLabels resource_labels = {std::string(name_parts[1]),
59+
std::string(name_parts[3]), "",
60+
"" /*=cluster*/, "" /*=zone*/};
61+
return resource_labels;
62+
}
63+
5164
} // namespace
5265
#endif // GOOGLE_CLOUD_CPP_BIGTABLE_WITH_OTEL_METRICS
5366

@@ -82,6 +95,16 @@ std::shared_ptr<OperationContext> OperationContextFactory::ReadModifyWriteRow(
8295
return std::make_shared<OperationContext>();
8396
}
8497

98+
std::shared_ptr<OperationContext> OperationContextFactory::PrepareQuery(
99+
std::string const&, std::string const&) {
100+
return std::make_shared<OperationContext>();
101+
}
102+
103+
std::shared_ptr<OperationContext> OperationContextFactory::ExecuteQuery(
104+
std::string const&, std::string const&) {
105+
return std::make_shared<OperationContext>();
106+
}
107+
85108
std::shared_ptr<OperationContext> SimpleOperationContextFactory::ReadRow(
86109
std::string const&, std::string const&) {
87110
return std::make_shared<OperationContext>();
@@ -113,6 +136,16 @@ SimpleOperationContextFactory::ReadModifyWriteRow(std::string const&,
113136
return std::make_shared<OperationContext>();
114137
}
115138

139+
std::shared_ptr<OperationContext> SimpleOperationContextFactory::PrepareQuery(
140+
std::string const&, std::string const&) {
141+
return std::make_shared<OperationContext>();
142+
}
143+
144+
std::shared_ptr<OperationContext> SimpleOperationContextFactory::ExecuteQuery(
145+
std::string const&, std::string const&) {
146+
return std::make_shared<OperationContext>();
147+
}
148+
116149
#ifdef GOOGLE_CLOUD_CPP_BIGTABLE_WITH_OTEL_METRICS
117150

118151
MetricsOperationContextFactory::MetricsOperationContextFactory(
@@ -154,6 +187,12 @@ MetricsOperationContextFactory::MetricsOperationContextFactory(
154187
absl::call_once(read_modify_write_row_metrics_.once, [this, metric]() {
155188
read_modify_write_row_metrics_.metrics.push_back(metric);
156189
});
190+
absl::call_once(prepare_query_metrics_.once, [this, metric]() {
191+
prepare_query_metrics_.metrics.push_back(metric);
192+
});
193+
absl::call_once(execute_query_metrics_.once, [this, metric]() {
194+
execute_query_metrics_.metrics.push_back(metric);
195+
});
157196
}
158197

159198
void MetricsOperationContextFactory::InitializeProvider(
@@ -424,6 +463,57 @@ MetricsOperationContextFactory::ReadModifyWriteRow(
424463
clock_);
425464
}
426465

466+
std::shared_ptr<OperationContext> MetricsOperationContextFactory::PrepareQuery(
467+
std::string const& instance_name, std::string const& app_profile) {
468+
auto constexpr kRpc = "PrepareQuery";
469+
absl::call_once(prepare_query_metrics_.once, [this, kRpc]() {
470+
std::vector<std::shared_ptr<Metric const>> v;
471+
v.emplace_back(std::make_shared<OperationLatency>(kRpc, provider_));
472+
v.emplace_back(std::make_shared<AttemptLatency>(kRpc, provider_));
473+
v.emplace_back(std::make_shared<RetryCount>(kRpc, provider_));
474+
v.emplace_back(std::make_shared<ServerLatency>(kRpc, provider_));
475+
v.emplace_back(std::make_shared<ConnectivityErrorCount>(kRpc, provider_));
476+
swap(prepare_query_metrics_.metrics, v);
477+
});
478+
479+
auto resource_labels = ResourceLabelsFromInstanceName(instance_name);
480+
DataLabels data_labels = {kRpc,
481+
"false", /*=streaming*/
482+
"cpp.Bigtable/" + version_string(),
483+
client_uid_,
484+
app_profile,
485+
"" /*=status*/};
486+
487+
return std::make_shared<OperationContext>(
488+
resource_labels, data_labels, prepare_query_metrics_.metrics, clock_);
489+
}
490+
491+
std::shared_ptr<OperationContext> MetricsOperationContextFactory::ExecuteQuery(
492+
std::string const& instance_name, std::string const& app_profile) {
493+
auto constexpr kRpc = "ExecuteQuery";
494+
absl::call_once(execute_query_metrics_.once, [this, kRpc]() {
495+
std::vector<std::shared_ptr<Metric const>> v;
496+
v.emplace_back(std::make_shared<OperationLatency>(kRpc, provider_));
497+
v.emplace_back(std::make_shared<AttemptLatency>(kRpc, provider_));
498+
v.emplace_back(std::make_shared<RetryCount>(kRpc, provider_));
499+
v.emplace_back(std::make_shared<FirstResponseLatency>(kRpc, provider_));
500+
v.emplace_back(std::make_shared<ServerLatency>(kRpc, provider_));
501+
v.emplace_back(std::make_shared<ConnectivityErrorCount>(kRpc, provider_));
502+
swap(execute_query_metrics_.metrics, v);
503+
});
504+
505+
auto resource_labels = ResourceLabelsFromInstanceName(instance_name);
506+
DataLabels data_labels = {kRpc,
507+
"true", /*=streaming*/
508+
"cpp.Bigtable/" + version_string(),
509+
client_uid_,
510+
app_profile,
511+
"" /*=status*/};
512+
513+
return std::make_shared<OperationContext>(
514+
resource_labels, data_labels, execute_query_metrics_.metrics, clock_);
515+
}
516+
427517
#endif // GOOGLE_CLOUD_CPP_BIGTABLE_WITH_OTEL_METRICS
428518

429519
GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END

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

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,10 @@ class OperationContextFactory {
5454
std::string const& name, std::string const& app_profile);
5555
virtual std::shared_ptr<OperationContext> ReadModifyWriteRow(
5656
std::string const& name, std::string const& app_profile);
57+
virtual std::shared_ptr<OperationContext> PrepareQuery(
58+
std::string const& instance_name, std::string const& app_profile);
59+
virtual std::shared_ptr<OperationContext> ExecuteQuery(
60+
std::string const& instance_name, std::string const& app_profile);
5761
};
5862

5963
class SimpleOperationContextFactory : public OperationContextFactory {
@@ -72,6 +76,12 @@ class SimpleOperationContextFactory : public OperationContextFactory {
7276
std::string const& name, std::string const& app_profile) override;
7377
std::shared_ptr<OperationContext> ReadModifyWriteRow(
7478
std::string const& name, std::string const& app_profile) override;
79+
std::shared_ptr<OperationContext> PrepareQuery(
80+
std::string const& instance_name,
81+
std::string const& app_profile) override;
82+
std::shared_ptr<OperationContext> ExecuteQuery(
83+
std::string const& instance_name,
84+
std::string const& app_profile) override;
7585
};
7686

7787
#ifdef GOOGLE_CLOUD_CPP_BIGTABLE_WITH_OTEL_METRICS
@@ -115,6 +125,13 @@ class MetricsOperationContextFactory : public OperationContextFactory {
115125
std::shared_ptr<OperationContext> ReadModifyWriteRow(
116126
std::string const& table_name, std::string const& app_profile) override;
117127

128+
std::shared_ptr<OperationContext> PrepareQuery(
129+
std::string const& instance_name,
130+
std::string const& app_profile) override;
131+
std::shared_ptr<OperationContext> ExecuteQuery(
132+
std::string const& instance_name,
133+
std::string const& app_profile) override;
134+
118135
private:
119136
void InitializeProvider(
120137
std::shared_ptr<monitoring_v3::MetricServiceConnection> conn,
@@ -137,6 +154,8 @@ class MetricsOperationContextFactory : public OperationContextFactory {
137154
MetricHolder check_and_mutate_row_metrics_;
138155
MetricHolder sample_row_keys_metrics_;
139156
MetricHolder read_modify_write_row_metrics_;
157+
MetricHolder prepare_query_metrics_;
158+
MetricHolder execute_query_metrics_;
140159
};
141160

142161
#endif // GOOGLE_CLOUD_CPP_BIGTABLE_WITH_OTEL_METRICS

0 commit comments

Comments
 (0)