Skip to content

Commit

Permalink
fix: Allow setting framework tags on cloud pubsub shim. (#571)
Browse files Browse the repository at this point in the history
  • Loading branch information
dpcollins-google committed Apr 1, 2021
1 parent a2ff14c commit 2231fbd
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 11 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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.
Expand All @@ -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());
Expand All @@ -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.
Expand All @@ -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 {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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) {
Expand Down

0 comments on commit 2231fbd

Please sign in to comment.