Tag Archives: AOP

Annotation for Spring Data Source

If you’ve ever written a Spring application that uses multiple datasources, you know that things can get a bit messy trying to figure out what datasource is used where. In this post, I will share simple Annotation-based solution I’ve used in production applications to manage data sources in DAOs declaratively.

To start with, let’s assume we have a Spring application which has three data sources, oracle, mysql, and warehouse.

To start, let’s create an implementation of Spring’s AbstractRoutingDataSource.  We’ll call it simply, Database.

Here’s what it looks like:

public class Database extends AbstractRoutingDataSource
{
   private final static String DATA_SOURCE_KEY = "currentDataSource";

   public enum Source
   {
      oracle, mySql, warehouse;  //names must match names in springapp-servlet.xml
   }

   public static void setDataSource( Source dataSource )
   {
      ThreadStorage.getInstance().put( DATA_SOURCE_KEY, dataSource );
   }

   public static Source getDataSource()
   {
      return (Source) ThreadStorage.getInstance().get( DATA_SOURCE_KEY );
   }

   @Override
   protected Object determineCurrentLookupKey()
   {
      return getDataSource().name();
   }
}

Our springapp-servlet.xml would look like this:

   <bean class="com.springapp.Database" id="dataSource">
      <property name="targetDataSources">
         <map key-type="java.lang.String">
         <!--These keys must match the names in Database.Source enum-->
            <entry key="oracle" value-ref="oracle"/>
            <entry key="mySql" value-ref="mySql"/>
            <entry key="warehouse" value-ref="warehouse"/>
         </map>
      </property>
      <property name="defaultTargetDataSource" ref="oracle"/>
   </bean>

As you can see, our implementation of AbstractRoutingDataSource works by storing the name the Source enum in ThreadStorage. That way, once it is set, any other section of code within the thread will have access to that key, which will be used by Spring to select the proper Data Source.

Now let’s look at how this can be used. Let’s say we have a DAO which provides access to User data:

@Repository
public class UserDaoImpl implements UserDao
{
   @Autowired
   private SimpleJdbcTemplate simpleJdbcTemplate;

   public String getUserName( String userId ) 
   {
      String sql = "select user_name_first_last from users where user_id = :userId";
      MapSqlParameterSource source = new MapSqlParameterSource();
      source.addValue( "userId", userId);
      return simpleJdbcTemplate.queryForObject( sql, String.class, source );
   }

   public List<Purchase> getUserPurchaseHistory( String userId )
   {
      String sql = "select * from member_purchase_history where userId = :userId";
      MapSqlParameterSource source = new MapSqlParameterSource();
      source.addValue( "userId", userId );
      return simpleJdbcTemplate.query( sql, new PurchaseRowMapper(), source );
   }
}

Now let’s say that getUserName() pulls data from an oracle database, but getUserPurchaseHistory() uses a warehouse connection. We could specify that like so:

   public String getUserName( String userId )
   {
      Database.setDatabaseSource( Database.Source.oracle );
      String sql = "select user_name_first_last from users where user_id = :userId";
      MapSqlParameterSource source = new MapSqlParameterSource();
      source.addValue( "userId", userId);
      return simpleJdbcTemplate.queryForObject( sql, String.class, source );
   }

   public List<Purchase> getUserPurchaseHistory( String userId )
   {
      Database.setDatabaseSource( Database.Source.warehouse );
      String sql = "select * from member_purchase_history where userId = :userId";
      MapSqlParameterSource source = new MapSqlParameterSource();
      source.addValue( "userId", userId );
      return simpleJdbcTemplate.query( sql, new PurchaseRowMapper(), source );
   }

This works fine – each method sets its Data Source and uses it appropriately. Now let’s suppose we have an Email Service which uses our UserDao:

@Service
public class EmailServiceImpl implements EmailService
{
   @Autowired
   SimpleJdbcTemplate simpleJdbcTemplate;
   
   @Autowired 
   UserDao userDao;

   public void sendEmail( String userId, String emailTypeCode, String emailAddress )
   {
      Database.setDatabaseSource( Database.Source.mySql );
      
      String userName = userDao.getUserName( userId );
            
      MapSqlParameterSource source = new MapSqlParameterSource();
      source.addValue(  "code", emailTypeCode );
      String emailBody = simpleJdbcTemplate.queryForObject( "select email_body from emailTypes where email_code = :code", String.class, source );
      
      sendEmail( userName, emailBody, emailAddress );
   }

Do you see the problem? Since the data source is stored in the Thread, when we call our userDao.getUserName(), it has the unintended consequence of changing our data source for all subsequent connections! We could fix this by saving off the Previous Data Source and setting it back at the end of each of our methods in UserDao, like this:

