From 8070a6d7de56f384fc0ad0d382bebf7d678f789e Mon Sep 17 00:00:00 2001
From: dfahlander <david.fahlander@gmail.com>
Date: Mon, 12 Nov 2018 11:12:38 +0100
Subject: [PATCH] Resolving issue #612 Calling upgrader the same way as it is
 done in transaction scope

---
 src/classes/version/schema-helpers.ts | 20 +++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

diff --git a/src/classes/version/schema-helpers.ts b/src/classes/version/schema-helpers.ts
index 6d3c01be8..8a9c53eff 100644
--- a/src/classes/version/schema-helpers.ts
+++ b/src/classes/version/schema-helpers.ts
@@ -3,7 +3,7 @@ import { DbSchema } from '../../public/types/db-schema';
 import { setProp, keys, slice, _global, isArray, shallowClone } from '../../functions/utils';
 import { Transaction } from '../transaction';
 import { Version } from './version';
-import Promise, { PSD, newScope } from '../../helpers/promise';
+import Promise, { PSD, newScope, NativePromise, decrementExpectedAwaits, AsyncFunction, incrementExpectedAwaits } from '../../helpers/promise';
 import { exceptions } from '../../errors';
 import { TableSchema } from '../../public/types/table-schema';
 import { IndexSpec } from '../../public/types/index-spec';
@@ -138,9 +138,23 @@ export function updateTablesAndIndexes(
         setApiOnPlace(db, [db.Transaction.prototype], keys(upgradeSchema), upgradeSchema);
         trans.schema = upgradeSchema;
 
-        return Promise.follow(() => {
-          contentUpgrade(trans);
+        // Support for native async await.
+        if (contentUpgrade.constructor === AsyncFunction) {
+          incrementExpectedAwaits();
+        }
+        let returnValue: any;
+        const promiseFollowed = Promise.follow(() => {
+          // Finally, call the scope function with our table and transaction arguments.
+          returnValue = contentUpgrade(trans);
+          if (returnValue) {
+            if (returnValue.constructor === NativePromise) {
+              var decrementor = decrementExpectedAwaits.bind(null, null);
+              returnValue.then(decrementor, decrementor);
+            }
+          }
         });
+        return (returnValue && typeof returnValue.then === 'function' ?
+          Promise.resolve(returnValue) : promiseFollowed.then(()=>returnValue));
       }
     });
     queue.push(idbtrans => {