merge cockroachdb

7 files changed, 86 insertions, 44 deletions

diff --git a/tox.ini b/tox.ini
index f9d9dbc..046d0d9 100755
--- a/tox.ini
+++ b/tox.ini
@@ -40,3 +40,6 @@ ignore =
     E203
     W503
 max-line-length = 88
+
+[pytest]
+addopts = -Werror::UserWarning
diff --git a/yoyo/backends/base.py b/yoyo/backends/base.py
index 3c1f9a7..8ca956e 100644
--- a/yoyo/backends/base.py
+++ b/yoyo/backends/base.py
@@ -45,9 +45,9 @@ class TransactionManager:
     when the context manager block closes
     """
 
-    def __init__(self, backend):
+    def __init__(self, backend, rollback_on_exit=False):
         self.backend = backend
-        self._rollback = False
+        self.rollback_on_exit = rollback_on_exit
 
     def __enter__(self):
         self._do_begin()
@@ -58,18 +58,11 @@ class TransactionManager:
             self._do_rollback()
             return None
 
-        if self._rollback:
+        if self.rollback_on_exit:
             self._do_rollback()
         else:
             self._do_commit()
 
-    def rollback(self):
-        """
-        Flag that the transaction will be rolled back when the with statement
-        exits
-        """
-        self._rollback = True
-
     def _do_begin(self):
         """
         Instruct the backend to begin a transaction
@@ -237,9 +230,12 @@ class DatabaseBackend:
         table_name = "yoyo_tmp_{}".format(utils.get_random_string(10))
         table_name_quoted = self.quote_identifier(table_name)
         sql = self.create_test_table_sql.format(table_name_quoted=table_name_quoted)
-        with self.transaction() as t:
-            self.execute(sql)
-            t.rollback()
+        try:
+            with self.transaction(rollback_on_exit=True):
+                self.execute(sql)
+        except self.DatabaseError:
+            return False
+
         try:
             with self.transaction():
                 self.execute("DROP TABLE {}".format(table_name_quoted))
@@ -259,12 +255,12 @@ class DatabaseBackend:
         )
         return [row[0] for row in cursor.fetchall()]
 
-    def transaction(self):
+    def transaction(self, rollback_on_exit=False):
         if not self._in_transaction:
-            return TransactionManager(self)
+            return TransactionManager(self, rollback_on_exit=rollback_on_exit)
 
         else:
-            return SavepointTransactionManager(self)
+            return SavepointTransactionManager(self, rollback_on_exit=rollback_on_exit)
 
     def cursor(self):
         return self.connection.cursor()
@@ -282,6 +278,7 @@ class DatabaseBackend:
         """
         Begin a new transaction
         """
+        assert not self._in_transaction
         self._in_transaction = True
         self.execute("BEGIN")
 
diff --git a/yoyo/backends/core/postgresql.py b/yoyo/backends/core/postgresql.py
index 0400ce9..09d0f06 100644
--- a/yoyo/backends/core/postgresql.py
+++ b/yoyo/backends/core/postgresql.py
@@ -12,12 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import warnings
 from contextlib import contextmanager
 
 from yoyo.backends.base import DatabaseBackend
 
 
 class PostgresqlBackend(DatabaseBackend):
+    """
+    Backend for PostgreSQL and PostgreSQL compatible databases.
+
+    This backend uses psycopg2. See
+    :class:`yoyo.backends.core.postgresql.PostgresqlPsycopgBackend`
+    if you need psycopg3.
+    """
 
     driver_module = "psycopg2"
     schema = None
@@ -26,8 +34,21 @@ class PostgresqlBackend(DatabaseBackend):
         "WHERE table_schema = :schema"
     )
 
+    @property
+    def TRANSACTION_STATUS_IDLE(self):
+        from psycopg2.extensions import TRANSACTION_STATUS_IDLE
+
+        return TRANSACTION_STATUS_IDLE
+
     def connect(self, dburi):
-        kwargs = {"dbname": dburi.database}
+        kwargs = {"dbname": dburi.database, "autocommit": True}
+
+        # Default to autocommit mode: without this psycopg sends a BEGIN before
+        # every query, causing a warning when we then explicitly start a
+        # transaction. This warning becomes an error in CockroachDB. See
+        # https://todo.sr.ht/~olly/yoyo/71
+        kwargs["autocommit"] = True
+
         kwargs.update(dburi.args)
         if dburi.username is not None:
             kwargs["user"] = dburi.username
@@ -38,7 +59,10 @@ class PostgresqlBackend(DatabaseBackend):
         if dburi.hostname is not None:
             kwargs["host"] = dburi.hostname
         self.schema = kwargs.pop("schema", None)
-        return self.driver.connect(**kwargs)
+        autocommit = bool(kwargs.pop("autocommit"))
+        connection = self.driver.connect(**kwargs)
+        connection.autocommit = autocommit
+        return connection
 
     @contextmanager
     def disable_transactions(self):
@@ -57,6 +81,25 @@ class PostgresqlBackend(DatabaseBackend):
         current_schema = self.execute("SELECT current_schema").fetchone()[0]
         return super(PostgresqlBackend, self).list_tables(schema=current_schema)
 
