Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(postgres): Remove SqlAlchemy dependency from postgres container #445

Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 22 additions & 4 deletions INDEX.rst
Original file line number Diff line number Diff line change
Expand Up @@ -45,15 +45,33 @@ Getting Started
>>> from testcontainers.postgres import PostgresContainer
>>> import sqlalchemy

>>> with PostgresContainer("postgres:9.5") as postgres:
... engine = sqlalchemy.create_engine(postgres.get_connection_url())
>>> with PostgresContainer("postgres:latest") as postgres:
... psql_url = postgres.get_connection_url()
... engine = sqlalchemy.create_engine(psql_url)
... with engine.begin() as connection:
... result = connection.execute(sqlalchemy.text("select version()"))
... version, = result.fetchone()
>>> version
'PostgreSQL 9.5...'
'PostgreSQL ...'

The snippet above will spin up the current latest version of a postgres database in a container. The :code:`get_connection_url()` convenience method returns a :code:`sqlalchemy` compatible url (using the :code:`psycopg2` driver per default) to connect to the database and retrieve the database version.

.. doctest::

>>> import asyncpg
>>> from testcontainers.postgres import PostgresContainer

>>> with PostgresContainer("postgres:16", driver=None) as postgres:
... psql_url = container.get_connection_url()
... with asyncpg.create_pool(dsn=psql_url,server_settings={"jit": "off"}) as pool:
... conn = await pool.acquire()
... ret = await conn.fetchval("SELECT 1")
... assert ret == 1

This snippet does the same, however using a specific version and the driver is set to None, to influence the :code:`get_connection_url()` convenience method to not include a driver in the URL (e.g. for compatibility with :code:`psycopg` v3).

Note, that the :code:`sqlalchemy` and :code:`psycopg2` packages are no longer a dependency of :code:`testcontainers[postgres]` and not needed to launch the Postgres container. Your project therefore needs to declare a dependency on the used driver and db access methods you use in your code.

The snippet above will spin up a postgres database in a container. The :code:`get_connection_url()` convenience method returns a :code:`sqlalchemy` compatible url we use to connect to the database and retrieve the database version.

Installation
------------
Expand Down
4 changes: 2 additions & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -12,13 +12,13 @@ For more information, see [the docs][readthedocs].
>>> from testcontainers.postgres import PostgresContainer
>>> import sqlalchemy

