feat(common): QuotaUserOption for gRPC-based libs

googleapis · Apr 6, 2024 · 404cd19 · 404cd19
1 parent 18f522e
commit 404cd19
Show file tree

Hide file tree

Showing 4 changed files with 75 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"
@@ -38,6 +39,12 @@ void SetMetadata(grpc::ClientContext& context, Options const& options,
   for (auto const& h : options.get<CustomHeadersOption>()) {
     context.AddMetadata(h.first, h.second);
   }
+  if (options.has<UserIpOption>() && !options.has<QuotaUserOption>()) {
+    context.AddMetadata("x-goog-user-ip", options.get<UserIpOption>());
+  }
+  if (options.has<QuotaUserOption>()) {
+    context.AddMetadata("x-goog-quota-user", options.get<QuotaUserOption>());
+  }
 }
 
 void ConfigureContext(grpc::ClientContext& context, Options const& opts) {

diff --git a/google/cloud/grpc_options_test.cc b/google/cloud/grpc_options_test.cc
@@ -328,6 +328,7 @@ TEST(GrpcSetMetadata, Full) {
       context,
       Options{}
           .set<UserProjectOption>("user-project")
+          .set<QuotaUserOption>("quota-user")
           .set<AuthorityOption>("authority.googleapis.com")
           .set<CustomHeadersOption>(
               {{"custom-header-1", "v1"}, {"custom-header-2", "v2"}}),
@@ -338,6 +339,7 @@ TEST(GrpcSetMetadata, Full) {
   EXPECT_THAT(metadata,
               UnorderedElementsAre(
                   Pair("x-goog-user-project", "user-project"),
+                  Pair("x-goog-quota-user", "quota-user"),
                   Pair("fixed-header-1", "v1"), Pair("fixed-header-2", "v2"),
                   Pair("custom-header-1", "v1"), Pair("custom-header-2", "v2"),
                   Pair("x-goog-api-client", "api-client-header")));
@@ -354,6 +356,44 @@ TEST(GrpcSetMetadata, Authority) {
   EXPECT_EQ(authority, "authority.googleapis.com");
 }
 
+TEST(GrpcSetMetadata, ConfigureUserIp) {
+  grpc::ClientContext context;
+  internal::SetMetadata(context, Options{}.set<UserIpOption>("1234"), {},
+                        "api-client-header");
+
+  ValidateMetadataFixture fixture;
+  auto const metadata = fixture.GetMetadata(context);
+  EXPECT_THAT(metadata, UnorderedElementsAre(
+                            Pair("x-goog-user-ip", "1234"),
+                            Pair("x-goog-api-client", "api-client-header")));
+}
+
+TEST(GrpcSetMetadata, ConfigureQuotaUser) {
+  grpc::ClientContext context;
+  internal::SetMetadata(context, Options{}.set<QuotaUserOption>("my-user"), {},
+                        "api-client-header");
+
+  ValidateMetadataFixture fixture;
+  auto const metadata = fixture.GetMetadata(context);
+  EXPECT_THAT(metadata, UnorderedElementsAre(
+                            Pair("x-goog-quota-user", "my-user"),
+                            Pair("x-goog-api-client", "api-client-header")));
+}
+
+TEST(GrpcSetMetadata, QuotaUserOverridesUserIp) {
+  grpc::ClientContext context;
+  internal::SetMetadata(
+      context,
+      Options{}.set<QuotaUserOption>("my-user").set<UserIpOption>("1234"), {},
+      "api-client-header");
+
+  ValidateMetadataFixture fixture;
+  auto const metadata = fixture.GetMetadata(context);
+  EXPECT_THAT(metadata, UnorderedElementsAre(
+                            Pair("x-goog-quota-user", "my-user"),
+                            Pair("x-goog-api-client", "api-client-header")));
+}
+
 TEST(GrpcClientContext, Configure) {
   auto setup = [](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.