Differential D7934 Diff 28586 swh/core/db/pytest_plugin.py

Changeset View

Standalone View

swh/core/db/pytest_plugin.py

# Copyright (C) 2020 The Software Heritage developers		# Copyright (C) 2020 The Software Heritage developers
# See the AUTHORS file at the top-level directory of this distribution		# See the AUTHORS file at the top-level directory of this distribution
# License: GNU General Public License version 3, or any later version		# License: GNU General Public License version 3, or any later version
# See top-level LICENSE file for more information		# See top-level LICENSE file for more information

import glob		import glob
from importlib import import_module		from importlib import import_module
import logging		import logging
import subprocess		import subprocess
from typing import Callable, Iterable, Iterator, List, Optional, Sequence, Set, Union		from typing import Callable, Iterable, Iterator, List, Optional, Sequence, Set, Union
		import warnings

from _pytest.fixtures import FixtureRequest		from _pytest.fixtures import FixtureRequest
		from deprecated import deprecated
import psycopg2		import psycopg2
import pytest		import pytest
from pytest_postgresql.compat import check_for_psycopg2, connection		from pytest_postgresql.compat import check_for_psycopg2, connection
from pytest_postgresql.executor import PostgreSQLExecutor		from pytest_postgresql.executor import PostgreSQLExecutor
from pytest_postgresql.executor_noop import NoopExecutor		from pytest_postgresql.executor_noop import NoopExecutor
from pytest_postgresql.janitor import DatabaseJanitor		from pytest_postgresql.janitor import DatabaseJanitor

from swh.core.db.db_utils import (		from swh.core.db.db_utils import initialize_database_for_module
init_admin_extensions,
populate_database_for_package,
swh_set_db_version,
)
from swh.core.utils import basename_sortkey		from swh.core.utils import basename_sortkey

# to keep mypy happy regardless pytest-postgresql version		# to keep mypy happy regardless pytest-postgresql version
try:		try:
_pytest_pgsql_get_config_module = import_module("pytest_postgresql.config")		_pytest_pgsql_get_config_module = import_module("pytest_postgresql.config")
except ImportError:		except ImportError:
# pytest_postgresql < 3.0.0		# pytest_postgresql < 3.0.0
_pytest_pgsql_get_config_module = import_module("pytest_postgresql.factories")		_pytest_pgsql_get_config_module = import_module("pytest_postgresql.factories")

_pytest_postgresql_get_config = getattr(_pytest_pgsql_get_config_module, "get_config")		_pytest_postgresql_get_config = getattr(_pytest_pgsql_get_config_module, "get_config")


logger = logging.getLogger(__name__)		logger = logging.getLogger(__name__)

		initialize_database_for_module = deprecated(
		version="2.10",
		reason="Use swh.core.db.db_utils.initialize_database_for_module instead.",
		)(initialize_database_for_module)

		warnings.warn(
		"This pytest plugin is deprecated, it should not be used any more.",
		category=DeprecationWarning,
		)


class SWHDatabaseJanitor(DatabaseJanitor):		class SWHDatabaseJanitor(DatabaseJanitor):
"""SWH database janitor implementation with a a different setup/teardown policy than		"""SWH database janitor implementation with a a different setup/teardown policy than
than the stock one. Instead of dropping, creating and initializing the database for		than the stock one. Instead of dropping, creating and initializing the database for
each test, it creates and initializes the db once, then truncates the tables (and		each test, it creates and initializes the db once, then truncates the tables (and
sequences) in between tests.		sequences) in between tests.

This is needed to have acceptable test performances.		This is needed to have acceptable test performances.
▲ Show 20 Lines • Show All 117 Lines • ▼ Show 20 Lines	def drop(self) -> None:
self._terminate_connection(cur, self.dbname)		self._terminate_connection(cur, self.dbname)
else:		else:
super().drop()		super().drop()


# the postgres_fact factory fixture below is mostly a copy of the code		# the postgres_fact factory fixture below is mostly a copy of the code
# from pytest-postgresql. We need a custom version here to be able to		# from pytest-postgresql. We need a custom version here to be able to
# specify our version of the DBJanitor we use.		# specify our version of the DBJanitor we use.
		@deprecated(version="2.10", reason="Use stock pytest_postgresql factory instead")
def postgresql_fact(		def postgresql_fact(
process_fixture_name: str,		process_fixture_name: str,
dbname: Optional[str] = None,		dbname: Optional[str] = None,
load: Optional[Sequence[Union[Callable, str]]] = None,		load: Optional[Sequence[Union[Callable, str]]] = None,
isolation_level: Optional[int] = None,		isolation_level: Optional[int] = None,
modname: Optional[str] = None,		modname: Optional[str] = None,
dump_files: Optional[Union[str, List[str]]] = None,		dump_files: Optional[Union[str, List[str]]] = None,
no_truncate_tables: Set[str] = {"dbversion"},		no_truncate_tables: Set[str] = {"dbversion"},
▲ Show 20 Lines • Show All 67 Lines • ▼ Show 20 Lines	def postgresql_factory(request: FixtureRequest) -> Iterator[connection]:
try:		try:
yield db_connection		yield db_connection
finally:		finally:
db_connection.close()		db_connection.close()

return postgresql_factory		return postgresql_factory


def initialize_database_for_module(modname, version, **kwargs):
conninfo = psycopg2.connect(**kwargs).dsn
init_admin_extensions(modname, conninfo)
populate_database_for_package(modname, conninfo)
try:
swh_set_db_version(conninfo, version)
except psycopg2.errors.UniqueViolation:
logger.warn(
"Version already set by db init scripts. "
"This generally means the swh.{modname} package needs to be "
"updated for swh.core>=1.2"
)


def gen_dump_files(dump_files: Union[str, Iterable[str]]) -> Iterator[str]:		def gen_dump_files(dump_files: Union[str, Iterable[str]]) -> Iterator[str]:
"""Generate files potentially resolving glob patterns if any"""		"""Generate files potentially resolving glob patterns if any"""
if isinstance(dump_files, str):		if isinstance(dump_files, str):
dump_files = [dump_files]		dump_files = [dump_files]
for dump_file in dump_files:		for dump_file in dump_files:
if glob.has_magic(dump_file):		if glob.has_magic(dump_file):
# if the dump_file is a glob pattern one, resolve it		# if the dump_file is a glob pattern one, resolve it
yield from (		yield from (
fname for fname in sorted(glob.glob(dump_file), key=basename_sortkey)		fname for fname in sorted(glob.glob(dump_file), key=basename_sortkey)
)		)
else:		else:
# otherwise, just return the filename		# otherwise, just return the filename
yield dump_file		yield dump_file