From 292188ad7c8abe9cf5e8f1814dc0c48f470cac5c Mon Sep 17 00:00:00 2001
From: Mike Lang <mikelang3000@gmail.com>
Date: Fri, 24 May 2019 19:19:20 -0700
Subject: [PATCH] database: Remove retry_on_conflict helper and default to
 autocommit

All our usage was of a single query anyway, so autocommit is easier to handle.
You can still opt into a longer transaction using the transaction() helper.
---
 common/common/database.py | 38 +++++++++++++++++++-------------------
 1 file changed, 19 insertions(+), 19 deletions(-)

diff --git a/common/common/database.py b/common/common/database.py
index ce3cfd8..1140857 100644
--- a/common/common/database.py
+++ b/common/common/database.py
@@ -5,6 +5,8 @@ Note that this code requires psycopg2 and psycogreen, but the common module
 as a whole does not to avoid needing to install them for components that don't need it.
 """
 
+from contextlib import contextmanager
+
 import psycopg2
 import psycopg2.extensions
 import psycopg2.extras
@@ -70,16 +72,15 @@ class DBManager(object):
 	the initial conn used to apply the schema, and you should use a real conn pool for
 	any non-trivial use.
 
-	Returned conns are set to seralizable isolation level, and use NamedTupleCursor cursors.
+	Returned conns are set to seralizable isolation level, autocommit, and use NamedTupleCursor cursors.
 	"""
 	def __init__(self, **connect_kwargs):
 		patch_psycopg()
 		self.conns = []
 		self.connect_kwargs = connect_kwargs
 		conn = self.get_conn()
-		with conn:
-			with conn.cursor() as cur:
-				cur.execute(SCHEMA)
+		with transaction(conn):
+			query(conn, SCHEMA)
 		self.put_conn(conn)
 
 	def put_conn(self, conn):
@@ -93,24 +94,23 @@ class DBManager(object):
 		# we don't care about the performance concerns and everything we do is easily retryable.
 		# This shouldn't matter in practice anyway since everything we're doing is either read-only
 		# searches or targetted single-row updates.
-		conn.set_session(psycopg2.extensions.ISOLATION_LEVEL_SERIALIZABLE)
+		conn.isolation_level = psycopg2.extensions.ISOLATION_LEVEL_SERIALIZABLE
+		conn.autocommit = True
 		return conn
 
 
-def retry_on_conflict(conn, func, *args, **kwargs):
-	"""Run func(conn, *args, **kwargs) in a transaction up to max_tries (given as kwarg) times
-	(at least once, default 5), retrying if it raises an error indicating a transaction conflict.
-	After max_tries, raises TransactionRollbackError.
-	"""
-	max_tries = kwargs.pop('max_tries', 5)
-	for _ in range(max_tries - 1):
-		try:
-			with conn:
-				return func(conn)
-		except psycopg2.extensions.TransactionRollbackError:
-			pass
-	with conn:
-		return func(conn)
+@contextmanager
+def transaction(conn):
+	"""Helper context manager that runs the code block as a single database transaction
+	instead of in autocommit mode. The only difference between this and "with conn" is
+	that we explicitly disable then re-enable autocommit."""
+	old_autocommit = conn.autocommit
+	conn.autocommit = False
+	try:
+		with conn:
+			yield
+	finally:
+		conn.autocommit = old_autocommit
 
 
 def query(conn, query, *args, **kwargs):