   public String getUserName( String userId )
   {
      Database2.Source previousDataSource = Database2.getDatabaseSource();
      Database.setDatabaseSource( Database.Source.oracle );
      String sql = "select user_name_first_last from users where user_id = :userId";
      MapSqlParameterSource source = new MapSqlParameterSource();
      source.addValue( "userId", userId);
      return simpleJdbcTemplate.queryForObject( sql, String.class, source );
      Database2.setDatabaseSource( previousDataSource );
   }

This works, but it a little verbose and also adds several lines which a Developer must Remember (not my strong suit) to write each time he writes a DB I/O method. We can do better.

First, let’s add a little bit of functionality to our Database class:

public class Database extends AbstractRoutingDataSource
{
   private final static String DATA_SOURCE_KEY = "currentDataSource";
   private final static String DATA_PREVIOUS_SOURCE_KEY = "previousDataSource";

   public enum Source
   {
      oracle, mySql, warehouse;  //names must match names in springapp-servlet.xml
   }

   public static void setPreviousDataSource( Source dataSource )
   {
      ThreadStorage.getInstance().put( DATA_PREVIOUS_SOURCE_KEY, dataSource );
   }

   public static Source getPreviousDataSource()
   {
      return (Source) ThreadStorage.getInstance().get( DATA_PREVIOUS_SOURCE_KEY );
   }

   public static void setDataSource( Source dataSource )
   {
      ThreadStorage.getInstance().put( DATA_SOURCE_KEY, dataSource );
   }

   public static Source getDataSource()
   {
      return (Source) ThreadStorage.getInstance().get( DATA_SOURCE_KEY );
   }

   @Override
   protected Object determineCurrentLookupKey()
   {
      return getDataSource().name();
   }
}

Now comes the fun part. We create an Annotation called DataSource:

public @interface DataSource 
{
   Database.Source value();
}

Now, let’s create a Spring AOP / AspectJ Aspect to advise methods which have a @DataSource annotation:

@Component
@Aspect
public class DataSourceAOP
{
   @Before("execution(* *.*(..)) && @annotation(dataSource)")
   public void setDatabaseSourceForMethods( DataSource dataSource )
   {
      setDataSource( dataSource );
   }

   @Before("execution(* *.*(..)) && @within(dataSource)")
   public void setDatabaseSourceForMethodsWithinClass( JoinPoint joinPoint, DataSource dataSource )
   {
      //If the method has its own method-level DataSource Annotation set, then do NOT set it to the class-level default
      if( !hasMethodLevelAnnotation( joinPoint ) )
      {
         setDataSource( dataSource );
      }
   }

   private boolean hasMethodLevelAnnotation( JoinPoint joinPoint )
   {
      boolean result = false;
      Signature signature = joinPoint.getStaticPart().getSignature();
      if ( signature instanceof MethodSignature )
      {
         Method method = ((MethodSignature)signature).getMethod(); //this will pull the interface method
         Object target = joinPoint.getTarget();
         Method implementingMethod = null;
         try
         {
            //this gets the implementing method, which is where any overriding DataSource annotations should be.
            implementingMethod = target.getClass().getMethod( method.getName(), method.getParameterTypes() );  
         }
         catch( Exception e )
         {
            e.printStackTrace( System.out );
         }

         if ( implementingMethod != null && implementingMethod.isAnnotationPresent( DataSource.class ) )
         {
            result = true;
         }
      }
      return result;
   }

   @After("execution(* *.*(..)) && @within(dataSource)")
   public void resetDatabaseSourceForMethodsWithinClass( JoinPoint joinPoint,DataSource dataSource )
   {
      if ( !hasMethodLevelAnnotation( joinPoint ) )
      {
         Database.setDatabaseSource( Database.getPreviousDatabaseSource() );
      }
   }

   @After("execution(* *.*(..)) && @annotation(dataSource)")
   public void resetDatabaseSourceForMethods( DataSource dataSource )
   {
      Database.setDatabaseSource( Database.getPreviousDatabaseSource() );
   }
   
