feat: Add cleanup_database option to postgres fixture.

Note, this could **potentially** reveal preexisting issues in code
under test, that could be perceived as a "breaking" change.

By deleting the database under test at the end of the test's execution,
any database connections left connected to the database might cause the
`DELETE DATABASE` command to fail.

PMR will **try** to use the `WITH FORCE` option on database versions >=
13.0, but that option does not exist on prior versions of postgres.

In any case, if this happens, it **is** ultimately revealing a "bug"
in the code it is testing. Additionally, you can simply turn off
database cleanup in one of various ways.

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pytest-mock-resources"
-version = "2.12.1"
+version = "2.13.0"
 description = "A pytest plugin for easily instantiating reproducible mock resources."
 authors = [
     "Omar Khan <[email protected]>",

src/pytest_mock_resources/fixture/postgresql.py

@@ -17,6 +17,7 @@
     PostgresConfig,
 )
 from pytest_mock_resources.fixture.base import asyncio_fixture, generate_fixture_id, Scope
+from pytest_mock_resources.hooks import should_cleanup
 from pytest_mock_resources.sqlalchemy import (
     bifurcate_actions,
     EngineManager,
@@ -114,6 +115,7 @@ def create_postgres_fixture(
     engine_kwargs=None,
     template_database=True,
     actions_share_transaction=None,
+    cleanup_database: Optional[bool] = None,
 ):
     """Produce a Postgres fixture.
 
@@ -141,6 +143,10 @@ def create_postgres_fixture(
             fixtures for backwards compatibility; and disabled by default for
             asynchronous fixtures (the way v2-style/async features work in SQLAlchemy can lead
             to bad default behavior).
+        cleanup_database: Whether to clean up the database after the test completes. Defaults to `None`,
+            which defers the decision to the pmr_cleanup/--pmr-cleanup pytest options (which default to True).
+            Note this does not currently clean up any "template" databases produced in service
+            of the fixture.
     """
     fixture_id = generate_fixture_id(enabled=template_database, name="pg")
 
@@ -155,13 +161,23 @@ def create_postgres_fixture(
     }
 
     @pytest.fixture(scope=scope)
-    def _sync(*_, pmr_postgres_container, pmr_postgres_config):
-        fixture = _sync_fixture(pmr_postgres_config, engine_manager_kwargs, engine_kwargs_)
+    def _sync(*_, pmr_postgres_container, pmr_postgres_config, pytestconfig):
+        fixture = _sync_fixture(
+            pmr_postgres_config,
+            engine_manager_kwargs,
+            engine_kwargs_,
+            cleanup_database=should_cleanup(pytestconfig, cleanup_database),
+        )
         for _, conn in fixture:
             yield conn
 
-    async def _async(*_, pmr_postgres_container, pmr_postgres_config):
-        fixture = _async_fixture(pmr_postgres_config, engine_manager_kwargs, engine_kwargs_)
+    async def _async(*_, pmr_postgres_container, pmr_postgres_config, pytestconfig):
+        fixture = _async_fixture(
+            pmr_postgres_config,
+            engine_manager_kwargs,
+            engine_kwargs_,
+            cleanup_database=should_cleanup(pytestconfig, cleanup_database),
+        )
         async for _, conn in fixture:
             yield conn
 
@@ -170,7 +186,14 @@ async def _async(*_, pmr_postgres_container, pmr_postgres_config):
     return _sync
 
 
-def _sync_fixture(pmr_config, engine_manager_kwargs, engine_kwargs, *, fixture="postgres"):
+def _sync_fixture(
+    pmr_config,
+    engine_manager_kwargs,
+    engine_kwargs,
+    *,
+    fixture="postgres",
+    cleanup_database: bool = True,
+):
     root_engine = cast(Engine, get_sqlalchemy_engine(pmr_config, pmr_config.root_database))
     conn = retry(root_engine.connect, retries=DEFAULT_RETRIES)
     conn.close()
@@ -216,10 +239,27 @@ def _sync_fixture(pmr_config, engine_manager_kwargs, engine_kwargs, *, fixture="
     engine = get_sqlalchemy_engine(pmr_config, database_name, **engine_kwargs)
     yield from engine_manager.manage_sync(engine)
 
+    if cleanup_database:
+        with root_engine.connect() as root_conn:
+            with root_conn.begin() as trans:
+                _drop_database(root_conn, database_name)
+                trans.commit()
+        root_engine.dispose()
+
+
+async def _async_fixture(
+    pmr_config,
+    engine_manager_kwargs,
+    engine_kwargs,
+    *,
+    fixture="postgres",
+    cleanup_database: bool = True,
+):
+    from sqlalchemy.ext.asyncio import AsyncEngine
 
-async def _async_fixture(pmr_config, engine_manager_kwargs, engine_kwargs, *, fixture="postgres"):
-    root_engine = get_sqlalchemy_engine(
-        pmr_config, pmr_config.root_database, async_=True, autocommit=True
+    root_engine = cast(
+        AsyncEngine,
+        get_sqlalchemy_engine(pmr_config, pmr_config.root_database, async_=True, autocommit=True),
     )
 
     root_conn = await async_retry(root_engine.connect, retries=DEFAULT_RETRIES)
@@ -239,7 +279,10 @@ async def _async_fixture(pmr_config, engine_manager_kwargs, engine_kwargs, *, fi
     if template_manager:
         assert template_database
 
-        engine = get_sqlalchemy_engine(pmr_config, template_database, **engine_kwargs, async_=True)
+        engine = cast(
+            AsyncEngine,
+            get_sqlalchemy_engine(pmr_config, template_database, **engine_kwargs, async_=True),
+        )
         async with engine.begin() as conn:
             await conn.run_sync(template_manager.run_static_actions)
             await conn.commit()
@@ -260,6 +303,13 @@ async def _async_fixture(pmr_config, engine_manager_kwargs, engine_kwargs, *, fi
     async for engine, conn in engine_manager.manage_async(engine):
         yield engine, conn
 
+    if cleanup_database:
+        async with root_engine.connect() as root_conn:
+            async with root_conn.begin() as trans:
+                await root_conn.run_sync(_drop_database, database_name)
+                await trans.commit()
+        await root_engine.dispose()
+
 
 def create_engine_manager(
     root_connection,
@@ -356,5 +406,15 @@ def _generate_database_name(conn):
     return f"pytest_mock_resource_db_{id_}"
 
 
+def _drop_database(root_conn: Connection, database_name: str):
+    with_force = ""
+
+    assert root_conn.dialect.server_version_info
+    if root_conn.dialect.server_version_info >= (13, 0):
+        with_force = " WITH FORCE"
+
+    root_conn.execute(text(f"DROP DATABASE {database_name}{with_force}"))
+
+
 pmr_postgres_config = create_postgres_config_fixture()
 pmr_postgres_container = create_postgres_container_fixture()

src/pytest_mock_resources/hooks.py

@@ -1,5 +1,6 @@
 import os
 import warnings
+from typing import Optional
 
 _resource_kinds = ["postgres", "redshift", "mongo", "redis", "mysql", "moto"]
 
@@ -17,6 +18,12 @@ def pytest_addoption(parser):
         type="bool",
         default=True,
     )
+    parser.addini(
+        "pmr_cleanup",
+        "Optionally cleanup created fixture resources.",
+        type="bool",
+        default=True,
+    )
     parser.addini(
         "pmr_docker_client",
         "Optional docker client name to use: docker, podman, nerdctl",
@@ -28,17 +35,38 @@ def pytest_addoption(parser):
     group.addoption(
         "--pmr-multiprocess-safe",
         action="store_true",
-        default=False,
+        default=None,
         help="Enable multiprocess-safe mode",
         dest="pmr_multiprocess_safe",
     )
     group.addoption(
         "--pmr-cleanup-container",
         action="store_true",
-        default=True,
+        default=None,
+        help="Optionally disable attempts to cleanup created containers",
+        dest="pmr_cleanup_container",
+    )
+    group.addoption(
+        "--no-pmr-cleanup-container",
+        action="store_false",
+        default=None,
         help="Optionally disable attempts to cleanup created containers",
         dest="pmr_cleanup_container",
     )
+    group.addoption(
+        "--pmr-cleanup",
+        action="store_true",
+        default=None,
+        help="Optionally cleanup created fixture resources.",
+        dest="pmr_cleanup",
+    )
+    group.addoption(
+        "--no-pmr-cleanup",
+        action="store_false",
+        default=None,
+        help="Optionally cleanup created fixture resources.",
+        dest="pmr_cleanup",
+    )
     group.addoption(
         "--pmr-docker-client",
         default=None,
@@ -49,7 +77,7 @@ def pytest_addoption(parser):
 
 def get_pytest_flag(config, name, *, default=None):
     value = getattr(config.option, name, default)
-    if value:
+    if value is not None:
         return value
 
     return config.getini(name)
@@ -59,6 +87,12 @@ def use_multiprocess_safe_mode(config):
     return bool(get_pytest_flag(config, "pmr_multiprocess_safe"))
 
 
+def should_cleanup(config, value: Optional[bool] = None) -> bool:
+    if value is None:
+        return bool(get_pytest_flag(config, "pmr_cleanup"))
+    return value
+
+
 def get_docker_client_name(config) -> str:
     pmr_docker_client = os.getenv("PMR_DOCKER_CLIENT")
     if pmr_docker_client:

tests/examples/test_postgres_no_cleanup/test_cleanup.py

@@ -0,0 +1,45 @@
+from sqlalchemy import Column, String, text
+
+from pytest_mock_resources import create_postgres_fixture, Rows
+from pytest_mock_resources.compat.sqlalchemy import declarative_base
+
+Base = declarative_base()
+
+
+class User(Base):
+    __tablename__ = "user"
+
+    name = Column(String, primary_key=True)
+
+
+rows = Rows(User(name="Harold"), User(name="Gump"))
+
+
+pg_no_clean = create_postgres_fixture(Base, rows, session=True, cleanup_database=False)
+
+
+non_cleaned_database_name = None
+
+
+def test_not_to_be_cleaned_up(pg_no_clean):
+    global non_cleaned_database_name
+    non_cleaned_database_name = pg_no_clean.pmr_credentials.database
+
+    names = [u.name for u in pg_no_clean.query(User).all()]
+    assert names == ["Harold", "Gump"]
+
+    names = pg_no_clean.execute(text("SELECT datname FROM pg_database")).all()
+    unique_names = {name for (name,) in names}
+    assert non_cleaned_database_name in unique_names
+
+
+def test_database_is_not_cleaned_up(pg_no_clean):
+    global non_cleaned_database_name
+
+    assert non_cleaned_database_name is not None
+
+    assert non_cleaned_database_name != pg_no_clean.pmr_credentials.database
+
+    names = pg_no_clean.execute(text("SELECT datname FROM pg_database")).all()
+    unique_names = {name for (name,) in names}
+    assert non_cleaned_database_name in unique_names

tests/fixture/redshift/utils.py

@@ -140,6 +140,7 @@ def copy_fn_to_test_create_engine_patch(redshift):
                     )
                 )
             )
+        engine.dispose()
 
         fetch_values_from_table_and_assert(engine)
 
@@ -201,6 +202,7 @@ def unload_fn_to_test_create_engine_patch(redshift):
                     )
                 )
             )
+        engine.dispose()
 
         fetch_values_from_s3_and_assert(engine, is_gzipped=False)
 

tests/fixture/test_database.py

@@ -91,6 +91,7 @@ def test_create_custom_connection(postgres_3):
 
     with engine.connect() as conn:
         conn.execute(text("select 1"))
+    engine.dispose()
 
 
 def test_create_custom_connection_from_dict(postgres_3):
@@ -103,13 +104,15 @@ def test_create_custom_connection_from_dict(postgres_3):
 
     with engine.connect() as conn:
         conn.execute(text("select 1"))
+    engine.dispose()
 
 
 def test_create_custom_connection_url(postgres_3):
     url = compat.sqlalchemy.URL(**postgres_3.pmr_credentials.as_sqlalchemy_url_kwargs())
     engine = create_engine(url, isolation_level="AUTOCOMMIT")
     with engine.connect() as conn:
         conn.execute(text("select 1"))
+    engine.dispose()
 
 
 def test_bad_actions(postgres):

tests/fixture/test_pmr_credentials.py

@@ -119,3 +119,4 @@ def verify_relational(connection, credentials, session=False):
     manual_engine = create_engine(credentials.as_url())
     with manual_engine.connect() as conn:
         conn.execute(text("select * from foo"))
+    manual_engine.dispose()

tests/test_examples.py

@@ -19,3 +19,12 @@ def test_multiprocess_container_cleanup_race_condition(pytester):
     args = ["-vv", "-n", "2", "--pmr-multiprocess-safe", "test_split.py"]
     result = pytester.inline_run(*args)
     result.assertoutcome(passed=2, skipped=0, failed=0)
+
+
+@pytest.mark.postgres
+def test_postgres_no_cleanup(pytester):
+    pytester.copy_example()
+
+    args = ["-vv", "--pmr-multiprocess-safe", "test_cleanup.py"]
+    result = pytester.inline_run(*args)
+    result.assertoutcome(passed=2, skipped=0, failed=0)