ITEndToEndTracingTest.java

/*
 * Copyright 2025 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.spanner.it;

import static com.google.common.truth.Truth.assertThat;
import static org.junit.Assert.assertTrue;
import static org.junit.Assume.assumeTrue;

import com.google.api.gax.core.FixedCredentialsProvider;
import com.google.api.gax.rpc.ApiException;
import com.google.api.gax.rpc.ResourceExhaustedException;
import com.google.api.gax.rpc.StatusCode;
import com.google.cloud.spanner.Database;
import com.google.cloud.spanner.DatabaseClient;
import com.google.cloud.spanner.IntegrationTestEnv;
import com.google.cloud.spanner.IntegrationTestEnv.TestEnvOptions;
import com.google.cloud.spanner.ParallelIntegrationTest;
import com.google.cloud.spanner.ResultSet;
import com.google.cloud.spanner.SpannerOptions;
import com.google.cloud.spanner.SpannerOptionsHelper;
import com.google.cloud.spanner.Statement;
import com.google.cloud.spanner.Struct;
import com.google.cloud.spanner.Type;
import com.google.cloud.spanner.Type.StructField;
import com.google.cloud.spanner.connection.ConnectionOptions;
import com.google.cloud.trace.v1.TraceServiceClient;
import com.google.cloud.trace.v1.TraceServiceSettings;
import com.google.common.base.Stopwatch;
import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.Tracer;
import io.opentelemetry.context.Scope;
import java.io.IOException;
import java.util.Arrays;
import java.util.Collection;
import java.util.concurrent.TimeUnit;
import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.junit.ClassRule;
import org.junit.Test;
import org.junit.experimental.categories.Category;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;

/** Integration tests for End to End Tracing. */
@Category(ParallelIntegrationTest.class)
@RunWith(JUnit4.class)
public class ITEndToEndTracingTest {
  public static Collection<TestEnvOptions> testEnvOptions =
      Arrays.asList(TestEnvOptions.USE_END_TO_END_TRACING);
  @ClassRule public static IntegrationTestEnv env = new IntegrationTestEnv(testEnvOptions);
  private static DatabaseClient googleStandardSQLClient;

  static {
    SpannerOptionsHelper.resetActiveTracingFramework();
    SpannerOptions.enableOpenTelemetryTraces();
  }

  private static String selectValueQuery = "SELECT @p1 + @p1";

  @BeforeClass
  public static void setUp() {
    setUpDatabase();
  }

  public static void setUpDatabase() {
    // Empty database.
    Database googleStandardSQLDatabase = env.getTestHelper().createTestDatabase();
    googleStandardSQLClient = env.getTestHelper().getDatabaseClient(googleStandardSQLDatabase);
  }

  @AfterClass
  public static void teardown() {
    ConnectionOptions.closeSpanner();
  }

  private void assertTrace(String traceId) throws IOException, InterruptedException {
    TraceServiceSettings settings =
        env.getTestHelper().getOptions().getCredentials() == null
            ? TraceServiceSettings.newBuilder().build()
            : TraceServiceSettings.newBuilder()
                .setCredentialsProvider(
                    FixedCredentialsProvider.create(
                        env.getTestHelper().getOptions().getCredentials()))
                .build();
    try (TraceServiceClient client = TraceServiceClient.create(settings)) {
      boolean foundTrace = false;
      Stopwatch metricsPollingStopwatch = Stopwatch.createStarted();
      while (!foundTrace && metricsPollingStopwatch.elapsed(TimeUnit.SECONDS) < 30) {
        // Try every 5 seconds
        Thread.sleep(5000);
        try {
          foundTrace =
              client.getTrace(env.getTestHelper().getInstanceId().getProject(), traceId)
                  .getSpansList().stream()
                  .anyMatch(span -> "Spanner.ExecuteStreamingSql".equals(span.getName()));
        } catch (ApiException apiException) {
          assumeTrue(
              apiException.getStatusCode() != null
                  && StatusCode.Code.NOT_FOUND.equals(apiException.getStatusCode().getCode()));
          System.out.println("Trace NOT_FOUND error ignored");
        }
      }
      assertTrue(foundTrace);
    } catch (ResourceExhaustedException resourceExhaustedException) {
      if (resourceExhaustedException
          .getMessage()
          .contains("Quota exceeded for quota metric 'Read requests (free)'")) {
        // Ignore and allow the test to succeed.
        System.out.println("RESOURCE_EXHAUSTED error ignored");
      } else {
        throw resourceExhaustedException;
      }
    }
  }

  private Struct executeWithRowResultType(Statement statement, Type expectedRowType) {
    ResultSet resultSet = statement.executeQuery(googleStandardSQLClient.singleUse());
    assertThat(resultSet.next()).isTrue();
    assertThat(resultSet.getType()).isEqualTo(expectedRowType);
    Struct row = resultSet.getCurrentRowAsStruct();
    assertThat(resultSet.next()).isFalse();
    return row;
  }

  @Test
  public void simpleSelect() throws IOException, InterruptedException {
    Tracer tracer =
        env.getTestHelper()
            .getOptions()
            .getOpenTelemetry()
            .getTracer(ITEndToEndTracingTest.class.getName());
    Span span = tracer.spanBuilder("simpleSelect").startSpan();
    Scope scope = span.makeCurrent();
    Type rowType = Type.struct(StructField.of("", Type.int64()));
    Struct row =
        executeWithRowResultType(
            Statement.newBuilder(selectValueQuery).bind("p1").to(1234).build(), rowType);
    assertThat(row.isNull(0)).isFalse();
    assertThat(row.getLong(0)).isEqualTo(2468);
    scope.close();
    span.end();
    assertTrace(span.getSpanContext().getTraceId());
  }
}