diff --git a/ci/abi-dumps/google_cloud_cpp_spanner.expected.abi.dump.gz b/ci/abi-dumps/google_cloud_cpp_spanner.expected.abi.dump.gz index 9960ffb9ac98..e27b5c0a26fb 100644 Binary files a/ci/abi-dumps/google_cloud_cpp_spanner.expected.abi.dump.gz and b/ci/abi-dumps/google_cloud_cpp_spanner.expected.abi.dump.gz differ diff --git a/google/cloud/spanner/connection.h b/google/cloud/spanner/connection.h index 6c471ba075cf..29a70211cbb6 100644 --- a/google/cloud/spanner/connection.h +++ b/google/cloud/spanner/connection.h @@ -78,6 +78,7 @@ class Connection { std::vector columns; ReadOptions read_options; absl::optional partition_token; + bool partition_data_boost = false; // when partition_token }; /// Wrap the arguments to `PartitionRead()`. @@ -93,6 +94,7 @@ class Connection { SqlStatement statement; QueryOptions query_options; absl::optional partition_token; + bool partition_data_boost = false; // when partition_token }; /// Wrap the arguments to `ExecutePartitionedDml()`. diff --git a/google/cloud/spanner/internal/connection_impl.cc b/google/cloud/spanner/internal/connection_impl.cc index d44cb5f0131d..720d48486d81 100644 --- a/google/cloud/spanner/internal/connection_impl.cc +++ b/google/cloud/spanner/internal/connection_impl.cc @@ -426,6 +426,9 @@ spanner::RowStream ConnectionImpl::ReadImpl( request->set_limit(params.read_options.limit); if (params.partition_token) { request->set_partition_token(*std::move(params.partition_token)); + if (params.partition_data_boost) { + request->set_data_boost_enabled(true); + } } request->mutable_request_options()->set_priority( ProtoRequestPriority(params.read_options.request_priority)); @@ -552,12 +555,17 @@ StatusOr> ConnectionImpl::PartitionReadImpl( return status; } + // Note: This is not `ReadParams::data_boost`. While that is the ultimate + // destination for the `PartitionOptions::data_boost` value, first it is + // encapsulated in a `ReadPartition`. + bool const data_boost = partition_options.data_boost; + std::vector read_partitions; for (auto const& partition : response->partitions()) { read_partitions.push_back(MakeReadPartition( response->transaction().id(), ctx.route_to_leader, ctx.tag, session->session_name(), partition.partition_token(), params.table, - params.keys, params.columns, params.read_options)); + params.keys, params.columns, data_boost, params.read_options)); } return read_partitions; @@ -589,6 +597,9 @@ StatusOr ConnectionImpl::ExecuteSqlImpl( request.set_query_mode(query_mode); if (params.partition_token) { request.set_partition_token(*std::move(params.partition_token)); + if (params.partition_data_boost) { + request.set_data_boost_enabled(true); + } } if (params.query_options.optimizer_version()) { request.mutable_query_options()->set_optimizer_version( @@ -811,8 +822,7 @@ ConnectionImpl::PartitionQueryImpl( *request.mutable_params() = std::move(*sql_statement.mutable_params()); *request.mutable_param_types() = std::move(*sql_statement.mutable_param_types()); - *request.mutable_partition_options() = - ToProto(std::move(params.partition_options)); + *request.mutable_partition_options() = ToProto(params.partition_options); auto stub = session_pool_->GetStub(*session); for (;;) { @@ -850,10 +860,10 @@ ConnectionImpl::PartitionQueryImpl( std::vector query_partitions; for (auto const& partition : response->partitions()) { - query_partitions.push_back( - MakeQueryPartition(response->transaction().id(), ctx.route_to_leader, - ctx.tag, session->session_name(), - partition.partition_token(), params.statement)); + query_partitions.push_back(MakeQueryPartition( + response->transaction().id(), ctx.route_to_leader, ctx.tag, + session->session_name(), partition.partition_token(), + params.partition_options.data_boost, params.statement)); } return query_partitions; } diff --git a/google/cloud/spanner/internal/connection_impl_test.cc b/google/cloud/spanner/internal/connection_impl_test.cc index b1f9ed6853a8..a37153fd70aa 100644 --- a/google/cloud/spanner/internal/connection_impl_test.cc +++ b/google/cloud/spanner/internal/connection_impl_test.cc @@ -2529,6 +2529,37 @@ TEST(ConnectionImplTest, RollbackInvalidatedTransaction) { HasSubstr("constraint error"))); } +TEST(ConnectionImplTest, ReadPartition) { + auto mock = std::make_shared(); + auto db = spanner::Database("placeholder_project", "placeholder_instance", + "placeholder_database_id"); + + EXPECT_CALL(*mock, BatchCreateSessions(_, HasDatabase(db))) + .WillOnce(Return(MakeSessionsResponse({"test-session-name"}))); + EXPECT_CALL(*mock, BeginTransaction).Times(0); + EXPECT_CALL(*mock, StreamingRead) + .WillOnce([](grpc::ClientContext&, + google::spanner::v1::ReadRequest const& request) { + EXPECT_EQ("test-session-name", request.session()); + EXPECT_EQ("Table", request.table()); + EXPECT_EQ("DEADBEEF", request.partition_token()); + EXPECT_TRUE(request.data_boost_enabled()); + return MakeReader( + {R"pb(metadata: { transaction: { id: " ABCDEF00 " } })pb"}); + }); + + auto conn = MakeConnectionImpl(db, mock); + internal::OptionsSpan span(MakeLimitedTimeOptions()); + auto rows = conn->Read({spanner::MakeReadOnlyTransaction(), + "Table", + spanner::KeySet::All(), + {"Column"}, + {}, + "DEADBEEF", + true}); + EXPECT_TRUE(ContainsNoRows(rows)); +} + TEST(ConnectionImplTest, PartitionReadSuccess) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", @@ -2572,12 +2603,14 @@ TEST(ConnectionImplTest, PartitionReadSuccess) { read_options.index_name = "index"; read_options.limit = 21; read_options.request_priority = spanner::RequestPriority::kLow; + bool data_boost = true; StatusOr> result = conn->PartitionRead({{txn, "table", spanner::KeySet::All(), {"UserId", "UserName"}, - read_options}}); + read_options}, + {absl::nullopt, absl::nullopt, data_boost}}); ASSERT_STATUS_OK(result); EXPECT_THAT(txn, HasSessionAndTransaction("test-session-name", "CAFEDEAD", false, "")); @@ -2585,10 +2618,12 @@ TEST(ConnectionImplTest, PartitionReadSuccess) { std::vector expected_read_partitions = { spanner_internal::MakeReadPartition( "CAFEDEAD", false, "", "test-session-name", "BADDECAF", "table", - spanner::KeySet::All(), {"UserId", "UserName"}, read_options), + spanner::KeySet::All(), {"UserId", "UserName"}, data_boost, + read_options), spanner_internal::MakeReadPartition( "CAFEDEAD", false, "", "test-session-name", "DEADBEEF", "table", - spanner::KeySet::All(), {"UserId", "UserName"}, read_options)}; + spanner::KeySet::All(), {"UserId", "UserName"}, data_boost, + read_options)}; EXPECT_THAT(*result, UnorderedPointwise(Eq(), expected_read_partitions)); } @@ -2649,6 +2684,35 @@ TEST(ConnectionImplTest, PartitionReadTooManyTransientFailures) { StatusIs(StatusCode::kUnavailable, HasSubstr("try-again"))); } +TEST(ConnectionImplTest, QueryPartition) { + auto mock = std::make_shared(); + auto db = spanner::Database("placeholder_project", "placeholder_instance", + "placeholder_database_id"); + + EXPECT_CALL(*mock, BatchCreateSessions(_, HasDatabase(db))) + .WillOnce(Return(MakeSessionsResponse({"test-session-name"}))); + EXPECT_CALL(*mock, BeginTransaction).Times(0); + EXPECT_CALL(*mock, ExecuteStreamingSql) + .WillOnce([](grpc::ClientContext&, + google::spanner::v1::ExecuteSqlRequest const& request) { + EXPECT_EQ("test-session-name", request.session()); + EXPECT_EQ("SELECT * FROM Table", request.sql()); + EXPECT_EQ("DEADBEEF", request.partition_token()); + EXPECT_TRUE(request.data_boost_enabled()); + return MakeReader( + {R"pb(metadata: { transaction: { id: " ABCDEF00 " } })pb"}); + }); + + auto conn = MakeConnectionImpl(db, mock); + internal::OptionsSpan span(MakeLimitedTimeOptions()); + auto rows = conn->ExecuteQuery({spanner::MakeReadOnlyTransaction(), + spanner::SqlStatement("SELECT * FROM Table"), + {}, + "DEADBEEF", + true}); + EXPECT_TRUE(ContainsNoRows(rows)); +} + TEST(ConnectionImplTest, PartitionQuerySuccess) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", @@ -2683,18 +2747,20 @@ TEST(ConnectionImplTest, PartitionQuerySuccess) { auto conn = MakeConnectionImpl(db, mock); internal::OptionsSpan span(MakeLimitedTimeOptions()); spanner::SqlStatement sql_statement("SELECT * FROM Table"); + bool data_boost = true; StatusOr> result = conn->PartitionQuery( {MakeReadOnlyTransaction(spanner::Transaction::ReadOnlyOptions()), - sql_statement, spanner::PartitionOptions()}); + sql_statement, + {absl::nullopt, absl::nullopt, data_boost}}); ASSERT_STATUS_OK(result); std::vector expected_query_partitions = { spanner_internal::MakeQueryPartition("CAFEDEAD", false, "", "test-session-name", "BADDECAF", - sql_statement), + data_boost, sql_statement), spanner_internal::MakeQueryPartition("CAFEDEAD", false, "", "test-session-name", "DEADBEEF", - sql_statement)}; + data_boost, sql_statement)}; EXPECT_THAT(*result, UnorderedPointwise(Eq(), expected_query_partitions)); } diff --git a/google/cloud/spanner/options.h b/google/cloud/spanner/options.h index b5df4f9ec7a1..dc9244f84da9 100644 --- a/google/cloud/spanner/options.h +++ b/google/cloud/spanner/options.h @@ -299,6 +299,20 @@ struct PartitionsMaximumOption { using Type = std::int64_t; }; +/** + * Option for `google::cloud::Options` to use "data boost" in the + * partitions returned from `Client::PartitionRead()` or `PartitionQuery()`. + * + * If true, the requests from the subsequent partitioned `Client::Read()` + * and `Client::ExecuteQuery()` requests will be executed via Spanner + * independent compute resources. + * + * @ingroup spanner-options + */ +struct PartitionDataBoostOption { + using Type = bool; +}; + /** * Option for `google::cloud::Options` to set a per-transaction tag. * diff --git a/google/cloud/spanner/partition_options.cc b/google/cloud/spanner/partition_options.cc index 1cca8794ea50..b96c3f0ad648 100644 --- a/google/cloud/spanner/partition_options.cc +++ b/google/cloud/spanner/partition_options.cc @@ -28,6 +28,9 @@ Options ToOptions(PartitionOptions const& po) { if (po.max_partitions) { opts.set(*po.max_partitions); } + if (po.data_boost) { + opts.set(true); + } return opts; } @@ -39,6 +42,9 @@ PartitionOptions ToPartitionOptions(Options const& opts) { if (opts.has()) { po.max_partitions = opts.get(); } + if (opts.has()) { + po.data_boost = opts.get(); + } return po; } @@ -55,6 +61,10 @@ google::spanner::v1::PartitionOptions ToProto( if (po.partition_size_bytes) { proto.set_partition_size_bytes(*po.partition_size_bytes); } + // Note: PartitionOptions.data_boost sets the data_boost_enabled + // field of the ReadRequest and ExecuteSqlRequest messages in + // ConnectionImpl::ReadImpl() and ConnectionImpl::ExecuteSqlImpl() + // respectively. return proto; } diff --git a/google/cloud/spanner/partition_options.h b/google/cloud/spanner/partition_options.h index 4f658d2a0d40..a7bb6574343d 100644 --- a/google/cloud/spanner/partition_options.h +++ b/google/cloud/spanner/partition_options.h @@ -54,11 +54,20 @@ struct PartitionOptions { * returned may be smaller or larger than this maximum count request. */ absl::optional max_partitions; + + /** + * Use "data boost" in the returned partitions. + * + * If true, the requests from the subsequent partitioned `Client::Read()` + * and `Client::ExecuteQuery()` requests will be executed via Spanner + * independent compute resources. + */ + bool data_boost = false; }; inline bool operator==(PartitionOptions const& a, PartitionOptions const& b) { return a.partition_size_bytes == b.partition_size_bytes && - a.max_partitions == b.max_partitions; + a.max_partitions == b.max_partitions && a.data_boost == b.data_boost; } inline bool operator!=(PartitionOptions const& a, PartitionOptions const& b) { diff --git a/google/cloud/spanner/partition_options_test.cc b/google/cloud/spanner/partition_options_test.cc index 32c7b6946814..35781a29b780 100644 --- a/google/cloud/spanner/partition_options_test.cc +++ b/google/cloud/spanner/partition_options_test.cc @@ -41,17 +41,23 @@ TEST(PartitionOptionsTest, Regular) { } TEST(PartitionOptionsTest, Proto) { - PartitionOptions po{1, 2}; + PartitionOptions po{1, 2, true}; auto proto = spanner_internal::ToProto(po); EXPECT_EQ(*po.partition_size_bytes, proto.partition_size_bytes()); EXPECT_EQ(*po.max_partitions, proto.max_partitions()); } TEST(PartitionOptionsTest, OptionsRoundTrip) { - for (auto const& po : - {PartitionOptions{absl::nullopt, absl::nullopt}, - PartitionOptions{42, absl::nullopt}, - PartitionOptions{absl::nullopt, 42}, PartitionOptions{32, 64}}) { + for (auto const& po : { + PartitionOptions{absl::nullopt, absl::nullopt, false}, + PartitionOptions{42, absl::nullopt, false}, + PartitionOptions{absl::nullopt, 42, false}, + PartitionOptions{32, 64, false}, + PartitionOptions{absl::nullopt, absl::nullopt, true}, + PartitionOptions{42, absl::nullopt, true}, + PartitionOptions{absl::nullopt, 42, true}, + PartitionOptions{32, 64, true}, + }) { EXPECT_EQ(po, ToPartitionOptions(ToOptions(po))); } } diff --git a/google/cloud/spanner/query_partition.cc b/google/cloud/spanner/query_partition.cc index 676ea48e1945..84c280e6cc98 100644 --- a/google/cloud/spanner/query_partition.cc +++ b/google/cloud/spanner/query_partition.cc @@ -27,13 +27,14 @@ constexpr int kRouteToLeaderFieldNumber = 511037314; QueryPartition::QueryPartition(std::string transaction_id, bool route_to_leader, std::string transaction_tag, std::string session_id, - std::string partition_token, + std::string partition_token, bool data_boost, SqlStatement sql_statement) : transaction_id_(std::move(transaction_id)), route_to_leader_(route_to_leader), transaction_tag_(std::move(transaction_tag)), session_id_(std::move(session_id)), partition_token_(std::move(partition_token)), + data_boost_(data_boost), sql_statement_(std::move(sql_statement)) {} bool operator==(QueryPartition const& a, QueryPartition const& b) { @@ -42,7 +43,7 @@ bool operator==(QueryPartition const& a, QueryPartition const& b) { a.transaction_tag_ == b.transaction_tag_ && a.session_id_ == b.session_id_ && a.partition_token_ == b.partition_token_ && - a.sql_statement_ == b.sql_statement_; + a.data_boost_ == b.data_boost_ && a.sql_statement_ == b.sql_statement_; } StatusOr SerializeQueryPartition( @@ -59,6 +60,9 @@ StatusOr SerializeQueryPartition( (*proto.mutable_params()->mutable_fields())[param_name] = type_value.second; (*proto.mutable_param_types())[param_name] = type_value.first; } + if (query_partition.data_boost()) { + proto.set_data_boost_enabled(true); + } // QueryOptions are not serialized, but are instead applied on the remote // side during the Client::ExecuteQuery(QueryPartition, QueryOptions) call. @@ -119,6 +123,7 @@ StatusOr DeserializeQueryPartition( QueryPartition query_partition(proto.transaction().id(), route_to_leader, proto.request_options().transaction_tag(), proto.session(), proto.partition_token(), + proto.data_boost_enabled(), SqlStatement(proto.sql(), sql_parameters)); return query_partition; } diff --git a/google/cloud/spanner/query_partition.h b/google/cloud/spanner/query_partition.h index e58404c20c47..fdbd280237c5 100644 --- a/google/cloud/spanner/query_partition.h +++ b/google/cloud/spanner/query_partition.h @@ -83,6 +83,8 @@ StatusOr DeserializeQueryPartition( * Instances of `QueryPartition` are created by `Client::PartitionQuery`. Once * created, `QueryPartition` objects can be serialized, transmitted to separate * processes, and used to read data in parallel using `Client::ExecuteQuery`. + * If `data_boost` is set, those requests will be executed via Spanner + * independent compute resources. */ class QueryPartition { public: @@ -123,7 +125,8 @@ class QueryPartition { QueryPartition(std::string transaction_id, bool route_to_leader, std::string transaction_tag, std::string session_id, - std::string partition_token, SqlStatement sql_statement); + std::string partition_token, bool data_boost, + SqlStatement sql_statement); // Accessor methods for use by friends. std::string const& transaction_id() const { return transaction_id_; } @@ -131,12 +134,14 @@ class QueryPartition { std::string const& transaction_tag() const { return transaction_tag_; } std::string const& session_id() const { return session_id_; } std::string const& partition_token() const { return partition_token_; } + bool data_boost() const { return data_boost_; } std::string transaction_id_; bool route_to_leader_; std::string transaction_tag_; std::string session_id_; std::string partition_token_; + bool data_boost_; SqlStatement sql_statement_; }; @@ -151,11 +156,11 @@ struct QueryPartitionInternals { static spanner::QueryPartition MakeQueryPartition( std::string const& transaction_id, bool route_to_leader, std::string const& transaction_tag, std::string const& session_id, - std::string const& partition_token, + std::string const& partition_token, bool data_boost, spanner::SqlStatement const& sql_statement) { return spanner::QueryPartition(transaction_id, route_to_leader, transaction_tag, session_id, partition_token, - sql_statement); + data_boost, sql_statement); } static spanner::Connection::SqlParams MakeSqlParams( @@ -166,18 +171,18 @@ struct QueryPartitionInternals { query_partition.route_to_leader(), query_partition.transaction_tag()), query_partition.sql_statement(), query_options, - query_partition.partition_token()}; + query_partition.partition_token(), query_partition.data_boost()}; } }; inline spanner::QueryPartition MakeQueryPartition( std::string const& transaction_id, bool route_to_leader, std::string const& transaction_tag, std::string const& session_id, - std::string const& partition_token, + std::string const& partition_token, bool data_boost, spanner::SqlStatement const& sql_statement) { return QueryPartitionInternals::MakeQueryPartition( transaction_id, route_to_leader, transaction_tag, session_id, - partition_token, sql_statement); + partition_token, data_boost, sql_statement); } inline spanner::Connection::SqlParams MakeSqlParams( diff --git a/google/cloud/spanner/query_partition_test.cc b/google/cloud/spanner/query_partition_test.cc index 8624124f1128..662db0ace61e 100644 --- a/google/cloud/spanner/query_partition_test.cc +++ b/google/cloud/spanner/query_partition_test.cc @@ -31,6 +31,7 @@ class QueryPartitionTester { spanner::SqlStatement const& Statement() const { return partition_.sql_statement(); } + bool DataBoost() const { return partition_.data_boost(); } std::string const& PartitionToken() const { return partition_.partition_token(); } @@ -62,6 +63,7 @@ using ::testing::Not; TEST(QueryPartitionTest, MakeQueryPartition) { std::string stmt("SELECT * FROM foo WHERE name = @name"); SqlStatement::ParamType params = {{"name", Value("Bob")}}; + bool data_boost = true; std::string partition_token("token"); std::string session_id("session"); std::string transaction_id("txn-id"); @@ -70,7 +72,7 @@ TEST(QueryPartitionTest, MakeQueryPartition) { QueryPartitionTester actual_partition(spanner_internal::MakeQueryPartition( transaction_id, route_to_leader, transaction_tag, session_id, - partition_token, SqlStatement(stmt, params))); + partition_token, data_boost, SqlStatement(stmt, params))); EXPECT_EQ(stmt, actual_partition.Statement().sql()); EXPECT_EQ(params, actual_partition.Statement().params()); EXPECT_EQ(partition_token, actual_partition.PartitionToken()); @@ -78,11 +80,13 @@ TEST(QueryPartitionTest, MakeQueryPartition) { EXPECT_EQ(route_to_leader, actual_partition.RouteToLeader()); EXPECT_EQ(transaction_tag, actual_partition.TransactionTag()); EXPECT_EQ(session_id, actual_partition.SessionId()); + EXPECT_EQ(data_boost, actual_partition.DataBoost()); } TEST(QueryPartitionTest, RegularSemantics) { std::string stmt("SELECT * FROM foo WHERE name = @name"); SqlStatement::ParamType params = {{"name", Value("Bob")}}; + bool data_boost = true; std::string partition_token("token"); std::string session_id("session"); std::string transaction_id("txn-id"); @@ -91,7 +95,7 @@ TEST(QueryPartitionTest, RegularSemantics) { QueryPartition query_partition(spanner_internal::MakeQueryPartition( transaction_id, route_to_leader, transaction_tag, session_id, - partition_token, SqlStatement(stmt, params))); + partition_token, data_boost, SqlStatement(stmt, params))); EXPECT_NE(query_partition, QueryPartition()); @@ -108,7 +112,7 @@ TEST(QueryPartitionTest, RegularSemantics) { TEST(QueryPartitionTest, SerializeDeserialize) { QueryPartitionTester expected_partition(spanner_internal::MakeQueryPartition( - "txn-id", true, "tag", "session", "token", + "txn-id", true, "tag", "session", "token", false, SqlStatement("SELECT * FROM foo WHERE name = @name", {{"name", Value("Bob")}}))); StatusOr partition = DeserializeQueryPartition( @@ -129,6 +133,7 @@ TEST(QueryPartitionTest, SerializeDeserialize) { actual_partition.Statement().sql()); EXPECT_EQ(expected_partition.Statement().params(), actual_partition.Statement().params()); + EXPECT_EQ(expected_partition.DataBoost(), actual_partition.DataBoost()); } TEST(QueryPartitionTest, FailedDeserialize) { @@ -140,7 +145,7 @@ TEST(QueryPartitionTest, FailedDeserialize) { TEST(QueryPartitionTest, MakeSqlParams) { QueryPartitionTester expected_partition(spanner_internal::MakeQueryPartition( - "txn-id", true, "tag", "session", "token", + "txn-id", true, "tag", "session", "token", true, SqlStatement("SELECT * FROM foo WHERE name = @name", {{"name", Value("Bob")}}))); @@ -152,6 +157,7 @@ TEST(QueryPartitionTest, MakeSqlParams) { SqlStatement("SELECT * FROM foo WHERE name = @name", {{"name", Value("Bob")}})); EXPECT_EQ(*params.partition_token, "token"); + EXPECT_TRUE(params.partition_data_boost); EXPECT_THAT(params.transaction, HasSessionAndTransaction("session", "txn-id", true, "tag")); EXPECT_EQ(*params.query_options.request_tag(), "request_tag"); diff --git a/google/cloud/spanner/read_partition.cc b/google/cloud/spanner/read_partition.cc index f05b6c46be68..a095502bf41e 100644 --- a/google/cloud/spanner/read_partition.cc +++ b/google/cloud/spanner/read_partition.cc @@ -31,6 +31,7 @@ ReadPartition::ReadPartition(std::string transaction_id, bool route_to_leader, std::string table_name, google::cloud::spanner::KeySet key_set, std::vector column_names, + bool data_boost, google::cloud::spanner::ReadOptions read_options) { proto_.set_session(std::move(session_id)); proto_.mutable_transaction()->set_id(std::move(transaction_id)); @@ -42,6 +43,9 @@ ReadPartition::ReadPartition(std::string transaction_id, bool route_to_leader, *proto_.mutable_key_set() = spanner_internal::ToProto(std::move(key_set)); proto_.set_limit(read_options.limit); proto_.set_partition_token(std::move(partition_token)); + if (data_boost) { + proto_.set_data_boost_enabled(true); + } if (read_options.request_priority) { auto* request_options = proto_.mutable_request_options(); switch (*read_options.request_priority) { diff --git a/google/cloud/spanner/read_partition.h b/google/cloud/spanner/read_partition.h index 5942330ab9dc..aa65388e1a9b 100644 --- a/google/cloud/spanner/read_partition.h +++ b/google/cloud/spanner/read_partition.h @@ -81,9 +81,11 @@ StatusOr DeserializeReadPartition( * The `ReadPartition` class is a regular type that represents a single * slice of a parallel Read operation. * - * Instances of `ReadPartition` are created by `Client::PartitionRead`. Once - * created, `ReadPartition` objects can be serialized, transmitted to separate - * processes, and used to read data in parallel using `Client::Read`. + * Instances of `ReadPartition` are created by `Client::PartitionRead`. + * Once created, `ReadPartition` objects can be serialized, transmitted to + * separate processes, and used to read data in parallel using `Client::Read`. + * If `data_boost` is set, those requests will be executed via Spanner + * independent compute resources. */ class ReadPartition { public: @@ -132,7 +134,7 @@ class ReadPartition { std::string transaction_tag, std::string session_id, std::string partition_token, std::string table_name, google::cloud::spanner::KeySet key_set, - std::vector column_names, + std::vector column_names, bool data_boost, google::cloud::spanner::ReadOptions read_options); // Accessor methods for use by friends. @@ -144,6 +146,7 @@ class ReadPartition { std::string SessionId() const { return proto_.session(); } std::string PartitionToken() const { return proto_.partition_token(); } google::spanner::v1::KeySet KeySet() const { return proto_.key_set(); } + bool DataBoost() const { return proto_.data_boost_enabled(); } google::spanner::v1::ReadRequest proto_; }; @@ -161,12 +164,12 @@ struct ReadPartitionInternals { std::string transaction_tag, std::string session_id, std::string partition_token, std::string table_name, spanner::KeySet key_set, std::vector column_names, - spanner::ReadOptions read_options) { + bool data_boost, spanner::ReadOptions read_options) { return spanner::ReadPartition( std::move(transaction_id), route_to_leader, std::move(transaction_tag), std::move(session_id), std::move(partition_token), std::move(table_name), std::move(key_set), std::move(column_names), - std::move(read_options)); + data_boost, std::move(read_options)); } static spanner::Connection::ReadParams MakeReadParams( @@ -179,7 +182,8 @@ struct ReadPartitionInternals { FromProto(read_partition.KeySet()), read_partition.ColumnNames(), read_partition.ReadOptions(), - read_partition.PartitionToken()}; + read_partition.PartitionToken(), + read_partition.DataBoost()}; } }; @@ -188,11 +192,12 @@ inline spanner::ReadPartition MakeReadPartition( std::string transaction_tag, std::string session_id, std::string partition_token, std::string table_name, spanner::KeySet key_set, std::vector column_names, - spanner::ReadOptions read_options) { + bool data_boost, spanner::ReadOptions read_options) { return ReadPartitionInternals::MakeReadPartition( std::move(transaction_id), route_to_leader, std::move(transaction_tag), std::move(session_id), std::move(partition_token), std::move(table_name), - std::move(key_set), std::move(column_names), std::move(read_options)); + std::move(key_set), std::move(column_names), data_boost, + std::move(read_options)); } inline spanner::Connection::ReadParams MakeReadParams( diff --git a/google/cloud/spanner/read_partition_test.cc b/google/cloud/spanner/read_partition_test.cc index faa37ebbb27a..a90a00208599 100644 --- a/google/cloud/spanner/read_partition_test.cc +++ b/google/cloud/spanner/read_partition_test.cc @@ -28,7 +28,7 @@ class ReadPartitionTester { ReadPartitionTester() = default; explicit ReadPartitionTester(spanner::ReadPartition partition) : partition_(std::move(partition)) {} - spanner::ReadPartition Partition() const { return partition_; } + spanner::ReadPartition const& Partition() const { return partition_; } std::string PartitionToken() const { return partition_.PartitionToken(); } std::string SessionId() const { return partition_.SessionId(); } std::string TransactionId() const { return partition_.TransactionId(); } @@ -39,6 +39,7 @@ class ReadPartitionTester { std::vector ColumnNames() const { return partition_.ColumnNames(); } + bool DataBoost() const { return partition_.DataBoost(); } google::cloud::spanner::ReadOptions ReadOptions() const { return partition_.ReadOptions(); } @@ -67,6 +68,7 @@ TEST(ReadPartitionTest, MakeReadPartition) { std::string transaction_tag("tag"); std::string table_name("Students"); std::vector column_names = {"LastName", "FirstName"}; + bool data_boost = true; ReadOptions read_options; read_options.index_name = "secondary"; read_options.limit = 42; @@ -75,7 +77,8 @@ TEST(ReadPartitionTest, MakeReadPartition) { ReadPartitionTester actual_partition(spanner_internal::MakeReadPartition( transaction_id, route_to_leader, transaction_tag, session_id, - partition_token, table_name, KeySet::All(), column_names, read_options)); + partition_token, table_name, KeySet::All(), column_names, data_boost, + read_options)); EXPECT_EQ(partition_token, actual_partition.PartitionToken()); EXPECT_EQ(transaction_id, actual_partition.TransactionId()); @@ -86,6 +89,7 @@ TEST(ReadPartitionTest, MakeReadPartition) { EXPECT_THAT(actual_partition.KeySet(), IsProtoEqual(spanner_internal::ToProto(KeySet::All()))); EXPECT_EQ(column_names, actual_partition.ColumnNames()); + EXPECT_TRUE(actual_partition.DataBoost()); EXPECT_EQ(read_options, actual_partition.ReadOptions()); } @@ -97,6 +101,7 @@ TEST(ReadPartitionTest, RegularSemantics) { std::string transaction_tag("tag"); std::string table_name("Students"); std::vector column_names = {"LastName", "FirstName"}; + bool data_boost = true; ReadOptions read_options; read_options.index_name = "secondary"; read_options.limit = 42; @@ -104,7 +109,8 @@ TEST(ReadPartitionTest, RegularSemantics) { ReadPartition read_partition(spanner_internal::MakeReadPartition( transaction_id, route_to_leader, transaction_tag, session_id, - partition_token, table_name, KeySet::All(), column_names, read_options)); + partition_token, table_name, KeySet::All(), column_names, data_boost, + read_options)); EXPECT_NE(read_partition, ReadPartition()); @@ -126,6 +132,7 @@ TEST(ReadPartitionTest, RouteToLeaderInequality) { std::string transaction_tag("tag"); std::string table_name("Students"); std::vector column_names = {"LastName", "FirstName"}; + bool data_boost = true; ReadOptions read_options; read_options.index_name = "secondary"; read_options.limit = 42; @@ -135,10 +142,12 @@ TEST(ReadPartitionTest, RouteToLeaderInequality) { // considers route_to_leader, which is stored in an "unknown field". auto read_partition_route_to_leader = spanner_internal::MakeReadPartition( transaction_id, /*route_to_leader=*/true, transaction_tag, session_id, - partition_token, table_name, KeySet::All(), column_names, read_options); + partition_token, table_name, KeySet::All(), column_names, data_boost, + read_options); auto read_partition_not_route_to_leader = spanner_internal::MakeReadPartition( transaction_id, /*route_to_leader=*/false, transaction_tag, session_id, - partition_token, table_name, KeySet::All(), column_names, read_options); + partition_token, table_name, KeySet::All(), column_names, data_boost, + read_options); EXPECT_NE(read_partition_route_to_leader, read_partition_not_route_to_leader); } @@ -149,7 +158,7 @@ TEST(ReadPartitionTest, SerializeDeserialize) { read_options.request_priority = RequestPriority::kMedium; ReadPartitionTester expected_partition(spanner_internal::MakeReadPartition( "txn-id", true, "tag", "session", "token", "Students", KeySet::All(), - std::vector{"LastName", "FirstName"}, read_options)); + std::vector{"LastName", "FirstName"}, true, read_options)); StatusOr partition = DeserializeReadPartition( *(SerializeReadPartition(expected_partition.Partition()))); @@ -169,6 +178,7 @@ TEST(ReadPartitionTest, SerializeDeserialize) { EXPECT_THAT(actual_partition.KeySet(), IsProtoEqual(expected_partition.KeySet())); EXPECT_EQ(expected_partition.ColumnNames(), actual_partition.ColumnNames()); + EXPECT_EQ(expected_partition.DataBoost(), actual_partition.DataBoost()); EXPECT_EQ(expected_partition.ReadOptions(), actual_partition.ReadOptions()); } @@ -187,13 +197,14 @@ TEST(ReadPartitionTest, MakeReadParams) { read_options.request_priority = RequestPriority::kMedium; ReadPartitionTester expected_partition(spanner_internal::MakeReadPartition( "txn-id", true, "tag", "session", "token", "Students", KeySet::All(), - columns, read_options)); + columns, false, read_options)); Connection::ReadParams params = spanner_internal::MakeReadParams(expected_partition.Partition()); EXPECT_EQ(*params.partition_token, "token"); EXPECT_EQ(params.read_options, read_options); + EXPECT_FALSE(params.partition_data_boost); EXPECT_EQ(params.columns, columns); EXPECT_EQ(params.keys, KeySet::All()); EXPECT_EQ(params.table, "Students"); diff --git a/google/cloud/spanner/samples/samples.cc b/google/cloud/spanner/samples/samples.cc index 9e1217e1a974..5d85aa59fc24 100644 --- a/google/cloud/spanner/samples/samples.cc +++ b/google/cloud/spanner/samples/samples.cc @@ -2489,8 +2489,9 @@ void UsePartitionQuery(google::cloud::spanner::Client client) { "SELECT SingerId, FirstName, LastName FROM Singers"); using RowType = std::tuple; - auto partitions = - client.PartitionQuery(std::move(txn), std::move(select), {}); + auto partitions = client.PartitionQuery( + std::move(txn), std::move(select), + google::cloud::Options{}.set(true)); if (!partitions) throw std::move(partitions).status(); // You would probably choose to execute these partitioned queries in @@ -3682,8 +3683,11 @@ void PartitionRead(google::cloud::spanner::Client client) { //! [partition-read] spanner::Transaction ro_transaction = spanner::MakeReadOnlyTransaction(); google::cloud::StatusOr> partitions = - client.PartitionRead(ro_transaction, "Singers", key_set, - {"SingerId", "FirstName", "LastName"}); + client.PartitionRead( + ro_transaction, "Singers", key_set, + {"SingerId", "FirstName", "LastName"}, + google::cloud::Options{}.set( + true)); if (!partitions) throw std::move(partitions).status(); for (auto& partition : *partitions) { remote_connection.SendPartitionToRemoteMachine(partition); @@ -3730,7 +3734,9 @@ void PartitionQuery(google::cloud::spanner::Client client) { client.PartitionQuery( spanner::MakeReadOnlyTransaction(), spanner::SqlStatement( - "SELECT SingerId, FirstName, LastName FROM Singers")); + "SELECT SingerId, FirstName, LastName FROM Singers"), + google::cloud::Options{}.set( + true)); if (!partitions) throw std::move(partitions).status(); for (auto& partition : *partitions) { remote_connection.SendPartitionToRemoteMachine(partition);