Make DatabaseBackend autodetect if DDL statements are transactional

2 files changed, 27 insertions, 1 deletions

diff --git a/yoyo/backends.py b/yoyo/backends.py
index 6f431c9..7ec5938 100644
--- a/yoyo/backends.py
+++ b/yoyo/backends.py
@@ -17,7 +17,7 @@ from importlib import import_module
 from itertools import count
 from logging import getLogger
 
-from . import exceptions
+from . import exceptions, utils
 from .migrations import topological_sort
 
 logger = getLogger('yoyo.migrations')
@@ -107,6 +107,7 @@ class DatabaseBackend(object):
                             "VALUES (?, ?)")
     delete_migration_sql = "DELETE FROM {0.migration_table} WHERE id=?"
     applied_ids_sql = "SELECT id FROM {0.migration_table} ORDER by ctime"
+    create_test_table_sql = "CREATE TABLE {table_name} (id INT PRIMARY KEY)"
 
     _driver = None
     _in_transaction = False
@@ -116,6 +117,7 @@ class DatabaseBackend(object):
         self._connection = self.connect(dburi)
         self.migration_table = migration_table
         self.create_migrations_table()
+        self.has_transactional_ddl = self._check_transactional_ddl()
 
     def _load_driver_module(self):
         """
@@ -136,6 +138,23 @@ class DatabaseBackend(object):
     def connection(self):
         return self._connection
 
+    def _check_transactional_ddl(self):
+        """
+        Return True if the database supports committing/rolling back
+        DDL statements within a transaction
+        """
+        table_name = '_yoyo_tmp_{}'.format(utils.get_random_string(10))
+        sql = self.create_test_table_sql.format(table_name=table_name)
+        with self.transaction() as t:
+            self.execute(sql)
+            t.rollback()
+        try:
+            with self.transaction():
+                self.execute("DROP TABLE {}".format(table_name))
+        except tuple(exceptions.DatabaseErrors):
+            return True
+        return False
+
     def transaction(self):
         if not self._in_transaction:
             return TransactionManager(self)
diff --git a/yoyo/tests/test_backends.py b/yoyo/tests/test_backends.py
index 6a6568f..6c2b43f 100644
--- a/yoyo/tests/test_backends.py
+++ b/yoyo/tests/test_backends.py
@@ -68,3 +68,10 @@ class TestTransactionHandling(object):
         with backend.transaction():
             rows = list(backend.execute("SELECT * FROM t").fetchall())
             assert rows == [('A',), ('C',)]
+
+    def test_backend_detects_transactional_ddl(self, backend):
+        expected = {backends.PostgresqlBackend: True,
+                    backends.SQLiteBackend: True,
+                    backends.MySQLBackend: False}
+        if backend.__class__ in expected:
+            assert backend.has_transactional_ddl is expected[backend.__class__]