killbill-memoizeit

Changes

Details

diff --git a/api/src/main/java/com/ning/billing/junction/api/Blockable.java b/api/src/main/java/com/ning/billing/junction/api/Blockable.java
index cb5b1a0..daab418 100644
--- a/api/src/main/java/com/ning/billing/junction/api/Blockable.java
+++ b/api/src/main/java/com/ning/billing/junction/api/Blockable.java
@@ -21,6 +21,7 @@ import java.util.UUID;
 import com.ning.billing.account.api.Account;
 import com.ning.billing.entitlement.api.user.Subscription;
 import com.ning.billing.entitlement.api.user.SubscriptionBundle;
+import com.ning.billing.util.dao.ObjectType;
 
 public interface Blockable {
 
@@ -51,6 +52,23 @@ public interface Blockable {
             throw new IllegalStateException("Unsupported type of blockable " + type);
         }
 
+        public static ObjectType getObjectType(final Blockable o) {
+            final Type type = get(o);
+            return getObjectType(type);
+        }
+
+        public static ObjectType getObjectType(final Type type) {
+            switch (type) {
+                case ACCOUNT:
+                    return ObjectType.ACCOUNT;
+                case SUBSCRIPTION_BUNDLE:
+                    return ObjectType.BUNDLE;
+                case SUBSCRIPTION:
+                    return ObjectType.SUBSCRIPTION;
+                default:
+                    throw new IllegalStateException("Unsupported type of blockable " + type);
+            }
+        }
     }
 
     public UUID getId();
diff --git a/invoice/src/main/java/com/ning/billing/invoice/api/invoice/DefaultInvoicePaymentApi.java b/invoice/src/main/java/com/ning/billing/invoice/api/invoice/DefaultInvoicePaymentApi.java
index 2414a36..88b2425 100644
--- a/invoice/src/main/java/com/ning/billing/invoice/api/invoice/DefaultInvoicePaymentApi.java
+++ b/invoice/src/main/java/com/ning/billing/invoice/api/invoice/DefaultInvoicePaymentApi.java
@@ -35,7 +35,6 @@ import com.ning.billing.invoice.api.InvoicePaymentApi;
 import com.ning.billing.invoice.dao.InvoiceDao;
 import com.ning.billing.invoice.model.DefaultInvoicePayment;
 import com.ning.billing.util.callcontext.CallContext;
-import com.ning.billing.util.callcontext.InternalCallContext;
 import com.ning.billing.util.callcontext.InternalCallContextFactory;
 import com.ning.billing.util.callcontext.InternalTenantContext;
 import com.ning.billing.util.callcontext.TenantContext;
@@ -107,7 +106,7 @@ public class DefaultInvoicePaymentApi implements InvoicePaymentApi {
         final InvoicePayment invoicePayment = new DefaultInvoicePayment(InvoicePaymentType.ATTEMPT, paymentId, invoiceId, paymentDate, amount, currency);
 
         // Retrieve the account id for the internal call context
-        final Invoice invoice = dao.getById(invoiceId, internalCallContextFactory.createInternalCallContext(context));
+        final Invoice invoice = dao.getById(invoiceId, internalCallContextFactory.createInternalTenantContext(context));
         final UUID accountId = invoice.getAccountId();
 
         dao.notifyOfPayment(invoicePayment, internalCallContextFactory.createInternalCallContext(accountId, context));
@@ -152,7 +151,7 @@ public class DefaultInvoicePaymentApi implements InvoicePaymentApi {
                 }
 
                 // Retrieve the account id for the internal call context
-                final InternalCallContext internalCallContextNoAccountId = internalCallContextFactory.createInternalCallContext(context);
+                final InternalTenantContext internalCallContextNoAccountId = internalCallContextFactory.createInternalTenantContext(context);
                 final List<InvoicePayment> invoicePayments = dao.getInvoicePayments(paymentId, internalCallContextNoAccountId);
                 final UUID accountId = dao.getAccountIdFromInvoicePaymentId(invoicePayments.get(0).getId(), internalCallContextNoAccountId);
 
@@ -173,7 +172,7 @@ public class DefaultInvoicePaymentApi implements InvoicePaymentApi {
             @Override
             public InvoicePayment doHandle() throws InvoiceApiException {
                 // Retrieve the account id for the internal call context
-                final UUID accountId = dao.getAccountIdFromInvoicePaymentId(invoicePaymentId, internalCallContextFactory.createInternalCallContext(context));
+                final UUID accountId = dao.getAccountIdFromInvoicePaymentId(invoicePaymentId, internalCallContextFactory.createInternalTenantContext(context));
                 return dao.postChargeback(invoicePaymentId, amount, internalCallContextFactory.createInternalCallContext(accountId, context));
             }
         });
diff --git a/invoice/src/main/java/com/ning/billing/invoice/api/user/DefaultInvoiceUserApi.java b/invoice/src/main/java/com/ning/billing/invoice/api/user/DefaultInvoiceUserApi.java
index f1c2ed5..759c867 100644
--- a/invoice/src/main/java/com/ning/billing/invoice/api/user/DefaultInvoiceUserApi.java
+++ b/invoice/src/main/java/com/ning/billing/invoice/api/user/DefaultInvoiceUserApi.java
@@ -26,6 +26,8 @@ import javax.annotation.Nullable;
 
 import org.joda.time.DateTime;
 import org.joda.time.LocalDate;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import com.ning.billing.ErrorCode;
 import com.ning.billing.account.api.Account;
@@ -42,7 +44,10 @@ import com.ning.billing.invoice.model.CreditAdjInvoiceItem;
 import com.ning.billing.invoice.model.ExternalChargeInvoiceItem;
 import com.ning.billing.invoice.template.HtmlInvoiceGenerator;
 import com.ning.billing.util.api.TagApiException;
+import com.ning.billing.util.bus.Bus;
+import com.ning.billing.util.bus.Bus.EventBusException;
 import com.ning.billing.util.callcontext.CallContext;
+import com.ning.billing.util.callcontext.InternalCallContext;
 import com.ning.billing.util.callcontext.InternalCallContextFactory;
 import com.ning.billing.util.callcontext.InternalTenantContext;
 import com.ning.billing.util.callcontext.TenantContext;