+    def commit(self):
+        # The connection is in autocommit mode and ignores calls to
+        # ``commit()`` and ``rollback()``, so we have to issue the SQL directly
+        self.execute("COMMIT")
+        super().commit()
+
+    def rollback(self):
+        self.execute("ROLLBACK")
+        super().rollback()
+
+    def begin(self):
+        if self.connection.info.transaction_status != self.TRANSACTION_STATUS_IDLE:
+            warnings.warn(
+                "Nested transaction requested; "
+                "this will raise an exception in some "
+                "PostgreSQL-compatible databases"
+            )
+        return super().begin()
+
 
 class PostgresqlPsycopgBackend(PostgresqlBackend):
     """
@@ -64,3 +107,9 @@ class PostgresqlPsycopgBackend(PostgresqlBackend):
     """
 
     driver_module = "psycopg"
+
+    @property
+    def TRANSACTION_STATUS_IDLE(self):
+        from psycopg.pq import TransactionStatus
+
+        return TransactionStatus.IDLE
diff --git a/yoyo/internalmigrations/__init__.py b/yoyo/internalmigrations/__init__.py
index 5bcb0b5..3164dfa 100644
--- a/yoyo/internalmigrations/__init__.py
+++ b/yoyo/internalmigrations/__init__.py
@@ -48,15 +48,12 @@ def get_current_version(backend):
         return 0
     if version_table not in tables:
         return 1
-    with backend.transaction():
-        cursor = backend.execute(
-            "SELECT max(version) FROM {} ".format(
-                backend.quote_identifier(version_table)
-            )
-        )
-        version = cursor.fetchone()[0]
-        assert version in schema_versions
-        return version
+    cursor = backend.execute(
+        f"SELECT max(version) FROM {backend.quote_identifier(version_table)}"
+    )
+    version = cursor.fetchone()[0]
+    assert version in schema_versions
+    return version
 
 
 def mark_schema_version(backend, version):
diff --git a/yoyo/migrations.py b/yoyo/migrations.py
index 5450d0a..5b30049 100755
--- a/yoyo/migrations.py
+++ b/yoyo/migrations.py
@@ -312,16 +312,14 @@ class TransactionWrapper(StepBase):
         return "<TransactionWrapper {!r}>".format(self.step)
 
     def apply(self, backend, force=False, direction="apply"):
-        with backend.transaction() as transaction:
-            try:
+        try:
+            with backend.transaction():
                 getattr(self.step, direction)(backend, force)
-            except backend.DatabaseError:
-                if force or self.ignore_errors in (direction, "all"):
-                    logger.exception("Ignored error in %r", self.step)
-                    transaction.rollback()
-                    return
-                else:
-                    raise
+        except backend.DatabaseError:
+            if force or self.ignore_errors in (direction, "all"):
+                logger.exception("Ignored error in %r", self.step)
+            else:
+                raise
 
     def rollback(self, backend, force=False):
         self.apply(backend, force, "rollback")
diff --git a/yoyo/tests/conftest.py b/yoyo/tests/conftest.py
index 1bf8d4a..f24509d 100644
--- a/yoyo/tests/conftest.py
+++ b/yoyo/tests/conftest.py
@@ -56,9 +56,9 @@ def drop_all_tables(backend):
 
 
 def drop_yoyo_tables(backend):
-    for table in backend.list_tables():
-        if table.startswith("yoyo") or table.startswith("_yoyo"):
-            with backend.transaction():
+    with backend.transaction():
+        for table in backend.list_tables():
+            if table.startswith("yoyo") or table.startswith("_yoyo"):
                 backend.execute("DROP TABLE {}".format(table))
 
 
diff --git a/yoyo/tests/test_backends.py b/yoyo/tests/test_backends.py
index 8d99009..e4b9727 100644
--- a/yoyo/tests/test_backends.py
+++ b/yoyo/tests/test_backends.py
@@ -44,11 +44,10 @@ class TestTransactionHandling(object):
         with backend.transaction():
             backend.execute("INSERT INTO yoyo_t values ('A')")
 
-            with backend.transaction() as trans:
+            with backend.transaction(rollback_on_exit=True):
                 backend.execute("INSERT INTO yoyo_t values ('B')")
-                trans.rollback()
 
-            with backend.transaction() as trans:
+            with backend.transaction():
                 backend.execute("INSERT INTO yoyo_t values ('C')")
 
         with backend.transaction():
@@ -95,12 +94,11 @@ class TestTransactionHandling(object):
         if backend.has_transactional_ddl:
             return
 
-        with backend.transaction() as trans:
+        with backend.transaction(rollback_on_exit=True):
             backend.execute("CREATE TABLE yoyo_a (id INT)")  # implicit commit
             backend.execute("INSERT INTO yoyo_a VALUES (1)")
             backend.execute("CREATE TABLE yoyo_b (id INT)")  # implicit commit
             backend.execute("INSERT INTO yoyo_b VALUES (1)")
-            trans.rollback()
 
         count_a = backend.execute("SELECT COUNT(1) FROM yoyo_a").fetchall()[0][0]
         assert count_a == 1