killbill-aplcache

diff --git a/invoice/src/main/java/org/killbill/billing/invoice/dao/DefaultInvoiceDao.java b/invoice/src/main/java/org/killbill/billing/invoice/dao/DefaultInvoiceDao.java
index 3f76659..7742de0 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/dao/DefaultInvoiceDao.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/dao/DefaultInvoiceDao.java
@@ -351,24 +351,22 @@ public class DefaultInvoiceDao extends EntityDaoBase<InvoiceModelDao, Invoice, I
         return transactionalSqlDao.execute(false, new EntitySqlDaoTransactionWrapper<List<InvoiceItemModelDao>>() {
             @Override
             public List<InvoiceItemModelDao> inTransaction(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory) throws Exception {
+                final InvoiceSqlDao invoiceSqlDao = entitySqlDaoWrapperFactory.become(InvoiceSqlDao.class);
+                final InvoiceItemSqlDao transInvoiceItemSqlDao = entitySqlDaoWrapperFactory.become(InvoiceItemSqlDao.class);
+
                 final ExistingInvoiceMetadata existingInvoiceMetadata;
                 if (existingInvoiceMetadataOrNull == null) {
-                    final List<InvoiceModelDao> allInvoicesForAccount = invoiceDaoHelper.getAllInvoicesByAccountFromTransaction(true, invoicesTags, entitySqlDaoWrapperFactory, context);
-                    invoiceDaoHelper.populateChildren(allInvoicesForAccount, invoicesTags, entitySqlDaoWrapperFactory, context);
-                    existingInvoiceMetadata = new ExistingInvoiceMetadata(allInvoicesForAccount);
+                    existingInvoiceMetadata = new ExistingInvoiceMetadata(invoiceSqlDao, transInvoiceItemSqlDao);
                 } else {
                     existingInvoiceMetadata = existingInvoiceMetadataOrNull;
                 }
 
-                final InvoiceSqlDao invoiceSqlDao = entitySqlDaoWrapperFactory.become(InvoiceSqlDao.class);
-                final InvoiceItemSqlDao transInvoiceItemSqlDao = entitySqlDaoWrapperFactory.become(InvoiceItemSqlDao.class);
-
                 final List<InvoiceItemModelDao> createdInvoiceItems = new LinkedList<InvoiceItemModelDao>();
                 for (final InvoiceModelDao invoiceModelDao : invoices) {
                     invoiceByInvoiceId.put(invoiceModelDao.getId(), invoiceModelDao);
                     final boolean isNotShellInvoice = invoiceIdsReferencedFromItems.remove(invoiceModelDao.getId());
 
-                    final InvoiceModelDao invoiceOnDisk = existingInvoiceMetadata.getExistingInvoice(invoiceModelDao.getId());
+                    final InvoiceModelDao invoiceOnDisk = existingInvoiceMetadata.getExistingInvoice(invoiceModelDao.getId(), context);
                     if (isNotShellInvoice) {
                         // Create the invoice if this is not a shell invoice and it does not already exist
                         if (invoiceOnDisk == null) {
@@ -382,7 +380,7 @@ public class DefaultInvoiceDao extends EntityDaoBase<InvoiceModelDao, Invoice, I
 
                     // Create the invoice items if needed (note: they may not necessarily belong to that invoice)
                     for (final InvoiceItemModelDao invoiceItemModelDao : invoiceModelDao.getInvoiceItems()) {
-                        final InvoiceItemModelDao existingInvoiceItem = existingInvoiceMetadata.getExistingInvoiceItem(invoiceItemModelDao.getId());
+                        final InvoiceItemModelDao existingInvoiceItem = existingInvoiceMetadata.getExistingInvoiceItem(invoiceItemModelDao.getId(), context);
                         // Because of AUTO_INVOICING_REUSE_DRAFT we expect an invoice were items might already exist.
                         // Also for ALLOWED_INVOICE_ITEM_TYPES, we expect plugins to potentially modify the amount
                         if (existingInvoiceItem == null) {

diff --git a/invoice/src/main/java/org/killbill/billing/invoice/dao/ExistingInvoiceMetadata.java b/invoice/src/main/java/org/killbill/billing/invoice/dao/ExistingInvoiceMetadata.java
index 58d8a3e..1147442 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/dao/ExistingInvoiceMetadata.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/dao/ExistingInvoiceMetadata.java
@@ -22,6 +22,7 @@ import java.util.HashMap;
 import java.util.Map;
 import java.util.UUID;
 
+import org.killbill.billing.callcontext.InternalTenantContext;
 import org.killbill.billing.invoice.api.Invoice;
 import org.killbill.billing.invoice.api.InvoiceItem;
 
@@ -30,6 +31,9 @@ public class ExistingInvoiceMetadata {
     private final Map<UUID, InvoiceModelDao> invoicesCache = new HashMap<UUID, InvoiceModelDao>();
     private final Map<UUID, InvoiceItemModelDao> invoiceItemsCache = new HashMap<UUID, InvoiceItemModelDao>();
 
+    private InvoiceSqlDao invoiceSqlDao;
+    private InvoiceItemSqlDao invoiceItemSqlDao;
+
     public ExistingInvoiceMetadata(final Iterable<Invoice> existingInvoices) {
         for (final Invoice invoice : existingInvoices) {
             invoicesCache.put(invoice.getId(), new InvoiceModelDao(invoice));
@@ -48,11 +52,24 @@ public class ExistingInvoiceMetadata {
         }
     }
 
-    public InvoiceModelDao getExistingInvoice(final UUID invoiceId) {
-        return invoicesCache.get(invoiceId);
+    public ExistingInvoiceMetadata(final InvoiceSqlDao invoiceSqlDao, final InvoiceItemSqlDao invoiceItemSqlDao) {
+        this.invoiceSqlDao = invoiceSqlDao;
+        this.invoiceItemSqlDao = invoiceItemSqlDao;
     }
 
-    public InvoiceItemModelDao getExistingInvoiceItem(final UUID invoiceItemId) {
-        return invoiceItemsCache.get(invoiceItemId);
+    public InvoiceModelDao getExistingInvoice(final UUID invoiceId, final InternalTenantContext context) {
+        if (invoiceSqlDao != null) {
+            return invoiceSqlDao.getById(invoiceId.toString(), context);
+        } else {
+            return invoicesCache.get(invoiceId);
+        }
+    }
+
+    public InvoiceItemModelDao getExistingInvoiceItem(final UUID invoiceItemId, final InternalTenantContext context) {
+        if (invoiceItemSqlDao != null) {
+            return invoiceItemSqlDao.getById(invoiceItemId.toString(), context);
+        } else {
+            return invoiceItemsCache.get(invoiceItemId);
+        }
     }
 }