com.undefinedlabs.scope.rules.StatementScopeAgentRule Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of scope-rule-java-sql Show documentation
Scope is a APM for tests to give engineering teams unprecedented visibility into their CI process to quickly identify, troubleshoot and fix failed builds. This artifact contains the classes to instrument the Java SQL package.
There is a newer version: 0.15.1-beta.2
Show newest version
package com.undefinedlabs.scope.rules;

import com.undefinedlabs.scope.ScopeGlobalTracer;
import com.undefinedlabs.scope.events.EventFieldsFactory;
import com.undefinedlabs.scope.events.exception.ThrowableEvent;
import com.undefinedlabs.scope.rules.sql.model.ConnectionInfo;
import com.undefinedlabs.scope.rules.sql.model.PreparedStatementQuery;
import com.undefinedlabs.scope.rules.sql.model.PreparedStatementQueryParameter;
import com.undefinedlabs.scope.rules.sql.provider.ConnectionInfoProviderRegistry;
import com.undefinedlabs.scope.rules.sql.provider.PreparedStatementQueryProviderRegistry;
import com.undefinedlabs.scope.rules.sql.provider.internal.PreparedStatementQueryUtils;
import com.undefinedlabs.scope.rules.transformer.ScopeAgentAdvicedTransformer;
import com.undefinedlabs.scope.settings.ScopeSettings;
import com.undefinedlabs.scope.settings.ScopeSettingsResolver;
import com.undefinedlabs.scope.utils.SpanUtils;
import com.undefinedlabs.scope.utils.event.EventValues;
import com.undefinedlabs.scope.utils.sourcecode.ExceptionSourceCodeFactory;
import com.undefinedlabs.scope.utils.sourcecode.ExceptionSourceCodeFrame;
import com.undefinedlabs.scope.utils.tag.TagKeys;
import com.undefinedlabs.scope.utils.tag.TagValues;
import io.opentracing.Span;
import io.opentracing.Tracer;
import io.opentracing.tag.Tags;
import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.asm.Advice;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.Collections;
import java.util.HashMap;

import static net.bytebuddy.matcher.ElementMatchers.*;

public class StatementScopeAgentRule extends AbstractScopeAgentRule {

    @Override
    protected Iterable transformers() {
        final Class instrumentedClass = lookUp("java.sql.Statement");
        return (instrumentedClass != null) ?
                Collections.singleton(newAgentBuilder()
                        .type(isSubTypeOf(instrumentedClass)).transform(new ScopeAgentAdvicedTransformer(StatementExecuteNativeAdvice.class, named("executeQuery").and(takesArgument(0, String.class))))
                        .type(isSubTypeOf(instrumentedClass)).transform(new ScopeAgentAdvicedTransformer(StatementExecuteNativeAdvice.class, named("executeUpdate").and(takesArgument(0, String.class))))
                        .type(isSubTypeOf(instrumentedClass)).transform(new ScopeAgentAdvicedTransformer(PreparedStatementExecuteAdvice.class, named("execute").and(takesArguments(0))))
                        .type(isSubTypeOf(instrumentedClass)).transform(new ScopeAgentAdvicedTransformer(PreparedStatementExecuteAdvice.class, named("executeQuery").and(takesArguments(0))))
                        .type(isSubTypeOf(instrumentedClass)).transform(new ScopeAgentAdvicedTransformer(PreparedStatementExecuteAdvice.class, named("executeUpdate").and(takesArguments(0)))))
                : emptyCollection();
    }

    public static class StatementExecuteNativeAdvice {

        @Advice.OnMethodEnter
        public static Span enter(@Advice.This Object thiz, @Advice.AllArguments Object[] args) {
            if (!(thiz instanceof Statement) && args.length == 0) {
                return null;
            }

            final Tracer tracer = ScopeGlobalTracer.get();
            final Span previousSpan = tracer.activeSpan();
            if (previousSpan == null) {
                return null;
            }

            try {
                final String sql = (String) args[0];

                final Statement statement = (Statement) thiz;
                final Connection connection = statement.getConnection();
                final ConnectionInfo connectionInfo = ConnectionInfoProviderRegistry.INSTANCE.getProvider(connection).extractInfo(connection);

                final Span span = tracer.buildSpan(connectionInfo.getProductName() + ":" + PreparedStatementQueryUtils.INSTANCE.extractSqlMethod(sql))
                        .withTag(TagKeys.COMPONENT, TagValues.Component.SQL)
                        .withTag(TagKeys.SPAN_KIND, Tags.SPAN_KIND_CLIENT)
                        .withTag(TagKeys.DB.DB_CONNECTION, connectionInfo.getUrl())
                        .withTag(TagKeys.DB.DB_TYPE, TagValues.DB.TYPE_SQL)
                        .withTag(TagKeys.DB.DB_USER, connectionInfo.getUserName())
                        .withTag(TagKeys.DB.DB_INSTANCE, connection.getCatalog())
                        .withTag(TagKeys.DB.DB_PRODUCT_NAME, connectionInfo.getProductName())
                        .withTag(TagKeys.DB.DB_PRODUCT_VERSION, connectionInfo.getProductVersion())
                        .withTag(TagKeys.DB.DB_DRIVER_NAME, connectionInfo.getDriverName())
                        .withTag(TagKeys.DB.DB_DRIVER_VERSION, connectionInfo.getDriverVersion())
                        .withTag(TagKeys.Network.PEER_SERVICE, connectionInfo.getPeerService())
                        .start();

                if ((boolean) ScopeSettingsResolver.INSTANCE.get().getSetting(ScopeSettings.SCOPE_INSTRUMENTATION_DB_STATEMENT_VALUES)) {
                    span.setTag(TagKeys.DB.DB_STATEMENT, sql);
                    SpanUtils.INSTANCE.setTagObject(span, TagKeys.DB.DB_PARAMS, Collections.unmodifiableMap(new HashMap()));
                } else {
                    span.setTag(TagKeys.DB.DB_STATEMENT_UNAVAILABLE, TagValues.DB.DB_STATEMENT_DISABLED);
                }


                return span;

            } catch (Exception e) {
                throw new RuntimeException(e);
            }


        }