   private void setDataSource( DataSource dataSource )
   {
      Database.setPreviousDatabaseSource( Database.getDatabaseSource() );
      Database.setDatabaseSource( dataSource.value() );
   }
}

There is the magic. This code will allow you to specify a Data Source for a method declaratively, like so:

@DataSource( Database.Source.mySql )
public void doSomething()
{
   //do stuff
}

Since our Aspect fires before each method call and sets the previous data source, then fires at the end of each call and resets the data source to what it was previously, you can call methods that use one data source inside of methods that use a different one, without any unexpected consequences. Our example from before becomes this:

   @DataSource( Database.Source.mySql )
   public void sendEmail( String userId, String emailTypeCode, String emailAddress )
   {
      Database.setDatabaseSource( Database.Source.mySql );
      
      String userName = userDao.getUserName( userId );
            
      MapSqlParameterSource source = new MapSqlParameterSource();
      source.addValue(  "code", emailTypeCode );
      String emailBody = simpleJdbcTemplate.queryForObject( "select email_body from emailTypes where email_code = :code", String.class, source );
      
      sendEmail( userName, emailBody, emailAddress );
   }

Then our UserDao looks like this:

   @DataSource( Database.Source.oracle )
   public String getUserName( String userId )
   {
      String sql = "select user_name_first_last from users where user_id = :userId";
      MapSqlParameterSource source = new MapSqlParameterSource();
      source.addValue( "userId", userId);
      return simpleJdbcTemplate.queryForObject( sql, String.class, source );
   }

   @DataSource( Database.Source.warehouse )
   public List<Purchase> getUserPurchaseHistory( String userId )
   {
      String sql = "select * from member_purchase_history where userId = :userId";
      MapSqlParameterSource source = new MapSqlParameterSource();
      source.addValue( "userId", userId );
      return simpleJdbcTemplate.query( sql, new PurchaseRowMapper(), source );
   }

Much better! For just one final touch, let’s think about DAOs which almost exclusively deal with the same data source. Do we have to specify @DataSource( Database.Source.mySql ) for each method? No! DataSourceAOP also lets you use our DataSource annotation at the class level, like so:

@Repository
@DataSource( Database.Source.mySql )
public class PurchaseDaoImpl implements PurchaseDao
{
   public void purchaseSomething( String userId, String productId )
   {
   }
}

Thanks to our handy-dandy DataSource annotation and DataSourceAOP, when purchaseSomething() gets called, his data source will already by set to mySql. If, however you specify a Data Source on a method, the method-level DataSource will take precedence. In other words, you can do this:

@Repository
@DataSource( Database.Source.oracle)
public class UserDaoImpl implements UserDao
{
   @Autowired
   private SimpleJdbcTemplate simpleJdbcTemplate;

   public String getUserName( String userId ) /*This method will get mySql as its data source!*/
   {
      String sql = "select user_name_first_last from users where user_id = :userId";
      MapSqlParameterSource source = new MapSqlParameterSource();
      source.addValue( "userId", userId);
      return simpleJdbcTemplate.queryForObject( sql, String.class, source );
   }

   @DataSource( Database.Source.mySql )
   public List<Purchase> getUserPurchaseHistory( String userId )
   {
      String sql = "select * from member_purchase_history where userId = :userId";
      MapSqlParameterSource source = new MapSqlParameterSource();
      source.addValue( "userId", userId );
      return simpleJdbcTemplate.query( sql, new PurchaseRowMapper(), source );
   }
}

Before we go, I want to make a few disclaimers. Since our solution is powered by AOP, it has some of the inherent weaknesses and pitfalls associated with AOP. Most notably, if you call one of our advised methods internally (i.e. from the same class), our advice will not be executed and data source will not be set.

In other words, this works:

public class SomeOtherClass
{
   @Annotation
   UserDao userDao;
   
   public void printUserName( String userId )
   {
      System.out.println( userDao.getUserName( userId );
   }
}

but this doesn’t:

@Repository
@DataSource( Database.Source.oracle)
public class UserDaoImpl implements UserDao
{
   /* stuff from above */

   public void printUserName( String userId )
   {
      System.out.println( getUserName( userId ); /*getUserName()'s DataSource doesn't get set*/
   }

I recently found an *awesome* workaround that can let you get around the AOP-doesn’t-work-on-internal-calls issue, which I will soon be implementing in my code.

If you have any questions or suggestions, please leave a comment!