BigQueryJdbcOpenTelemetry.java

/*
 * Copyright 2026 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.google.cloud.bigquery.jdbc;

import com.google.api.gax.rpc.HeaderProvider;
import com.google.auth.Credentials;
import com.google.auth.oauth2.GoogleCredentials;
import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException;
import com.google.cloud.logging.Logging;
import com.google.cloud.logging.LoggingOptions;
import com.google.common.hash.Hashing;
import io.opentelemetry.api.GlobalOpenTelemetry;
import io.opentelemetry.api.OpenTelemetry;
import io.opentelemetry.api.baggage.Baggage;
import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.SpanKind;
import io.opentelemetry.api.trace.StatusCode;
import io.opentelemetry.api.trace.Tracer;
import io.opentelemetry.context.Context;
import io.opentelemetry.context.Scope;
import io.opentelemetry.exporter.otlp.http.trace.OtlpHttpSpanExporter;
import io.opentelemetry.exporter.otlp.http.trace.OtlpHttpSpanExporterBuilder;
import io.opentelemetry.exporter.otlp.trace.OtlpGrpcSpanExporter;
import io.opentelemetry.sdk.OpenTelemetrySdk;
import io.opentelemetry.sdk.autoconfigure.AutoConfiguredOpenTelemetrySdk;
import io.opentelemetry.sdk.common.export.ProxyOptions;
import java.net.InetSocketAddress;
import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.sql.SQLException;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.Callable;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.logging.Handler;
import java.util.logging.Logger;

class BigQueryJdbcOpenTelemetry {

  static final String INSTRUMENTATION_SCOPE_NAME = "com.google.cloud.bigquery.jdbc";
  static final String BIGQUERY_NAMESPACE = "com.google.cloud.bigquery";
  public static final String CONNECTION_ID_BAGGAGE_KEY = "jdbc.connection_id";
  public static final String DB_SYSTEM_KEY = "db.system";
  public static final String DB_SYSTEM_VALUE = "bigquery";
  public static final String DB_CONNECTION_ID_KEY = "db.connection_id";
  public static final String DB_APPLICATION_KEY = "db.application";
  public static final String DEFAULT_APPLICATION_NAME = "Google-BigQuery-JDBC-Driver";
  public static final String DB_STATEMENT_KEY = "db.statement";
  public static final String DB_STATEMENT_COUNT_KEY = "db.statement.count";
  public static final String DB_BATCH_STATEMENTS_KEY = "db.batch.statements";
  private static final String OTEL_TRACES_EXPORTER = "otel.traces.exporter";
  private static final String OTEL_EXPORTER_OTLP_ENDPOINT = "otel.exporter.otlp.endpoint";
  private static final String OTEL_LOGS_EXPORTER = "otel.logs.exporter";
  private static final String OTEL_METRICS_EXPORTER = "otel.metrics.exporter";
  private static final String GOOGLE_CLOUD_PROJECT = "google.cloud.project";
  private static final String OTLP_ENDPOINT_VALUE = "https://telemetry.googleapis.com:443";
  private static final URI OTLP_ENDPOINT_URI = URI.create(OTLP_ENDPOINT_VALUE);
  private static final String EXPORTER_NONE = "none";
  private static final String EXPORTER_OTLP = "otlp";
  private static final String OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT =
      "otel.span.attribute.value.length.limit";
  private static final String OTEL_ATTRIBUTE_VALUE_LENGTH_LIMIT =
      "otel.attribute.value.length.limit";
  private static final String DEFAULT_ATTRIBUTE_LENGTH_LIMIT = "32768";
  private static final BigQueryJdbcCustomLogger LOG =
      new BigQueryJdbcCustomLogger("BigQueryJdbcOpenTelemetry");

  private static final class SdkCacheKey {
    private final String projectId;
    private final String credentialsHashOrPath;
    private final boolean enableTrace;
    private final String proxyHost;
    private final String proxyPort;

    SdkCacheKey(
        String projectId,
        String credentialsHashOrPath,
        boolean enableTrace,
        Map<String, String> proxyProperties) {
      this.projectId = projectId;
      this.credentialsHashOrPath = credentialsHashOrPath;
      this.enableTrace = enableTrace;
      this.proxyHost =
          proxyProperties != null
              ? proxyProperties.get(BigQueryJdbcUrlUtility.PROXY_HOST_PROPERTY_NAME)
              : null;
      this.proxyPort =
          proxyProperties != null
              ? proxyProperties.get(BigQueryJdbcUrlUtility.PROXY_PORT_PROPERTY_NAME)
              : null;
    }

    @Override
    public boolean equals(Object o) {
      if (this == o) return true;
      if (o == null || getClass() != o.getClass()) return false;
      SdkCacheKey that = (SdkCacheKey) o;
      return enableTrace == that.enableTrace
          && Objects.equals(projectId, that.projectId)
          && Objects.equals(credentialsHashOrPath, that.credentialsHashOrPath)
          && Objects.equals(proxyHost, that.proxyHost)
          && Objects.equals(proxyPort, that.proxyPort);
    }

    @Override
    public int hashCode() {
      return Objects.hash(projectId, credentialsHashOrPath, enableTrace, proxyHost, proxyPort);
    }
  }

  private static final class CachedSdk {
    final OpenTelemetrySdk sdk;
    final AtomicInteger refCount = new AtomicInteger(1);

    CachedSdk(OpenTelemetrySdk sdk) {
      this.sdk = sdk;
    }
  }

  private static final ConcurrentHashMap<SdkCacheKey, CachedSdk> sdkCache =
      new ConcurrentHashMap<>();

  static class TelemetryConfig {
    final OpenTelemetry openTelemetry;
    final Logging loggingClient;
    final boolean useDirectGcpLogging;

    TelemetryConfig(
        OpenTelemetry openTelemetry, Logging loggingClient, Boolean useDirectGcpLogging) {
      this.openTelemetry = openTelemetry;
      this.loggingClient = loggingClient;
      this.useDirectGcpLogging = useDirectGcpLogging != null ? useDirectGcpLogging : false;
    }
  }

  private static final ConcurrentHashMap<String, TelemetryConfig> connectionConfigs =
      new ConcurrentHashMap<>();

  private BigQueryJdbcOpenTelemetry() {}

  static {
    ensureGlobalHandlerAttached();
  }

  public static synchronized void ensureGlobalHandlerAttached() {
    Logger logger = Logger.getLogger(BIGQUERY_NAMESPACE);
    boolean present = false;
    for (Handler h : logger.getHandlers()) {
      if (h instanceof OpenTelemetryJulHandler) {
        present = true;
        break;
      }
    }
    if (!present) {
      logger.addHandler(new OpenTelemetryJulHandler());
    }
  }

  public static void registerConnection(
      String connectionId,
      OpenTelemetry openTelemetry,
      Logging loggingClient,
      Boolean useDirectGcpLogging) {
    connectionConfigs.put(
        connectionId, new TelemetryConfig(openTelemetry, loggingClient, useDirectGcpLogging));
  }

  public static void unregisterConnection(String connectionId) {
    TelemetryConfig config = connectionConfigs.remove(connectionId);
    if (config != null && config.loggingClient != null) {
      try {
        config.loggingClient.close();
      } catch (Exception e) {
        LOG.warning("Failed to close Logging client during unregister: %s", e.getMessage());
      }
    }
  }

  public static Logging createLoggingClient(
      boolean enableGcpLogExporter,
      OpenTelemetry customOpenTelemetry,
      String effectiveCredentials,
      String effectiveProjectId,
      Credentials fallbackCredentials,
      HeaderProvider headerProvider) {

    if (!enableGcpLogExporter || customOpenTelemetry != null) {
      return null;
    }

    try {
      Credentials credentials;
      if (effectiveCredentials != null) {
        credentials = resolveCredentialsFromString(effectiveCredentials);
      } else {
        credentials = fallbackCredentials;
      }

      LoggingOptions.Builder loggingOptionsBuilder =
          LoggingOptions.newBuilder().setProjectId(effectiveProjectId);
      if (credentials != null) {
        loggingOptionsBuilder.setCredentials(credentials);
      }
      if (headerProvider != null) {
        loggingOptionsBuilder.setHeaderProvider(headerProvider);
      }
      return loggingOptionsBuilder.build().getService();
    } catch (Exception e) {
      throw new BigQueryJdbcRuntimeException("Failed to initialize Logging client", e);
    }
  }

  public static void releaseSdk(OpenTelemetry openTelemetry) {
    if (openTelemetry == null || openTelemetry == OpenTelemetry.noop()) {
      return;
    }

    AtomicBoolean shouldClose = new AtomicBoolean(false);
    for (Map.Entry<SdkCacheKey, CachedSdk> entry : sdkCache.entrySet()) {
      if (entry.getValue().sdk == openTelemetry) {
        sdkCache.computeIfPresent(
            entry.getKey(),
            (k, cachedSdk) -> {
              if (cachedSdk.sdk != openTelemetry) {
                return cachedSdk;
              }
              if (cachedSdk.refCount.decrementAndGet() <= 0) {
                shouldClose.set(true);
                return null;
              }
              return cachedSdk;
            });
        break;
      }
    }

    if (shouldClose.get() && openTelemetry instanceof OpenTelemetrySdk) {
      try {
        ((OpenTelemetrySdk) openTelemetry).close();
      } catch (Exception e) {
        // Swallow exceptions so that telemetry shutdown failures (e.g. flush timeouts)
        // do not propagate and disrupt the core BigQueryConnection.close() logic.
        // Throwing here could prevent the core connection from cleaning up correctly.
        LOG.warning("Failed to close OpenTelemetry SDK: %s", e.getMessage());
      }
    }
  }

  private static Credentials resolveCredentialsFromString(String credsString) {
    Map<String, String> authProperties = new java.util.HashMap<>();
    authProperties.put(
        BigQueryJdbcUrlUtility.OAUTH_TYPE_PROPERTY_NAME,
        BigQueryJdbcOAuthUtility.AuthType.GOOGLE_SERVICE_ACCOUNT.name()); // Service Account

    if (!BigQueryJdbcOAuthUtility.isFileExists(credsString)) {
      authProperties.put(BigQueryJdbcUrlUtility.OAUTH_PVT_KEY_PROPERTY_NAME, credsString);
    } else {
      authProperties.put(BigQueryJdbcUrlUtility.OAUTH_PVT_KEY_PATH_PROPERTY_NAME, credsString);
    }

    return BigQueryJdbcOAuthUtility.getCredentials(
        authProperties,
        new java.util.HashMap<>(),
        false,
        null,
        BigQueryJdbcOpenTelemetry.class.getName());
  }

  public static TelemetryConfig getConnectionConfig(String connectionId) {
    return connectionConfigs.get(connectionId);
  }

  public static Collection<TelemetryConfig> getRegisteredConfigs() {
    return connectionConfigs.values();
  }

  private static Map<String, String> getAuthHeaders(Credentials credentials, String projectId) {
    try {

      Map<String, List<String>> metadata = credentials.getRequestMetadata();
      Map<String, String> headers = new HashMap<>();
      metadata.forEach(
          (headerKey, headerValues) -> {
            if (!headerValues.isEmpty()) {
              headers.put(headerKey, headerValues.get(0));
            }
          });
      if (projectId != null && !projectId.isEmpty()) {
        headers.put("x-goog-user-project", projectId);
        headers.put("x-goog-request-params", "project=" + projectId);
        headers.put("x-goog-request-reason", "bigquery-jdbc");
      }
      return headers;
    } catch (Exception e) {
      // We log the warning and return an empty map, allowing the exporter to fail gracefully
      // with a standard OTLP response code (e.g., 401 Unauthorized) handled by OTel.
      LOG.warning(e, "Failed to get auth headers");
      return new HashMap<>();
    }
  }

  private static String getCredentialsIdentifier(String credentials) {
    if (credentials == null) {
      return "";
    }
    byte[] credsBytes = credentials.getBytes(StandardCharsets.UTF_8);
    if (BigQueryJdbcOAuthUtility.isJson(credsBytes)) {
      return Hashing.sha256().hashString(credentials, StandardCharsets.UTF_8).toString();
    }
    return credentials;
  }

  /**
   * Initializes or returns the OpenTelemetry instance based on hybrid logic. Prefer
   * customOpenTelemetry if provided; fallback to an auto-configured GCP exporter if requested.
   */
  public static OpenTelemetry getOpenTelemetry(
      boolean useGlobalOpenTelemetry,
      boolean enableGcpTraceExporter,
      boolean enableGcpLogExporter,
      OpenTelemetry customOpenTelemetry,
      String gcpTelemetryCredentials,
      String gcpTelemetryProjectId,
      Credentials fallbackCredentials,
      Map<String, String> proxyProperties) {

    if (customOpenTelemetry != null) {
      return customOpenTelemetry;
    }

    if (useGlobalOpenTelemetry) {
      return GlobalOpenTelemetry.get();
    }

    if (!enableGcpTraceExporter && !enableGcpLogExporter) {
      return OpenTelemetry.noop();
    }

    SdkCacheKey key =
        new SdkCacheKey(
            gcpTelemetryProjectId,
            getCredentialsIdentifier(gcpTelemetryCredentials),
            enableGcpTraceExporter,
            proxyProperties);
    CachedSdk fastCheck = sdkCache.get(key);
    if (fastCheck != null) {
      CachedSdk result =
          sdkCache.computeIfPresent(
              key,
              (k, cachedSdk) -> {
                cachedSdk.refCount.incrementAndGet();
                return cachedSdk;
              });
      if (result != null) {
        return result.sdk;
      }
    }

    Map<String, String> props = new HashMap<>();

    if (enableGcpTraceExporter) {
      props.put(OTEL_TRACES_EXPORTER, EXPORTER_OTLP);
      props.put(OTEL_EXPORTER_OTLP_ENDPOINT, OTLP_ENDPOINT_VALUE);
    } else {
      props.put(OTEL_TRACES_EXPORTER, EXPORTER_NONE);
    }

    // Logs are handled directly via GCP logging
    props.put(OTEL_LOGS_EXPORTER, EXPORTER_NONE);
    // Metrics are deferred to a future phase
    props.put(OTEL_METRICS_EXPORTER, EXPORTER_NONE);

    if (gcpTelemetryProjectId != null) {
      props.put(GOOGLE_CLOUD_PROJECT, gcpTelemetryProjectId);
    }

    // Set safe, generous default limits on attribute value lengths (32KB) to protect
    // customers from GCP Cloud Trace 64KB span ingestion failures when logging massive
    // exception stack traces or database schema metadata.
    // Respect any existing user configuration overrides.
    if (!props.containsKey(OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT)) {
      props.put(OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT, DEFAULT_ATTRIBUTE_LENGTH_LIMIT);
    }
    if (!props.containsKey(OTEL_ATTRIBUTE_VALUE_LENGTH_LIMIT)) {
      props.put(OTEL_ATTRIBUTE_VALUE_LENGTH_LIMIT, DEFAULT_ATTRIBUTE_LENGTH_LIMIT);
    }

    AutoConfiguredOpenTelemetrySdk autoConfigured =
        AutoConfiguredOpenTelemetrySdk.builder()
            .addPropertiesSupplier(() -> props)
            .addResourceCustomizer(
                (resource, configProperties) -> {
                  if (gcpTelemetryProjectId != null) {
                    return resource.merge(
                        io.opentelemetry.sdk.resources.Resource.builder()
                            .put("gcp.project_id", gcpTelemetryProjectId)
                            .build());
                  }
                  return resource;
                })
            .addSpanExporterCustomizer(
                (spanExporter, configProperties) -> {
                  try {
                    Credentials credentials;
                    if (gcpTelemetryCredentials != null) {
                      credentials = resolveCredentialsFromString(gcpTelemetryCredentials);
                    } else {
                      credentials = fallbackCredentials;
                    }

                    if (credentials instanceof GoogleCredentials) {
                      GoogleCredentials googleCredentials = (GoogleCredentials) credentials;
                      if (googleCredentials.createScopedRequired()) {
                        credentials =
                            googleCredentials.createScoped(
                                Collections.singletonList(
                                    "https://www.googleapis.com/auth/cloud-platform"));
                      }
                    }

                    final Credentials finalCredentials = credentials;

                    if (spanExporter instanceof OtlpHttpSpanExporter) {
                      OtlpHttpSpanExporterBuilder builder =
                          ((OtlpHttpSpanExporter) spanExporter).toBuilder();
                      builder.setHeaders(
                          () -> getAuthHeaders(finalCredentials, gcpTelemetryProjectId));

                      ProxyOptions proxyOptions = createProxyOptions(proxyProperties);
                      if (proxyOptions != null) {
                        builder.setProxy(proxyOptions);
                      }

                      return builder.build();
                    }
                    if (spanExporter instanceof OtlpGrpcSpanExporter) {
                      return ((OtlpGrpcSpanExporter) spanExporter)
                          .toBuilder()
                              .setHeaders(
                                  () -> getAuthHeaders(finalCredentials, gcpTelemetryProjectId))
                              .build();
                    }
                  } catch (Exception e) {
                    LOG.warning(
                        e,
                        "Failed to resolve telemetry credentials. Telemetry will be exported using default OpenTelemetry configuration (custom authentication headers will not be injected).");
                  }
                  return spanExporter;
                })
            .build();

    OpenTelemetrySdk newSdk = autoConfigured.getOpenTelemetrySdk();

    return sdkCache.compute(
            key,
            (k, cachedSdk) -> {
              if (cachedSdk != null) {
                cachedSdk.refCount.incrementAndGet();
                newSdk.close(); // Clean up the duplicate we just made
                return cachedSdk;
              }
              return new CachedSdk(newSdk);
            })
        .sdk;
  }

  /** Gets a Tracer for the JDBC driver instrumentation scope. */
  public static Tracer getTracer(OpenTelemetry openTelemetry) {
    return openTelemetry.getTracer(INSTRUMENTATION_SCOPE_NAME);
  }

  public static <T> T withTracing(
      String spanName, BigQueryConnection connection, String sql, Callable<T> operation)
      throws SQLException {

    Tracer tracer = connection.getTracer();
    Span span = tracer.spanBuilder(spanName).setSpanKind(SpanKind.CLIENT).startSpan();

    span.setAttribute(DB_SYSTEM_KEY, DB_SYSTEM_VALUE);
    span.setAttribute(DB_CONNECTION_ID_KEY, connection.getConnectionId());

    String appName = connection.getPartnerToken();
    if (appName == null || appName.isEmpty()) {
      appName = DEFAULT_APPLICATION_NAME;
    }
    span.setAttribute(DB_APPLICATION_KEY, appName);

    if (sql != null) {
      span.setAttribute(DB_STATEMENT_KEY, sql);
    }

    Baggage updatedBaggage =
        Baggage.fromContext(Context.current()).toBuilder()
            .put(CONNECTION_ID_BAGGAGE_KEY, connection.getConnectionId())
            .build();

    // Create full context with new span and updated baggage
    Context fullContext = Context.current().with(span).with(updatedBaggage);

    try (Scope scope = fullContext.makeCurrent()) {
      return operation.call();
    } catch (Exception ex) {
      span.recordException(ex);
      span.setStatus(StatusCode.ERROR, ex.getMessage());

      if (ex instanceof SQLException) {
        throw (SQLException) ex;
      }
      if (ex instanceof RuntimeException) {
        throw (RuntimeException) ex;
      }
      if (ex instanceof InterruptedException) {
        Thread.currentThread().interrupt();
        throw new BigQueryJdbcRuntimeException("Operation interrupted", ex);
      }
      throw new BigQueryJdbcRuntimeException(ex);
    } finally {
      span.end();
    }
  }

  private static ProxyOptions createProxyOptions(Map<String, String> proxyProperties) {
    if (proxyProperties == null) {
      return null;
    }

    final String host = proxyProperties.get(BigQueryJdbcUrlUtility.PROXY_HOST_PROPERTY_NAME);
    final String portStr = proxyProperties.get(BigQueryJdbcUrlUtility.PROXY_PORT_PROPERTY_NAME);
    if (host == null || host.isEmpty() || portStr == null || portStr.isEmpty()) {
      return null;
    }

    int port;
    try {
      port = Integer.parseInt(portStr);
    } catch (NumberFormatException e) {
      throw new BigQueryJdbcRuntimeException("Invalid proxy port number: " + portStr, e);
    }

    return ProxyOptions.create(InetSocketAddress.createUnresolved(host, port));
  }
}