Skip to content

feat: add Pub/Sub emulator support for source connector #429

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
dsaheb merged 3 commits into from
Jun 26, 2026
Merged

feat: add Pub/Sub emulator support for source connector #429

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
704e379
feat: add Pub/Sub emulator support for source connector
AdalbertMemSQL May 26, 2026
f73091d
Removed unused import
AdalbertMemSQL May 27, 2026
5f868df
Fixed lint failure
AdalbertMemSQL Jun 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -173,6 +173,7 @@ configurations:
| cps.subscription | String | REQUIRED (No default) | The Pub/Sub subscription ID, e.g. "baz" for subscription "/projects/bar/subscriptions/baz". |
| cps.project | String | REQUIRED (No default) | The project containing the Pub/Sub subscription, e.g. "bar" from above. |
| cps.endpoint | String | "pubsub.googleapis.com:443" | The [Pub/Sub endpoint](https://cloud.google.com/pubsub/docs/reference/service_apis_overview#service_endpoints) to use. |
| cps.useEmulator | Boolean | false | When true, use the Pub/Sub emulator instead of the production service. The emulator endpoint will be determined by the PUBSUB_EMULATOR_HOST environment variable, or fallback to the cps.endpoint configuration. |
| kafka.topic | String | REQUIRED (No default) | The Kafka topic which will receive messages from the Pub/Sub subscription. |
| cps.maxBatchSize | Integer | 100 | The maximum number of messages per batch in a pull request to Pub/Sub. |
| cps.makeOrderingKeyAttribute | Boolean | false | When true, copy the ordering key to the set of attributes set in the Kafka message. |
Expand Down
13 changes: 13 additions & 0 deletions src/main/java/com/google/pubsub/kafka/common/ConnectorUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,4 +58,17 @@ public static synchronized ScheduledExecutorService getSystemExecutor() {
}
return SYSTEM_EXECUTOR.get();
}

