From 3bc6b554cd50fe8a5413c5f3ae9ccadb706e60f5 Mon Sep 17 00:00:00 2001
From: Pat Riehecky <riehecky@fnal.gov>
Date: Mon, 11 Jan 2021 11:32:12 -0600
Subject: [PATCH] Fixes #357 add option to alter isolation level

---
 README.rst                         | 3 +++
 src/pytest_postgresql/factories.py | 7 +++++--
 src/pytest_postgresql/janitor.py   | 8 ++++++--
 3 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/README.rst b/README.rst
index 74aa907d..0bc9d48f 100644
--- a/README.rst
+++ b/README.rst
@@ -297,6 +297,9 @@ or use it as a context manager:
     DatabaseJanitor manages the state of the database, but you'll have to create
     connection to use in test code yourself.
 
+    You can optionally pass in a recognized postgresql ISOLATION_LEVEL for
+    additional control.
+
 Package resources
 -----------------
 
diff --git a/src/pytest_postgresql/factories.py b/src/pytest_postgresql/factories.py
index fafc1908..8ffe4aa6 100644
--- a/src/pytest_postgresql/factories.py
+++ b/src/pytest_postgresql/factories.py
@@ -213,7 +213,8 @@ def postgresql_noproc_fixture(request: FixtureRequest) -> NoopExecutor:
 
 
 def postgresql(
-        process_fixture_name: str, db_name: str = None, load: List[str] = None
+        process_fixture_name: str, db_name: str = None, load: List[str] = None,
+        isolation_level: int = 0,
 ) -> Callable[[FixtureRequest], connection]:
     """
     Return connection fixture factory for PostgreSQL.
@@ -221,6 +222,8 @@ def postgresql(
     :param process_fixture_name: name of the process fixture
     :param db_name: database name
     :param load: SQL to automatically load into our test database
+    :param isolation_level: optional postgresql isolation level
+                            defaults to ISOLATION_LEVEL_AUTOCOMMIT
     :returns: function which makes a connection to postgresql
     """
 
@@ -252,7 +255,7 @@ def postgresql_factory(request: FixtureRequest) -> connection:
 
         with DatabaseJanitor(
                 pg_user, pg_host, pg_port, pg_db, proc_fixture.version,
-                pg_password
+                pg_password, isolation_level
         ):
             db_connection: connection = psycopg2.connect(
                 dbname=pg_db,
diff --git a/src/pytest_postgresql/janitor.py b/src/pytest_postgresql/janitor.py
index 5ec63819..a1ce3894 100644
--- a/src/pytest_postgresql/janitor.py
+++ b/src/pytest_postgresql/janitor.py
@@ -23,7 +23,8 @@ def __init__(
             port: str,
             db_name: str,
             version: Union[str, float, Version],
-            password: str = None
+            password: str = None,
+            isolation_level: int = 0,
     ) -> None:
         """
         Initialize janitor.
@@ -34,12 +35,15 @@ def __init__(
         :param db_name: database name
         :param version: postgresql version number
         :param password: optional postgresql password
+        :param isolation_level: optional postgresql isolation level
+                                defaults to ISOLATION_LEVEL_AUTOCOMMIT
         """
         self.user = user
         self.password = password
         self.host = host
         self.port = port
         self.db_name = db_name
+        self.isolation_level = isolation_level
         if not isinstance(version, Version):
             self.version = parse_version(str(version))
         else:
@@ -79,7 +83,7 @@ def cursor(self) -> cursor:
             host=self.host,
             port=self.port,
         )
-        conn.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
+        conn.set_isolation_level(self.isolation_level)
         cur = conn.cursor()
         try:
             yield cur