Skip to content

Commit f30eb1d

Browse files
jankatinsJason Turim
and
Jason Turim
authoredMar 10, 2024··
feat(postgres): Remove SqlAlchemy dependency from postgres container (#445)
Updates the pg testcontainer implementation to not use (and not install) SQLAlchemy nor psycopg2. Closes: #340 Closes: #336 Closes: #320 --------- Co-authored-by: Jason Turim <jason@opscanvas.com>
1 parent 2c4f171 commit f30eb1d

File tree

6 files changed

+100
-25
lines changed

6 files changed

+100
-25
lines changed
 

‎INDEX.rst

+22-4
Original file line numberDiff line numberDiff line change
@@ -45,15 +45,33 @@ Getting Started
4545
>>> from testcontainers.postgres import PostgresContainer
4646
>>> import sqlalchemy
4747

48-
>>> with PostgresContainer("postgres:9.5") as postgres:
49-
... engine = sqlalchemy.create_engine(postgres.get_connection_url())
48+
>>> with PostgresContainer("postgres:latest") as postgres:
49+
... psql_url = postgres.get_connection_url()
50+
... engine = sqlalchemy.create_engine(psql_url)
5051
... with engine.begin() as connection:
5152
... result = connection.execute(sqlalchemy.text("select version()"))
5253
... version, = result.fetchone()
5354
>>> version
54-
'PostgreSQL 9.5...'
55+
'PostgreSQL ...'
56+
57+
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.
58+
59+
.. doctest::
60+
61+
>>> import asyncpg
62+
>>> from testcontainers.postgres import PostgresContainer
63+
64+
>>> with PostgresContainer("postgres:16", driver=None) as postgres:
65+
... psql_url = container.get_connection_url()
66+
... with asyncpg.create_pool(dsn=psql_url,server_settings={"jit": "off"}) as pool:
67+
... conn = await pool.acquire()
68+
... ret = await conn.fetchval("SELECT 1")
69+
... assert ret == 1
70+
71+
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).
72+
73+
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.
5574

56-
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.
5775

5876
Installation
5977
------------

‎README.md

+2-2
Original file line numberDiff line numberDiff line change
@@ -12,13 +12,13 @@ For more information, see [the docs][readthedocs].
1212
>>> from testcontainers.postgres import PostgresContainer
1313
>>> import sqlalchemy
1414

