Skip to content

Commit

Permalink
feat: Improve default postgres driver selection heuristics. (#206)
Browse files Browse the repository at this point in the history
  • Loading branch information
DanCardin authored Apr 11, 2024
1 parent 978190f commit f013626
Show file tree
Hide file tree
Showing 4 changed files with 37 additions and 13 deletions.
12 changes: 6 additions & 6 deletions docs/source/postgres.rst
Original file line number Diff line number Diff line change
Expand Up @@ -7,13 +7,13 @@ Postgres
and `asyncpg` for async fixtures. If you want to use a different driver, you
can configure the `drivername` field using the `pmr_postgres_config` fixture:

```python
from pytest_mock_resources import PostgresConfig
.. code-block:: python
@pytest.fixture
def pmr_postgres_config():
PostgresConfig(drivername='postgresql+psycopg2') # but whatever drivername you require.
```
from pytest_mock_resources import PostgresConfig
@pytest.fixture(scope='session')
def pmr_postgres_config():
return PostgresConfig(drivername='postgresql+psycopg2') # but whatever driver you require.
Note however, that the `asyncpg` driver **only** works with the async fixture, and the
`psycopg2` driver **only** works with the synchronous fixture. These are inherent
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[tool.poetry]
name = "pytest-mock-resources"
version = "2.10.1"
version = "2.10.2"
description = "A pytest plugin for easily instantiating reproducible mock resources."
authors = [
"Omar Khan <oakhan3@gmail.com>",
Expand Down
34 changes: 29 additions & 5 deletions src/pytest_mock_resources/container/postgres.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from typing import ClassVar, Iterable
import sys
from typing import ClassVar, Iterable, Optional

import sqlalchemy
import sqlalchemy.exc
Expand All @@ -7,6 +8,11 @@
from pytest_mock_resources.config import DockerContainerConfig, fallback
from pytest_mock_resources.container.base import ContainerCheckFailed

if sys.version_info < (3, 8):
from importlib_metadata import Distribution
else:
from importlib.metadata import Distribution


class PostgresConfig(DockerContainerConfig):
"""Define the configuration object for postgres.
Expand Down Expand Up @@ -48,7 +54,7 @@ class PostgresConfig(DockerContainerConfig):
"username": "user",
"password": "password",
"root_database": "dev",
"drivername": "postgresql+psycopg2",
"drivername": None,
}

@fallback
Expand Down Expand Up @@ -94,9 +100,7 @@ def check_fn(self):
def get_sqlalchemy_engine(config, database_name, async_=False, autocommit=False, **engine_kwargs):
# For backwards compatibility, our hardcoded default is psycopg2, and async fixtures
# will not work with psycopg2, so we instead swap the default to the preferred async driver.
drivername = config.drivername
if async_ and drivername.endswith("psycopg2"):
drivername = drivername.replace("psycopg2", "asyncpg")
drivername = detect_driver(config.drivername, async_=async_)

url = URL(
drivername=drivername,
Expand All @@ -118,3 +122,23 @@ def get_sqlalchemy_engine(config, database_name, async_=False, autocommit=False,
engine = sqlalchemy.create_engine(url, **engine_kwargs)

return engine


def detect_driver(drivername: Optional[str] = None, async_: bool = False) -> str:
if drivername:
return drivername

if any(Distribution.discover(name="psycopg")):
return "postgresql+psycopg"

if async_:
if any(Distribution.discover(name="asyncpg")):
return "postgresql+asyncpg"
else:
if any(Distribution.discover(name="psycopg2")):
return "postgresql+psycopg2"

raise ValueError( # pragma: no cover
"No suitable driver found for Postgres. Please install `psycopg`, `psycopg2`, "
"`asyncpg`, or explicitly configure the `drivername=` field of `PostgresConfig`."
)
2 changes: 1 addition & 1 deletion src/pytest_mock_resources/container/redshift.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,5 +43,5 @@ class RedshiftConfig(PostgresConfig):
"username": "user",
"password": "password",
"root_database": "dev",
"drivername": "postgresql+psycopg2",
"drivername": None,
}

0 comments on commit f013626

Please sign in to comment.