From f1f0b069345fbf4d7170e6dba1fd3b0d111eb6cc Mon Sep 17 00:00:00 2001
From: Olly Cope <olly@ollycope.com>
Date: Sat, 29 Oct 2022 16:27:36 +0000
Subject: postgresql: put connection into autocommit mode, handle transactions
 manually

Both psycopg2 and psycopg3 inject a BEGIN statement before the first SQL
statement. But yoyo attempts to manage its own transactions, sending a BEGIN
statement, duplicating the one already sent by psycopg. This duplicate BEGIN
raises a warning in PostgreSQL and an error in CockroachDB, making yoyo
unusable with CockroachDB.

Setting `connection.autocommit` connection stops psycopg from injecting BEGIN
statements, giving yoyo full control over the emitted SQL.

https://www.psycopg.org/docs/usage.html#transactions-control
https://www.psycopg.org/psycopg3/docs/basic/transactions.html
---
 yoyo/backends/core/postgresql.py | 31 +++++++++++++++++++++++++++++--
 1 file changed, 29 insertions(+), 2 deletions(-)

diff --git a/yoyo/backends/core/postgresql.py b/yoyo/backends/core/postgresql.py
index d879928..17e9963 100644
--- a/yoyo/backends/core/postgresql.py
+++ b/yoyo/backends/core/postgresql.py
@@ -19,6 +19,13 @@ 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
@@ -34,7 +41,14 @@ class PostgresqlBackend(DatabaseBackend):
         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
@@ -45,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
 
     def transaction(self, rollback_on_exit=False):
 
@@ -74,6 +91,16 @@ 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()
+
 
 class PostgresqlPsycopgBackend(PostgresqlBackend):
     """
-- 
cgit v1.2.1