15-
>>> with PostgresContainer("postgres:9.5") as postgres:
15+
>>> with PostgresContainer("postgres:16") as postgres:
1616
... engine = sqlalchemy.create_engine(postgres.get_connection_url())
1717
... with engine.begin() as connection:
1818
... result = connection.execute(sqlalchemy.text("select version()"))
1919
... version, = result.fetchone()
2020
>>> version
21-
'PostgreSQL 9.5...'
21+
'PostgreSQL 16...'
2222
```
2323

2424
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.

‎modules/postgres/testcontainers/postgres/__init__.py

+37-8
Original file line numberDiff line numberDiff line change
@@ -11,16 +11,23 @@
1111
# License for the specific language governing permissions and limitations
1212
# under the License.
1313
import os
14+
from time import sleep
1415
from typing import Optional
1516

17+
from testcontainers.core.config import MAX_TRIES, SLEEP_TIME
1618
from testcontainers.core.generic import DbContainer
1719
from testcontainers.core.utils import raise_for_deprecated_parameter
20+
from testcontainers.core.waiting_utils import wait_container_is_ready, wait_for_logs
21+
22+
_UNSET = object()
1823

1924

2025
class PostgresContainer(DbContainer):
2126
"""
2227
Postgres database container.
2328
29+
To get a URL without a driver, pass in :code:`driver=None`.
30+
2431
Example:
2532
2633
The example spins up a Postgres database and connects to it using the :code:`psycopg`
@@ -31,7 +38,7 @@ class PostgresContainer(DbContainer):
3138
>>> from testcontainers.postgres import PostgresContainer
3239
>>> import sqlalchemy
3340
34-
>>> postgres_container = PostgresContainer("postgres:9.5")
41+
>>> postgres_container = PostgresContainer("postgres:16")
3542
>>> with postgres_container as postgres:
3643
... engine = sqlalchemy.create_engine(postgres.get_connection_url())
3744
... with engine.begin() as connection:
@@ -48,16 +55,16 @@ def __init__(
4855
username: Optional[str] = None,
4956
password: Optional[str] = None,
5057
dbname: Optional[str] = None,
51-
driver: str = "psycopg2",
58+
driver: Optional[str] = "psycopg2",
5259
**kwargs,
5360
) -> None:
5461
raise_for_deprecated_parameter(kwargs, "user", "username")
5562
super().__init__(image=image, **kwargs)
56-
self.username = username or os.environ.get("POSTGRES_USER", "test")
57-
self.password = password or os.environ.get("POSTGRES_PASSWORD", "test")
58-
self.dbname = dbname or os.environ.get("POSTGRES_DB", "test")
63+
self.username: str = username or os.environ.get("POSTGRES_USER", "test")
64+
self.password: str = password or os.environ.get("POSTGRES_PASSWORD", "test")
65+
self.dbname: str = dbname or os.environ.get("POSTGRES_DB", "test")
5966
self.port = port
60-
self.driver = driver
67+
self.driver = f"+{driver}" if driver else ""
6168

6269
self.with_exposed_ports(self.port)
6370

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

69-
def get_connection_url(self, host=None) -> str:
76+
def get_connection_url(self, host: Optional[str] = None, driver: Optional[str] = _UNSET) -> str:
77+
"""Get a DB connection URL to connect to the PG DB.
78+
79+
If a driver is set in the constructor (defaults to psycopg2!), the URL will contain the
80+
driver. The optional driver argument to :code:`get_connection_url` overwrites the constructor
81+
set value. Pass :code:`driver=None` to get URLs without a driver.
82+
"""
83+
driver_str = self.driver if driver is _UNSET else f"+{driver}"
7084
return super()._create_connection_url(
71-
dialect=f"postgresql+{self.driver}",
85+
dialect=f"postgresql{driver_str}",
7286
username=self.username,
7387
password=self.password,
7488
dbname=self.dbname,
7589
host=host,
7690
port=self.port,
7791
)
92+
93+
@wait_container_is_ready()
94+
def _connect(self) -> None:
95+
wait_for_logs(self, ".*database system is ready to accept connections.*", MAX_TRIES, SLEEP_TIME)
96+
97+
count = 0
98+
while count < MAX_TRIES:
99+
status, _ = self.exec(f"pg_isready -hlocalhost -p{self.port} -U{self.username}")
100+
if status == 0:
101+
return
102+
103+
sleep(SLEEP_TIME)
104+
count += 1
105+
106+
raise RuntimeError("Postgres could not get into a ready state")

‎modules/postgres/tests/test_postgres.py

+27-2
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,34 @@
1-
import sqlalchemy
1+
import sys
2+
3+
import pytest
24

35
from testcontainers.postgres import PostgresContainer
6+
import sqlalchemy
7+
8+
9+
# https://www.postgresql.org/support/versioning/
10+
@pytest.mark.parametrize("version", ["12", "13", "14", "15", "16", "latest"])
11+
def test_docker_run_postgres(version: str, monkeypatch):
12+
def fail(*args, **kwargs):
13+
raise AssertionError("SQLA was called during PG container setup")
14+
15+
monkeypatch.setattr(sqlalchemy, "create_engine", fail)
16+
postgres_container = PostgresContainer(f"postgres:{version}")
17+
with postgres_container as postgres:
18+
status, msg = postgres.exec(f"pg_isready -hlocalhost -p{postgres.port} -U{postgres.username}")
19+
20+
assert msg.decode("utf-8").endswith("accepting connections\n")
21+
assert status == 0
22+
23+
status, msg = postgres.exec(
24+
f"psql -hlocalhost -p{postgres.port} -U{postgres.username} -c 'select 2*3*5*7*11*13*17 as a;' "
25+
)
26+
assert "510510" in msg.decode("utf-8")
27+
assert "(1 row)" in msg.decode("utf-8")
28+
assert status == 0
429

530

6-
def test_docker_run_postgres():
31+
def test_docker_run_postgres_with_sqlalchemy():
732
postgres_container = PostgresContainer("postgres:9.5")
833
with postgres_container as postgres:
934
engine = sqlalchemy.create_engine(postgres.get_connection_url())

‎poetry.lock

+6-6
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎pyproject.toml

+6-3
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,6 @@ pymysql = { version = "*", extras = ["rsa"], optional = true }
7979
neo4j = { version = "*", optional = true }
8080
opensearch-py = { version = "*", optional = true }
8181
cx_Oracle = { version = "*", optional = true }
82-
psycopg2-binary = { version = "*", optional = true }
8382
pika = { version = "*", optional = true }
8483
redis = { version = "*", optional = true }
8584
selenium = { version = "*", optional = true }
@@ -102,20 +101,24 @@ neo4j = ["neo4j"]
102101
nginx = []
103102
opensearch = ["opensearch-py"]
104103
oracle = ["sqlalchemy", "cx_Oracle"]
105-
postgres = ["sqlalchemy", "psycopg2-binary"]
104+
postgres = []
106105
rabbitmq = ["pika"]
107106
redis = ["redis"]
108107
selenium = ["selenium"]
109108

110109
[tool.poetry.group.dev.dependencies]
111110
mypy = "1.7.1"
112111
pre-commit = "^3.6"
113-
pg8000 = "*"
114112
pytest = "7.4.3"
115113
pytest-cov = "4.1.0"
116114
sphinx = "^7.2.6"
117115
twine = "^4.0.2"
118116
anyio = "^4.3.0"
117+
# for tests only
118+
psycopg2-binary = "*"
119+
pg8000 = "*"
120+
sqlalchemy = "*"
121+
119122

120123
[[tool.poetry.source]]
121124
name = "PyPI"

0 commit comments

Comments
 (0)
Please sign in to comment.