content: /*
   * Copyright The OpenTelemetry Authors
   * SPDX-License-Identifier: Apache-2.0
   */
  
  package io.opentelemetry.javaagent.instrumentation.cassandra.v4_0;
  
  import static io.opentelemetry.javaagent.instrumentation.cassandra.v4_0.CassandraSingletons.instrumenter;
  
  import com.datastax.oss.driver.api.core.CqlSession;
  import com.datastax.oss.driver.api.core.DriverException;
  import com.datastax.oss.driver.api.core.cql.AsyncResultSet;
  import com.datastax.oss.driver.api.core.cql.BoundStatement;
  import com.datastax.oss.driver.api.core.cql.ExecutionInfo;
  import com.datastax.oss.driver.api.core.cql.ResultSet;
  import com.datastax.oss.driver.api.core.cql.SimpleStatement;
  import com.datastax.oss.driver.api.core.cql.Statement;
  import io.opentelemetry.context.Context;
  import io.opentelemetry.context.Scope;
  import java.lang.reflect.Proxy;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.List;
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.CompletionStage;
  import java.util.function.Supplier;
  import javax.annotation.Nullable;
  
  final class TracingCqlSession {
    private TracingCqlSession() {}
  
    static CqlSession wrapSession(CqlSession session) {
      if (session == null) {
        return null;
      }
  
      List<Class<?>> interfaces = new ArrayList<>();
      Class<?> clazz = session.getClass();
      while (clazz != Object.class) {
        interfaces.addAll(Arrays.asList(clazz.getInterfaces()));
        clazz = clazz.getSuperclass();
      }
      return (CqlSession)
          Proxy.newProxyInstance(
              session.getClass().getClassLoader(),
              interfaces.toArray(new Class<?>[0]),
              (proxy, method, args) -> {
                if ("execute".equals(method.getName()) && method.getParameterCount() == 1) {
                  if (method.getParameterTypes()[0] == String.class) {
                    String query = (String) args[0];
                    return execute(session, query);
                  }
                  if (method.getParameterTypes()[0] == Statement.class) {
                    Statement<?> statement = (Statement<?>) args[0];
                    return execute(session, statement);
                  }
                } else if ("executeAsync".equals(method.getName())
                    && method.getParameterCount() == 1) {
                  if (method.getParameterTypes()[0] == String.class) {
                    String query = (String) args[0];
                    return executeAsync(session, query);
                  }
                  if (method.getParameterTypes()[0] == Statement.class) {
                    Statement<?> statement = (Statement<?>) args[0];
                    return executeAsync(session, statement);
                  }
                }
  
                return method.invoke(session, args);
              });
    }
  
    private static ResultSet execute(CqlSession session, String query) {
      CassandraRequest request = CassandraRequest.create(session, query);
      Context context = instrumenter().start(Context.current(), request);
      ResultSet resultSet;
      try (Scope ignored = context.makeCurrent()) {
        resultSet = session.execute(query);
      } catch (Throwable exception) {
        instrumenter().end(context, request, getExecutionInfo(exception), exception);
        throw exception;
      }
      instrumenter().end(context, request, resultSet.getExecutionInfo(), null);
      return resultSet;
    }
  
    private static ResultSet execute(CqlSession session, Statement<?> statement) {
      String query = getQuery(statement);
      CassandraRequest request = CassandraRequest.create(session, query);
      Context context = instrumenter().start(Context.current(), request);
      ResultSet resultSet;
      try (Scope ignored = context.makeCurrent()) {
        resultSet = session.execute(statement);
      } catch (Throwable exception) {
        instrumenter().end(context, request, getExecutionInfo(exception), exception);
        throw exception;
      }
      instrumenter().end(context, request, resultSet.getExecutionInfo(), null);
      return resultSet;
    }
  
    private static CompletionStage<AsyncResultSet> executeAsync(
        CqlSession session, Statement<?> statement) {
      String query = getQuery(statement);
      CassandraRequest request = CassandraRequest.create(session, query);
      return executeAsync(request, () -> session.executeAsync(statement));
    }
  
    private static CompletionStage<AsyncResultSet> executeAsync(CqlSession session, String query) {
      CassandraRequest request = CassandraRequest.create(session, query);
      return executeAsync(request, () -> session.executeAsync(query));
    }
  
    private static CompletionStage<AsyncResultSet> executeAsync(
        CassandraRequest request, Supplier<CompletionStage<AsyncResultSet>> query) {
      Context parentContext = Context.current();
      Context context = instrumenter().start(parentContext, request);
      try (Scope ignored = context.makeCurrent()) {
        CompletionStage<AsyncResultSet> stage = query.get();
        return wrap(
            stage.whenComplete(
                (asyncResultSet, throwable) ->
                    instrumenter()
                        .end(
                            context,
                            request,
                            getExecutionInfo(asyncResultSet, throwable),
                            throwable)),
            parentContext);
      }
    }
  
    private static <T> CompletableFuture<T> wrap(CompletionStage<T> future, Context context) {
      CompletableFuture<T> result = new CompletableFuture<>();
      future.whenComplete(
          (T value, Throwable throwable) -> {
            try (Scope ignored = context.makeCurrent()) {
              if (throwable != null) {
                result.completeExceptionally(throwable);
              } else {
                result.complete(value);
              }
            }
          });
  
      return result;
    }
  
    private static String getQuery(Statement<?> statement) {
      String query = null;
      if (statement instanceof SimpleStatement) {
        query = ((SimpleStatement) statement).getQuery();
      } else if (statement instanceof BoundStatement) {
        query = ((BoundStatement) statement).getPreparedStatement().getQuery();
      }
  
      return query == null ? "" : query;
    }
  
    private static ExecutionInfo getExecutionInfo(
        @Nullable AsyncResultSet asyncResultSet, @Nullable Throwable throwable) {
      if (asyncResultSet != null) {
        return asyncResultSet.getExecutionInfo();
      } else {
        return getExecutionInfo(throwable);
      }
    }
  
    private static ExecutionInfo getExecutionInfo(@Nullable Throwable throwable) {
      if (throwable instanceof DriverException) {
        return ((DriverException) throwable).getExecutionInfo();
      } else if (throwable != null && throwable.getCause() instanceof DriverException) {
        // TODO (trask) find out if this is needed and if so add comment explaining
        return ((DriverException) throwable.getCause()).getExecutionInfo();
      } else {
        return null;
      }
    }
  }
  