Skip to content

Commit

Permalink
feat(postgres): Remove SqlAlchemy dependency from postgres container
Browse files Browse the repository at this point in the history
Remove test that was testing sqlalchemy support for driver types
Add tests for supported versions of Postgres
Modify the `get_connection_url` convenience method to support a
driverless url

Co-authored-by: Jason Turim <jason@opscanvas.com>
Co-authored-by: Jan Katins <jasc@gmx.net>
  • Loading branch information
Jason Turim and jankatins committed Mar 8, 2024
1 parent ed3b9fa commit cfc3b27
Show file tree
Hide file tree
Showing 6 changed files with 92 additions and 23 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.
48 changes: 40 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,37 @@ 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.
"""
if driver is _UNSET:
driver_str = self.driver
else:
driver_str = 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")
22 changes: 20 additions & 2 deletions modules/postgres/tests/test_postgres.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,27 @@
import sqlalchemy

from testcontainers.postgres import PostgresContainer


def test_docker_run_postgres():
# https://www.postgresql.org/support/versioning/
supported_versions = ["12", "13", "14", "15", "16", "latest"]

for version in supported_versions:
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_with_sqlalchemy():
import sqlalchemy
postgres_container = PostgresContainer("postgres:9.5")
with postgres_container as postgres:
engine = sqlalchemy.create_engine(postgres.get_connection_url())
Expand Down
7 changes: 3 additions & 4 deletions poetry.lock

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

8 changes: 5 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,23 @@ 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 = "*"


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

0 comments on commit cfc3b27

Please sign in to comment.