Skip to content

Commit d799e61

Browse files
jsrinnnjsrinnn
authored
impl(bigquery): Stub and Decorator changes for GetQueryResults api (#12261)
1 parent 8660f8a commit d799e61

10 files changed

Lines changed: 211 additions & 0 deletions

File tree

‎google/cloud/bigquery/v2/minimal/internal/job_logging.cc‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,20 @@ StatusOr<QueryResponse> BigQueryJobLogging::Query(
9797
tracing_options_);
9898
}
9999

100+
StatusOr<GetQueryResultsResponse> BigQueryJobLogging::GetQueryResults(
101+
rest_internal::RestContext& rest_context,
102+
GetQueryResultsRequest const& request) {
103+
return LogWrapper(
104+
[this](rest_internal::RestContext& rest_context,
105+
GetQueryResultsRequest const& request) {
106+
return child_->GetQueryResults(rest_context, request);
107+
},
108+
rest_context, request, __func__,
109+
"google.cloud.bigquery.v2.minimal.internal.GetQueryResultsRequest",
110+
"google.cloud.bigquery.v2.minimal.internal.GetQueryResultsResponse",
111+
tracing_options_);
112+
}
113+
100114
GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
101115
} // namespace bigquery_v2_minimal_internal
102116
} // namespace cloud

