package com.hubspot.rosetta.jdbi;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.hubspot.rosetta.Rosetta;
import com.hubspot.rosetta.internal.RosettaModule;
import org.skife.jdbi.v2.DBI;
import org.skife.jdbi.v2.Handle;
import org.skife.jdbi.v2.SQLStatement;
import org.skife.jdbi.v2.StatementContext;

public class RosettaObjectMapperOverride {
  public static final String ATTRIBUTE_NAME = "_rosetta_object_mapper";

  private final ObjectMapper objectMapper;

  public RosettaObjectMapperOverride(ObjectMapper objectMapper) {
    this.objectMapper = objectMapper.copy().registerModule(new RosettaModule());
  }

  public void override(DBI dbi) {
    dbi.define(ATTRIBUTE_NAME, objectMapper);
  }

  public void override(Handle handle) {
    handle.define(ATTRIBUTE_NAME, objectMapper);
  }

  public void override(SQLStatement<?> statement) {
    statement.define(ATTRIBUTE_NAME, objectMapper);
  }

  public static ObjectMapper resolve(StatementContext context) {
    Object override = context.getAttribute(ATTRIBUTE_NAME);
    return override == null ? Rosetta.getMapper() : (ObjectMapper) override;
  }
}

@@ -32,147 +32,8 @@
 i;%0A%0A
-import com.fasterxml.jackson.databind.ObjectMapper;%0Aimport com.hubspot.rosetta.Rosetta;%0Aimport com.hubspot.rosetta.internal.RosettaModule;%0A
 impo
@@ -58,16 +58,16 @@
 v2.DBI;%0A
+
 import o
@@ -91,16 +91,47 @@
 Handle;%0A
+import org.skife.jdbi.v2.IDBI;%0A
 import o
@@ -205,16 +205,156 @@
 ntext;%0A%0A
+import com.fasterxml.jackson.databind.ObjectMapper;%0Aimport com.hubspot.rosetta.Rosetta;%0Aimport com.hubspot.rosetta.internal.RosettaModule;%0A%0A
 public c
@@ -649,32 +649,32 @@
 Module());%0A  %7D%0A%0A
-
   public void ov
@@ -687,24 +687,90 @@
 (DBI dbi) %7B%0A
+    override((IDBI) dbi);%0A  %7D%0A%0A  public void override(IDBI dbi) %7B%0A
     dbi.defi

Add RosettaObjectMapperOverride#override(IDBI)