// Resolve the endpoint. When using the emulator, prefer PUBSUB_EMULATOR_HOST and fall back to
// the configured cps.endpoint.
public static String getPubsubEndpoint(boolean useEmulator, String cpsEndpoint) {
if (useEmulator) {
String emulatorHost = System.getenv(PUBSUB_EMULATOR_HOST);
if (emulatorHost != null && !emulatorHost.isEmpty()) {
return emulatorHost;
}
}

return cpsEndpoint;
}
}
8 changes: 2 additions & 6 deletions src/main/java/com/google/pubsub/kafka/sink/CloudPubSubSinkTask.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -418,11 +418,8 @@ private void createPublisher() {

// Configure endpoint, credentials and channel based on whether we're using emulator or
// production
String endpoint = ConnectorUtils.getPubsubEndpoint(useEmulator, cpsEndpoint);
if (useEmulator) {
// For emulator: use PUBSUB_EMULATOR_HOST env var, fallback to configured cps.endpoint, then
// default
String emulatorHost = System.getenv(ConnectorUtils.PUBSUB_EMULATOR_HOST);
String endpoint = emulatorHost != null ? emulatorHost : cpsEndpoint;
builder
.setCredentialsProvider(com.google.api.gax.core.NoCredentialsProvider.create())
.setChannelProvider(
Expand All @@ -431,8 +428,7 @@ private void createPublisher() {
.setChannelConfigurator(channel -> channel.usePlaintext())
.build());
} else {
// For production: use configured credentials and endpoint
builder.setCredentialsProvider(gcpCredentialsProvider).setEndpoint(cpsEndpoint);
builder.setCredentialsProvider(gcpCredentialsProvider).setEndpoint(endpoint);
}
if (orderingKeySource != OrderingKeySource.NONE) {
builder.setEnableMessageOrdering(true);
Expand Down
42 changes: 32 additions & 10 deletions src/main/java/com/google/pubsub/kafka/source/CloudPubSubGRPCSubscriber.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@
import com.google.api.core.ApiFuture;
import com.google.api.core.ApiFutures;
import com.google.api.gax.core.CredentialsProvider;
import com.google.api.gax.core.NoCredentialsProvider;
import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
import com.google.cloud.pubsub.v1.stub.GrpcSubscriberStub;
import com.google.cloud.pubsub.v1.stub.SubscriberStubSettings;
import com.google.common.util.concurrent.MoreExecutors;
Expand Down Expand Up @@ -48,16 +50,19 @@ public class CloudPubSubGRPCSubscriber implements CloudPubSubSubscriber {
private final String endpoint;
private final ProjectSubscriptionName subscriptionName;
private final int cpsMaxBatchSize;
private final boolean useEmulator;

CloudPubSubGRPCSubscriber(
CredentialsProvider gcpCredentialsProvider,
String endpoint,
ProjectSubscriptionName subscriptionName,
int cpsMaxBatchSize) {
int cpsMaxBatchSize,
boolean useEmulator) {
this.gcpCredentialsProvider = gcpCredentialsProvider;
this.endpoint = endpoint;
this.subscriptionName = subscriptionName;
this.cpsMaxBatchSize = cpsMaxBatchSize;
this.useEmulator = useEmulator;
makeSubscriber();
}

Expand Down Expand Up @@ -103,15 +108,32 @@ private void makeSubscriber() {
subscriber.close();
}
log.info("Creating subscriber.");
SubscriberStubSettings subscriberStubSettings =
SubscriberStubSettings.newBuilder()
.setTransportChannelProvider(
SubscriberStubSettings.defaultGrpcTransportProviderBuilder()
.setMaxInboundMessageSize(20 << 20) // 20MB
.build())
.setCredentialsProvider(gcpCredentialsProvider)
.setEndpoint(endpoint)
.build();

// Configure endpoint, credentials and channel based on whether we're using emulator or
// production
SubscriberStubSettings subscriberStubSettings;
if (useEmulator) {
subscriberStubSettings =
SubscriberStubSettings.newBuilder()
.setCredentialsProvider(NoCredentialsProvider.create())
.setTransportChannelProvider(
InstantiatingGrpcChannelProvider.newBuilder()
.setMaxInboundMessageSize(20 << 20) // 20MB
.setEndpoint(endpoint)
.setChannelConfigurator(channel -> channel.usePlaintext())
.build())
.build();
} else {
subscriberStubSettings =
SubscriberStubSettings.newBuilder()
.setTransportChannelProvider(
SubscriberStubSettings.defaultGrpcTransportProviderBuilder()
.setMaxInboundMessageSize(20 << 20) // 20MB
.build())
.setCredentialsProvider(gcpCredentialsProvider)
.setEndpoint(endpoint)
.build();
}
subscriber = GrpcSubscriberStub.create(subscriberStubSettings);
// We change the subscriber every 25 - 35 minutes in order to avoid GOAWAY errors.
nextSubscriberResetTime =
Expand Down
5 changes: 3 additions & 2 deletions src/main/java/com/google/pubsub/kafka/source/CloudPubSubRoundRobinSubscriber.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,12 +38,13 @@ public CloudPubSubRoundRobinSubscriber(
CredentialsProvider gcpCredentialsProvider,
String endpoint,
ProjectSubscriptionName subscriptionName,
int cpsMaxBatchSize) {
int cpsMaxBatchSize,
boolean useEmulator) {
subscribers = new ArrayList<>();
for (int i = 0; i < subscriberCount; ++i) {
subscribers.add(
new CloudPubSubGRPCSubscriber(
gcpCredentialsProvider, endpoint, subscriptionName, cpsMaxBatchSize));
gcpCredentialsProvider, endpoint, subscriptionName, cpsMaxBatchSize, useEmulator));
}
}

Expand Down
56 changes: 43 additions & 13 deletions src/main/java/com/google/pubsub/kafka/source/CloudPubSubSourceConnector.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@
package com.google.pubsub.kafka.source;

import com.google.api.gax.core.CredentialsProvider;
import com.google.api.gax.core.NoCredentialsProvider;
import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
import com.google.cloud.pubsub.v1.stub.GrpcSubscriberStub;
import com.google.cloud.pubsub.v1.stub.SubscriberStubSettings;
import com.google.common.annotations.VisibleForTesting;
Expand Down Expand Up @@ -140,10 +142,13 @@ public void start(Map<String, String> props) {
Map<String, Object> validated = config().parse(props);
String cpsProject = validated.get(ConnectorUtils.CPS_PROJECT_CONFIG).toString();
String cpsSubscription = validated.get(CPS_SUBSCRIPTION_CONFIG).toString();
String cpsEndpoint = (String) validated.get(ConnectorUtils.CPS_ENDPOINT);
boolean useEmulator = (Boolean) validated.get(ConnectorUtils.CPS_USE_EMULATOR);
String endpoint = ConnectorUtils.getPubsubEndpoint(useEmulator, cpsEndpoint);
ConnectorCredentialsProvider credentialsProvider =
ConnectorCredentialsProvider.fromConfig(validated);

verifySubscription(cpsProject, cpsSubscription, credentialsProvider);
verifySubscription(cpsProject, cpsSubscription, credentialsProvider, endpoint, useEmulator);
this.props = props;
log.info("Started the CloudPubSubSourceConnector");
}
Expand Down Expand Up @@ -297,7 +302,13 @@ public ConfigDef config() {
Type.STRING,
ConnectorUtils.CPS_DEFAULT_ENDPOINT,
Importance.LOW,
"The Pub/Sub endpoint to use.");
"The Pub/Sub endpoint to use.")
.define(
ConnectorUtils.CPS_USE_EMULATOR,
Type.BOOLEAN,
false,
Importance.LOW,
"When true, use the Pub/Sub emulator instead of the production service.");
}