‎google/cloud/bigquery/v2/minimal/internal/job_logging.h‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,10 @@ class BigQueryJobLogging : public BigQueryJobRestStub {
4949
StatusOr<QueryResponse> Query(rest_internal::RestContext& rest_context,
5050
PostQueryRequest const& request) override;
5151

52+
StatusOr<GetQueryResultsResponse> GetQueryResults(
53+
rest_internal::RestContext& rest_context,
54+
GetQueryResultsRequest const& request) override;
55+
5256
private:
5357
std::shared_ptr<BigQueryJobRestStub> child_;
5458
TracingOptions tracing_options_;

‎google/cloud/bigquery/v2/minimal/internal/job_logging_test.cc‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,8 @@ namespace cloud {
2828
namespace bigquery_v2_minimal_internal {
2929
GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
3030

31+
using ::google::cloud::bigquery_v2_minimal_testing::
32+
MakeGetQueryResultsResponsePayload;
3133
using ::google::cloud::bigquery_v2_minimal_testing::MakePartialJob;
3234
using ::google::cloud::bigquery_v2_minimal_testing::MakeQueryRequest;
3335
using ::google::cloud::bigquery_v2_minimal_testing::MakeQueryResponsePayload;
@@ -323,6 +325,52 @@ TEST(JobLoggingClientTest, Query) {
323325
EXPECT_THAT(actual_lines, Contains(HasSubstr(R"(value: "value-2")")));
324326
}
325327

328+
TEST(JobLoggingClientTest, GetQueryResults) {
329+
ScopedLog log;
330+
331+
auto mock_stub = std::make_shared<MockBigQueryJobRestStub>();
332+
auto expected_payload = MakeGetQueryResultsResponsePayload();
333+
334+
EXPECT_CALL(*mock_stub, GetQueryResults)
335+
.WillOnce([&](rest_internal::RestContext&,
336+
GetQueryResultsRequest const& request)
337+
-> StatusOr<GetQueryResultsResponse> {
338+
EXPECT_THAT(request.project_id(), Not(IsEmpty()));
339+
EXPECT_THAT(request.job_id(), Not(IsEmpty()));
340+
BigQueryHttpResponse http_response;
341+
http_response.payload = expected_payload;
342+
return GetQueryResultsResponse::BuildFromHttpResponse(
343+
std::move(http_response));
344+
});
345+
346+
auto client = CreateMockJobLogging(std::move(mock_stub));
347+
GetQueryResultsRequest request("p123", "j123");
348+
349+
rest_internal::RestContext context;
350+
context.AddHeader("header-1", "value-1");
351+
context.AddHeader("header-2", "value-2");
352+
353+
client->GetQueryResults(context, request);
354+
355+
auto actual_lines = log.ExtractLines();
356+
357+
EXPECT_THAT(actual_lines, Contains(HasSubstr(" << ")));
358+
EXPECT_THAT(actual_lines, Contains(HasSubstr(R"(GetQueryResultsRequest)")));
359+
EXPECT_THAT(actual_lines, Contains(HasSubstr(R"(job_id: "j123")")).Times(2));
360+
EXPECT_THAT(actual_lines,
361+
Contains(HasSubstr(R"(project_id: "p123")")).Times(2));
362+
EXPECT_THAT(actual_lines, Contains(HasSubstr(R"(GetQueryResultsResponse)")));
363+
EXPECT_THAT(actual_lines, Contains(HasSubstr(R"(kind: "query-kind")")));
364+
EXPECT_THAT(actual_lines, Contains(HasSubstr(R"(etag: "query-etag")")));
365+
EXPECT_THAT(actual_lines, Contains(HasSubstr(R"(job_id: "j123")")));
366+
EXPECT_THAT(actual_lines, Contains(HasSubstr(R"(page_token: "np123")")));
367+
EXPECT_THAT(actual_lines, Contains(HasSubstr(R"(Context)")));
368+
EXPECT_THAT(actual_lines, Contains(HasSubstr(R"(name: "header-1")")));
369+
EXPECT_THAT(actual_lines, Contains(HasSubstr(R"(value: "value-1")")));
370+
EXPECT_THAT(actual_lines, Contains(HasSubstr(R"(name: "header-2")")));
371+
EXPECT_THAT(actual_lines, Contains(HasSubstr(R"(value: "value-2")")));
372+
}
373+
326374
GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
327375
} // namespace bigquery_v2_minimal_internal
328376
} // namespace cloud

‎google/cloud/bigquery/v2/minimal/internal/job_metadata.cc‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,13 @@ StatusOr<QueryResponse> BigQueryJobMetadata::Query(
6363
return child_->Query(context, request);
6464
}
6565

66+
StatusOr<GetQueryResultsResponse> BigQueryJobMetadata::GetQueryResults(
67+
rest_internal::RestContext& context,
68+
GetQueryResultsRequest const& request) {
69+
SetMetadata(context);
70+
return child_->GetQueryResults(context, request);
71+
}
72+
6673
void BigQueryJobMetadata::SetMetadata(rest_internal::RestContext& rest_context,
6774
std::vector<std::string> const& params) {
6875
rest_context.AddHeader("x-goog-api-client", api_client_header_);

‎google/cloud/bigquery/v2/minimal/internal/job_metadata.h‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,10 @@ class BigQueryJobMetadata : public BigQueryJobRestStub {
4646
StatusOr<QueryResponse> Query(rest_internal::RestContext& rest_context,
4747
PostQueryRequest const& request) override;
4848

49+
StatusOr<GetQueryResultsResponse> GetQueryResults(
50+
rest_internal::RestContext& rest_context,
51+
GetQueryResultsRequest const& request) override;
52+
4953
private:
5054
void SetMetadata(rest_internal::RestContext& context,
5155
std::vector<std::string> const& params = {});

‎google/cloud/bigquery/v2/minimal/internal/job_metadata_test.cc‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@ namespace bigquery_v2_minimal_internal {
2929
GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
3030

3131
using ::google::cloud::bigquery_v2_minimal_testing::GetMetadataOptions;
32+
using ::google::cloud::bigquery_v2_minimal_testing::
33+
MakeGetQueryResultsResponsePayload;
3234
using ::google::cloud::bigquery_v2_minimal_testing::MakePartialJob;
3335
using ::google::cloud::bigquery_v2_minimal_testing::MakeQueryRequest;
3436
using ::google::cloud::bigquery_v2_minimal_testing::MakeQueryResponsePayload;
@@ -232,6 +234,35 @@ TEST(JobMetadataTest, Query) {
232234
VerifyMetadataContext(context, "bigquery_v2_job");
233235
}
234236

237+
TEST(JobMetadataTest, GetQueryResults) {
238+
auto mock_stub = std::make_shared<MockBigQueryJobRestStub>();
239+
auto expected_payload = MakeGetQueryResultsResponsePayload();
240+
241+
EXPECT_CALL(*mock_stub, GetQueryResults)
242+
.WillOnce([&](rest_internal::RestContext&,
243+
GetQueryResultsRequest const& request)
244+
-> StatusOr<GetQueryResultsResponse> {
245+
EXPECT_THAT(request.project_id(), Not(IsEmpty()));
246+
EXPECT_THAT(request.job_id(), Not(IsEmpty()));
247+
BigQueryHttpResponse http_response;
248+
http_response.payload = expected_payload;
249+
return GetQueryResultsResponse::BuildFromHttpResponse(
250+
std::move(http_response));
251+
});
252+
253+
auto metadata = CreateMockJobMetadata(std::move(mock_stub));
254+
GetQueryResultsRequest request;
255+
rest_internal::RestContext context;
256+
request.set_project_id("test-project-id");
257+
request.set_job_id("test-job-id");
258+
259+
internal::OptionsSpan span(GetMetadataOptions());
260+
261+
auto result = metadata->GetQueryResults(context, request);
262+
ASSERT_STATUS_OK(result);
263+
VerifyMetadataContext(context, "bigquery_v2_job");
264+
}
265+
235266
GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
236267
} // namespace bigquery_v2_minimal_internal
237268
} // namespace cloud

‎google/cloud/bigquery/v2/minimal/internal/job_rest_stub.cc‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,19 @@ StatusOr<QueryResponse> DefaultBigQueryJobRestStub::Query(
111111
{absl::MakeConstSpan(json_payload.dump())}));
112112
}
113113

