OracleConnectionMetricsSettingAspect

OracleConnectionMetricsSettingAspect.java

package my.package;

import static oracle.jdbc.OracleConnection.END_TO_END_CLIENTID_INDEX;
import static oracle.jdbc.OracleConnection.END_TO_END_MODULE_INDEX;
import static oracle.jdbc.OracleConnection.END_TO_END_STATE_INDEX_MAX;
import static org.springframework.util.StringUtils.hasText;
import oracle.jdbc.OracleConnection;

import org.aspectj.lang.annotation.AfterReturning;
import org.aspectj.lang.annotation.Aspect;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.context.SecurityContext;
import org.springframework.stereotype.Component;

import java.sql.Connection;
import java.sql.SQLException;

@Aspect
@Component
public class OracleConnectionMetricsSettingAspect {

  private static final String APPLICATION_NAME = "MY_APP";

  private SecurityContext securityContext;

  @AfterReturning(
      pointcut = "execution(* javax.sql.DataSource.getConnection(..))",
      returning = "connection")
  public void setMetrics(Connection connection) throws SQLException {

    Connection metaDataConnection = connection.getMetaData().getConnection();

    if (!(metaDataConnection instanceof OracleConnection)) {
      return;
    }

    String[] metrics = new String[END_TO_END_STATE_INDEX_MAX];

    if (hasText(getLoggedInUser())) {
      metrics[END_TO_END_CLIENTID_INDEX] = getLoggedInUser();
    }

    metrics[END_TO_END_MODULE_INDEX] = APPLICATION_NAME;

    ((OracleConnection) metaDataConnection).setEndToEndMetrics(metrics, (short) 0);
  }

  private String getLoggedInUser() {
    Authentication authentication = securityContext.getAuthentication();
    if (authentication == null) {
      return null;
    }
    return authentication.getName();
  }

  @Autowired
  public void setSecurityContext(SecurityContext securityContext) {
    this.securityContext = securityContext;
  }
}

OracleConnectionMetricsSettingAspectTest.java

package my.package;

import static org.mockito.AdditionalMatchers.aryEq;
import static org.mockito.Matchers.eq;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import my.package.Unit4MockitoTestCase;

import oracle.jdbc.OracleConnection;

import org.junit.Test;
import org.mockito.InjectMocks;
import org.mockito.Mock;
import org.springframework.jdbc.datasource.ConnectionProxy;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.context.SecurityContext;

import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.SQLException;

public class OracleConnectionMetricsSettingAspectTest extends Unit4MockitoTestCase {

  @Mock private Connection connection;
  @Mock private DatabaseMetaData dbMeta;
  @Mock private OracleConnection oraConnection;
  @Mock private ConnectionProxy otherConnection;

  @Mock private SecurityContext securityContext;
  @Mock private Authentication authentication;

  @InjectMocks
  private OracleConnectionMetricsSettingAspect aspect = new OracleConnectionMetricsSettingAspect();

  @Test
  public void doesNothingWhenNotOracleConnection() throws SQLException {
    when(connection.getMetaData()).thenReturn(dbMeta);
    when(dbMeta.getConnection()).thenReturn(otherConnection);
    aspect.setMetrics(connection);
  }

  @Test
  public void setConnectionMetricsWithoutLoggedInUser() throws SQLException {
    when(connection.getMetaData()).thenReturn(dbMeta);
    when(dbMeta.getConnection()).thenReturn(oraConnection);
    aspect.setMetrics(connection);
    verify(oraConnection).setEndToEndMetrics(aryEq(new String[] { null, null, null, "MY_APP" }), eq((short) 0));
  }

  @Test
  public void setConnectionMetricsWithLoggedInUser() throws SQLException {
    when(connection.getMetaData()).thenReturn(dbMeta);
    when(dbMeta.getConnection()).thenReturn(oraConnection);
    when(securityContext.getAuthentication()).thenReturn(authentication);
    when(authentication.getName()).thenReturn("TheName");
    aspect.setMetrics(connection);
    verify(oraConnection).setEndToEndMetrics(aryEq(new String[] { null, "TheName", null, "MY_APP" }), eq((short) 0));
  }
}