Skip to content

Commit

Permalink
feat(postgres): Remove SqlAlchemy dependency from postgres container (t…
Browse files Browse the repository at this point in the history
…estcontainers#445)

Updates the pg testcontainer implementation to not use (and not install)
SQLAlchemy nor psycopg2.

Closes: testcontainers#340
Closes: testcontainers#336
Closes: testcontainers#320

---------

Co-authored-by: Jason Turim <jason@opscanvas.com>
  • Loading branch information
2 people authored and bstrausser committed Mar 30, 2024
1 parent 656a2bd commit 0f79144
Show file tree
Hide file tree
Showing 6 changed files with 99 additions and 24 deletions.
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
10 changes: 5 additions & 5 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 @@ -80,7 +80,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 @@ -104,20 +103,24 @@ neo4j = ["neo4j"]
nginx = []
opensearch = ["opensearch-py"]
oracle = ["sqlalchemy", "cx_Oracle"]
postgres = ["sqlalchemy", "psycopg2-binary"]
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

0 comments on commit 0f79144

Please sign in to comment.