114+
StatusOr<GetQueryResultsResponse> DefaultBigQueryJobRestStub::GetQueryResults(
115+
rest_internal::RestContext& rest_context,
116+
GetQueryResultsRequest const& request) {
117+
// Prepare the RestRequest from GetQueryResultsRequest.
118+
auto rest_request =
119+
PrepareRestRequest<GetQueryResultsRequest>(rest_context, request);
120+
121+
// Call the rest stub and parse the RestResponse.
122+
rest_internal::RestContext context;
123+
return ParseFromRestResponse<GetQueryResultsResponse>(
124+
rest_stub_->Get(context, std::move(*rest_request)));
125+
}
126+
114127
GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
115128
} // namespace bigquery_v2_minimal_internal
116129
} // namespace cloud

‎google/cloud/bigquery/v2/minimal/internal/job_rest_stub.h‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,9 @@ class BigQueryJobRestStub {
4848
virtual StatusOr<QueryResponse> Query(
4949
rest_internal::RestContext& rest_context,
5050
PostQueryRequest const& request) = 0;
51+
virtual StatusOr<GetQueryResultsResponse> GetQueryResults(
52+
rest_internal::RestContext& rest_context,
53+
GetQueryResultsRequest const& request) = 0;
5154
};
5255

5356
class DefaultBigQueryJobRestStub : public BigQueryJobRestStub {
@@ -71,6 +74,10 @@ class DefaultBigQueryJobRestStub : public BigQueryJobRestStub {
7174
StatusOr<QueryResponse> Query(rest_internal::RestContext& rest_context,
7275
PostQueryRequest const& request) override;
7376

77+
StatusOr<GetQueryResultsResponse> GetQueryResults(
78+
rest_internal::RestContext& rest_context,
79+
GetQueryResultsRequest const& request) override;
80+
7481
private:
7582
std::unique_ptr<rest_internal::RestClient> rest_stub_;
7683
};

‎google/cloud/bigquery/v2/minimal/internal/job_rest_stub_test.cc‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,11 @@ GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
3232

3333
namespace rest = ::google::cloud::rest_internal;
3434

35+
using ::google::cloud::bigquery_v2_minimal_testing::
36+
MakeFullGetQueryResultsRequest;
37+
using ::google::cloud::bigquery_v2_minimal_testing::MakeGetQueryResults;
38+
using ::google::cloud::bigquery_v2_minimal_testing::
39+
MakeGetQueryResultsResponsePayload;
3540
using ::google::cloud::bigquery_v2_minimal_testing::MakePartialJob;
3641
using ::google::cloud::bigquery_v2_minimal_testing::MakeQueryRequest;
3742
using ::google::cloud::bigquery_v2_minimal_testing::MakeQueryResponsePayload;
@@ -495,6 +500,78 @@ TEST(BigQueryJobStubTest, QueryRestResponseError) {
495500
EXPECT_THAT(status, StatusIs(InvalidArgumentError().code()));
496501
}
497502

503+
TEST(BigQueryJobStubTest, GetQueryResultsSuccess) {
504+
std::string response_payload = MakeGetQueryResultsResponsePayload();
505+
auto mock_response = std::make_unique<MockRestResponse>();
506+
507+
EXPECT_CALL(*mock_response, StatusCode)
508+
.WillRepeatedly(Return(HttpStatusCode::kOk));
509+
EXPECT_CALL(*mock_response, Headers)
510+
.WillRepeatedly(Return(std::multimap<std::string, std::string>()));
511+
EXPECT_CALL(std::move(*mock_response), ExtractPayload)
512+
.WillOnce(Return(ByMove(MakeMockHttpPayloadSuccess(response_payload))));
513+
514+
auto mock_rest_client = std::make_unique<MockRestClient>();
515+
EXPECT_CALL(*mock_rest_client, Get(_, An<rest::RestRequest const&>()))
516+
.WillOnce(Return(ByMove(
517+
std::unique_ptr<rest::RestResponse>(std::move(mock_response)))));
518+
519+
GetQueryResultsRequest request = MakeFullGetQueryResultsRequest();
520+
521+
rest_internal::RestContext context;
522+
DefaultBigQueryJobRestStub rest_stub(std::move(mock_rest_client));
523+
524+
auto expected = MakeGetQueryResults();
525+
526+
auto actual_result = rest_stub.GetQueryResults(context, std::move(request));
527+
ASSERT_STATUS_OK(actual_result);
528+
EXPECT_THAT(actual_result->http_response.http_status_code,
529+
Eq(HttpStatusCode::kOk));
530+
EXPECT_THAT(actual_result->http_response.payload, Eq(response_payload));
531+
bigquery_v2_minimal_testing::AssertEquals(expected,
532+
actual_result->get_query_results);
533+
}
534+
535+
TEST(BigQueryJobStubTest, GetQueryResultsRestClientError) {
536+
// Get() fails.
537+
auto mock_rest_client = std::make_unique<MockRestClient>();
538+
EXPECT_CALL(*mock_rest_client, Get(_, An<rest::RestRequest const&>()))
539+
.WillOnce(
540+
Return(rest::AsStatus(HttpStatusCode::kInternalServerError, "")));
541+
542+
rest_internal::RestContext context;
543+
DefaultBigQueryJobRestStub rest_stub(std::move(mock_rest_client));
544+
545+
GetQueryResultsRequest request = MakeFullGetQueryResultsRequest();
546+
547+
auto response = rest_stub.GetQueryResults(context, request);
548+
EXPECT_THAT(response, StatusIs(StatusCode::kUnavailable));
549+
}
550+
551+
TEST(BigQueryJobStubTest, GetQueryResultsRestResponseError) {
552+
// Invalid Rest response.
553+
auto mock_payload = std::make_unique<MockHttpPayload>();
554+
auto mock_response = std::make_unique<MockRestResponse>();
555+
EXPECT_CALL(*mock_response, StatusCode)
556+
.WillRepeatedly(Return(HttpStatusCode::kBadRequest));
557+
EXPECT_CALL(std::move(*mock_response), ExtractPayload)
558+
.WillOnce(Return(std::move(mock_payload)));
559+
560+
// Get() is successful.
561+
auto mock_rest_client = std::make_unique<MockRestClient>();
562+
EXPECT_CALL(*mock_rest_client, Get(_, An<rest::RestRequest const&>()))
563+
.WillOnce(Return(ByMove(
564+
std::unique_ptr<rest::RestResponse>(std::move(mock_response)))));
565+
566+
rest_internal::RestContext context;
567+
DefaultBigQueryJobRestStub rest_stub(std::move(mock_rest_client));
568+
569+
GetQueryResultsRequest request = MakeFullGetQueryResultsRequest();
570+
571+
auto response = rest_stub.GetQueryResults(context, std::move(request));
572+
EXPECT_THAT(response, StatusIs(StatusCode::kInvalidArgument));
573+
}
574+
498575
GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
499576
} // namespace bigquery_v2_minimal_internal
500577
} // namespace cloud

‎google/cloud/bigquery/v2/minimal/testing/mock_job_rest_stub.h‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,12 @@ class MockBigQueryJobRestStub
4949
(rest_internal::RestContext & rest_context,
5050
bigquery_v2_minimal_internal::PostQueryRequest const& request),
5151
(override));
52+
MOCK_METHOD(
53+
StatusOr<bigquery_v2_minimal_internal::GetQueryResultsResponse>,
54+
GetQueryResults,
55+
(rest_internal::RestContext & rest_context,
56+
bigquery_v2_minimal_internal::GetQueryResultsRequest const& request),
57+
(override));
5258
};
5359

5460
GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END

0 commit comments

Comments
 (0)