killbill-aplcache

util: return rehydrated entity after write operations In

3/30/2017 3:42:42 AM

Details

diff --git a/account/src/main/java/org/killbill/billing/account/dao/AccountSqlDao.java b/account/src/main/java/org/killbill/billing/account/dao/AccountSqlDao.java
index 2eccd3c..77bea77 100644
--- a/account/src/main/java/org/killbill/billing/account/dao/AccountSqlDao.java
+++ b/account/src/main/java/org/killbill/billing/account/dao/AccountSqlDao.java
@@ -1,7 +1,9 @@
 /*
  * Copyright 2010-2013 Ning, Inc.
+ * Copyright 2014-2017 Groupon, Inc
+ * Copyright 2014-2017 The Billing Project, LLC
  *
- * Ning licenses this file to you under the Apache License, version 2.0
+ * The Billing Project 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:
  *
@@ -44,7 +46,7 @@ public interface AccountSqlDao extends EntitySqlDao<AccountModelDao, Account> {
 
     @SqlQuery
     public Integer getBCD(@Bind("id") String accountId,
-                       @BindBean final InternalTenantContext context);
+                          @BindBean final InternalTenantContext context);
 
     @SqlUpdate
     @Audited(ChangeType.UPDATE)
@@ -53,9 +55,9 @@ public interface AccountSqlDao extends EntitySqlDao<AccountModelDao, Account> {
 
     @SqlUpdate
     @Audited(ChangeType.UPDATE)
-    public void updatePaymentMethod(@Bind("id") String accountId,
-                                    @Bind("paymentMethodId") String paymentMethodId,
-                                    @BindBean final InternalCallContext context);
+    public Object updatePaymentMethod(@Bind("id") String accountId,
+                                      @Bind("paymentMethodId") String paymentMethodId,
+                                      @BindBean final InternalCallContext context);
 
     @SqlQuery
     List<AccountModelDao> getAccountsByParentId(@Bind("parentAccountId") UUID parentAccountId,
diff --git a/account/src/main/java/org/killbill/billing/account/dao/DefaultAccountDao.java b/account/src/main/java/org/killbill/billing/account/dao/DefaultAccountDao.java
index 171e157..4103bf8 100644
--- a/account/src/main/java/org/killbill/billing/account/dao/DefaultAccountDao.java
+++ b/account/src/main/java/org/killbill/billing/account/dao/DefaultAccountDao.java
@@ -215,9 +215,8 @@ public class DefaultAccountDao extends EntityDaoBase<AccountModelDao, Account, A
                 }
 
                 final String thePaymentMethodId = paymentMethodId != null ? paymentMethodId.toString() : null;
-                transactional.updatePaymentMethod(accountId.toString(), thePaymentMethodId, context);
+                final AccountModelDao account = (AccountModelDao) transactional.updatePaymentMethod(accountId.toString(), thePaymentMethodId, context);
 
-                final AccountModelDao account = transactional.getById(accountId.toString(), context);
                 final AccountChangeInternalEvent changeEvent = new DefaultAccountChangeEvent(accountId, currentAccount, account,
                                                                                              context.getAccountRecordId(),
                                                                                              context.getTenantRecordId(),
@@ -245,7 +244,7 @@ public class DefaultAccountDao extends EntityDaoBase<AccountModelDao, Account, A
                     throw new AccountApiException(ErrorCode.ACCOUNT_EMAIL_ALREADY_EXISTS, email.getId());
                 }
 
-                transactional.create(email, context);
+                createAndRefresh(transactional, email, context);
                 return null;
             }
         });
diff --git a/entitlement/src/main/java/org/killbill/billing/entitlement/dao/DefaultBlockingStateDao.java b/entitlement/src/main/java/org/killbill/billing/entitlement/dao/DefaultBlockingStateDao.java
index 164a9b8..29c193c 100644
--- a/entitlement/src/main/java/org/killbill/billing/entitlement/dao/DefaultBlockingStateDao.java
+++ b/entitlement/src/main/java/org/killbill/billing/entitlement/dao/DefaultBlockingStateDao.java
@@ -230,7 +230,7 @@ public class DefaultBlockingStateDao extends EntityDaoBase<BlockingStateModelDao
                     boolean inserted = false;
                     // Create the state, if needed
                     if (!blockingStatesToRemove.contains(newBlockingStateModelDao.getId())) {
-                        sqlDao.create(newBlockingStateModelDao, context);
+                        createAndRefresh(sqlDao, newBlockingStateModelDao, context);
                         inserted = true;
                     }
 
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 b9743f6..0afa55c 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
@@ -326,7 +326,7 @@ public class DefaultInvoiceDao extends EntityDaoBase<InvoiceModelDao, Invoice, I
                         // We only want to insert that invoice if there are real invoiceItems associated to it -- if not, this is just
                         // a shell invoice and we only need to insert the invoiceItems -- for the already existing invoices
                         if (isRealInvoice) {
-                            invoiceSqlDao.create(invoiceModelDao, context);
+                            createAndRefresh(invoiceSqlDao, invoiceModelDao, context);
                             createdInvoiceIds.add(invoiceModelDao.getId());
                         }
                     }
@@ -334,8 +334,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()) {
                         if (transInvoiceItemSqlDao.getById(invoiceItemModelDao.getId().toString(), context) == null) {
-                            createInvoiceItemFromTransaction(transInvoiceItemSqlDao, invoiceItemModelDao, context);
-                            createdInvoiceItems.add(transInvoiceItemSqlDao.getById(invoiceItemModelDao.getId().toString(), context));
+                            createdInvoiceItems.add(createInvoiceItemFromTransaction(transInvoiceItemSqlDao, invoiceItemModelDao, context));
 
                             adjustedInvoiceIds.add(invoiceItemModelDao.getInvoiceId());
                         }
@@ -551,14 +550,13 @@ public class DefaultInvoiceDao extends EntityDaoBase<InvoiceModelDao, Invoice, I
                                                                                  payment.getInvoiceId(), paymentId,
                                                                                  context.getCreatedDate(), requestedPositiveAmount.negate(),
                                                                                  payment.getCurrency(), payment.getProcessedCurrency(), transactionExternalKey, payment.getId(), true);
-                transactional.create(refund, context);
+                createAndRefresh(transactional, refund, context);
 
                 // Retrieve invoice after the Refund
                 final InvoiceModelDao invoice = transInvoiceDao.getById(payment.getInvoiceId().toString(), context);
                 Preconditions.checkState(invoice != null, "Invoice shouldn't be null for payment " + payment.getId());
                 invoiceDaoHelper.populateChildren(invoice, invoicesTags, entitySqlDaoWrapperFactory, context);
 
-                final BigDecimal invoiceBalanceAfterRefund = InvoiceModelDaoHelper.getBalance(invoice);
                 final InvoiceItemSqlDao transInvoiceItemDao = entitySqlDaoWrapperFactory.become(InvoiceItemSqlDao.class);
 
                 // At this point, we created the refund which made the invoice balance positive and applied any existing
@@ -631,7 +629,7 @@ public class DefaultInvoiceDao extends EntityDaoBase<InvoiceModelDao, Invoice, I
                                                                                      payment.getInvoiceId(), payment.getPaymentId(), context.getCreatedDate(),
                                                                                      requestedChargedBackAmount.negate(), payment.getCurrency(), payment.getProcessedCurrency(),
                                                                                      chargebackTransactionExternalKey, payment.getId(), true);
-                transactional.create(chargeBack, context);
+                createAndRefresh(transactional, chargeBack, context);
 
                 // Notify the bus since the balance of the invoice changed
                 final UUID accountId = transactional.getAccountIdFromInvoicePaymentId(chargeBack.getId().toString(), context);
@@ -817,7 +815,7 @@ public class DefaultInvoiceDao extends EntityDaoBase<InvoiceModelDao, Invoice, I
                     }).orNull();
 
                     if (existingAttempt == null) {
-                        transactional.create(invoicePayment, context);
+                        createAndRefresh(transactional, invoicePayment, context);
                     } else {
                         transactional.updateAttempt(existingAttempt.getRecordId(),
                                                     invoicePayment.getPaymentId().toString(),
@@ -1033,7 +1031,7 @@ public class DefaultInvoiceDao extends EntityDaoBase<InvoiceModelDao, Invoice, I
         }
     }
 
-    private void createInvoiceItemFromTransaction(final InvoiceItemSqlDao invoiceItemSqlDao, final InvoiceItemModelDao invoiceItemModelDao, final InternalCallContext context) throws EntityPersistenceException, InvoiceApiException {
+    private InvoiceItemModelDao createInvoiceItemFromTransaction(final InvoiceItemSqlDao invoiceItemSqlDao, final InvoiceItemModelDao invoiceItemModelDao, final InternalCallContext context) throws EntityPersistenceException, InvoiceApiException {
         // There is no efficient way to retrieve an invoice item given an ID today (and invoice plugins can put item adjustments
         // on a different invoice than the original item), so it's easier to do the check in the DAO rather than in the API layer
         // See also https://github.com/killbill/killbill/issues/7
@@ -1041,7 +1039,7 @@ public class DefaultInvoiceDao extends EntityDaoBase<InvoiceModelDao, Invoice, I
             validateInvoiceItemToBeAdjusted(invoiceItemSqlDao, invoiceItemModelDao, context);
         }
 
-        invoiceItemSqlDao.create(invoiceItemModelDao, context);
+        return createAndRefresh(invoiceItemSqlDao, invoiceItemModelDao, context);
     }
 
     private void validateInvoiceItemToBeAdjusted(final InvoiceItemSqlDao invoiceItemSqlDao, final InvoiceItemModelDao invoiceItemModelDao, final InternalCallContext context) throws InvoiceApiException {
@@ -1115,7 +1113,7 @@ public class DefaultInvoiceDao extends EntityDaoBase<InvoiceModelDao, Invoice, I
             @Override
             public Void inTransaction(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory) throws Exception {
                 final InvoiceParentChildrenSqlDao transactional = entitySqlDaoWrapperFactory.become(InvoiceParentChildrenSqlDao.class);
-                transactional.create(invoiceRelation, context);
+                createAndRefresh(transactional, invoiceRelation, context);
                 return null;
             }
         });
@@ -1231,14 +1229,14 @@ public class DefaultInvoiceDao extends EntityDaoBase<InvoiceModelDao, Invoice, I
 
                 // save invoices and invoice items
                 final InvoiceModelDao childInvoice = new InvoiceModelDao(invoiceForExternalCharge);
-                invoiceSqlDao.create(childInvoice, childAccountContext);
+                createAndRefresh(invoiceSqlDao, childInvoice, childAccountContext);
                 final InvoiceItemModelDao childExternalChargeItem = new InvoiceItemModelDao(externalChargeItem);
                 createInvoiceItemFromTransaction(transInvoiceItemSqlDao, childExternalChargeItem, childAccountContext);
                 // Keep invoice up-to-date for CBA below
                 childInvoice.addInvoiceItem(childExternalChargeItem);
 
                 final InvoiceModelDao parentInvoice = new InvoiceModelDao(invoiceForCredit);
-                invoiceSqlDao.create(parentInvoice, parentAccountContext);
+                createAndRefresh(invoiceSqlDao, parentInvoice, parentAccountContext);
                 final InvoiceItemModelDao parentCreditItem = new InvoiceItemModelDao(creditItem);
                 createInvoiceItemFromTransaction(transInvoiceItemSqlDao, parentCreditItem, parentAccountContext);
                 // Keep invoice up-to-date for CBA below
diff --git a/payment/src/main/java/org/killbill/billing/payment/dao/DefaultPaymentDao.java b/payment/src/main/java/org/killbill/billing/payment/dao/DefaultPaymentDao.java
index 4d00d40..22e43e7 100644
--- a/payment/src/main/java/org/killbill/billing/payment/dao/DefaultPaymentDao.java
+++ b/payment/src/main/java/org/killbill/billing/payment/dao/DefaultPaymentDao.java
@@ -1,7 +1,7 @@
 /*
  * Copyright 2010-2013 Ning, Inc.
- * Copyright 2014-2016 Groupon, Inc
- * Copyright 2014-2016 The Billing Project, LLC
+ * Copyright 2014-2017 Groupon, Inc
+ * Copyright 2014-2017 The Billing Project, LLC
  *
  * The Billing Project 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
@@ -107,9 +107,7 @@ public class DefaultPaymentDao extends EntityDaoBase<PaymentModelDao, Payment, P
             @Override
             public PaymentAttemptModelDao inTransaction(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory) throws Exception {
                 final PaymentAttemptSqlDao transactional = entitySqlDaoWrapperFactory.become(PaymentAttemptSqlDao.class);
-                transactional.create(attempt, context);
-                final PaymentAttemptModelDao result = transactional.getById(attempt.getId().toString(), context);
-                return result;
+                return createAndRefresh(transactional, attempt, context);
             }
         });
     }
@@ -297,9 +295,12 @@ public class DefaultPaymentDao extends EntityDaoBase<PaymentModelDao, Payment, P
             @Override
             public PaymentModelDao inTransaction(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory) throws Exception {
                 final PaymentSqlDao paymentSqlDao = entitySqlDaoWrapperFactory.become(PaymentSqlDao.class);
-                paymentSqlDao.create(payment, context);
-                entitySqlDaoWrapperFactory.become(TransactionSqlDao.class).create(paymentTransaction, context);
-                return paymentSqlDao.getById(payment.getId().toString(), context);
+                final PaymentModelDao paymentModelDao = createAndRefresh(paymentSqlDao, payment, context);
+
+                final TransactionSqlDao transactionSqlDao = entitySqlDaoWrapperFactory.become(TransactionSqlDao.class);
+                createAndRefresh(transactionSqlDao, paymentTransaction, context);
+
+                return paymentModelDao;
             }
         });
     }
@@ -310,8 +311,7 @@ public class DefaultPaymentDao extends EntityDaoBase<PaymentModelDao, Payment, P
             @Override
             public PaymentTransactionModelDao inTransaction(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory) throws Exception {
                 final TransactionSqlDao transactional = entitySqlDaoWrapperFactory.become(TransactionSqlDao.class);
-                transactional.create(paymentTransaction, context);
-                final PaymentTransactionModelDao paymentTransactionModelDao = transactional.getById(paymentTransaction.getId().toString(), context);
+                final PaymentTransactionModelDao paymentTransactionModelDao = createAndRefresh(transactional, paymentTransaction, context);
 
                 entitySqlDaoWrapperFactory.become(PaymentSqlDao.class).updatePaymentForNewTransaction(paymentId.toString(), contextWithUpdatedDate(context));
 
@@ -428,9 +428,7 @@ public class DefaultPaymentDao extends EntityDaoBase<PaymentModelDao, Payment, P
     private PaymentMethodModelDao insertPaymentMethodInTransaction(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory, final PaymentMethodModelDao paymentMethod, final InternalCallContext context)
             throws EntityPersistenceException {
         final PaymentMethodSqlDao transactional = entitySqlDaoWrapperFactory.become(PaymentMethodSqlDao.class);
-        transactional.create(paymentMethod, context);
-
-        return transactional.getById(paymentMethod.getId().toString(), context);
+        return createAndRefresh(transactional, paymentMethod, context);
     }
 
     @Override
diff --git a/subscription/src/main/java/org/killbill/billing/subscription/engine/dao/DefaultSubscriptionDao.java b/subscription/src/main/java/org/killbill/billing/subscription/engine/dao/DefaultSubscriptionDao.java
index af140fb..332d497 100644
--- a/subscription/src/main/java/org/killbill/billing/subscription/engine/dao/DefaultSubscriptionDao.java
+++ b/subscription/src/main/java/org/killbill/billing/subscription/engine/dao/DefaultSubscriptionDao.java
@@ -1,7 +1,7 @@
 /*
  * Copyright 2010-2013 Ning, Inc.
- * Copyright 2014-2016 Groupon, Inc
- * Copyright 2014-2016 The Billing Project, LLC
+ * Copyright 2014-2017 Groupon, Inc
+ * Copyright 2014-2017 The Billing Project, LLC
  *
  * The Billing Project 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
@@ -258,10 +258,9 @@ public class DefaultSubscriptionDao extends EntityDaoBase<SubscriptionBundleMode
         return transactionalSqlDao.execute(new EntitySqlDaoTransactionWrapper<SubscriptionBaseBundle>() {
             @Override
             public SubscriptionBaseBundle inTransaction(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory) throws EntityPersistenceException {
-
                 final SubscriptionBundleModelDao model = new SubscriptionBundleModelDao(bundle);
-                entitySqlDaoWrapperFactory.become(BundleSqlDao.class).create(model, context);
-                final SubscriptionBundleModelDao result = entitySqlDaoWrapperFactory.become(BundleSqlDao.class).getById(bundle.getId().toString(), context);
+                final BundleSqlDao bundleSqlDao = entitySqlDaoWrapperFactory.become(BundleSqlDao.class);
+                final SubscriptionBundleModelDao result = createAndRefresh(bundleSqlDao, model, context);
                 return SubscriptionBundleModelDao.toSubscriptionbundle(result);
             }
         });
@@ -425,7 +424,7 @@ public class DefaultSubscriptionDao extends EntityDaoBase<SubscriptionBundleMode
                 final SubscriptionEventSqlDao transactional = entitySqlDaoWrapperFactory.become(SubscriptionEventSqlDao.class);
                 final UUID subscriptionId = subscription.getId();
                 cancelNextPhaseEventFromTransaction(subscriptionId, entitySqlDaoWrapperFactory, context);
-                transactional.create(new SubscriptionEventModelDao(nextPhaseEvent), context);
+                createAndRefresh(transactional, new SubscriptionEventModelDao(nextPhaseEvent), context);
                 recordFutureNotificationFromTransaction(entitySqlDaoWrapperFactory,
                                                         nextPhaseEvent.getEffectiveDate(),
                                                         new SubscriptionNotificationKey(nextPhaseEvent.getId()), context);
@@ -503,11 +502,11 @@ public class DefaultSubscriptionDao extends EntityDaoBase<SubscriptionBundleMode
             @Override
             public Void inTransaction(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory) throws Exception {
                 final SubscriptionSqlDao transactional = entitySqlDaoWrapperFactory.become(SubscriptionSqlDao.class);
-                transactional.create(new SubscriptionModelDao(subscription), context);
+                createAndRefresh(transactional, new SubscriptionModelDao(subscription), context);
 
                 final SubscriptionEventSqlDao eventsDaoFromSameTransaction = entitySqlDaoWrapperFactory.become(SubscriptionEventSqlDao.class);
                 for (final SubscriptionBaseEvent cur : initialEvents) {
-                    eventsDaoFromSameTransaction.create(new SubscriptionEventModelDao(cur), context);
+                    createAndRefresh(eventsDaoFromSameTransaction, new SubscriptionEventModelDao(cur), context);
 
                     final boolean isBusEvent = cur.getEffectiveDate().compareTo(clock.getUTCNow()) <= 0 && (cur.getType() == EventType.API_USER);
                     recordBusOrFutureNotificationFromTransaction(subscription, cur, entitySqlDaoWrapperFactory, isBusEvent, 0, context);
@@ -534,12 +533,11 @@ public class DefaultSubscriptionDao extends EntityDaoBase<SubscriptionBundleMode
                     for (final SubscriptionBase subscriptionBase : subscription.getSubscriptionBaseList()) {
                         // Safe cast
                         final DefaultSubscriptionBase defaultSubscriptionBase = (DefaultSubscriptionBase) subscriptionBase;
-
-                        transactional.create(new SubscriptionModelDao(defaultSubscriptionBase), context);
+                        createAndRefresh(transactional, new SubscriptionModelDao(defaultSubscriptionBase), context);
 
                         final List<SubscriptionBaseEvent> initialEvents = initialEventsMap.get(defaultSubscriptionBase.getId());
                         for (final SubscriptionBaseEvent cur : initialEvents) {
-                            eventsDaoFromSameTransaction.create(new SubscriptionEventModelDao(cur), context);
+                            createAndRefresh(eventsDaoFromSameTransaction, new SubscriptionEventModelDao(cur), context);
 
                             final boolean isBusEvent = cur.getEffectiveDate().compareTo(clock.getUTCNow()) <= 0 && (cur.getType() == EventType.API_USER);
                             recordBusOrFutureNotificationFromTransaction(defaultSubscriptionBase, cur, entitySqlDaoWrapperFactory, isBusEvent, 0, context);
@@ -654,7 +652,7 @@ public class DefaultSubscriptionDao extends EntityDaoBase<SubscriptionBundleMode
                 cancelFutureEventsFromTransaction(subscriptionId, changeEvents.get(0).getEffectiveDate(), entitySqlDaoWrapperFactory, false, context);
 
                 for (final SubscriptionBaseEvent cur : changeEvents) {
-                    transactional.create(new SubscriptionEventModelDao(cur), context);
+                    createAndRefresh(transactional, new SubscriptionEventModelDao(cur), context);
 
                     final boolean isBusEvent = cur.getEffectiveDate().compareTo(clock.getUTCNow()) <= 0 && (cur.getType() == EventType.API_USER);
                     recordBusOrFutureNotificationFromTransaction(subscription, cur, entitySqlDaoWrapperFactory, isBusEvent, 0, context);
@@ -701,7 +699,8 @@ public class DefaultSubscriptionDao extends EntityDaoBase<SubscriptionBundleMode
             throws EntityPersistenceException {
         final UUID subscriptionId = subscription.getId();
         cancelFutureEventsFromTransaction(subscriptionId, cancelEvent.getEffectiveDate(), entitySqlDaoWrapperFactory, true, context);
-        entitySqlDaoWrapperFactory.become(SubscriptionEventSqlDao.class).create(new SubscriptionEventModelDao(cancelEvent), context);
+        final SubscriptionEventSqlDao subscriptionEventSqlDao = entitySqlDaoWrapperFactory.become(SubscriptionEventSqlDao.class);
+        createAndRefresh(subscriptionEventSqlDao, new SubscriptionEventModelDao(cancelEvent), context);
 
         final boolean isBusEvent = cancelEvent.getEffectiveDate().compareTo(clock.getUTCNow()) <= 0;
         recordBusOrFutureNotificationFromTransaction(subscription, cancelEvent, entitySqlDaoWrapperFactory, isBusEvent, seqId, context);
@@ -957,7 +956,7 @@ public class DefaultSubscriptionDao extends EntityDaoBase<SubscriptionBundleMode
             @Override
             public Void inTransaction(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory) throws Exception {
                 final SubscriptionEventSqlDao transactional = entitySqlDaoWrapperFactory.become(SubscriptionEventSqlDao.class);
-                transactional.create(new SubscriptionEventModelDao(bcdEvent), context);
+                createAndRefresh(transactional, new SubscriptionEventModelDao(bcdEvent), context);
 
                 // Notify the Bus
                 notifyBusOfRequestedChange(entitySqlDaoWrapperFactory, subscription, bcdEvent, SubscriptionBaseTransitionType.BCD_CHANGE, context);
@@ -1072,20 +1071,20 @@ public class DefaultSubscriptionDao extends EntityDaoBase<SubscriptionBundleMode
         for (final SubscriptionTransferData curSubscription : bundleTransferData.getSubscriptions()) {
             final DefaultSubscriptionBase subData = curSubscription.getData();
             for (final SubscriptionBaseEvent curEvent : curSubscription.getInitialEvents()) {
-                transactional.create(new SubscriptionEventModelDao(curEvent), context);
+                createAndRefresh(transactional, new SubscriptionEventModelDao(curEvent), context);
                 recordFutureNotificationFromTransaction(entitySqlDaoWrapperFactory,
                                                         curEvent.getEffectiveDate(),
                                                         new SubscriptionNotificationKey(curEvent.getId()),
                                                         context);
             }
-            transSubDao.create(new SubscriptionModelDao(subData), context);
+            createAndRefresh(transSubDao, new SubscriptionModelDao(subData), context);
 
             // Notify the Bus of the latest requested change
             final SubscriptionBaseEvent finalEvent = curSubscription.getInitialEvents().get(curSubscription.getInitialEvents().size() - 1);
             notifyBusOfRequestedChange(entitySqlDaoWrapperFactory, subData, finalEvent, SubscriptionBaseTransitionType.TRANSFER, context);
         }
 
-        transBundleDao.create(new SubscriptionBundleModelDao(bundleData), context);
+        createAndRefresh(transBundleDao, new SubscriptionBundleModelDao(bundleData), context);
     }
 
     //
diff --git a/tenant/src/main/java/org/killbill/billing/tenant/dao/DefaultTenantDao.java b/tenant/src/main/java/org/killbill/billing/tenant/dao/DefaultTenantDao.java
index 81477ce..18f2238 100644
--- a/tenant/src/main/java/org/killbill/billing/tenant/dao/DefaultTenantDao.java
+++ b/tenant/src/main/java/org/killbill/billing/tenant/dao/DefaultTenantDao.java
@@ -1,7 +1,7 @@
 /*
  * Copyright 2010-2013 Ning, Inc.
- * Copyright 2014-2016 Groupon, Inc
- * Copyright 2014-2016 The Billing Project, LLC
+ * Copyright 2014-2017 Groupon, Inc
+ * Copyright 2014-2017 The Billing Project, LLC
  *
  * The Billing Project 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
@@ -98,7 +98,8 @@ public class DefaultTenantDao extends EntityDaoBase<TenantModelDao, Tenant, Tena
                 final TenantModelDao tenantModelDaoWithSecret = new TenantModelDao(entity.getId(), context.getCreatedDate(), context.getUpdatedDate(),
                                                                                    entity.getExternalKey(), entity.getApiKey(),
                                                                                    hashedPasswordBase64, salt.toBase64());
-                entitySqlDaoWrapperFactory.become(TenantSqlDao.class).create(tenantModelDaoWithSecret, context);
+                final TenantSqlDao tenantSqlDao = entitySqlDaoWrapperFactory.become(TenantSqlDao.class);
+                createAndRefresh(tenantSqlDao, tenantModelDaoWithSecret, context);
                 return null;
             }
         });
@@ -145,8 +146,7 @@ public class DefaultTenantDao extends EntityDaoBase<TenantModelDao, Tenant, Tena
                 if (uniqueKey) {
                     deleteFromTransaction(key, entitySqlDaoWrapperFactory, context);
                 }
-                tenantKVSqlDao.create(tenantKVModelDao, context);
-                final TenantKVModelDao rehydrated = tenantKVSqlDao.getById(tenantKVModelDao.getId().toString(), context);
+                final TenantKVModelDao rehydrated = createAndRefresh(tenantKVSqlDao, tenantKVModelDao, context);
                 broadcastConfigurationChangeFromTransaction(rehydrated.getRecordId(), key, entitySqlDaoWrapperFactory, context);
                 return null;
             }
@@ -164,14 +164,14 @@ public class DefaultTenantDao extends EntityDaoBase<TenantModelDao, Tenant, Tena
                 // Retrieve all values for key ordered with recordId (last at the end)
                 final List<TenantKVModelDao> tenantKV = tenantKVSqlDao.getTenantValueForKey(key, context);
                 final String id;
+                final TenantKVModelDao rehydrated;
                 if (!tenantKV.isEmpty()) {
                     id = tenantKV.get(tenantKV.size() - 1).getId().toString();
-                    tenantKVSqlDao.updateTenantValueKey(id, value, context);
+                    rehydrated = (TenantKVModelDao) tenantKVSqlDao.updateTenantValueKey(id, value, context);
                 } else {
                     id = tenantKVModelDao.getId().toString();
-                    tenantKVSqlDao.create(tenantKVModelDao, context);
+                    rehydrated = createAndRefresh(tenantKVSqlDao, tenantKVModelDao, context);
                 }
-                final TenantKVModelDao rehydrated = tenantKVSqlDao.getById(id, context);
                 broadcastConfigurationChangeFromTransaction(rehydrated.getRecordId(), key, entitySqlDaoWrapperFactory, context);
                 return null;
             }
@@ -225,7 +225,8 @@ public class DefaultTenantDao extends EntityDaoBase<TenantModelDao, Tenant, Tena
                                                              final InternalCallContext context) throws EntityPersistenceException {
         if (isSystemKey(key)) {
             final TenantBroadcastModelDao broadcast = new TenantBroadcastModelDao(kvRecordId, key, context.getUserToken());
-            entitySqlDaoWrapperFactory.become(TenantBroadcastSqlDao.class).create(broadcast, context);
+            final TenantBroadcastSqlDao tenantBroadcastSqlDao = entitySqlDaoWrapperFactory.become(TenantBroadcastSqlDao.class);
+            createAndRefresh(tenantBroadcastSqlDao, broadcast, context);
         }
     }
 
diff --git a/tenant/src/main/java/org/killbill/billing/tenant/dao/TenantKVSqlDao.java b/tenant/src/main/java/org/killbill/billing/tenant/dao/TenantKVSqlDao.java
index fd523cb..c0cac34 100644
--- a/tenant/src/main/java/org/killbill/billing/tenant/dao/TenantKVSqlDao.java
+++ b/tenant/src/main/java/org/killbill/billing/tenant/dao/TenantKVSqlDao.java
@@ -49,9 +49,7 @@ public interface TenantKVSqlDao extends EntitySqlDao<TenantKVModelDao, TenantKV>
 
     @SqlUpdate
     @Audited(ChangeType.UPDATE)
-    public void updateTenantValueKey(@Bind("id") final String id,
-                                     @Bind("tenantValue") final String tenantValue,
-                                     @BindBean final InternalCallContext context);
-
-
+    public Object updateTenantValueKey(@Bind("id") final String id,
+                                       @Bind("tenantValue") final String tenantValue,
+                                       @BindBean final InternalCallContext context);
 }
diff --git a/util/src/main/java/org/killbill/billing/util/entity/dao/EntityDaoBase.java b/util/src/main/java/org/killbill/billing/util/entity/dao/EntityDaoBase.java
index c152b58..72287d4 100644
--- a/util/src/main/java/org/killbill/billing/util/entity/dao/EntityDaoBase.java
+++ b/util/src/main/java/org/killbill/billing/util/entity/dao/EntityDaoBase.java
@@ -1,7 +1,9 @@
 /*
- * Copyright 2010-2012 Ning, Inc.
+ * Copyright 2010-2013 Ning, Inc.
+ * Copyright 2014-2017 Groupon, Inc
+ * Copyright 2014-2017 The Billing Project, LLC
  *
- * Ning licenses this file to you under the Apache License, version 2.0
+ * The Billing Project 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:
  *
@@ -22,6 +24,7 @@ import java.util.UUID;
 import org.killbill.billing.BillingExceptionBase;
 import org.killbill.billing.callcontext.InternalCallContext;
 import org.killbill.billing.callcontext.InternalTenantContext;
+import org.killbill.billing.entity.EntityPersistenceException;
 import org.killbill.billing.util.audit.ChangeType;
 import org.killbill.billing.util.entity.DefaultPagination;
 import org.killbill.billing.util.entity.Entity;
@@ -56,9 +59,7 @@ public abstract class EntityDaoBase<M extends EntityModelDao<E>, E extends Entit
                 if (checkEntityAlreadyExists(transactional, entity, context)) {
                     throw generateAlreadyExistsException(entity, context);
                 }
-                transactional.create(entity, context);
-
-                final M refreshedEntity = transactional.getById(entity.getId().toString(), context);
+                final M refreshedEntity = createAndRefresh(transactional, entity, context);
 
                 postBusEventFromTransaction(entity, refreshedEntity, ChangeType.INSERT, entitySqlDaoWrapperFactory, context);
                 return null;
@@ -66,6 +67,11 @@ public abstract class EntityDaoBase<M extends EntityModelDao<E>, E extends Entit
         };
     }
 
+    protected <F extends EntityModelDao> F createAndRefresh(final EntitySqlDao transactional, final F entity, final InternalCallContext context) throws EntityPersistenceException {
+        // We have overridden the jDBI return type in EntitySqlDaoWrapperInvocationHandler
+        return (F) transactional.create(entity, context);
+    }
+
     protected boolean checkEntityAlreadyExists(final EntitySqlDao<M, E> transactional, final M entity, final InternalCallContext context) {
         return transactional.getRecordId(entity.getId().toString(), context) != null;
     }
diff --git a/util/src/main/java/org/killbill/billing/util/entity/dao/EntitySqlDao.java b/util/src/main/java/org/killbill/billing/util/entity/dao/EntitySqlDao.java
index d20a5e7..ede762c 100644
--- a/util/src/main/java/org/killbill/billing/util/entity/dao/EntitySqlDao.java
+++ b/util/src/main/java/org/killbill/billing/util/entity/dao/EntitySqlDao.java
@@ -45,8 +45,8 @@ public interface EntitySqlDao<M extends EntityModelDao<E>, E extends Entity> ext
 
     @SqlUpdate
     @Audited(ChangeType.INSERT)
-    public void create(@BindBean final M entity,
-                       @BindBean final InternalCallContext context) throws EntityPersistenceException;
+    public Object create(@BindBean final M entity,
+                         @BindBean final InternalCallContext context) throws EntityPersistenceException;
 
     @SqlQuery
     public M getById(@Bind("id") final String id,
diff --git a/util/src/main/java/org/killbill/billing/util/entity/dao/EntitySqlDaoWrapperInvocationHandler.java b/util/src/main/java/org/killbill/billing/util/entity/dao/EntitySqlDaoWrapperInvocationHandler.java
index 465c1b6..efc6b37 100644
--- a/util/src/main/java/org/killbill/billing/util/entity/dao/EntitySqlDaoWrapperInvocationHandler.java
+++ b/util/src/main/java/org/killbill/billing/util/entity/dao/EntitySqlDaoWrapperInvocationHandler.java
@@ -94,7 +94,7 @@ public class EntitySqlDaoWrapperInvocationHandler<S extends EntitySqlDao<M, E>, 
     private final Clock clock;
     private final NonEntityDao nonEntityDao;
     private final InternalCallContextFactory internalCallContextFactory;
-    private final Profiling prof;
+    private final Profiling<Object, Throwable> prof;
 
     public EntitySqlDaoWrapperInvocationHandler(final Class<S> sqlDaoClass,
                                                 final S sqlDao,
@@ -117,7 +117,7 @@ public class EntitySqlDaoWrapperInvocationHandler<S extends EntitySqlDao<M, E>, 
     @Override
     public Object invoke(final Object proxy, final Method method, final Object[] args) throws Throwable {
         try {
-            return prof.executeWithProfiling(ProfilingFeatureType.DAO, sqlDaoClass.getSimpleName() + ":" + method.getName(), new WithProfilingCallback() {
+            return prof.executeWithProfiling(ProfilingFeatureType.DAO, sqlDaoClass.getSimpleName() + ":" + method.getName(), new WithProfilingCallback<Object, Throwable>() {
                 @Override
                 public Object execute() throws Throwable {
                     return invokeSafely(proxy, method, args);
@@ -204,7 +204,7 @@ public class EntitySqlDaoWrapperInvocationHandler<S extends EntitySqlDao<M, E>, 
     }
 
     private Object invokeRaw(final Method method, final Object[] args) throws Throwable {
-        return prof.executeWithProfiling(ProfilingFeatureType.DAO_DETAILS, sqlDaoClass.getSimpleName() + " (raw):" + method.getName(), new WithProfilingCallback() {
+        return prof.executeWithProfiling(ProfilingFeatureType.DAO_DETAILS, sqlDaoClass.getSimpleName() + " (raw):" + method.getName(), new WithProfilingCallback<Object, Throwable>() {
             @Override
             public Object execute() throws Throwable {
                 Object result = method.invoke(sqlDao, args);
@@ -327,17 +327,26 @@ public class EntitySqlDaoWrapperInvocationHandler<S extends EntitySqlDao<M, E>, 
         }
 
         // Real jdbc call
-        final Object obj = prof.executeWithProfiling(ProfilingFeatureType.DAO_DETAILS, sqlDaoClass.getSimpleName() + " (raw) :", new WithProfilingCallback() {
+        final Object obj = prof.executeWithProfiling(ProfilingFeatureType.DAO_DETAILS, sqlDaoClass.getSimpleName() + " (raw) :", new WithProfilingCallback<Object, Throwable>() {
             @Override
             public Object execute() throws Throwable {
                 return method.invoke(sqlDao, args);
             }
         });
 
+        M m = null;
         for (final String entityId : entityIds) {
-            updateHistoryAndAudit(entityId, deletedEntities.get(entityId), changeType, context);
+            m = updateHistoryAndAudit(entityId, deletedEntities.get(entityId), changeType, context);
+        }
+
+        // PERF: override the return value with the reHydrated entity to avoid an extra 'get' in the transaction,
+        // (see EntityDaoBase#createAndRefresh for an example, but it works for updates as well).
+        if (entityIds.size() == 1) {
+            return m;
+        } else {
+            // jDBI will return the number of rows modified otherwise
+            return obj;
         }
-        return obj;
     }
 
     private void populateCacheOnGetByIdInvocation(M model) {
@@ -365,10 +374,10 @@ public class EntitySqlDaoWrapperInvocationHandler<S extends EntitySqlDao<M, E>, 
                rawKey;
     }
 
-    private void updateHistoryAndAudit(final String entityId, @Nullable final M deletedEntity, final ChangeType changeType, final InternalCallContext context) throws Throwable {
-        prof.executeWithProfiling(ProfilingFeatureType.DAO_DETAILS, sqlDaoClass.getSimpleName() + " (history/audit) :", new WithProfilingCallback() {
+    private M updateHistoryAndAudit(final String entityId, @Nullable final M deletedEntity, final ChangeType changeType, final InternalCallContext context) throws Throwable {
+        final Object reHydratedEntity = prof.executeWithProfiling(ProfilingFeatureType.DAO_DETAILS, sqlDaoClass.getSimpleName() + " (history/audit) :", new WithProfilingCallback<Object, Throwable>() {
             @Override
-            public Object execute() {
+            public M execute() throws Throwable {
                 final M reHydratedEntity;
                 if (changeType == ChangeType.DELETE) {
                     reHydratedEntity = deletedEntity;
@@ -390,9 +399,11 @@ public class EntitySqlDaoWrapperInvocationHandler<S extends EntitySqlDao<M, E>, 
 
                 // Make sure to re-hydrate the object (especially needed for create calls)
                 insertAudits(tableName, reHydratedEntity, entityRecordId, historyRecordId, changeType, context);
-                return null;
+                return reHydratedEntity;
             }
         });
+        //noinspection unchecked
+        return (M) reHydratedEntity;
     }
 
     private List<String> retrieveEntityIdsFromArguments(final Method method, final Object[] args) {
diff --git a/util/src/main/java/org/killbill/billing/util/tag/dao/DefaultTagDefinitionDao.java b/util/src/main/java/org/killbill/billing/util/tag/dao/DefaultTagDefinitionDao.java
index 6c78854..90d2d75 100644
--- a/util/src/main/java/org/killbill/billing/util/tag/dao/DefaultTagDefinitionDao.java
+++ b/util/src/main/java/org/killbill/billing/util/tag/dao/DefaultTagDefinitionDao.java
@@ -155,7 +155,7 @@ public class DefaultTagDefinitionDao extends EntityDaoBase<TagDefinitionModelDao
 
                     // Create it
                     final TagDefinitionModelDao tagDefinition = new TagDefinitionModelDao(context.getCreatedDate(), definitionName, description);
-                    tagDefinitionSqlDao.create(tagDefinition, context);
+                    createAndRefresh(tagDefinitionSqlDao, tagDefinition, context);
 
                     // Post an event to the bus
                     final boolean isControlTag = TagModelDaoHelper.isControlTag(tagDefinition.getName());