/**
Expand All @@ -306,24 +317,43 @@ public ConfigDef config() {
*/
@VisibleForTesting
public void verifySubscription(
String cpsProject, String cpsSubscription, CredentialsProvider credentialsProvider) {
String cpsProject,
String cpsSubscription,
CredentialsProvider credentialsProvider,
String endpoint,
boolean useEmulator) {
try {
SubscriberStubSettings subscriberStubSettings =
SubscriberStubSettings.newBuilder()
.setTransportChannelProvider(
SubscriberStubSettings.defaultGrpcTransportProviderBuilder()
.setMaxInboundMessageSize(20 << 20) // 20MB
.build())
.setCredentialsProvider(credentialsProvider)
.build();
GrpcSubscriberStub stub = GrpcSubscriberStub.create(subscriberStubSettings);
// Configure endpoint, credentials and channel based on whether we're using emulator or
// production
SubscriberStubSettings.Builder settingsBuilder = SubscriberStubSettings.newBuilder();
if (useEmulator) {
settingsBuilder
.setCredentialsProvider(NoCredentialsProvider.create())
.setTransportChannelProvider(
InstantiatingGrpcChannelProvider.newBuilder()
.setMaxInboundMessageSize(20 << 20) // 20MB
.setEndpoint(endpoint)
.setChannelConfigurator(channel -> channel.usePlaintext())
.build());
} else {
settingsBuilder
.setTransportChannelProvider(
SubscriberStubSettings.defaultGrpcTransportProviderBuilder()
.setMaxInboundMessageSize(20 << 20) // 20MB
.build())
.setCredentialsProvider(credentialsProvider)
.setEndpoint(endpoint);
}

GetSubscriptionRequest request =
GetSubscriptionRequest.newBuilder()
.setSubscription(
String.format(
ConnectorUtils.CPS_SUBSCRIPTION_FORMAT, cpsProject, cpsSubscription))
.build();
stub.getSubscriptionCallable().call(request);
try (GrpcSubscriberStub stub = GrpcSubscriberStub.create(settingsBuilder.build())) {
stub.getSubscriptionCallable().call(request);
}
} catch (Exception e) {
throw new ConnectException(
"Error verifying the subscription " + cpsSubscription + " for project " + cpsProject, e);
Expand Down
24 changes: 20 additions & 4 deletions src/main/java/com/google/pubsub/kafka/source/CloudPubSubSourceTask.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,7 @@ public class CloudPubSubSourceTask extends SourceTask {
private CloudPubSubSubscriber subscriber;
private final Set<String> standardAttributes = new HashSet<>();
private boolean useKafkaHeaders;
private boolean useEmulator;

public CloudPubSubSourceTask() {}

Expand All @@ -99,6 +100,8 @@ public void start(Map<String, String> props) {
validatedProps.get(CloudPubSubSourceConnector.CPS_SUBSCRIPTION_CONFIG).toString())
.build();
String cpsEndpoint = (String) validatedProps.get(ConnectorUtils.CPS_ENDPOINT);
useEmulator = (Boolean) validatedProps.get(ConnectorUtils.CPS_USE_EMULATOR);
String endpoint = ConnectorUtils.getPubsubEndpoint(useEmulator, cpsEndpoint);
kafkaTopic = validatedProps.get(CloudPubSubSourceConnector.KAFKA_TOPIC_CONFIG).toString();
int cpsMaxBatchSize =
(Integer) validatedProps.get(CloudPubSubSourceConnector.CPS_MAX_BATCH_SIZE_CONFIG);
Expand Down Expand Up @@ -142,16 +145,28 @@ public void start(Map<String, String> props) {
receiver -> {
Subscriber.Builder builder =
Subscriber.newBuilder(cpsSubscription, receiver)
.setCredentialsProvider(gcpCredentialsProvider)
.setFlowControlSettings(
FlowControlSettings.newBuilder()
.setLimitExceededBehavior(LimitExceededBehavior.Block)
.setMaxOutstandingElementCount(streamingPullMessages)
.setMaxOutstandingRequestBytes(streamingPullBytes)
.build())
.setParallelPullCount(streamingPullParallelStreams)
.setEndpoint(cpsEndpoint)
.setExecutorProvider(FixedExecutorProvider.create(getSystemExecutor()));
// Configure endpoint, credentials and channel based on whether we're using
// emulator or production
if (useEmulator) {
builder
.setCredentialsProvider(
com.google.api.gax.core.NoCredentialsProvider.create())
.setChannelProvider(
com.google.api.gax.grpc.InstantiatingGrpcChannelProvider.newBuilder()
.setEndpoint(endpoint)
.setChannelConfigurator(channel -> channel.usePlaintext())
.build());
} else {
builder.setCredentialsProvider(gcpCredentialsProvider).setEndpoint(endpoint);
}
if (streamingPullMaxAckDeadlineMs > 0) {
builder.setMaxAckExtensionPeriod(
Duration.ofMillis(streamingPullMaxAckDeadlineMs));
Expand All @@ -168,9 +183,10 @@ public void start(Map<String, String> props) {
new CloudPubSubRoundRobinSubscriber(
NUM_CPS_SUBSCRIBERS,
gcpCredentialsProvider,
cpsEndpoint,
endpoint,
cpsSubscription,
cpsMaxBatchSize),
cpsMaxBatchSize,
useEmulator),
runnable ->
getSystemExecutor()
.scheduleAtFixedRate(runnable, 100, 100, TimeUnit.MILLISECONDS));
Expand Down
17 changes: 13 additions & 4 deletions src/test/java/com/google/pubsub/kafka/source/CloudPubSubSourceConnectorTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,7 @@
package com.google.pubsub.kafka.source;

import static org.junit.Assert.assertEquals;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.*;
import static org.mockito.Mockito.doNothing;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.spy;
Expand Down Expand Up @@ -56,7 +55,12 @@ public void setup() {
public void testStartWhenSubscriptionNonexistant() {
doThrow(new ConnectException(""))
.when(connector)
.verifySubscription(anyString(), anyString(), any(ConnectorCredentialsProvider.class));
.verifySubscription(
anyString(),
anyString(),
any(ConnectorCredentialsProvider.class),
anyString(),
anyBoolean());
connector.start(props);
}

Expand All @@ -69,7 +73,12 @@ public void testStartWhenRequiredConfigMissing() {
public void testTaskConfigs() {
doNothing()
.when(connector)
.verifySubscription(anyString(), anyString(), any(ConnectorCredentialsProvider.class));
.verifySubscription(
anyString(),
anyString(),
any(ConnectorCredentialsProvider.class),
anyString(),
anyBoolean());
connector.start(props);
List<Map<String, String>> taskConfigs = connector.taskConfigs(NUM_TASKS);
assertEquals(taskConfigs.size(), NUM_TASKS);
Expand Down
Loading
Loading