Add postgresql support to test suite

Author: jagerman

sogs/db.py

@@ -587,24 +587,17 @@ def create_admin_user(dbconn):
     )
 
 
-if config.DB_URL.startswith('postgresql'):
-    # rooms.token is a 'citext' (case-insensitive text), which sqlalchemy doesn't recognize out of
-    # the box.  Map it to a plain TEXT which is good enough for what we need (if we actually needed
-    # to generate this wouldn't suffice: we'd have to use something like the sqlalchemy-citext
-    # module).
-    from sqlalchemy.dialects.postgresql.base import ischema_names
-
-    if 'citext' not in ischema_names:
-        ischema_names['citext'] = ischema_names['text']
-
-
 engine, engine_initial_pid, metadata = None, None, None
 
 
 def _init_engine(*args, **kwargs):
     """
     Initializes or reinitializes db.engine.  (Only the test suite should be calling this externally
     to reinitialize).
+
+    Arguments:
+    sogs_preinit - a callable to invoke after setting up `engine` but before calling
+    `database_init()`.
     """
     global engine, engine_initial_pid, metadata, have_returning
 
@@ -616,9 +609,17 @@ def _init_engine(*args, **kwargs):
             return
         args = (config.DB_URL,)
 
-    # Disable *sqlalchemy*-level autocommit, which works so badly that it got completely removed in
-    # 2.0.  (We put the actual sqlite into driver-level autocommit mode below).
-    engine = sqlalchemy.create_engine(*args, **kwargs).execution_options(autocommit=False)
+    preinit = kwargs.pop('sogs_preinit', None)
+
+    exec_opts_args = {}
+    if args[0].startswith('postgresql'):
+        exec_opts_args['isolation_level'] = 'READ COMMITTED'
+    else:
+        # SQLite's Python code is seriously broken, so we have to force off autocommit mode and turn
+        # on driver-level autocommit (which we do below).
+        exec_opts_args['autocommit'] = False
+
+    engine = sqlalchemy.create_engine(*args, **kwargs).execution_options(**exec_opts_args)
     engine_initial_pid = os.getpid()
     metadata = sqlalchemy.MetaData()
 
@@ -651,6 +652,18 @@ def do_begin(conn):
     else:
         have_returning = True
 
+        # rooms.token is a 'citext' (case-insensitive text), which sqlalchemy doesn't recognize out
+        # of the box.  Map it to a plain TEXT which is good enough for what we need (if we actually
+        # needed to generate this wouldn't suffice: we'd have to use something like the
+        # sqlalchemy-citext module).
+        from sqlalchemy.dialects.postgresql.base import ischema_names
+
+        if 'citext' not in ischema_names:
+            ischema_names['citext'] = ischema_names['text']
+
+    if preinit:
+        preinit()
+
     database_init()
 
 

tests/conftest.py

@@ -14,10 +14,25 @@
 sogs.omq.test_suite = True
 
 
+def pgsql_url(arg):
+    if not arg.startswith('postgresql:'):
+        raise ValueError("Invalid postgresql url; see SQLAlchemy postgresql docs")
+    return arg
+
+
 def pytest_addoption(parser):
     parser.addoption(
         "--sql-tracing", action="store_true", default=False, help="Log all SQL queries"
     )
+    parser.addoption(
+        "--pgsql", type=pgsql_url, help="Use the given postgresql database url for testing"
+    )
+    parser.addoption(
+        "--pgsql-no-drop-schema",
+        action="store_true",
+        default=False,
+        help="Don't clean up the final test schema; typically used with --maxfail=1",
+    )
 
 
 @pytest.fixture
@@ -40,28 +55,63 @@ def db(request):
     """
 
     trace = request.config.getoption("--sql-tracing")
+    pgsql = request.config.getoption("--pgsql")
 
     from sogs import db as db_
 
     global db_counter_
     db_counter_ += 1
-    sqlite_uri = f'file:sogs_testdb{db_counter_}?mode=memory&cache=shared'
 
-    web.app.logger.warning(f"using sqlite {sqlite_uri}")
+    if pgsql:
+        web.app.logger.warning(f"using postgresql {pgsql}")
+
+        first = True
+
+        def pg_setup_schema():
+            # Run everything in a separate schema that we can easily drop when done
+            from sqlalchemy import event
+
+            @event.listens_for(db_.engine, "connect", insert=True)
+            def setup_schema(dbapi_connection, connection_record):
+                existing_autocommit = dbapi_connection.autocommit
+                dbapi_connection.autocommit = True
+
+                cursor = dbapi_connection.cursor()
+                nonlocal first
+                if first:
+                    cursor.execute("DROP SCHEMA IF EXISTS sogs_tests CASCADE")
+                    first = False
+                cursor.execute("CREATE SCHEMA IF NOT EXISTS sogs_tests")
+                cursor.execute("SET search_path TO sogs_tests")
+                cursor.close()
+
+                dbapi_connection.autocommit = existing_autocommit
+
+        db_._init_engine(pgsql, echo=trace, sogs_preinit=pg_setup_schema)
+
+    else:
+        sqlite_uri = f'file:sogs_testdb{db_counter_}?mode=memory&cache=shared'
+
+        web.app.logger.warning(f"using sqlite {sqlite_uri}")
+
+        def sqlite_connect():
+            import sqlite3
 
-    def sqlite_connect():
-        import sqlite3
+            web.app.logger.warning(f"connecting to {sqlite_uri}")
+            return sqlite3.connect(sqlite_uri, uri=True)
 
-        web.app.logger.warning(f"connecting to {sqlite_uri}")
-        return sqlite3.connect(sqlite_uri, uri=True)
+        db_._init_engine("sqlite://", creator=sqlite_connect, echo=trace)
 
-    db_._init_engine("sqlite://", creator=sqlite_connect, echo=trace)
+    db_.database_init()
 
     web.appdb = db_.get_conn()
 
     yield db_
 
     web.app.logger.warning("closing db")
+    if pgsql and not request.config.getoption("--pgsql-no-drop-schema"):
+        web.app.logger.critical("DROPPING SCHEMA")
+        db_.query("DROP SCHEMA sogs_tests CASCADE")
     web.appdb.close()