Introduce DatabaseJanitor - closes #226

* Introduce DatabaseJanitor

* streamline dependency installation

* add type hints to DatabaseJanitor

.travis.yml

@@ -11,7 +11,7 @@ branches:
   - requires-io-master
 install:
 - pip install -r requirements-test.txt
-- pip install -e .[tests] coveralls wheel
+- pip install coveralls wheel
 script:
 - py.test -n 0 --cov src/pytest_postgresql
 after_success:
@@ -25,7 +25,6 @@ jobs:
       python: 3.7
       install:
       - pip install -r requirements-lint.txt
-      - pip install .[tests] psycopg2-binary
       script:
       - pycodestyle
       - pydocstyle

CHANGES.rst

@@ -1,6 +1,13 @@
 CHANGELOG
 =========
 
+unreleased
+-------
+
+- [enhancement] Gather helper functions maintaining postgresql database in DatabaseJanitor class.
+- [deprecate] Deprecate ``init_postgresql_database`` in favour of ``DatabaseJanitor.init``
+- [deprecate] Deprecate ``drop_postgresql_database`` in favour of ``DatabaseJanitor.drop``
+
 2.0.0
 -------
 

README.rst

@@ -163,6 +163,39 @@ Example usage:
         [pytest]
         postgresql_port = 8888
 
+Maintaining database state outside of the fixtures
+--------------------------------------------------
+
+It is possible and appears it's used in other libraries for tests,
+to maintain database state with the use of the ``pytest-postgresql`` database
+managing functionality:
+
+For this import DatabaseJanitor and use it's init and drop methods:
+
+
+.. code-block:: python
+
+    from pytest_postgresql.factories import DatabaseJanitor
+
+    # variable definition
+
+    janitor = DatabaseJanitor(user, host, port, db_name, version)
+    janitor.init()
+    # your code, or yield
+    janitor.drop()
+    # at this moment you'll have clean database step
+
+or use it as a context manager:
+
+.. code-block:: python
+
+    from pytest_postgresql.factories import DatabaseJanitor
+
+    # variable definition
+
+    with DatabaseJanitor(user, host, port, db_name, version):
+        # do something here
+
 Package resources
 -----------------
 

docs/source/conf.py

@@ -121,7 +121,7 @@
 # html_theme_path = []
 
 # The name for this set of Sphinx documents.  If None, it defaults to
-# "<project> v<release> documentation".
+# "<project> version<release> documentation".
 # html_title = None
 
 # A shorter title for the navigation bar.  Default is the same as html_title.

requirements-lint.txt

@@ -4,3 +4,4 @@ pydocstyle==4.0.0
 pylint==2.3.1
 pygments
 pyroma==2.5
+-r requirements-test.txt

requirements-test.txt

@@ -7,3 +7,4 @@ psycopg2-binary==2.8.3; platform_python_implementation != "PyPy"
 psycopg2cffi==2.8.1; platform_python_implementation == "PyPy"
 port-for==0.4
 mirakuru==2.0.1
+-e .[tests]

src/pytest_postgresql/factories.py

@@ -21,15 +21,11 @@
 import platform
 import subprocess
 from tempfile import gettempdir
+from warnings import warn
 
 import pytest
-from pkg_resources import parse_version
-
-try:
-    import psycopg2
-except ImportError:
-    psycopg2 = False
 
+from pytest_postgresql.janitor import DatabaseJanitor, psycopg2
 from pytest_postgresql.executor import PostgreSQLExecutor
 from pytest_postgresql.port import get_port
 
@@ -58,12 +54,12 @@ def init_postgresql_database(user, host, port, db_name):
     :param str port: postgresql port
     :param str db_name: database name
     """
-    conn = psycopg2.connect(user=user, host=host, port=port)
-    conn.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
-    cur = conn.cursor()
-    cur.execute('CREATE DATABASE "{}";'.format(db_name))
-    cur.close()
-    conn.close()
+    warn(
+        'init_postgresql_database is deprecated, '
+        'use DatabaseJanitor.init istead.',
+        DeprecationWarning
+    )
+    DatabaseJanitor(user, host, port, db_name, 0.0).init()
 
 
 def drop_postgresql_database(user, host, port, db_name, version):
@@ -76,26 +72,12 @@ def drop_postgresql_database(user, host, port, db_name, version):
     :param str db_name: database name
     :param packaging.version.Version version: postgresql version number
     """
