feat(common): QuotaUserOption for gRPC-based libs

googleapis · Apr 5, 2024 · 6d428bb · 6d428bb
1 parent 4db2434
commit 6d428bb
Show file tree

Hide file tree

Showing 4 changed files with 72 additions and 27 deletions.
diff --git a/google/cloud/common_options.h b/google/cloud/common_options.h
@@ -91,6 +91,33 @@ struct UserProjectOption {
   using Type = std::string;
 };
 
+/**
+ * Configure the QuotaUser system parameter.
+ *
+ * A pseudo user identifier for charging per-user quotas. If not specified, the
+ * authenticated principal is used. If there is no authenticated principal, the
+ * client IP address will be used. When specified, a valid API key with service
+ * restrictions must be used to identify the quota project. Otherwise, this
+ * parameter is ignored.
+ *
+ * @ingroup rest-options
+ */
+struct QuotaUserOption {
+  using Type = std::string;
+};
+
+/**
+ * Configure the UserIp query parameter.
+ *
+ * This can be used to separate quota usage by source IP address.
+ *
+ * @deprecated prefer using `google::cloud::QuotaUser`.
+ * @ingroup rest-options
+ */
+struct UserIpOption {
+  using Type = std::string;
+};
+
 /**
  * Configure the "authority" attribute.
  *

diff --git a/google/cloud/grpc_options.cc b/google/cloud/grpc_options.cc
@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include "google/cloud/grpc_options.h"
+#include "google/cloud/common_options.h"
 #include "google/cloud/internal/absl_str_cat_quiet.h"
 #include "google/cloud/internal/absl_str_join_quiet.h"
 #include "google/cloud/internal/background_threads_impl.h"
@@ -31,6 +32,12 @@ void ConfigureContext(grpc::ClientContext& context, Options const& opts) {
     context.set_compression_algorithm(
         opts.get<GrpcCompressionAlgorithmOption>());
   }
+  if (opts.has<UserIpOption>() && !opts.has<QuotaUserOption>()) {
+    context.AddMetadata("x-goog-user-ip", opts.get<UserIpOption>());
+  }
+  if (opts.has<QuotaUserOption>()) {
+    context.AddMetadata("x-goog-quota-user", opts.get<QuotaUserOption>());
+  }
 }
 
 void ConfigurePollContext(grpc::ClientContext& context, Options const& opts) {

diff --git a/google/cloud/grpc_options_test.cc b/google/cloud/grpc_options_test.cc
@@ -16,6 +16,7 @@
 #include "google/cloud/common_options.h"
 #include "google/cloud/internal/background_threads_impl.h"
 #include "google/cloud/testing_util/scoped_log.h"
+#include "google/cloud/testing_util/validate_metadata.h"
 #include <gmock/gmock.h>
 #include <string>
 #include <utility>
@@ -26,11 +27,14 @@ GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
 
 namespace {
 
+using ::google::cloud::testing_util::ValidateMetadataFixture;
 using ::testing::Contains;
 using ::testing::ContainsRegex;
 using ::testing::IsEmpty;
 using ::testing::IsNull;
 using ::testing::NotNull;
+using ::testing::Pair;
+using ::testing::UnorderedElementsAre;
 using ms = std::chrono::milliseconds;
 using ThreadPool = internal::AutomaticallyCreatedBackgroundThreads;
 
@@ -322,6 +326,39 @@ TEST(GrpcClientContext, Configure) {
   EXPECT_EQ(GRPC_COMPRESS_DEFLATE, context.compression_algorithm());
 }
 
+TEST(GrpcClientContext, ConfigureUserIp) {
+  auto opts = Options{}.set<UserIpOption>("1234");
+  grpc::ClientContext context;
+  internal::ConfigureContext(context, std::move(opts));
+
+  ValidateMetadataFixture fixture;
+  auto const metadata = fixture.GetMetadata(context);
+  EXPECT_THAT(metadata, UnorderedElementsAre(Pair("x-goog-user-ip", "1234")));
+}
+
+TEST(GrpcClientContext, ConfigureQuotaUser) {
+  auto opts = Options{}.set<QuotaUserOption>("my-user");
+  grpc::ClientContext context;
+  internal::ConfigureContext(context, std::move(opts));
+
+  ValidateMetadataFixture fixture;
+  auto const metadata = fixture.GetMetadata(context);
+  EXPECT_THAT(metadata,
+              UnorderedElementsAre(Pair("x-goog-quota-user", "my-user")));
+}
+
+TEST(GrpcClientContext, QuotaUserOverridesUserIp) {
+  auto opts =
+      Options{}.set<QuotaUserOption>("my-user").set<UserIpOption>("1234");
+  grpc::ClientContext context;
+  internal::ConfigureContext(context, std::move(opts));
+
+  ValidateMetadataFixture fixture;
+  auto const metadata = fixture.GetMetadata(context);
+  EXPECT_THAT(metadata,
+              UnorderedElementsAre(Pair("x-goog-quota-user", "my-user")));
+}
+
 TEST(GrpcClientContext, ConfigurePoll) {
   auto setup_poll = [](grpc::ClientContext& context) {
     // This might not be the most useful setting, but it is the most easily

diff --git a/google/cloud/rest_options.h b/google/cloud/rest_options.h
@@ -15,6 +15,7 @@
 #ifndef GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_REST_OPTIONS_H
 #define GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_REST_OPTIONS_H
 
+#include "google/cloud/common_options.h"
 #include "google/cloud/options.h"
 #include "google/cloud/tracing_options.h"
 #include "google/cloud/version.h"
@@ -25,33 +26,6 @@ namespace google {
 namespace cloud {
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
 
-/**
- * Configure the QuotaUser system parameter.
- *
- * A pseudo user identifier for charging per-user quotas. If not specified, the
- * authenticated principal is used. If there is no authenticated principal, the
- * client IP address will be used. When specified, a valid API key with service
- * restrictions must be used to identify the quota project. Otherwise, this
- * parameter is ignored.
- *
- * @ingroup rest-options
- */
-struct QuotaUserOption {
-  using Type = std::string;
-};
-
-/**
- * Configure the UserIp query parameter.
- *
- * This can be used to separate quota usage by source IP address.
- *
- * @deprecated prefer using `google::cloud::QuotaUser`.
- * @ingroup rest-options
- */
-struct UserIpOption {
-  using Type = std::string;
-};
-
 /**
  * Timeout for the server to finish processing the request. This system param
  * only applies to REST APIs for which client-side timeout is not applicable.