@@ -56,22 +61,26 @@ import com.google.inject.Inject;
 
 public class DefaultInvoiceUserApi implements InvoiceUserApi {
 
+    private static final Logger log = LoggerFactory.getLogger(DefaultInvoiceUserApi.class);
+
     private final InvoiceDao dao;
     private final InvoiceDispatcher dispatcher;
     private final AccountInternalApi accountUserApi;
-    private final TagInternalApi tagUserApi;
+    private final TagInternalApi tagApi;
     private final HtmlInvoiceGenerator generator;
     private final InternalCallContextFactory internalCallContextFactory;
+    private final Bus eventBus;
 
     @Inject
-    public DefaultInvoiceUserApi(final InvoiceDao dao, final InvoiceDispatcher dispatcher, final AccountInternalApi accountUserApi,
-                                 final TagInternalApi tagUserApi, final HtmlInvoiceGenerator generator, final InternalCallContextFactory internalCallContextFactory) {
+    public DefaultInvoiceUserApi(final InvoiceDao dao, final InvoiceDispatcher dispatcher, final AccountInternalApi accountUserApi,  final Bus eventBus,
+                                 final TagInternalApi tagApi, final HtmlInvoiceGenerator generator, final InternalCallContextFactory internalCallContextFactory) {
         this.dao = dao;
         this.dispatcher = dispatcher;
         this.accountUserApi = accountUserApi;
-        this.tagUserApi = tagUserApi;
+        this.tagApi = tagApi;
         this.generator = generator;
         this.internalCallContextFactory = internalCallContextFactory;
+        this.eventBus = eventBus;
     }
 
     @Override
@@ -87,7 +96,7 @@ public class DefaultInvoiceUserApi implements InvoiceUserApi {
     @Override
     public void notifyOfPayment(final InvoicePayment invoicePayment, final CallContext context) throws InvoiceApiException {
         // Retrieve the account id for the internal call context
-        final UUID accountId = dao.getAccountIdFromInvoicePaymentId(invoicePayment.getId(), internalCallContextFactory.createInternalCallContext(context));
+        final UUID accountId = dao.getAccountIdFromInvoicePaymentId(invoicePayment.getId(), internalCallContextFactory.createInternalTenantContext(context));
         dao.notifyOfPayment(invoicePayment, internalCallContextFactory.createInternalCallContext(accountId, context));
     }
 
@@ -139,16 +148,26 @@ public class DefaultInvoiceUserApi implements InvoiceUserApi {
 
     @Override
     public void tagInvoiceAsWrittenOff(final UUID invoiceId, final CallContext context) throws TagApiException, InvoiceApiException {
-        // Retrieve the invoice for the internal call context
-        final Invoice invoice = dao.getById(invoiceId, internalCallContextFactory.createInternalCallContext(context));
-        dao.setWrittenOff(invoiceId, internalCallContextFactory.createInternalCallContext(invoice.getAccountId(), context));
+        // Note: the tagApi is audited
+        final InternalCallContext internalContext = internalCallContextFactory.createInternalCallContext(context);
+        tagApi.addTag(invoiceId, ObjectType.INVOICE, ControlTagType.WRITTEN_OFF.getId(), internalContext);
+
+        // Retrieve the invoice for the account id
+        final Invoice invoice = dao.getById(invoiceId, internalContext);
+        // This is for overdue
+        notifyBusOfInvoiceAdjustment(invoiceId, invoice.getAccountId(), context.getUserToken());
     }
 
     @Override
     public void tagInvoiceAsNotWrittenOff(final UUID invoiceId, final CallContext context) throws TagApiException, InvoiceApiException {
-        // Retrieve the invoice for the internal call context
-        final Invoice invoice = dao.getById(invoiceId, internalCallContextFactory.createInternalCallContext(context));
-        dao.removeWrittenOff(invoiceId, internalCallContextFactory.createInternalCallContext(invoice.getAccountId(), context));
+        // Note: the tagApi is audited
+        final InternalCallContext internalContext = internalCallContextFactory.createInternalCallContext(context);
+        tagApi.removeTag(invoiceId, ObjectType.INVOICE, ControlTagType.WRITTEN_OFF.getId(), internalContext);
+
+        // Retrieve the invoice for the account id
+        final Invoice invoice = dao.getById(invoiceId, internalContext);
+        // This is for overdue
+        notifyBusOfInvoiceAdjustment(invoiceId, invoice.getAccountId(), context.getUserToken());
     }
 
     @Override
@@ -253,7 +272,7 @@ public class DefaultInvoiceUserApi implements InvoiceUserApi {
 
         // Check if this account has the MANUAL_PAY system tag
         boolean manualPay = false;
-        final Map<String, Tag> accountTags = tagUserApi.getTags(account.getId(), ObjectType.ACCOUNT, internalContext);
+        final Map<String, Tag> accountTags = tagApi.getTags(account.getId(), ObjectType.ACCOUNT, internalContext);
         for (final Tag tag : accountTags.values()) {
             if (ControlTagType.MANUAL_PAY.getId().equals(tag.getTagDefinitionId())) {
                 manualPay = true;
@@ -264,4 +283,11 @@ public class DefaultInvoiceUserApi implements InvoiceUserApi {
         return generator.generateInvoice(account, invoice, manualPay);
     }
 
+    private void notifyBusOfInvoiceAdjustment(final UUID invoiceId, final UUID accountId, final UUID userToken) {
+        try {
+            eventBus.post(new DefaultInvoiceAdjustmentEvent(invoiceId, accountId, userToken));
+        } catch (EventBusException e) {
+            log.warn("Failed to post adjustment event for invoice " + invoiceId, e);
+        }
+    }
 }
diff --git a/invoice/src/main/java/com/ning/billing/invoice/dao/AuditedInvoiceDao.java b/invoice/src/main/java/com/ning/billing/invoice/dao/AuditedInvoiceDao.java
index 524ec45..61bc576 100644
--- a/invoice/src/main/java/com/ning/billing/invoice/dao/AuditedInvoiceDao.java
+++ b/invoice/src/main/java/com/ning/billing/invoice/dao/AuditedInvoiceDao.java
@@ -55,17 +55,13 @@ import com.ning.billing.invoice.model.RecurringInvoiceItem;
 import com.ning.billing.invoice.model.RefundAdjInvoiceItem;
 import com.ning.billing.invoice.notification.NextBillingDatePoster;
 import com.ning.billing.util.ChangeType;
-import com.ning.billing.util.api.TagApiException;
 import com.ning.billing.util.bus.Bus;
 import com.ning.billing.util.bus.Bus.EventBusException;
 import com.ning.billing.util.callcontext.InternalCallContext;
 import com.ning.billing.util.callcontext.InternalTenantContext;
 import com.ning.billing.util.clock.Clock;
 import com.ning.billing.util.dao.EntityAudit;
-import com.ning.billing.util.dao.ObjectType;
 import com.ning.billing.util.dao.TableName;
-import com.ning.billing.util.svcapi.tag.TagInternalApi;
-import com.ning.billing.util.tag.ControlTagType;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Objects;
@@ -80,7 +76,6 @@ public class AuditedInvoiceDao implements InvoiceDao {
 
     private final InvoiceSqlDao invoiceSqlDao;
     private final InvoicePaymentSqlDao invoicePaymentSqlDao;
-    private final TagInternalApi tagInternalApi;
     private final NextBillingDatePoster nextBillingDatePoster;
     private final InvoiceItemSqlDao invoiceItemSqlDao;
     private final Clock clock;
@@ -89,14 +84,12 @@ public class AuditedInvoiceDao implements InvoiceDao {
     @Inject
     public AuditedInvoiceDao(final IDBI dbi,
             final NextBillingDatePoster nextBillingDatePoster,
-            final TagInternalApi tagInternalApi,
             final Clock clock,
             final Bus eventBus) {
         this.invoiceSqlDao = dbi.onDemand(InvoiceSqlDao.class);
         this.invoicePaymentSqlDao = dbi.onDemand(InvoicePaymentSqlDao.class);
         this.invoiceItemSqlDao = dbi.onDemand(InvoiceItemSqlDao.class);
         this.nextBillingDatePoster = nextBillingDatePoster;
-        this.tagInternalApi = tagInternalApi;
         this.clock = clock;
         this.eventBus = eventBus;
     }
@@ -326,43 +319,7 @@ public class AuditedInvoiceDao implements InvoiceDao {
     }
 
     @Override
-    public void setWrittenOff(final UUID invoiceId, final InternalCallContext context) throws TagApiException {
 
-        tagInternalApi.addTag(invoiceId, ObjectType.INVOICE, ControlTagType.WRITTEN_OFF.getId(), context);
-
-        invoiceSqlDao.inTransaction(new Transaction<Void, InvoiceSqlDao>() {
-            @Override
-            public Void inTransaction(final InvoiceSqlDao transactional, final TransactionStatus status) throws Exception {
-
-
-                final Invoice invoice = transactional.getById(invoiceId.toString(), context);
-                notifyBusOfInvoiceAdjustment(transactional, invoiceId, invoice.getAccountId(), context.getUserToken());
-
-                return null;
-            }
-        });
-    }
-
-    @Override
-    public void removeWrittenOff(final UUID invoiceId, final InternalCallContext context) throws TagApiException {
-
-        // Note: the tagInternalApi is audited
-        tagInternalApi.removeTag(invoiceId, ObjectType.INVOICE, ControlTagType.WRITTEN_OFF.getId(), context);
-
-
-        invoiceSqlDao.inTransaction(new Transaction<Void, InvoiceSqlDao>() {
-            @Override
-            public Void inTransaction(final InvoiceSqlDao transactional, final TransactionStatus status) throws Exception {
-
-                final Invoice invoice = transactional.getById(invoiceId.toString(), context);
-                notifyBusOfInvoiceAdjustment(transactional, invoiceId, invoice.getAccountId(), context.getUserToken());
-
-                return null;
-            }
-        });
-    }
-
-    @Override
     public InvoicePayment createRefund(final UUID paymentId, final BigDecimal requestedRefundAmount, final boolean isInvoiceAdjusted,
                                        final Map<UUID, BigDecimal> invoiceItemIdsWithNullAmounts, final UUID paymentCookieId,
                                        final InternalCallContext context)
diff --git a/invoice/src/main/java/com/ning/billing/invoice/dao/InvoiceDao.java b/invoice/src/main/java/com/ning/billing/invoice/dao/InvoiceDao.java
index a97492e..f9e9b4d 100644
--- a/invoice/src/main/java/com/ning/billing/invoice/dao/InvoiceDao.java
+++ b/invoice/src/main/java/com/ning/billing/invoice/dao/InvoiceDao.java
@@ -30,7 +30,6 @@ import com.ning.billing.invoice.api.Invoice;
 import com.ning.billing.invoice.api.InvoiceApiException;
 import com.ning.billing.invoice.api.InvoiceItem;
 import com.ning.billing.invoice.api.InvoicePayment;
-import com.ning.billing.util.api.TagApiException;
 import com.ning.billing.util.callcontext.InternalCallContext;
 import com.ning.billing.util.callcontext.InternalTenantContext;
 
@@ -66,10 +65,6 @@ public interface InvoiceDao {
 
     List<Invoice> getAllInvoicesByAccount(UUID accountId, InternalTenantContext context);
 
-    void setWrittenOff(UUID invoiceId, InternalCallContext context) throws TagApiException;
-
-    void removeWrittenOff(UUID invoiceId, InternalCallContext context) throws TagApiException;
-
     InvoicePayment postChargeback(UUID invoicePaymentId, BigDecimal amount, InternalCallContext context) throws InvoiceApiException;
 
     /**
diff --git a/invoice/src/test/java/com/ning/billing/invoice/api/invoice/TestDefaultInvoicePaymentApi.java b/invoice/src/test/java/com/ning/billing/invoice/api/invoice/TestDefaultInvoicePaymentApi.java
index bc5c182..0fb1ae3 100644
--- a/invoice/src/test/java/com/ning/billing/invoice/api/invoice/TestDefaultInvoicePaymentApi.java
+++ b/invoice/src/test/java/com/ning/billing/invoice/api/invoice/TestDefaultInvoicePaymentApi.java
@@ -51,11 +51,6 @@ import com.ning.billing.util.bus.Bus;
 import com.ning.billing.util.callcontext.InternalCallContextFactory;
 import com.ning.billing.util.clock.Clock;
 import com.ning.billing.util.clock.ClockMock;
-import com.ning.billing.util.svcapi.tag.TagInternalApi;
-import com.ning.billing.util.tag.dao.MockTagDao;
-import com.ning.billing.util.tag.dao.MockTagDefinitionDao;
-import com.ning.billing.util.tag.dao.TagDao;
-import com.ning.billing.util.tag.dao.TagDefinitionDao;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -84,12 +79,8 @@ public class TestDefaultInvoicePaymentApi extends InvoiceTestSuiteWithEmbeddedDB
         invoiceItemSqlDao.test(internalCallContext);
 
         final NextBillingDatePoster nextBillingDatePoster = new MockNextBillingDatePoster();
-        final TagDefinitionDao tagDefinitionDao = new MockTagDefinitionDao();
-        final TagDao tagDao = new MockTagDao();
         internalCallContextFactory = new InternalCallContextFactory(dbi, clock);
-        // API_FIX
-        final TagInternalApi tagUserApi = null ; // new DefaultTagInternalApi(internalCallContextFactory, tagDefinitionDao, tagDao);
-        final InvoiceDao invoiceDao = new AuditedInvoiceDao(dbi, nextBillingDatePoster, tagUserApi, clock, Mockito.mock(Bus.class));
+        final InvoiceDao invoiceDao = new AuditedInvoiceDao(dbi, nextBillingDatePoster, clock, Mockito.mock(Bus.class));
         invoicePaymentApi = new DefaultInvoicePaymentApi(invoiceDao, internalCallContextFactory);
     }
 
diff --git a/invoice/src/test/java/com/ning/billing/invoice/api/migration/InvoiceApiTestBase.java b/invoice/src/test/java/com/ning/billing/invoice/api/migration/InvoiceApiTestBase.java
index 477a6ab..76db6d5 100644
--- a/invoice/src/test/java/com/ning/billing/invoice/api/migration/InvoiceApiTestBase.java
+++ b/invoice/src/test/java/com/ning/billing/invoice/api/migration/InvoiceApiTestBase.java
@@ -49,6 +49,7 @@ import com.ning.billing.invoice.generator.InvoiceGenerator;
 import com.ning.billing.invoice.notification.NullInvoiceNotifier;
 import com.ning.billing.invoice.tests.InvoicingTestBase;
 import com.ning.billing.mock.api.MockBillCycleDay;
+import com.ning.billing.util.api.TagUserApi;
 import com.ning.billing.util.bus.BusService;
 import com.ning.billing.util.bus.DefaultBusService;
 import com.ning.billing.util.callcontext.InternalCallContextFactory;
@@ -91,6 +92,9 @@ public abstract class InvoiceApiTestBase extends InvoicingTestBase {
     protected InvoiceDao invoiceDao;
 
     @Inject
+    protected TagUserApi tagUserApi;
+
+    @Inject
     protected GlobalLocker locker;
 
     @Inject
diff --git a/invoice/src/test/java/com/ning/billing/invoice/api/user/TestDefaultInvoiceUserApi.java b/invoice/src/test/java/com/ning/billing/invoice/api/user/TestDefaultInvoiceUserApi.java
index 7fa510f..cf256f0 100644
--- a/invoice/src/test/java/com/ning/billing/invoice/api/user/TestDefaultInvoiceUserApi.java
+++ b/invoice/src/test/java/com/ning/billing/invoice/api/user/TestDefaultInvoiceUserApi.java
@@ -17,6 +17,7 @@
 package com.ning.billing.invoice.api.user;
 
 import java.math.BigDecimal;
+import java.util.Map;
 import java.util.UUID;
 
 import javax.annotation.Nullable;
@@ -33,11 +34,17 @@ import com.ning.billing.invoice.api.InvoiceItem;
 import com.ning.billing.invoice.api.InvoiceItemType;
 import com.ning.billing.invoice.api.migration.InvoiceApiTestBase;
 import com.ning.billing.invoice.model.InvoicingConfiguration;
+import com.ning.billing.util.api.TagApiException;
 import com.ning.billing.util.callcontext.CallContext;
 import com.ning.billing.util.callcontext.CallOrigin;
 import com.ning.billing.util.callcontext.DefaultCallContextFactory;
 import com.ning.billing.util.callcontext.TenantContext;
 import com.ning.billing.util.callcontext.UserType;
+import com.ning.billing.util.dao.ObjectType;
+import com.ning.billing.util.tag.ControlTagType;
+import com.ning.billing.util.tag.Tag;
+
+import static org.testng.Assert.assertEquals;
 
 public class TestDefaultInvoiceUserApi extends InvoiceApiTestBase {
 
@@ -316,4 +323,17 @@ public class TestDefaultInvoiceUserApi extends InvoiceApiTestBase {
                                                                            .setScale(InvoicingConfiguration.getNumberOfDecimals(),
                                                                                      InvoicingConfiguration.getRoundingMode())), 0);
     }
+
+    @Test(groups = "slow")
+    public void testAddRemoveWrittenOffTag() throws InvoiceApiException, TagApiException {
+        invoiceUserApi.tagInvoiceAsWrittenOff(invoiceId, callContext);
+
+        Map<String, Tag> tags = tagUserApi.getTags(invoiceId, ObjectType.INVOICE, tenantContext);
+        assertEquals(tags.size(), 1);
+        assertEquals(tags.values().iterator().next().getTagDefinitionId(), ControlTagType.WRITTEN_OFF.getId());
+
+        invoiceUserApi.tagInvoiceAsNotWrittenOff(invoiceId, callContext);
+        tags = tagUserApi.getTags(invoiceId, ObjectType.INVOICE, tenantContext);
+        assertEquals(tags.size(), 0);
+    }
 }
diff --git a/invoice/src/test/java/com/ning/billing/invoice/dao/InvoiceDaoTestBase.java b/invoice/src/test/java/com/ning/billing/invoice/dao/InvoiceDaoTestBase.java
index 7224afe..f5e0046 100644
--- a/invoice/src/test/java/com/ning/billing/invoice/dao/InvoiceDaoTestBase.java
+++ b/invoice/src/test/java/com/ning/billing/invoice/dao/InvoiceDaoTestBase.java
@@ -39,12 +39,7 @@ import com.ning.billing.util.bus.Bus;
 import com.ning.billing.util.bus.InMemoryBus;
 import com.ning.billing.util.clock.Clock;
 import com.ning.billing.util.clock.ClockMock;
-import com.ning.billing.util.svcapi.tag.TagInternalApi;
 import com.ning.billing.util.tag.api.user.TagEventBuilder;
-import com.ning.billing.util.tag.dao.AuditedTagDao;
-import com.ning.billing.util.tag.dao.MockTagDefinitionDao;
-import com.ning.billing.util.tag.dao.TagDao;
-import com.ning.billing.util.tag.dao.TagDefinitionDao;
 
 public class InvoiceDaoTestBase extends InvoicingTestBase {
     protected final TagEventBuilder tagEventBuilder = new TagEventBuilder();
@@ -103,11 +98,7 @@ public class InvoiceDaoTestBase extends InvoicingTestBase {
         bus.start();
 
         final NextBillingDatePoster nextBillingDatePoster = new MockNextBillingDatePoster();
-        final TagDefinitionDao tagDefinitionDao = new MockTagDefinitionDao();
-        final TagDao tagDao = new AuditedTagDao(dbi, tagEventBuilder, bus);
-        // API_FIX
-        final TagInternalApi tagUserApi = null; // new DefaultTagInternalApi(new InternalCallContextFactory(dbi, clock), tagDefinitionDao, tagDao);
-        invoiceDao = new AuditedInvoiceDao(dbi, nextBillingDatePoster, tagUserApi, clock, bus);
+        invoiceDao = new AuditedInvoiceDao(dbi, nextBillingDatePoster, clock, bus);
         invoiceDao.test(internalCallContext);
 
         invoiceItemSqlDao = dbi.onDemand(InvoiceItemSqlDao.class);
diff --git a/invoice/src/test/java/com/ning/billing/invoice/dao/MockInvoiceDao.java b/invoice/src/test/java/com/ning/billing/invoice/dao/MockInvoiceDao.java
index 4065ab6..58f5990 100644
--- a/invoice/src/test/java/com/ning/billing/invoice/dao/MockInvoiceDao.java
+++ b/invoice/src/test/java/com/ning/billing/invoice/dao/MockInvoiceDao.java
@@ -222,16 +222,6 @@ public class MockInvoiceDao implements InvoiceDao {
     }
 
     @Override
-    public void setWrittenOff(final UUID objectId, final InternalCallContext context) {
-        throw new UnsupportedOperationException();
-    }
-
-    @Override
-    public void removeWrittenOff(final UUID objectId, final InternalCallContext context) {
-        throw new UnsupportedOperationException();
-    }
-
-    @Override
     public InvoicePayment postChargeback(final UUID invoicePaymentId, final BigDecimal amount, final InternalCallContext context) throws InvoiceApiException {
         throw new UnsupportedOperationException();
     }
diff --git a/invoice/src/test/java/com/ning/billing/invoice/dao/TestDefaultInvoiceDao.java b/invoice/src/test/java/com/ning/billing/invoice/dao/TestDefaultInvoiceDao.java
index 0fb91ad..3d860a1 100644
--- a/invoice/src/test/java/com/ning/billing/invoice/dao/TestDefaultInvoiceDao.java
+++ b/invoice/src/test/java/com/ning/billing/invoice/dao/TestDefaultInvoiceDao.java
@@ -38,10 +38,6 @@ import com.ning.billing.invoice.notification.NextBillingDatePoster;
 import com.ning.billing.util.bus.Bus;
 import com.ning.billing.util.callcontext.InternalTenantContext;
 import com.ning.billing.util.clock.Clock;
-import com.ning.billing.util.dao.ObjectType;
-import com.ning.billing.util.svcapi.tag.TagInternalApi;
-import com.ning.billing.util.tag.ControlTagType;
-import com.ning.billing.util.tag.Tag;
 import com.ning.billing.util.tag.dao.MockTagDao;
 import com.ning.billing.util.tag.dao.MockTagDefinitionDao;
 import com.ning.billing.util.tag.dao.TagDao;
@@ -52,7 +48,6 @@ import com.google.common.collect.ImmutableMap;
 public class TestDefaultInvoiceDao extends InvoiceTestSuite {
 
     private InvoiceSqlDao invoiceSqlDao;
-    private TagInternalApi tagInternalApi;
     private AuditedInvoiceDao dao;
 
     @BeforeMethod(groups = "fast")
@@ -77,9 +72,7 @@ public class TestDefaultInvoiceDao extends InvoiceTestSuite {
         final NextBillingDatePoster poster = Mockito.mock(NextBillingDatePoster.class);
         final TagDefinitionDao tagDefinitionDao = new MockTagDefinitionDao();
         final TagDao tagDao = new MockTagDao();
-        // API_FIX
-        tagInternalApi = null; //new DefaultTagInternalApi(new InternalCallContextFactory(Mockito.mock(IDBI.class), new ClockMock()), tagDefinitionDao, tagDao);
-        dao = new AuditedInvoiceDao(idbi, poster, tagInternalApi, Mockito.mock(Clock.class), Mockito.mock(Bus.class));
+        dao = new AuditedInvoiceDao(idbi, poster, Mockito.mock(Clock.class), Mockito.mock(Bus.class));
     }
 
     @Test(groups = "fast")
@@ -132,33 +125,4 @@ public class TestDefaultInvoiceDao extends InvoiceTestSuite {
         } catch (InvoiceApiException e) {
         }
     }
-
-    @Test(groups = "fast")
-    public void testSetWrittenOff() throws Exception {
-        final UUID invoiceId = UUID.randomUUID();
-
-        final Map<String, Tag> beforeTags = tagInternalApi.getTags(invoiceId, ObjectType.INVOICE, internalCallContext);
-        Assert.assertEquals(beforeTags.keySet().size(), 0);
-
-        dao.setWrittenOff(invoiceId, internalCallContext);
-
-        final Map<String, Tag> afterTags = tagInternalApi.getTags(invoiceId, ObjectType.INVOICE, internalCallContext);
-        Assert.assertEquals(afterTags.keySet().size(), 1);
-        final UUID tagDefinitionId = ControlTagType.WRITTEN_OFF.getId();
-        Assert.assertEquals(afterTags.values().iterator().next().getTagDefinitionId(), tagDefinitionId);
-    }
-
-    @Test(groups = "fast")
-    public void testRemoveWrittenOff() throws Exception {
-        final UUID invoiceId = UUID.randomUUID();
-
-        dao.setWrittenOff(invoiceId, internalCallContext);
-
-        final Map<String, Tag> beforeTags = tagInternalApi.getTags(invoiceId, ObjectType.INVOICE, internalCallContext);
-        Assert.assertEquals(beforeTags.keySet().size(), 1);
-        dao.removeWrittenOff(invoiceId, internalCallContext);
-
-        final Map<String, Tag> afterTags = tagInternalApi.getTags(invoiceId, ObjectType.INVOICE, internalCallContext);
-        Assert.assertEquals(afterTags.keySet().size(), 0);
-    }
 }
diff --git a/invoice/src/test/java/com/ning/billing/invoice/dao/TestInvoiceDao.java b/invoice/src/test/java/com/ning/billing/invoice/dao/TestInvoiceDao.java
index ddfb9e0..dd8c0ce 100644
--- a/invoice/src/test/java/com/ning/billing/invoice/dao/TestInvoiceDao.java
+++ b/invoice/src/test/java/com/ning/billing/invoice/dao/TestInvoiceDao.java
@@ -16,6 +16,11 @@
 
 package com.ning.billing.invoice.dao;
 
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertNull;
+import static org.testng.Assert.assertTrue;
+
 import java.math.BigDecimal;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -59,25 +64,14 @@ import com.ning.billing.invoice.model.DefaultInvoicePayment;
 import com.ning.billing.invoice.model.FixedPriceInvoiceItem;
 import com.ning.billing.invoice.model.RecurringInvoiceItem;
 import com.ning.billing.invoice.model.RepairAdjInvoiceItem;
-import com.ning.billing.util.api.TagApiException;
 import com.ning.billing.util.clock.ClockMock;
-import com.ning.billing.util.dao.ObjectType;
 import com.ning.billing.util.entity.EntityPersistenceException;
 import com.ning.billing.util.svcapi.junction.BillingEvent;
 import com.ning.billing.util.svcapi.junction.BillingEventSet;
 import com.ning.billing.util.svcapi.junction.BillingModeType;
-import com.ning.billing.util.tag.ControlTagType;
-import com.ning.billing.util.tag.Tag;
-import com.ning.billing.util.tag.dao.AuditedTagDao;
-import com.ning.billing.util.tag.dao.TagDao;
 
 import com.google.common.collect.ImmutableMap;
 
-import static org.testng.Assert.assertEquals;
-import static org.testng.Assert.assertNotNull;
-import static org.testng.Assert.assertNull;
-import static org.testng.Assert.assertTrue;
-
 public class TestInvoiceDao extends InvoiceDaoTestBase {
 
     @Test(groups = "slow")
@@ -1317,72 +1311,6 @@ public class TestInvoiceDao extends InvoiceDaoTestBase {
     }
 
     @Test(groups = "slow")
-    public void testAddingWrittenOffTag() throws InvoiceApiException, TagApiException {
-        final Subscription subscription = getZombieSubscription();
-
-        final Plan plan = Mockito.mock(Plan.class);
-        Mockito.when(plan.getName()).thenReturn("plan");
-
-        final PlanPhase phase1 = Mockito.mock(PlanPhase.class);
-        Mockito.when(phase1.getName()).thenReturn("plan-phase1");
-
-        final DateTime targetDate1 = clock.getUTCNow();
-        final Currency currency = Currency.USD;
-
-        // create pseudo-random invoice
-        final BillingEvent event1 = createMockBillingEvent(null, subscription, targetDate1, plan, phase1, null,
-                                                           TEN, currency,
-                                                           BillingPeriod.MONTHLY, 31, BillingModeType.IN_ADVANCE,
-                                                           "testEvent1", 1L, SubscriptionTransitionType.CHANGE);
-        final BillingEventSet events = new MockBillingEventSet();
-        events.add(event1);
-
-        final Invoice invoice = generator.generateInvoice(UUID.randomUUID(), events, null, new LocalDate(targetDate1), DateTimeZone.UTC, Currency.USD);
-        invoiceDao.create(invoice, invoice.getTargetDate().getDayOfMonth(), true, internalCallContext);
-        invoiceDao.setWrittenOff(invoice.getId(), internalCallContext);
-
-        final TagDao tagDao = new AuditedTagDao(dbi, tagEventBuilder, bus);
-        final Map<String, Tag> tags = tagDao.loadEntities(invoice.getId(), ObjectType.INVOICE, internalCallContext);
-        assertEquals(tags.size(), 1);
-        assertEquals(tags.values().iterator().next().getTagDefinitionId(), ControlTagType.WRITTEN_OFF.getId());
-    }
-
-    @Test(groups = "slow")
-    public void testRemoveWrittenOffTag() throws InvoiceApiException, TagApiException {
-        final Subscription subscription = getZombieSubscription();
-
-        final Plan plan = Mockito.mock(Plan.class);
-        Mockito.when(plan.getName()).thenReturn("plan");
-
-        final PlanPhase phase1 = Mockito.mock(PlanPhase.class);
-        Mockito.when(phase1.getName()).thenReturn("plan-phase1");
-
-        final DateTime targetDate1 = clock.getUTCNow();
-        final Currency currency = Currency.USD;
-
-        // create pseudo-random invoice
-        final BillingEvent event1 = createMockBillingEvent(null, subscription, targetDate1, plan, phase1, null,
-                                                           TEN, currency,
-                                                           BillingPeriod.MONTHLY, 31, BillingModeType.IN_ADVANCE,
-                                                           "testEvent1", 1L, SubscriptionTransitionType.CHANGE);
-        final BillingEventSet events = new MockBillingEventSet();
-        events.add(event1);
-
-        final Invoice invoice = generator.generateInvoice(UUID.randomUUID(), events, null, new LocalDate(targetDate1), DateTimeZone.UTC, Currency.USD);
-        invoiceDao.create(invoice, invoice.getTargetDate().getDayOfMonth(), true, internalCallContext);
-        invoiceDao.setWrittenOff(invoice.getId(), internalCallContext);
-
-        final TagDao tagDao = new AuditedTagDao(dbi, tagEventBuilder, bus);
-        Map<String, Tag> tags = tagDao.loadEntities(invoice.getId(), ObjectType.INVOICE, internalCallContext);
-        assertEquals(tags.size(), 1);
-        assertEquals(tags.values().iterator().next().getTagDefinitionId(), ControlTagType.WRITTEN_OFF.getId());
-
-        invoiceDao.removeWrittenOff(invoice.getId(), internalCallContext);
-        tags = tagDao.loadEntities(invoice.getId(), ObjectType.INVOICE, internalCallContext);
-        assertEquals(tags.size(), 0);
-    }
-
-    @Test(groups = "slow")
     public void testDeleteCBANotConsumed() throws Exception {
         final UUID accountId = UUID.randomUUID();
 
diff --git a/invoice/src/test/java/com/ning/billing/invoice/tests/InvoiceTestUtils.java b/invoice/src/test/java/com/ning/billing/invoice/tests/InvoiceTestUtils.java
index 0bace86..d572117 100644
--- a/invoice/src/test/java/com/ning/billing/invoice/tests/InvoiceTestUtils.java
+++ b/invoice/src/test/java/com/ning/billing/invoice/tests/InvoiceTestUtils.java
@@ -75,12 +75,12 @@ public class InvoiceTestUtils {
         final List<InvoiceItem> invoiceItems = new ArrayList<InvoiceItem>();
         for (final BigDecimal amount : amounts) {
             final InvoiceItem invoiceItem = createInvoiceItem(clock, invoiceId, accountId, amount, currency);
-            invoiceItemSqlDao.create(invoiceItem, internalCallContextFactory.createInternalCallContext(callContext));
+            invoiceItemSqlDao.create(invoiceItem, internalCallContextFactory.createInternalCallContext(accountId, callContext));
             invoiceItems.add(invoiceItem);
         }
         Mockito.when(invoice.getInvoiceItems()).thenReturn(invoiceItems);
 
-        invoiceSqlDao.create(invoice, internalCallContextFactory.createInternalCallContext(callContext));
+        invoiceSqlDao.create(invoice, internalCallContextFactory.createInternalCallContext(accountId, callContext));
 
         return invoice;
     }
diff --git a/invoice/src/test/java/com/ning/billing/invoice/tests/TestChargeBacks.java b/invoice/src/test/java/com/ning/billing/invoice/tests/TestChargeBacks.java
index b3009ed..01576b8 100644
--- a/invoice/src/test/java/com/ning/billing/invoice/tests/TestChargeBacks.java
+++ b/invoice/src/test/java/com/ning/billing/invoice/tests/TestChargeBacks.java
@@ -55,11 +55,6 @@ import com.ning.billing.util.bus.Bus;
 import com.ning.billing.util.callcontext.InternalCallContextFactory;
 import com.ning.billing.util.clock.Clock;
 import com.ning.billing.util.clock.ClockMock;
-import com.ning.billing.util.svcapi.tag.TagInternalApi;
-import com.ning.billing.util.tag.dao.MockTagDao;
-import com.ning.billing.util.tag.dao.MockTagDefinitionDao;
-import com.ning.billing.util.tag.dao.TagDao;
-import com.ning.billing.util.tag.dao.TagDefinitionDao;
 
 public class TestChargeBacks extends InvoiceTestSuiteWithEmbeddedDB {
 
@@ -87,12 +82,8 @@ public class TestChargeBacks extends InvoiceTestSuiteWithEmbeddedDB {
         invoiceItemSqlDao = dbi.onDemand(InvoiceItemSqlDao.class);
         invoiceItemSqlDao.test(internalCallContext);
         final NextBillingDatePoster nextBillingDatePoster = new MockNextBillingDatePoster();
-        final TagDefinitionDao tagDefinitionDao = new MockTagDefinitionDao();
-        final TagDao tagDao = new MockTagDao();
         internalCallContextFactory = new InternalCallContextFactory(dbi, clock);
-        // API_FIX
-        final TagInternalApi tagUserApi = null; //new DefaultTagInternalApi(internalCallContextFactory, tagDefinitionDao, tagDao);
-        final InvoiceDao invoiceDao = new AuditedInvoiceDao(dbi, nextBillingDatePoster, tagUserApi, clock, Mockito.mock(Bus.class));
+        final InvoiceDao invoiceDao = new AuditedInvoiceDao(dbi, nextBillingDatePoster, clock, Mockito.mock(Bus.class));
         invoicePaymentApi = new DefaultInvoicePaymentApi(invoiceDao, internalCallContextFactory);
     }
 
diff --git a/junction/src/test/java/com/ning/billing/junction/api/blocking/TestDefaultBlockingApi.java b/junction/src/test/java/com/ning/billing/junction/api/blocking/TestDefaultBlockingApi.java
new file mode 100644
index 0000000..17b9f6d
--- /dev/null
+++ b/junction/src/test/java/com/ning/billing/junction/api/blocking/TestDefaultBlockingApi.java
@@ -0,0 +1,90 @@
+/*
+ * Copyright 2010-2012 Ning, Inc.
+ *
+ * Ning licenses this file to you under the Apache License, version 2.0
+ * (the "License"); you may not use this file except in compliance with the
+ * License.  You may obtain a copy of the License at:
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package com.ning.billing.junction.api.blocking;
+
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+
+import org.skife.jdbi.v2.Handle;
+import org.skife.jdbi.v2.tweak.HandleCallback;
+import org.testng.Assert;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import com.ning.billing.junction.JunctionTestSuiteWithEmbeddedDB;
+import com.ning.billing.junction.api.Blockable.Type;
+import com.ning.billing.junction.api.BlockingState;
+import com.ning.billing.junction.dao.BlockingStateDao;
+import com.ning.billing.junction.dao.BlockingStateSqlDao;
+import com.ning.billing.util.clock.ClockMock;
+import com.ning.billing.util.svcapi.junction.DefaultBlockingState;
+
+public class TestDefaultBlockingApi extends JunctionTestSuiteWithEmbeddedDB {
+
+    private final ClockMock clock = new ClockMock();
+
+    private DefaultBlockingApi blockingApi;
+
+    @BeforeMethod(groups = "slow")
+    public void setUp() throws Exception {
+        final BlockingStateDao blockingStateDao = getMysqlTestingHelper().getDBI().onDemand(BlockingStateSqlDao.class);
+        blockingApi = new DefaultBlockingApi(blockingStateDao, clock);
+    }
+
+    @Test(groups = "slow")
+    public void testSetBlockingStateOnBundle() throws Exception {
+        final UUID bundleId = UUID.randomUUID();
+        final Long accountRecordId = 123049714L;
+        getMysqlTestingHelper().getDBI().withHandle(new HandleCallback<Void>() {
+            @Override
+            public Void withHandle(final Handle handle) throws Exception {
+                handle.execute("DROP TABLE IF EXISTS bundles;\n" +
+                               "CREATE TABLE bundles (\n" +
+                               "    record_id int(11) unsigned NOT NULL AUTO_INCREMENT,\n" +
+                               "    id char(36) NOT NULL,\n" +
+                               "    external_key varchar(64) NOT NULL,\n" +
+                               "    account_id char(36) NOT NULL,\n" +
+                               "    last_sys_update_date datetime,\n" +
+                               "    account_record_id int(11) unsigned default null,\n" +
+                               "    tenant_record_id int(11) unsigned default null,\n" +
+                               "    PRIMARY KEY(record_id)\n" +
+                               ") ENGINE=innodb;");
+                handle.execute("insert into bundles (id, external_key, account_id, account_record_id) values (?, 'foo', ?, ?)",
+                               bundleId.toString(), UUID.randomUUID().toString(), accountRecordId);
+                return null;
+            }
+        });
+
+        final BlockingState blockingState = new DefaultBlockingState(bundleId, "BLOCKED", Type.SUBSCRIPTION_BUNDLE, "myService", true, true, true, clock.getUTCToday().toDateTimeAtStartOfDay());
+        blockingApi.setBlockingState(blockingState, internalCallContext);
+
+        // Verify the blocking state was applied
+        Assert.assertEquals(blockingApi.getBlockingStateFor(bundleId, internalCallContext), blockingState);
+        // Verify the account_record_id was populated
+        getMysqlTestingHelper().getDBI().withHandle(new HandleCallback<Void>() {
+            @Override
+            public Void withHandle(final Handle handle) throws Exception {
+                final List<Map<String, Object>> values = handle.select("select account_record_id from blocking_states where id = ?", bundleId.toString());
+                Assert.assertEquals(values.size(), 1);
+                Assert.assertEquals(values.get(0).keySet().size(), 1);
+                Assert.assertEquals(values.get(0).get("account_record_id"), accountRecordId);
+                return null;
+            }
+        });
+    }
+}
diff --git a/overdue/src/main/java/com/ning/billing/overdue/api/DefaultOverdueUserApi.java b/overdue/src/main/java/com/ning/billing/overdue/api/DefaultOverdueUserApi.java
index 997ecb4..15b359d 100644
--- a/overdue/src/main/java/com/ning/billing/overdue/api/DefaultOverdueUserApi.java
+++ b/overdue/src/main/java/com/ning/billing/overdue/api/DefaultOverdueUserApi.java
@@ -22,6 +22,7 @@ import org.slf4j.LoggerFactory;
 import com.ning.billing.ErrorCode;
 import com.ning.billing.entitlement.api.user.SubscriptionBundle;
 import com.ning.billing.junction.api.Blockable;
+import com.ning.billing.junction.api.Blockable.Type;
 import com.ning.billing.overdue.OverdueApiException;
 import com.ning.billing.overdue.OverdueState;
 import com.ning.billing.overdue.OverdueUserApi;
@@ -32,8 +33,10 @@ import com.ning.billing.overdue.config.api.OverdueStateSet;
 import com.ning.billing.overdue.wrapper.OverdueWrapper;
 import com.ning.billing.overdue.wrapper.OverdueWrapperFactory;
 import com.ning.billing.util.callcontext.CallContext;
+import com.ning.billing.util.callcontext.InternalCallContext;
 import com.ning.billing.util.callcontext.InternalCallContextFactory;
 import com.ning.billing.util.callcontext.TenantContext;
+import com.ning.billing.util.dao.ObjectType;
 import com.ning.billing.util.svcapi.junction.BlockingApi;
 
 import com.google.inject.Inject;
@@ -78,8 +81,12 @@ public class DefaultOverdueUserApi implements OverdueUserApi {
     public <T extends Blockable> OverdueState<T> refreshOverdueStateFor(final T blockable, final CallContext context) throws OverdueException, OverdueApiException {
         log.info(String.format("Refresh of %s requested", blockable.getId()));
         final OverdueWrapper<T> wrapper = factory.createOverdueWrapperFor(blockable);
-        // TODO accountId?
-        return wrapper.refresh(internalCallContextFactory.createInternalCallContext(context));
+        return wrapper.refresh(createInternalCallContext(blockable, context));
+    }
+
+    private <T extends Blockable> InternalCallContext createInternalCallContext(final T blockable, final CallContext context) {
+        final ObjectType objectType = Type.getObjectType(blockable);
+        return internalCallContextFactory.createInternalCallContext(blockable.getId(), objectType, context);
     }
 
     @Override
diff --git a/overdue/src/test/java/com/ning/billing/overdue/applicator/TestOverdueStateApplicator.java b/overdue/src/test/java/com/ning/billing/overdue/applicator/TestOverdueStateApplicator.java
index 8152a14..7ce68f7 100644
--- a/overdue/src/test/java/com/ning/billing/overdue/applicator/TestOverdueStateApplicator.java
+++ b/overdue/src/test/java/com/ning/billing/overdue/applicator/TestOverdueStateApplicator.java
@@ -41,6 +41,7 @@ import com.ning.billing.util.svcapi.junction.DefaultBlockingState;
 import com.google.inject.Inject;
 
 public class TestOverdueStateApplicator extends OverdueTestBase {
+
     @Inject
     OverdueStateApplicator<SubscriptionBundle> applicator;
 
diff --git a/util/src/main/java/com/ning/billing/util/callcontext/CallContextSqlDao.java b/util/src/main/java/com/ning/billing/util/callcontext/CallContextSqlDao.java
index 5580273..eb05589 100644
--- a/util/src/main/java/com/ning/billing/util/callcontext/CallContextSqlDao.java
+++ b/util/src/main/java/com/ning/billing/util/callcontext/CallContextSqlDao.java
@@ -26,5 +26,4 @@ public interface CallContextSqlDao {
 
     @SqlQuery("select record_id from accounts where id = :accountId;")
     public Long getAccountRecordId(@Bind("accountId") final String accountId);
-
 }
diff --git a/util/src/main/java/com/ning/billing/util/callcontext/InternalCallContextFactory.java b/util/src/main/java/com/ning/billing/util/callcontext/InternalCallContextFactory.java
index cbadb35..fac8975 100644
--- a/util/src/main/java/com/ning/billing/util/callcontext/InternalCallContextFactory.java
+++ b/util/src/main/java/com/ning/billing/util/callcontext/InternalCallContextFactory.java
@@ -16,16 +16,22 @@
 
 package com.ning.billing.util.callcontext;
 
+import java.util.List;
+import java.util.Map;
 import java.util.UUID;
 
 import javax.annotation.Nullable;
 import javax.inject.Inject;
 
+import org.skife.jdbi.v2.Handle;
 import org.skife.jdbi.v2.IDBI;
+import org.skife.jdbi.v2.tweak.HandleCallback;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import com.ning.billing.util.clock.Clock;
+import com.ning.billing.util.dao.ObjectType;
+import com.ning.billing.util.dao.TableName;
 
 public class InternalCallContextFactory {
 
@@ -34,11 +40,13 @@ public class InternalCallContextFactory {
     public static final UUID INTERNAL_TENANT_ID = new UUID(0L, 0L);
     public static final long INTERNAL_TENANT_RECORD_ID = 0L;
 
+    private final IDBI dbi;
     private final CallContextSqlDao callContextSqlDao;
     private final Clock clock;
 
     @Inject
     public InternalCallContextFactory(final IDBI dbi, final Clock clock) {
+        this.dbi = dbi;
         this.callContextSqlDao = dbi.onDemand(CallContextSqlDao.class);
         this.clock = clock;
     }
@@ -48,7 +56,7 @@ public class InternalCallContextFactory {
         return createInternalTenantContext(INTERNAL_TENANT_RECORD_ID, null);
     }
 
-    // Used for r/o or update/delete operations - we don't need the account id in that case
+    // Used for r/o operations - we don't need the account id in that case
     public InternalTenantContext createInternalTenantContext(final TenantContext context) {
         return createInternalTenantContext(getTenantRecordId(context), null);
     }
@@ -63,8 +71,51 @@ public class InternalCallContextFactory {
         return createInternalCallContext(INTERNAL_TENANT_RECORD_ID, null, new DefaultCallContext(INTERNAL_TENANT_ID, userName, callOrigin, userType, userToken, clock));
     }
 
-    // Used for r/o or update/delete operations - we don't need the account id in that case
-    // TODO - more work is needed for this statement to hold (especially for junction, overdue, custom fields and tags)
+    /**
+     * Crate an internal call context from a call context, and retrieving the account_record_id from another table
+     *
+     * @param objectId   the id of the row in the table pointed by object type where to look for account_record_id
+     * @param objectType the object type pointed by this objectId
+     * @param context    original call context
+     * @return internal call context from context, with a non null account_record_id (if found)
+     */
+    public InternalCallContext createInternalCallContext(final UUID objectId, final ObjectType objectType, final CallContext context) {
+        final Long accountRecordId;
+
+        final TableName tableName = TableName.fromObjectType(objectType);
+        if (tableName != null) {
+            accountRecordId = dbi.withHandle(new HandleCallback<Long>() {
+                @Override
+                public Long withHandle(final Handle handle) throws Exception {
+                    final String columnName;
+                    if (TableName.TAG_DEFINITIONS.equals(tableName) || TableName.TAG_DEFINITION_HISTORY.equals(tableName)) {
+                        // Not tied to an account
+                        return null;
+                    } else if (TableName.ACCOUNT.equals(tableName) || TableName.ACCOUNT_HISTORY.equals(tableName)) {
+                        // Lookup the record_id directly
+                        columnName = "record_id";
+                    } else {
+                        // The table should have an account_record_id column
+                        columnName = "account_record_id";
+                    }
+
+                    final List<Map<String, Object>> values = handle.select(String.format("select %s from %s where id = ?;", columnName, tableName.getTableName()), objectId.toString());
+                    if (values.size() == 0) {
+                        return null;
+                    } else {
+                        return (Long) values.get(0).get(columnName);
+                    }
+                }
+            });
+        } else {
+            accountRecordId = null;
+        }
+
+        return createInternalCallContext(getTenantRecordId(context), accountRecordId, context);
+    }
+
+    // Used for update/delete operations - we don't need the account id in that case
+    // Used also when we don't have an account_record_id column (e.g. tenants, tag_definitions)
     public InternalCallContext createInternalCallContext(final CallContext context) {
         return createInternalCallContext(getTenantRecordId(context), null, context);
     }
diff --git a/util/src/main/java/com/ning/billing/util/customfield/api/DefaultCustomFieldUserApi.java b/util/src/main/java/com/ning/billing/util/customfield/api/DefaultCustomFieldUserApi.java
index ec836b8..d41285e 100644
--- a/util/src/main/java/com/ning/billing/util/customfield/api/DefaultCustomFieldUserApi.java
+++ b/util/src/main/java/com/ning/billing/util/customfield/api/DefaultCustomFieldUserApi.java
@@ -48,7 +48,6 @@ public class DefaultCustomFieldUserApi implements CustomFieldUserApi {
 
     @Override
     public void saveCustomFields(final UUID objectId, final ObjectType objectType, final List<CustomField> fields, final CallContext context) {
-        // TODO accountId?
-        customFieldDao.saveEntities(objectId, objectType, fields, internalCallContextFactory.createInternalCallContext(context));
+        customFieldDao.saveEntities(objectId, objectType, fields, internalCallContextFactory.createInternalCallContext(objectId, objectType, context));
     }
 }
diff --git a/util/src/main/java/com/ning/billing/util/dao/TableName.java b/util/src/main/java/com/ning/billing/util/dao/TableName.java
index 84e9710..e958798 100644
--- a/util/src/main/java/com/ning/billing/util/dao/TableName.java
+++ b/util/src/main/java/com/ning/billing/util/dao/TableName.java
@@ -39,7 +39,8 @@ public enum TableName {
     SUBSCRIPTION_EVENTS("subscription_events", ObjectType.SUBSCRIPTION_EVENT),
     REFUND_HISTORY("refund_history"),
     REFUNDS("refunds", ObjectType.REFUND, REFUND_HISTORY),
-    TAG_DEFINITIONS("tag_definitions", ObjectType.TAG_DEFINITION),
+    TAG_DEFINITION_HISTORY("tag_definition_history"),
+    TAG_DEFINITIONS("tag_definitions", ObjectType.TAG_DEFINITION, TAG_DEFINITION_HISTORY),
     TAG_HISTORY("tag_history"),
     TENANT("tenants", ObjectType.TENANT),
     TAG("tags", TAG_HISTORY);
@@ -66,6 +67,15 @@ public enum TableName {
         this(tableName, null, null);
     }
 
+    public static TableName fromObjectType(final ObjectType objectType) {
+        for (final TableName tableName : values()) {
+            if (tableName.getObjectType() != null && tableName.getObjectType().equals(objectType)) {
+                return tableName;
+            }
+        }
+        return null;
+    }
+
     public String getTableName() {
         return tableName;
     }
diff --git a/util/src/main/java/com/ning/billing/util/svcapi/junction/BlockingApi.java b/util/src/main/java/com/ning/billing/util/svcapi/junction/BlockingApi.java
index 65a2ae4..2dde101 100644
--- a/util/src/main/java/com/ning/billing/util/svcapi/junction/BlockingApi.java
+++ b/util/src/main/java/com/ning/billing/util/svcapi/junction/BlockingApi.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2011 Ning, Inc.
+ * Copyright 2010-2012 Ning, Inc.
  *
  * Ning licenses this file to you under the Apache License, version 2.0
  * (the "License"); you may not use this file except in compliance with the
diff --git a/util/src/main/java/com/ning/billing/util/svcapi/junction/DefaultBlockingState.java b/util/src/main/java/com/ning/billing/util/svcapi/junction/DefaultBlockingState.java
index 56e45d5..4d4f615 100644
--- a/util/src/main/java/com/ning/billing/util/svcapi/junction/DefaultBlockingState.java
+++ b/util/src/main/java/com/ning/billing/util/svcapi/junction/DefaultBlockingState.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2011 Ning, Inc.
+ * Copyright 2010-2012 Ning, Inc.
  *
  * Ning licenses this file to you under the Apache License, version 2.0
  * (the "License"); you may not use this file except in compliance with the
diff --git a/util/src/main/java/com/ning/billing/util/tag/api/DefaultTagUserApi.java b/util/src/main/java/com/ning/billing/util/tag/api/DefaultTagUserApi.java
index 87eba08..71b8274 100644
--- a/util/src/main/java/com/ning/billing/util/tag/api/DefaultTagUserApi.java
+++ b/util/src/main/java/com/ning/billing/util/tag/api/DefaultTagUserApi.java
@@ -79,29 +79,25 @@ public class DefaultTagUserApi implements TagUserApi {
     public void addTags(final UUID objectId, final ObjectType objectType, final Collection<UUID> tagDefinitionIds, final CallContext context) throws TagApiException {
         // TODO: consider making this batch
         for (final UUID tagDefinitionId : tagDefinitionIds) {
-            // TODO accountId?
-            tagDao.insertTag(objectId, objectType, tagDefinitionId, internalCallContextFactory.createInternalCallContext(context));
+            tagDao.insertTag(objectId, objectType, tagDefinitionId, internalCallContextFactory.createInternalCallContext(objectId, objectType, context));
         }
     }
 
     @Override
     public void addTag(final UUID objectId, final ObjectType objectType, final UUID tagDefinitionId, final CallContext context) throws TagApiException {
-        // TODO accountId?
-        tagDao.insertTag(objectId, objectType, tagDefinitionId, internalCallContextFactory.createInternalCallContext(context));
+        tagDao.insertTag(objectId, objectType, tagDefinitionId, internalCallContextFactory.createInternalCallContext(objectId, objectType, context));
     }
 
     @Override
     public void removeTag(final UUID objectId, final ObjectType objectType, final UUID tagDefinitionId, final CallContext context) throws TagApiException {
-        // TODO accountId?
-        tagDao.deleteTag(objectId, objectType, tagDefinitionId, internalCallContextFactory.createInternalCallContext(context));
+        tagDao.deleteTag(objectId, objectType, tagDefinitionId, internalCallContextFactory.createInternalCallContext(objectId, objectType, context));
     }
 
     @Override
     public void removeTags(final UUID objectId, final ObjectType objectType, final Collection<UUID> tagDefinitionIds, final CallContext context) throws TagApiException {
         // TODO: consider making this batch
         for (final UUID tagDefinitionId : tagDefinitionIds) {
-            // TODO accountId?
-            tagDao.deleteTag(objectId, objectType, tagDefinitionId, internalCallContextFactory.createInternalCallContext(context));
+            tagDao.deleteTag(objectId, objectType, tagDefinitionId, internalCallContextFactory.createInternalCallContext(objectId, objectType, context));
         }
     }
 
diff --git a/util/src/test/java/com/ning/billing/dbi/MysqlTestingHelper.java b/util/src/test/java/com/ning/billing/dbi/MysqlTestingHelper.java
index 02d090a..65bf7c9 100644
--- a/util/src/test/java/com/ning/billing/dbi/MysqlTestingHelper.java
+++ b/util/src/test/java/com/ning/billing/dbi/MysqlTestingHelper.java
@@ -191,6 +191,10 @@ public class MysqlTestingHelper {
         initDb("drop table if exists accounts; create table accounts(record_id int(11) unsigned not null auto_increment, id char(36) not null, primary key(record_id)) engine=innodb;");
         initDb("drop table if exists tenants; create table tenants(record_id int(11) unsigned not null auto_increment, id char(36) not null, primary key(record_id)) engine=innodb;");
 
+        // We always want the basic tables when we do account_record_id lookups (e.g. for custom fields, tags or junction)
+        initDb("drop table if exists bundles; create table bundles(record_id int(11) unsigned not null auto_increment, id char(36) not null, account_record_id int(11) unsigned not null, primary key(record_id)) engine=innodb;");
+        initDb("drop table if exists subscriptions; create table subscriptions(record_id int(11) unsigned not null auto_increment, id char(36) not null, account_record_id int(11) unsigned not null, primary key(record_id)) engine=innodb;");
+
         for (final String pack : new String[]{"account", "analytics", "entitlement", "util", "payment", "invoice", "junction", "tenant"}) {
             final String ddl;
             try {
diff --git a/util/src/test/java/com/ning/billing/util/callcontext/TestInternalCallContextFactory.java b/util/src/test/java/com/ning/billing/util/callcontext/TestInternalCallContextFactory.java
new file mode 100644
index 0000000..0ac9552
--- /dev/null
+++ b/util/src/test/java/com/ning/billing/util/callcontext/TestInternalCallContextFactory.java
@@ -0,0 +1,107 @@
+/*
+ * Copyright 2010-2012 Ning, Inc.
+ *
+ * Ning licenses this file to you under the Apache License, version 2.0
+ * (the "License"); you may not use this file except in compliance with the
+ * License.  You may obtain a copy of the License at:
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package com.ning.billing.util.callcontext;
+
+import java.util.UUID;
+
+import org.skife.jdbi.v2.Handle;
+import org.skife.jdbi.v2.tweak.HandleCallback;
+import org.testng.Assert;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import com.ning.billing.util.UtilTestSuiteWithEmbeddedDB;
+import com.ning.billing.util.clock.ClockMock;
+import com.ning.billing.util.dao.ObjectType;
+
+public class TestInternalCallContextFactory extends UtilTestSuiteWithEmbeddedDB {
+
+    private InternalCallContextFactory internalCallContextFactory;
+
+    @BeforeMethod(groups = "slow")
+    public void setUp() throws Exception {
+        internalCallContextFactory = new InternalCallContextFactory(getMysqlTestingHelper().getDBI(), new ClockMock());
+    }
+
+    @Test(groups = "slow")
+    public void testCreateInternalCallContextWithAccountRecordIdFromSimpleObjectType() throws Exception {
+        final UUID invoiceId = UUID.randomUUID();
+        final Long accountRecordId = 19384012L;
+
+        getMysqlTestingHelper().getDBI().withHandle(new HandleCallback<Void>() {
+            @Override
+            public Void withHandle(final Handle handle) throws Exception {
+                handle.execute("DROP TABLE IF EXISTS invoices;\n" +
+                               "CREATE TABLE invoices (\n" +
+                               "    record_id int(11) unsigned NOT NULL AUTO_INCREMENT,\n" +
+                               "    id char(36) NOT NULL,\n" +
+                               "    account_id char(36) NOT NULL,\n" +
+                               "    invoice_date date NOT NULL,\n" +
+                               "    target_date date NOT NULL,\n" +
+                               "    currency char(3) NOT NULL,\n" +
+                               "    migrated bool NOT NULL,\n" +
+                               "    created_by varchar(50) NOT NULL,\n" +
+                               "    created_date datetime NOT NULL,\n" +
+                               "    account_record_id int(11) unsigned default null,\n" +
+                               "    tenant_record_id int(11) unsigned default null,\n" +
+                               "    PRIMARY KEY(record_id)\n" +
+                               ") ENGINE=innodb;");
+                handle.execute("insert into invoices (id, account_id, invoice_date, target_date, currency, migrated, created_by, created_date, account_record_id) values " +
+                               "(?, ?, now(), now(), 'USD', 0, 'test', now(), ?)", invoiceId.toString(), UUID.randomUUID().toString(), accountRecordId);
+                return null;
+            }
+        });
+
+        final InternalCallContext context = internalCallContextFactory.createInternalCallContext(invoiceId, ObjectType.INVOICE, callContext);
+        // The account record id should have been looked up in the invoices table
+        Assert.assertEquals(context.getAccountRecordId(), accountRecordId);
+        verifyInternalCallContext(context);
+    }
+
+    @Test(groups = "slow")
+    public void testCreateInternalCallContextWithAccountRecordIdFromAccountObjectType() throws Exception {
+        final UUID accountId = UUID.randomUUID();
+        final Long accountRecordId = 19384012L;
+
+        getMysqlTestingHelper().getDBI().withHandle(new HandleCallback<Void>() {
+            @Override
+            public Void withHandle(final Handle handle) throws Exception {
+                // Note: we always create an accounts table, see MysqlTestingHelper
+                handle.execute("insert into accounts (record_id, id) values (?, ?)", accountRecordId, accountId.toString());
+                return null;
+            }
+        });
+
+        final InternalCallContext context = internalCallContextFactory.createInternalCallContext(accountId, ObjectType.ACCOUNT, callContext);
+        // The account record id should have been looked up in the accounts table
+        Assert.assertEquals(context.getAccountRecordId(), accountRecordId);
+        verifyInternalCallContext(context);
+    }
+
+    private void verifyInternalCallContext(final InternalCallContext context) {
+        Assert.assertEquals(context.getCallOrigin(), callContext.getCallOrigin());
+        Assert.assertEquals(context.getComment(), callContext.getComment());
+        Assert.assertEquals(context.getCreatedDate(), callContext.getCreatedDate());
+        Assert.assertEquals(context.getReasonCode(), callContext.getReasonCode());
+        Assert.assertEquals(context.getUpdatedDate(), callContext.getUpdatedDate());
+        Assert.assertEquals(context.getUserName(), callContext.getUserName());
+        Assert.assertEquals(context.getUserToken(), callContext.getUserToken());
+        Assert.assertEquals(context.getUserType(), callContext.getUserType());
+        // Our test callContext doesn't have a tenant id
+        Assert.assertEquals(context.getTenantRecordId(), (Long) InternalCallContextFactory.INTERNAL_TENANT_RECORD_ID);
+    }
+}
diff --git a/util/src/test/java/com/ning/billing/util/customfield/api/TestDefaultCustomFieldUserApi.java b/util/src/test/java/com/ning/billing/util/customfield/api/TestDefaultCustomFieldUserApi.java
new file mode 100644
index 0000000..76182f6
--- /dev/null
+++ b/util/src/test/java/com/ning/billing/util/customfield/api/TestDefaultCustomFieldUserApi.java
@@ -0,0 +1,84 @@
+/*
+ * Copyright 2010-2012 Ning, Inc.
+ *
+ * Ning licenses this file to you under the Apache License, version 2.0
+ * (the "License"); you may not use this file except in compliance with the
+ * License.  You may obtain a copy of the License at:
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package com.ning.billing.util.customfield.api;
+
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+
+import org.skife.jdbi.v2.Handle;
+import org.skife.jdbi.v2.tweak.HandleCallback;
+import org.testng.Assert;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import com.ning.billing.util.UtilTestSuiteWithEmbeddedDB;
+import com.ning.billing.util.callcontext.InternalCallContextFactory;
+import com.ning.billing.util.clock.ClockMock;
+import com.ning.billing.util.customfield.CustomField;
+import com.ning.billing.util.customfield.StringCustomField;
+import com.ning.billing.util.customfield.dao.AuditedCustomFieldDao;
+import com.ning.billing.util.customfield.dao.CustomFieldDao;
+import com.ning.billing.util.dao.ObjectType;
+
+import com.google.common.collect.ImmutableList;
+
+public class TestDefaultCustomFieldUserApi extends UtilTestSuiteWithEmbeddedDB {
+
+    private DefaultCustomFieldUserApi customFieldUserApi;
+
+    @BeforeMethod(groups = "slow")
+    public void setUp() throws Exception {
+        final InternalCallContextFactory internalCallContextFactory = new InternalCallContextFactory(getMysqlTestingHelper().getDBI(), new ClockMock());
+        final CustomFieldDao customFieldDao = new AuditedCustomFieldDao(getMysqlTestingHelper().getDBI());
+        customFieldUserApi = new DefaultCustomFieldUserApi(internalCallContextFactory, customFieldDao);
+    }
+
+    @Test(groups = "slow")
+    public void testSaveCustomFieldWithAccountRecordId() throws Exception {
+        final UUID accountId = UUID.randomUUID();
+        final Long accountRecordId = 19384012L;
+
+        getMysqlTestingHelper().getDBI().withHandle(new HandleCallback<Void>() {
+            @Override
+            public Void withHandle(final Handle handle) throws Exception {
+                // Note: we always create an accounts table, see MysqlTestingHelper
+                handle.execute("insert into accounts (record_id, id) values (?, ?)", accountRecordId, accountId.toString());
+                return null;
+            }
+        });
+
+        final CustomField customField = new StringCustomField(UUID.randomUUID().toString().substring(1, 4), UUID.randomUUID().toString().substring(1, 4));
+        customFieldUserApi.saveCustomFields(accountId, ObjectType.ACCOUNT, ImmutableList.<CustomField>of(customField), callContext);
+
+        // Verify the field was saved
+        final Map<String, CustomField> customFields = customFieldUserApi.getCustomFields(accountId, ObjectType.ACCOUNT, callContext);
+        Assert.assertEquals(customFields.keySet().size(), 1);
+        Assert.assertEquals(customFields.get(customField.getName()), customField);
+        // Verify the account_record_id was populated
+        getMysqlTestingHelper().getDBI().withHandle(new HandleCallback<Void>() {
+            @Override
+            public Void withHandle(final Handle handle) throws Exception {
+                final List<Map<String, Object>> values = handle.select("select account_record_id from custom_fields where object_id = ?", accountId.toString());
+                Assert.assertEquals(values.size(), 1);
+                Assert.assertEquals(values.get(0).keySet().size(), 1);
+                Assert.assertEquals(values.get(0).get("account_record_id"), accountRecordId);
+                return null;
+            }
+        });
+    }
+}