Merge pull request #6 from flea89/update-transaction

Newer version of gcloudc use django transaction

Author: adamalton

massmigration/utils/transaction.py

@@ -1,7 +1,6 @@
 # Third party
 from django.conf import settings
 from django.db import router, transaction as django_transaction
-from gcloudc.db import transaction as datastore_transaction
 
 # Mass Migration
 from massmigration.models import MigrationRecord
@@ -13,8 +12,15 @@
 # That said, should you be allowed to run the same migration on two different databases?
 
 
-def get_transaction(db_alias):
-    engine = settings.DATABASES[db_alias]["ENGINE"]
-    if engine == "gcloudc.db.backends.datastore":
-        return datastore_transaction
-    return django_transaction
+def get_transaction():
+    connection = router.db_for_write(MigrationRecord)
+    engine = settings.DATABASES[connection]["ENGINE"]
+    transaction = django_transaction
+    if engine.startswith("gcloudc.db.backends"):
+        try:
+            from gcloudc.db import transaction as gcloudc_transaction
+            transaction = gcloudc_transaction
+        except ImportError:
+            # Newer versions of gcloudc use Django transactions.
+            pass
+    return transaction