>>> with PostgresContainer("postgres:9.5") as postgres:
>>> with PostgresContainer("postgres:16") as postgres:
... engine = sqlalchemy.create_engine(postgres.get_connection_url())
... with engine.begin() as connection:
... result = connection.execute(sqlalchemy.text("select version()"))
... version, = result.fetchone()
>>> version
'PostgreSQL 9.5...'
'PostgreSQL 16...'
```

The snippet above will spin up a postgres database in a container. The `get_connection_url()` convenience method returns a `sqlalchemy` compatible url we use to connect to the database and retrieve the database version.
45 changes: 37 additions & 8 deletions modules/postgres/testcontainers/postgres/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,16 +11,23 @@
# License for the specific language governing permissions and limitations
# under the License.
import os
from time import sleep
from typing import Optional

from testcontainers.core.config import MAX_TRIES, SLEEP_TIME
from testcontainers.core.generic import DbContainer
from testcontainers.core.utils import raise_for_deprecated_parameter
from testcontainers.core.waiting_utils import wait_container_is_ready, wait_for_logs

_UNSET = object()


class PostgresContainer(DbContainer):
"""
Postgres database container.

To get a URL without a driver, pass in :code:`driver=None`.

Example:

The example spins up a Postgres database and connects to it using the :code:`psycopg`
Expand All @@ -31,7 +38,7 @@ class PostgresContainer(DbContainer):
>>> from testcontainers.postgres import PostgresContainer
>>> import sqlalchemy

>>> postgres_container = PostgresContainer("postgres:9.5")
>>> postgres_container = PostgresContainer("postgres:16")
>>> with postgres_container as postgres:
... engine = sqlalchemy.create_engine(postgres.get_connection_url())
... with engine.begin() as connection:
Expand All @@ -48,16 +55,16 @@ def __init__(
username: Optional[str] = None,
password: Optional[str] = None,
dbname: Optional[str] = None,
driver: str = "psycopg2",
driver: Optional[str] = "psycopg2",
**kwargs,
) -> None:
raise_for_deprecated_parameter(kwargs, "user", "username")
super().__init__(image=image, **kwargs)
self.username = username or os.environ.get("POSTGRES_USER", "test")
self.password = password or os.environ.get("POSTGRES_PASSWORD", "test")
self.dbname = dbname or os.environ.get("POSTGRES_DB", "test")
self.username: str = username or os.environ.get("POSTGRES_USER", "test")
self.password: str = password or os.environ.get("POSTGRES_PASSWORD", "test")
self.dbname: str = dbname or os.environ.get("POSTGRES_DB", "test")
self.port = port
self.driver = driver
self.driver = f"+{driver}" if driver else ""

self.with_exposed_ports(self.port)

Expand All @@ -66,12 +73,34 @@ def _configure(self) -> None:
self.with_env("POSTGRES_PASSWORD", self.password)
self.with_env("POSTGRES_DB", self.dbname)

def get_connection_url(self, host=None) -> str:
def get_connection_url(self, host: Optional[str] = None, driver: Optional[str] = _UNSET) -> str:
"""Get a DB connection URL to connect to the PG DB.

If a driver is set in the constructor (defaults to psycopg2!), the URL will contain the
driver. The optional driver argument to :code:`get_connection_url` overwrites the constructor
set value. Pass :code:`driver=None` to get URLs without a driver.
"""
driver_str = self.driver if driver is _UNSET else f"+{driver}"
return super()._create_connection_url(
dialect=f"postgresql+{self.driver}",
dialect=f"postgresql{driver_str}",
username=self.username,
password=self.password,
dbname=self.dbname,
host=host,
port=self.port,
)

@wait_container_is_ready()
def _connect(self) -> None:
wait_for_logs(self, ".*database system is ready to accept connections.*", MAX_TRIES, SLEEP_TIME)

count = 0
while count < MAX_TRIES:
status, _ = self.exec(f"pg_isready -hlocalhost -p{self.port} -U{self.username}")
if status == 0:
return

sleep(SLEEP_TIME)
count += 1

raise RuntimeError("Postgres could not get into a ready state")
29 changes: 27 additions & 2 deletions modules/postgres/tests/test_postgres.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,34 @@
import sqlalchemy
import sys

import pytest

from testcontainers.postgres import PostgresContainer
import sqlalchemy


# https://www.postgresql.org/support/versioning/
@pytest.mark.parametrize("version", ["12", "13", "14", "15", "16", "latest"])
def test_docker_run_postgres(version: str, monkeypatch):
def fail(*args, **kwargs):
raise AssertionError("SQLA was called during PG container setup")

monkeypatch.setattr(sqlalchemy, "create_engine", fail)
postgres_container = PostgresContainer(f"postgres:{version}")
with postgres_container as postgres:
status, msg = postgres.exec(f"pg_isready -hlocalhost -p{postgres.port} -U{postgres.username}")

assert msg.decode("utf-8").endswith("accepting connections\n")
assert status == 0

status, msg = postgres.exec(
f"psql -hlocalhost -p{postgres.port} -U{postgres.username} -c 'select 2*3*5*7*11*13*17 as a;' "
)
assert "510510" in msg.decode("utf-8")
assert "(1 row)" in msg.decode("utf-8")
assert status == 0


def test_docker_run_postgres():
def test_docker_run_postgres_with_sqlalchemy():
postgres_container = PostgresContainer("postgres:9.5")
with postgres_container as postgres:
engine = sqlalchemy.create_engine(postgres.get_connection_url())
Expand Down
12 changes: 6 additions & 6 deletions poetry.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

9 changes: 6 additions & 3 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,6 @@ pymysql = { version = "*", extras = ["rsa"], optional = true }
neo4j = { version = "*", optional = true }
opensearch-py = { version = "*", optional = true }
cx_Oracle = { version = "*", optional = true }
psycopg2-binary = { version = "*", optional = true }
pika = { version = "*", optional = true }
redis = { version = "*", optional = true }
selenium = { version = "*", optional = true }
Expand All @@ -102,20 +101,24 @@ neo4j = ["neo4j"]
nginx = []
opensearch = ["opensearch-py"]
oracle = ["sqlalchemy", "cx_Oracle"]
postgres = ["sqlalchemy", "psycopg2-binary"]
jankatins marked this conversation as resolved.
Show resolved Hide resolved
postgres = []
rabbitmq = ["pika"]
redis = ["redis"]
selenium = ["selenium"]

[tool.poetry.group.dev.dependencies]
mypy = "1.7.1"
pre-commit = "^3.6"
pg8000 = "*"
pytest = "7.4.3"
pytest-cov = "4.1.0"
sphinx = "^7.2.6"
twine = "^4.0.2"
anyio = "^4.3.0"
# for tests only
psycopg2-binary = "*"
pg8000 = "*"
sqlalchemy = "*"


[[tool.poetry.source]]
name = "PyPI"
Expand Down