test(bigtable): add conformance test for ExecuteQuery

Author: diegomarquezp

google/cloud/bigtable/test_proxy/cbt_test_proxy.cc

@@ -25,6 +25,8 @@
 #include <google/protobuf/util/time_util.h>
 #include <chrono>
 
+#include "google/cloud/bigtable/client.h"
+
 namespace google {
 namespace cloud {
 namespace bigtable {
@@ -34,6 +36,8 @@ namespace {
 namespace v2 = ::google::bigtable::v2;
 namespace testpb = ::google::bigtable::testproxy;
 
+using ms = std::chrono::milliseconds;
+
 // Convert `Status` to `grpc::Status`, discarding any `details`.
 ::grpc::Status ToGrpcStatus(Status const& status) {
   return ::grpc::Status(static_cast<grpc::StatusCode>(status.code()),
@@ -360,6 +364,66 @@ grpc::Status CbtTestProxy::ReadModifyWriteRow(
   return grpc::Status();
 }
 
+
+grpc::Status CbtTestProxy::ExecuteQuery(grpc::ServerContext*,
+    const google::bigtable::testproxy::ExecuteQueryRequest* request,
+    google::bigtable::testproxy::ExecuteQueryResult* response) {
+  // Client options
+  auto retry_policy_option = DataLimitedErrorCountRetryPolicy(0).clone();
+  auto backoff_policy_option =
+      google::cloud::internal::ExponentialBackoffPolicy(ms(0), ms(0), 2.0)
+          .clone();
+  auto query_refresh_option =
+      bigtable::experimental::QueryPlanRefreshLimitedErrorCountRetryPolicy(0)
+          .clone();
+  auto opts =
+      Options{}
+  .set<DataRetryPolicyOption>(std::move(retry_policy_option))
+  .set<DataBackoffPolicyOption>(std::move(backoff_policy_option))
+  .set<bigtable::experimental::QueryPlanRefreshRetryPolicyOption>(
+      std::move(query_refresh_option));
+
+  // Retrieve connection
+  auto const& conn = GetConnection(request->client_id());
+  if (!conn.ok()) return ToGrpcStatus(std::move(conn).status());
+  auto client = bigtable::Client(*conn, opts);
+  auto request_proto = request->request();
+
+  // Call prepare query
+  auto instance = MakeInstanceResource(request_proto.instance_name());
+  bigtable::SqlStatement sql_statement{request_proto.query()};
+  auto prepared_query = client.PrepareQuery(*std::move(instance), sql_statement);
+  if (!prepared_query.ok()) return ToGrpcStatus(std::move(prepared_query).status());
+
+  // Bind parameters
+  std::unordered_map<std::string, Value> params;
+  for (auto const& param : request_proto.params()) {
+    auto value = bigtable_internal::FromProto(param.second.type(), param.second);
+    params.insert(std::make_pair(param.first, std::move(value)));
+  }
+  auto bound_query = prepared_query->BindParameters(params);
+  RowStream result = client.ExecuteQuery(std::move(bound_query), {});
+
+  Status status;
+  for (auto const& row : result) {
+    if (!row.ok()) {
+      status = row.status();
+      GCP_LOG(INFO) << "Error reading row: " << row.status();
+      continue;
+    }
+    google::bigtable::testproxy::SqlRow proxy_row;
+    for (auto const& value : row->values()) {
+      *proxy_row.add_values() = bigtable_internal::ToProto(value).second;
+    }
+    *response->add_rows() = std::move(proxy_row);
+  }
+
+  // TODO: we need to set response->mutable_metadata()
+
+  *response->mutable_status() = ToRpcStatus(status);
+  return grpc::Status();
+}
+
 StatusOr<std::shared_ptr<DataConnection>> CbtTestProxy::GetConnection(
     std::string const& client_id) {
   std::lock_guard<std::mutex> lk(mu_);

google/cloud/bigtable/test_proxy/cbt_test_proxy.h

@@ -91,13 +91,18 @@ class CbtTestProxy final
       ::google::bigtable::testproxy::ReadModifyWriteRowRequest const* request,
       ::google::bigtable::testproxy::RowResult* response) override;
 
- private:
+  grpc::Status ExecuteQuery(grpc::ServerContext* context,
+      const google::bigtable::testproxy::ExecuteQueryRequest* request,
+      google::bigtable::testproxy::ExecuteQueryResult* response) override;
+
+    private:
   StatusOr<std::shared_ptr<DataConnection>> GetConnection(
       std::string const& client_id);
 
   StatusOr<Table> GetTableFromRequest(std::string const& client_id,
                                       std::string const& table_name);
 
+private:
   std::unordered_map<std::string, std::shared_ptr<DataConnection>> connections_;
   std::mutex mu_;
 };

protos/google/cloud/bigtable/test_proxy/test_proxy.proto

@@ -1,4 +1,4 @@
-// Copyright 2023 Google LLC
+// 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.
@@ -12,22 +12,56 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// Copied from https://github.com/googleapis/cndb-client-testing-protos/blob/16a532ea8a56164b921e5f015f04d39e531f8d4d/google/bigtable/testproxy/test_proxy.proto
+
 syntax = "proto3";
 
 package google.bigtable.testproxy;
 
 import "google/api/client.proto";
 import "google/bigtable/v2/bigtable.proto";
 import "google/bigtable/v2/data.proto";
+import "google/protobuf/descriptor.proto";
 import "google/protobuf/duration.proto";
 import "google/rpc/status.proto";
 
-option go_package = "./testproxypb";
+option go_package = "cloud.google.com/go/bigtable/testproxy/testproxypb;testproxypb";
 option java_multiple_files = true;
 option java_package = "com.google.cloud.bigtable.testproxy";
 
+// A config flag that dictates how the optional features should be enabled
+// during the client creation. The optional features customize how the client
+// interacts with the server, and are defined in
+// https://github.com/googleapis/googleapis/blob/master/google/bigtable/v2/feature_flags.proto
+enum OptionalFeatureConfig {
+  OPTIONAL_FEATURE_CONFIG_DEFAULT = 0;
+
+  OPTIONAL_FEATURE_CONFIG_ENABLE_ALL = 1;
+}
+
 // Request to test proxy service to create a client object.
 message CreateClientRequest {
+  message SecurityOptions {
+    // Access token to use for client credentials. If empty, the client will not
+    // use any call credentials. Certain implementations may require `use_ssl`
+    // to be set when using this.
+    string access_token = 1;
+
+    // Whether to use SSL channel credentials when connecting to the data
+    // endpoint.
+    bool use_ssl = 2;
+
+    // If using SSL channel credentials, override the SSL endpoint to match the
+    // host that is specified in the backend's certificate. Also sets the
+    // client's authority header value.
+    string ssl_endpoint_override = 3;
+
+    // PEM encoding of the server root certificates. If not set, the default
+    // root certs will be used instead. The default can be overridden via the
+    // GRPC_DEFAULT_SSL_ROOTS_FILE_PATH env var.
+    string ssl_root_certs_pem = 4;
+  }
+
   // A unique ID associated with the client object to be created.
   string client_id = 1;
 
@@ -52,6 +86,21 @@ message CreateClientRequest {
   // the created client. Otherwise, the default timeout from the client library
   // will be used. Note that the override applies to all the methods.
   google.protobuf.Duration per_operation_timeout = 6;
+
+  // Optional config that dictates how the optional features should be enabled
+  // during the client creation. Please check the enum type's docstring above.
+  OptionalFeatureConfig optional_feature_config = 7;
+
+  // Options to allow connecting to backends with channel and/or call
+  // credentials. This is needed internally by Cloud Bigtable's own testing
+  // frameworks.It is not necessary to support these fields for client
+  // conformance testing.
+  //
+  // WARNING: this allows the proxy to connect to a real production
+  // CBT backend with the right options, however, the proxy itself is insecure
+  // so it is not recommended to use it with real credentials or outside testing
+  // contexts.
+  SecurityOptions security_options = 8;
 }
 
 // Response from test proxy service for CreateClientRequest.
@@ -203,6 +252,44 @@ message ReadModifyWriteRowRequest {
   google.bigtable.v2.ReadModifyWriteRowRequest request = 2;
 }
 
+// Request to test proxy service to execute a query.
+message ExecuteQueryRequest {
+  // The ID of the target client object.
+  string client_id = 1;
+
+  // The raw request to the Bigtable server.
+  google.bigtable.v2.ExecuteQueryRequest request = 2;
+
+  // A collection of proto descriptor files for deserializing PROTO/ENUM columns
+  // in the query result. The descriptor files must be provided in dependency
+  // order.
+  google.protobuf.FileDescriptorSet file_descriptor_set = 3;
+}
+
+// Response from test proxy service for ExecuteQueryRequest.
+message ExecuteQueryResult {
+  // The RPC status from the client binding.
+  google.rpc.Status status = 1;
+
+  // Name and type information for the query result.
+  ResultSetMetadata metadata = 4;
+
+  // Encoded version of the ResultSet. Should not contain type information.
+  repeated SqlRow rows = 3;
+}
+
+// Schema information for the query result.
+message ResultSetMetadata {
+  // Column metadata for each column inthe query result.
+  repeated google.bigtable.v2.ColumnMetadata columns = 1;
+}
+
+// Representation of a single row in the query result.
+message SqlRow {
+  // Columnar values returned by the query.
+  repeated google.bigtable.v2.Value values = 1;
+}
+
 // Note that all RPCs are unary, even when the equivalent client binding call
 // may be streaming. This is an intentional simplification.
 //
@@ -265,4 +352,7 @@ service CloudBigtableV2TestProxy {
 
   // Performs a read-modify-write operation with the client.
   rpc ReadModifyWriteRow(ReadModifyWriteRowRequest) returns (RowResult) {}
+
+  // Executes a BTQL query with the client.
+  rpc ExecuteQuery(ExecuteQueryRequest) returns (ExecuteQueryResult) {}
 }