        @Advice.OnMethodExit(onThrowable = SQLException.class)
        public static void exit(@Advice.This Object thiz, @Advice.Enter Object spanObj, @Advice.Thrown Throwable throwable) {
            if (!(thiz instanceof Statement) || spanObj == null) {
                return;
            }

            final Span span = (Span) spanObj;

            if (throwable != null) {
                span.setTag(TagKeys.ERROR, true);

                final ExceptionSourceCodeFrame exceptionSourceCodeFrame = ExceptionSourceCodeFactory.INSTANCE.createFrame(throwable);
                final ThrowableEvent.Builder throwableEventBuilder = ThrowableEvent.newBuilder();
                throwableEventBuilder
                        .withEventType(EventValues.General.ERROR)
                        .withThrowable(exceptionSourceCodeFrame.getUserThrowable())
                        .withSource(exceptionSourceCodeFrame.getSourceCodeFrame().getLinkPathWithMethodLine());

                span.log(EventFieldsFactory.INSTANCE.createFields(throwableEventBuilder.build()));
            }

            span.finish();
        }

    }


    public static class PreparedStatementExecuteAdvice {

        @Advice.OnMethodEnter
        public static Span enter(@Advice.This Object thiz) {
            if (!(thiz instanceof PreparedStatement)) {
                return null;
            }

            final Tracer tracer = ScopeGlobalTracer.get();
            final Span previousSpan = tracer.activeSpan();
            if (previousSpan == null) {
                return null;
            }

            try {
                final PreparedStatement preparedStatement = (PreparedStatement) thiz;
                final Connection connection = preparedStatement.getConnection();
                final ConnectionInfo connectionInfo = ConnectionInfoProviderRegistry.INSTANCE.getProvider(connection).extractInfo(connection);
                final PreparedStatementQuery preparedStatementQuery = PreparedStatementQueryProviderRegistry.INSTANCE.getProvider(preparedStatement).create(preparedStatement);

                if (preparedStatementQuery.equals(PreparedStatementQuery.EMPTY)) {
                    return null;
                }

                final Span span = tracer.buildSpan(connectionInfo.getProductName() + ":" + preparedStatementQuery.getSqlMethod())
                        .withTag(TagKeys.COMPONENT, TagValues.Component.SQL)
                        .withTag(TagKeys.SPAN_KIND, Tags.SPAN_KIND_CLIENT)
                        .withTag(TagKeys.DB.DB_CONNECTION, connectionInfo.getUrl())
                        .withTag(TagKeys.DB.DB_TYPE, TagValues.DB.TYPE_SQL)
                        .withTag(TagKeys.DB.DB_USER, connectionInfo.getUserName())
                        .withTag(TagKeys.DB.DB_INSTANCE, connection.getCatalog())
                        .withTag(TagKeys.DB.DB_PREPARED_STATEMENT, preparedStatementQuery.getSqlPreparedStatement())
                        .withTag(TagKeys.DB.DB_PRODUCT_NAME, connectionInfo.getProductName())
                        .withTag(TagKeys.DB.DB_PRODUCT_VERSION, connectionInfo.getProductVersion())
                        .withTag(TagKeys.DB.DB_DRIVER_NAME, connectionInfo.getDriverName())
                        .withTag(TagKeys.DB.DB_DRIVER_VERSION, connectionInfo.getDriverVersion())
                        .withTag(TagKeys.Network.PEER_SERVICE, connectionInfo.getPeerService())
                        .start();


                if ((boolean) ScopeSettingsResolver.INSTANCE.get().getSetting(ScopeSettings.SCOPE_INSTRUMENTATION_DB_STATEMENT_VALUES)) {
                    span.setTag(TagKeys.DB.DB_STATEMENT, preparedStatementQuery.getSqlStatement());
                    SpanUtils.INSTANCE.setTagObject(span, TagKeys.DB.DB_PARAMS, preparedStatementQuery.getSqlParameterMap());
                } else {
                    span.setTag(TagKeys.DB.DB_STATEMENT_UNAVAILABLE, TagValues.DB.DB_STATEMENT_DISABLED);
                }

                return span;

            } catch (Exception e) {
                throw new RuntimeException(e);
            }

        }

        @Advice.OnMethodExit(onThrowable = SQLException.class)
        public static void exit(@Advice.This Object thiz, @Advice.Enter Object spanObj, @Advice.Thrown Throwable throwable) {
            if (!(thiz instanceof PreparedStatement) || spanObj == null) {
                return;
            }

            final Span span = (Span) spanObj;

            if (throwable != null) {
                span.setTag(TagKeys.ERROR, true);

                final ExceptionSourceCodeFrame exceptionSourceCodeFrame = ExceptionSourceCodeFactory.INSTANCE.createFrame(throwable);
                final ThrowableEvent.Builder throwableEventBuilder = ThrowableEvent.newBuilder();
                throwableEventBuilder
                        .withEventType(EventValues.General.ERROR)
                        .withThrowable(exceptionSourceCodeFrame.getUserThrowable())
                        .withSource(exceptionSourceCodeFrame.getSourceCodeFrame().getLinkPathWithMethodLine());

                span.log(EventFieldsFactory.INSTANCE.createFields(throwableEventBuilder.build()));
            }

            span.finish();
        }

    }
}