fix: Allow setting framework tags on cloud pubsub shim. (#571)

googleapis · Apr 1, 2021 · 2231fbd · 2231fbd
1 parent a2ff14c
commit 2231fbd
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 11 deletions.
diff --git a/...d-pubsublite/src/main/java/com/google/cloud/pubsublite/cloudpubsub/PublisherSettings.java b/...d-pubsublite/src/main/java/com/google/cloud/pubsublite/cloudpubsub/PublisherSettings.java
@@ -59,7 +59,6 @@ public abstract class PublisherSettings {
           .setRequestByteThreshold(Constants.MAX_PUBLISH_BATCH_BYTES)
           .setDelayThreshold(Duration.ofMillis(50))
           .build();
-  private static final Framework FRAMEWORK = Framework.of("CLOUD_PUBSUB_SHIM");
 
   // Required parameters.
 
@@ -79,6 +78,11 @@ public abstract class PublisherSettings {
   /** A provider for credentials. */
   abstract CredentialsProvider credentialsProvider();
 
+  /**
+   * A Framework tag for internal metrics. Please set this if integrating with a public framework!
+   */
+  abstract Framework framework();
+
   /**
    * A supplier for new PublisherServiceClients. Should return a new client each time. If present,
    * ignores CredentialsProvider.
@@ -94,6 +98,7 @@ public abstract class PublisherSettings {
   /** Get a new builder for a PublisherSettings. */
   public static Builder newBuilder() {
     return new AutoValue_PublisherSettings.Builder()
+        .setFramework(Framework.of("CLOUD_PUBSUB_SHIM"))
         .setCredentialsProvider(
             PublisherServiceSettings.defaultCredentialsProviderBuilder().build())
         .setUnderlyingBuilder(SinglePartitionPublisherBuilder.newBuilder());
@@ -120,6 +125,11 @@ public abstract Builder setMessageTransformer(
     /** A provider for credentials. */
     public abstract Builder setCredentialsProvider(CredentialsProvider credentialsProvider);
 
+    /**
+     * A Framework tag for internal metrics. Please set this if integrating with a public framework!
+     */
+    public abstract Builder setFramework(Framework framework);
+
     /**
      * A supplier for new PublisherServiceClients. Should return a new client each time. If present,
      * ignores CredentialsProvider.
@@ -143,7 +153,7 @@ private PublisherServiceClient newServiceClient(Partition partition) throws ApiE
     settingsBuilder = settingsBuilder.setCredentialsProvider(credentialsProvider());
     settingsBuilder =
         addDefaultMetadata(
-            PubsubContext.of(FRAMEWORK),
+            PubsubContext.of(framework()),
             RoutingMetadata.of(topicPath(), partition),
             settingsBuilder);
     try {

diff --git a/...-pubsublite/src/main/java/com/google/cloud/pubsublite/cloudpubsub/SubscriberSettings.java b/...-pubsublite/src/main/java/com/google/cloud/pubsublite/cloudpubsub/SubscriberSettings.java
@@ -60,9 +60,6 @@
  */
 @AutoValue
 public abstract class SubscriberSettings {
-
-  private static final Framework FRAMEWORK = Framework.of("CLOUD_PUBSUB_SHIM");
-
   // Required parameters.
   /**
    * The receiver which handles new messages sent by the Pub/Sub Lite system. Only one downcall from
@@ -97,6 +94,11 @@ public abstract class SubscriberSettings {
   /** A provider for credentials. */
   abstract CredentialsProvider credentialsProvider();
 
+  /**
+   * A Framework tag for internal metrics. Please set this if integrating with a public framework!
+   */
+  abstract Framework framework();
+
   /**
    * A supplier for new SubscriberServiceClients. Should return a new client each time. If present,
    * ignores CredentialsProvider.
@@ -124,16 +126,15 @@ public abstract class SubscriberSettings {
 
   public static Builder newBuilder() {
     return new AutoValue_SubscriberSettings.Builder()
+        .setFramework(Framework.of("CLOUD_PUBSUB_SHIM"))
         .setPartitions(ImmutableList.of())
         .setCredentialsProvider(
             SubscriberServiceSettings.defaultCredentialsProviderBuilder().build());
   }
 
   @AutoValue.Builder
   public abstract static class Builder {
-
     // Required parameters.
-
     /**
      * The receiver which handles new messages sent by the Pub/Sub Lite system. Only one downcall
      * from any connected partition will be outstanding at a time, and blocking in this receiver
@@ -168,6 +169,11 @@ public abstract Builder setTransformer(
     /** A provider for credentials. */
     public abstract Builder setCredentialsProvider(CredentialsProvider provider);
 
+    /**
+     * A Framework tag for internal metrics. Please set this if integrating with a public framework!
+     */
+    public abstract Builder setFramework(Framework framework);
+
     /**
      * A supplier for new SubscriberServiceClients. Should return a new client each time. If
      * present, ignores CredentialsProvider.
@@ -206,10 +212,11 @@ private SubscriberServiceClient newSubscriberServiceClient(Partition partition)
     try {
       SubscriberServiceSettings.Builder settingsBuilder =
           SubscriberServiceSettings.newBuilder().setCredentialsProvider(credentialsProvider());
-      addDefaultMetadata(
-          PubsubContext.of(FRAMEWORK),
-          RoutingMetadata.of(subscriptionPath(), partition),
-          settingsBuilder);
+      settingsBuilder =
+          addDefaultMetadata(
+              PubsubContext.of(framework()),
+              RoutingMetadata.of(subscriptionPath(), partition),
+              settingsBuilder);
       return SubscriberServiceClient.create(
           addDefaultSettings(subscriptionPath().location().region(), settingsBuilder));
     } catch (Throwable t) {