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

Go to download
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();
        }

    }
}