Skip to content

Commit

Permalink
feat(spanner): replaced setting channel provider with setting channel…
Browse files Browse the repository at this point in the history
… configurator
  • Loading branch information
sagnghos committed Jan 7, 2025
1 parent 03ce64f commit 7259b97
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 50 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -27,11 +27,9 @@
import com.google.api.gax.core.GaxProperties;
import com.google.api.gax.grpc.GrpcCallContext;
import com.google.api.gax.grpc.GrpcInterceptorProvider;
import com.google.api.gax.grpc.GrpcTransportChannel;
import com.google.api.gax.longrunning.OperationTimedPollAlgorithm;
import com.google.api.gax.retrying.RetrySettings;
import com.google.api.gax.rpc.ApiCallContext;
import com.google.api.gax.rpc.FixedTransportChannelProvider;
import com.google.api.gax.rpc.TransportChannelProvider;
import com.google.api.gax.tracing.ApiTracerFactory;
import com.google.api.gax.tracing.BaseApiTracerFactory;
Expand Down Expand Up @@ -71,19 +69,17 @@
import io.grpc.CompressorRegistry;
import io.grpc.Context;
import io.grpc.ExperimentalApi;
import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import io.grpc.MethodDescriptor;
import io.grpc.netty.shaded.io.grpc.netty.GrpcSslContexts;
import io.grpc.netty.shaded.io.grpc.netty.NettyChannelBuilder;
import io.grpc.netty.shaded.io.netty.handler.ssl.SslContext;
import io.opentelemetry.api.GlobalOpenTelemetry;
import io.opentelemetry.api.OpenTelemetry;
import io.opentelemetry.api.common.Attributes;
import java.io.File;
import java.io.IOException;
import java.net.MalformedURLException;
import java.net.URI;
import java.net.URISyntaxException;
import java.net.URL;
import java.time.Duration;
import java.util.ArrayList;
Expand All @@ -98,8 +94,6 @@
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.logging.Level;
import java.util.logging.Logger;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import javax.annotation.concurrent.GuardedBy;
Expand Down Expand Up @@ -952,7 +946,6 @@ public static class Builder
private CloseableExecutorProvider asyncExecutorProvider;
private String compressorName;
private String emulatorHost = System.getenv("SPANNER_EMULATOR_HOST");
private ManagedChannel managedChannel;
private boolean leaderAwareRoutingEnabled = true;
private boolean attemptDirectPath = true;
private DirectedReadOptions directedReadOptions;
Expand All @@ -963,6 +956,7 @@ public static class Builder
private boolean enableEndToEndTracing = SpannerOptions.environment.isEnableEndToEndTracing();
private boolean enableBuiltInMetrics = SpannerOptions.environment.isEnableBuiltInMetrics();
private String monitoringHost = SpannerOptions.environment.getMonitoringHost();
private SslContext mTLSContext = null;

private static String createCustomClientLibToken(String token) {
return token + " " + ServiceOptions.getGoogApiClientLibName();
Expand Down Expand Up @@ -1496,34 +1490,27 @@ public Builder setEmulatorHost(String emulatorHost) {
return this;
}

public Builder useClientCert(String host, String clientCertificate, String clientKey) {
/**
* Configures mTLS authentication using the provided client certificate and key files. mTLS is
* only supported for external spanner hosts.
*
* @param clientCertificate Path to the client certificate file.
* @param clientCertificateKey Path to the client private key file.
* @throws SpannerException If an error occurs while configuring the mTLS context
*/
@ExperimentalApi("https://github.com/googleapis/java-spanner/pull/3574")
public Builder useClientCert(String clientCertificate, String clientCertificateKey) {
try {
URI uri = new URI(host);
managedChannel =
NettyChannelBuilder.forAddress(uri.getHost(), uri.getPort())
.sslContext(
GrpcSslContexts.forClient()
.keyManager(new File(clientCertificate), new File(clientKey))
.build())
this.mTLSContext =
GrpcSslContexts.forClient()
.keyManager(new File(clientCertificate), new File(clientCertificateKey))
.build();

setChannelProvider(
FixedTransportChannelProvider.create(GrpcTransportChannel.create(managedChannel)));
} catch (URISyntaxException e) {
throw new IllegalArgumentException(
"Invalid host format. Expected format: 'protocol://host[:port]'.", e);
} catch (Exception e) {
throw new RuntimeException("Unexpected error during mTLS setup.", e);
throw SpannerExceptionFactory.asSpannerException(e);
}
return this;
}

public Builder usePlainText() {
this.setChannelConfigurator(ManagedChannelBuilder::usePlaintext);
this.setCredentials(NoCredentials.getInstance());
return this;
}

/**
* Sets OpenTelemetry object to be used for Spanner Metrics and Traces. GlobalOpenTelemetry will
* be used as fallback if this options is not set.
Expand Down Expand Up @@ -1632,24 +1619,15 @@ public SpannerOptions build() {
this.setChannelConfigurator(ManagedChannelBuilder::usePlaintext);
// As we are using plain text, we should never send any credentials.
this.setCredentials(NoCredentials.getInstance());
} else if (managedChannel != null) {
// Add shutdown hook for the ManagedChannel if created to prevent resource leak
Runtime.getRuntime()
.addShutdownHook(
new Thread(
() -> {
final Logger logger = Logger.getLogger(SpannerOptions.class.getName());
try {
managedChannel.shutdown();
logger.log(
Level.INFO, "[SpannerOptions] ManagedChannel shut down successfully.");
} catch (Exception e) {
logger.log(
Level.WARNING,
"[SpannerOptions] Failed to shut down ManagedChannel.",
e);
}
}));
}
if (mTLSContext != null) {
this.setChannelConfigurator(
builder -> {
if (builder instanceof NettyChannelBuilder) {
((NettyChannelBuilder) builder).sslContext(mTLSContext);
}
return builder;
});
}
if (this.numChannels == null) {
this.numChannels =
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,8 @@ static class SpannerPoolKey {
private final Boolean enableExtendedTracing;
private final Boolean enableApiTracing;
private final boolean enableEndToEndTracing;
private final String clientCertificate;
private final String clientCertificateKey;

@VisibleForTesting
static SpannerPoolKey of(ConnectionOptions options) {
Expand Down Expand Up @@ -192,6 +194,8 @@ private SpannerPoolKey(ConnectionOptions options) throws IOException {
this.enableExtendedTracing = options.isEnableExtendedTracing();
this.enableApiTracing = options.isEnableApiTracing();
this.enableEndToEndTracing = options.isEndToEndTracingEnabled();
this.clientCertificate = options.getClientCertificate();
this.clientCertificateKey = options.getClientCertificateKey();
}

@Override
Expand Down Expand Up @@ -393,9 +397,8 @@ Spanner createSpanner(SpannerPoolKey key, ConnectionOptions options) {
// Set a custom channel configurator to allow http instead of https.
builder.setChannelConfigurator(ManagedChannelBuilder::usePlaintext);
}
if (options.getClientCertificate() != null && options.getClientCertificateKey() != null) {
builder.useClientCert(
options.getHost(), options.getClientCertificate(), options.getClientCertificateKey());
if (key.clientCertificate != null && key.clientCertificateKey != null) {
builder.useClientCert(key.clientCertificate, key.clientCertificateKey);
}
if (options.getConfigurator() != null) {
options.getConfigurator().configure(builder);
Expand Down

0 comments on commit 7259b97

Please sign in to comment.