-    conn = psycopg2.connect(user=user, host=host, port=port)
-    conn.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
-    cur = conn.cursor()
-    # We cannot drop the database while there are connections to it, so we
-    # terminate all connections first while not allowing new connections.
-    if version >= parse_version('9.2'):
-        pid_column = 'pid'
-    else:
-        pid_column = 'procpid'
-    cur.execute(
-        'UPDATE pg_database SET datallowconn=false WHERE datname = %s;',
-        (db_name,))
-    cur.execute(
-        'SELECT pg_terminate_backend(pg_stat_activity.{})'
-        'FROM pg_stat_activity WHERE pg_stat_activity.datname = %s;'.format(
-            pid_column),
-        (db_name,))
-    cur.execute('DROP DATABASE IF EXISTS "{}";'.format(db_name))
-    cur.close()
-    conn.close()
+    warn(
+        'drop_postgresql_database is deprecated, '
+        'use DatabaseJanitor.drop istead.',
+        DeprecationWarning
+    )
+    DatabaseJanitor(user, host, port, db_name, version).init()
 
 
 def postgresql_proc(
@@ -213,24 +195,18 @@ def postgresql_factory(request):
         pg_options = proc_fixture.options
         pg_db = db_name or config['dbname']
 
-        init_postgresql_database(pg_user, pg_host, pg_port, pg_db)
-        connection = psycopg2.connect(
-            dbname=pg_db,
-            user=pg_user,
-            host=pg_host,
-            port=pg_port,
-            options=pg_options
-        )
-
-        def drop_database():
-            connection.close()
-            drop_postgresql_database(
+        with DatabaseJanitor(
                 pg_user, pg_host, pg_port, pg_db, proc_fixture.version
+        ):
+            connection = psycopg2.connect(
+                dbname=pg_db,
+                user=pg_user,
+                host=pg_host,
+                port=pg_port,
+                options=pg_options
             )
-
-        request.addfinalizer(drop_database)
-
-        return connection
+            yield connection
+            connection.close()
 
     return postgresql_factory
 

src/pytest_postgresql/janitor.py

@@ -0,0 +1,97 @@
+"""Database Janitor."""
+from contextlib import contextmanager
+from types import TracebackType
+from typing import TypeVar, Union, Optional, Type, Any
+
+from pkg_resources import parse_version
+Version = type(parse_version('1'))  # pylint:disable=invalid-name
+
+try:
+    import psycopg2
+    try:
+        from psycopg2._psycopg import cursor
+    except ImportError:
+        from psycopg2cffi._impl.cursor import Cursor as cursor
+except ImportError:
+    psycopg2 = False
+    # if there's no postgres, just go with the flow.
+    cursor = Any  # pylint:disable=invalid-name
+
+DatabaseJanitorType = TypeVar("DatabaseJanitorType", bound="DatabaseJanitor")
+
+
+class DatabaseJanitor:
+    """Manage database state for specific tasks."""
+
+    def __init__(
+            self,
+            user: str,
+            host: str,
+            port: str,
+            db_name: str,
+            version: Union[str, float, Version]
+    ) -> None:
+        """
+        Initialize janitor.
+
+        :param user: postgresql username
+        :param host: postgresql host
+        :param port: postgresql port
+        :param db_name: database name
+        :param version: postgresql version number
+        """
+        self.user = user
+        self.host = host
+        self.port = port
+        self.db_name = db_name
+        if not isinstance(version, Version):
+            self.version = parse_version(str(version))
+        else:
+            self.version = version
+
+    def init(self) -> None:
+        """Create database in postgresql."""
+        with self.cursor() as cur:
+            cur.execute('CREATE DATABASE "{}";'.format(self.db_name))
+
+    def drop(self) -> None:
+        """Drop database in postgresql."""
+        # We cannot drop the database while there are connections to it, so we
+        # terminate all connections first while not allowing new connections.
+        if self.version >= parse_version('9.2'):
+            pid_column = 'pid'
+        else:
+            pid_column = 'procpid'
+        with self.cursor() as cur:
+            cur.execute(
+                'UPDATE pg_database SET datallowconn=false WHERE datname = %s;',
+                (self.db_name,))
+            cur.execute(
+                'SELECT pg_terminate_backend(pg_stat_activity.{})'
+                'FROM pg_stat_activity '
+                'WHERE pg_stat_activity.datname = %s;'.format(pid_column),
+                (self.db_name,))
+            cur.execute('DROP DATABASE IF EXISTS "{}";'.format(self.db_name))
+
+    @contextmanager
+    def cursor(self) -> cursor:
+        """Return postgresql cursor."""
+        conn = psycopg2.connect(user=self.user, host=self.host, port=self.port)
+        conn.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
+        cur = conn.cursor()
+        try:
+            yield cur
+        finally:
+            cur.close()
+            conn.close()
+
+    def __enter__(self: DatabaseJanitorType) -> DatabaseJanitorType:
+        self.init()
+        return self
+
+    def __exit__(
+            self: DatabaseJanitorType,
+            exc_type: Optional[Type[BaseException]],
+            exc_val: Optional[BaseException],
+            exc_tb: Optional[TracebackType]) -> None:
+        self.drop()

tests/test_janitor.py

@@ -0,0 +1,14 @@
+"""Database Janitor tests."""
+import pytest
+from pkg_resources import parse_version
+
+from pytest_postgresql.janitor import DatabaseJanitor
+
+VERSION = parse_version('9.2')
+
+
+@pytest.mark.parametrize('version', (VERSION, 9.2, '9.2'))
+def test_version_cast(version):
+    """Test that version is cast to Version object."""
+    janitor = DatabaseJanitor(None, None, None, None, version)
+    assert janitor.version == VERSION