Skip to content

Commit 56f8081

Browse files
logicbombjankatins
andcommitted
feat(postgres): Remove SqlAlchemy dependency from postgres container
Create a new generic database subclass - DependencyFreeDbContainer 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 <[email protected]> Co-authored-by: Jan Katins <[email protected]>
1 parent ed3b9fa commit 56f8081

File tree

5 files changed

+94
-27
lines changed

5 files changed

+94
-27
lines changed

INDEX.rst

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,15 +46,31 @@ Getting Started
4646
>>> import sqlalchemy
4747

4848
>>> with PostgresContainer("postgres:9.5") as postgres:
49-
... engine = sqlalchemy.create_engine(postgres.get_connection_url())
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 ......'
5556

5657
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.
5758

59+
.. doctest::
60+
61+
>>> import asyncpg
62+
>>> from testcontainers.postgres import PostgresContainer
63+
64+
>>> with PostgresContainer("postgres:9.5", 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 the driver is set to None, to influence the :code:`get_connection_url()` convenience method. Note, that the :code:`sqlalchemy` package is no longer a dependency to launch the Postgres container, so your project must provide support for the specified driver.
72+
73+
5874
Installation
5975
------------
6076

core/testcontainers/core/generic.py

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,46 @@
2626
pass
2727

2828

29+
class DependencyFreeDbContainer(DockerContainer):
30+
"""
31+
A generic database without any package dependencies
32+
"""
33+
34+
def start(self) -> "DbContainer":
35+
self._configure()
36+
super().start()
37+
self._verify_status()
38+
return self
39+
40+
def _verify_status(self) -> "DependencyFreeDbContainer":
41+
"""override this method to ensure the database is running and accepting connections"""
42+
raise NotImplementedError
43+
44+
def _configure(self) -> None:
45+
raise NotImplementedError
46+
47+
def _create_connection_url(
48+
self,
49+
dialect: str,
50+
username: str,
51+
password: str,
52+
host: Optional[str] = None,
53+
port: Optional[int] = None,
54+
dbname: Optional[str] = None,
55+
**kwargs,
56+
) -> str:
57+
if raise_for_deprecated_parameter(kwargs, "db_name", "dbname"):
58+
raise ValueError(f"Unexpected arguments: {','.join(kwargs)}")
59+
if self._container is None:
60+
raise ContainerStartException("container has not been started")
61+
host = host or self.get_container_host_ip()
62+
port = self.get_exposed_port(port)
63+
url = f"{dialect}://{username}:{password}@{host}:{port}"
64+
if dbname:
65+
url = f"{url}/{dbname}"
66+
return url
67+
68+
2969
class DbContainer(DockerContainer):
3070
"""
3171
Generic database container.

modules/postgres/testcontainers/postgres/__init__.py

Lines changed: 27 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -11,13 +11,16 @@
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

16-
from testcontainers.core.generic import DbContainer
17+
from testcontainers.core.config import MAX_TRIES, SLEEP_TIME
18+
from testcontainers.core.generic import DependencyFreeDbContainer
1719
from testcontainers.core.utils import raise_for_deprecated_parameter
20+
from testcontainers.core.waiting_utils import wait_container_is_ready, wait_for_logs
1821

1922

20-
class PostgresContainer(DbContainer):
23+
class PostgresContainer(DependencyFreeDbContainer):
2124
"""
2225
Postgres database container.
2326
@@ -31,7 +34,7 @@ class PostgresContainer(DbContainer):
3134
>>> from testcontainers.postgres import PostgresContainer
3235
>>> import sqlalchemy
3336
34-
>>> postgres_container = PostgresContainer("postgres:9.5")
37+
>>> postgres_container = PostgresContainer("postgres:16")
3538
>>> with postgres_container as postgres:
3639
... engine = sqlalchemy.create_engine(postgres.get_connection_url())
3740
... with engine.begin() as connection:
@@ -48,16 +51,16 @@ def __init__(
4851
username: Optional[str] = None,
4952
password: Optional[str] = None,
5053
dbname: Optional[str] = None,
51-
driver: str = "psycopg2",
54+
driver: str | None = "psycopg2",
5255
**kwargs,
5356
) -> None:
5457
raise_for_deprecated_parameter(kwargs, "user", "username")
5558
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")
59+
self.username: str = username or os.environ.get("POSTGRES_USER", "test")
60+
self.password: str = password or os.environ.get("POSTGRES_PASSWORD", "test")
61+
self.dbname: str = dbname or os.environ.get("POSTGRES_DB", "test")
5962
self.port = port
60-
self.driver = driver
63+
self.driver = f"+{driver}" if driver else ""
6164

6265
self.with_exposed_ports(self.port)
6366

@@ -68,10 +71,25 @@ def _configure(self) -> None:
6871

6972
def get_connection_url(self, host=None) -> str:
7073
return super()._create_connection_url(
71-
dialect=f"postgresql+{self.driver}",
74+
dialect=f"postgresql{self.driver}",
7275
username=self.username,
7376
password=self.password,
7477
dbname=self.dbname,
7578
host=host,
7679
port=self.port,
7780
)
81+
82+
@wait_container_is_ready()
83+
def _verify_status(self) -> None:
84+
wait_for_logs(self, ".*database system is ready to accept connections.*", MAX_TRIES, SLEEP_TIME)
85+
86+
count = 0
87+
while count < MAX_TRIES:
88+
status, _ = self.exec(f"pg_isready -hlocalhost -p{self.port} -U{self.username}")
89+
if status == 0:
90+
return
91+
92+
sleep(SLEEP_TIME)
93+
count += 1
94+
95+
raise RuntimeError("Postgres could not get into a ready state")
Lines changed: 8 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,14 @@
1-
import sqlalchemy
2-
31
from testcontainers.postgres import PostgresContainer
42

53

64
def test_docker_run_postgres():
7-
postgres_container = PostgresContainer("postgres:9.5")
8-
with postgres_container as postgres:
9-
engine = sqlalchemy.create_engine(postgres.get_connection_url())
10-
with engine.begin() as connection:
11-
result = connection.execute(sqlalchemy.text("select version()"))
12-
for row in result:
13-
assert row[0].lower().startswith("postgresql 9.5")
5+
# https://www.postgresql.org/support/versioning/
6+
supported_versions = ["12", "13", "14", "15", "16", "latest"]
147

8+
for version in supported_versions:
9+
postgres_container = PostgresContainer(f"postgres:{version}")
10+
with postgres_container as postgres:
11+
status, msg = postgres.exec(f"pg_isready -hlocalhost -p{postgres.port} -U{postgres.username}")
1512

16-
def test_docker_run_postgres_with_driver_pg8000():
17-
postgres_container = PostgresContainer("postgres:9.5", driver="pg8000")
18-
with postgres_container as postgres:
19-
engine = sqlalchemy.create_engine(postgres.get_connection_url())
20-
with engine.begin() as connection:
21-
connection.execute(sqlalchemy.text("select 1=1"))
13+
assert msg.decode("utf-8").endswith("accepting connections\n")
14+
assert status == 0

pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,7 @@ neo4j = ["neo4j"]
102102
nginx = []
103103
opensearch = ["opensearch-py"]
104104
oracle = ["sqlalchemy", "cx_Oracle"]
105-
postgres = ["sqlalchemy", "psycopg2-binary"]
105+
postgres = []
106106
rabbitmq = ["pika"]
107107
redis = ["redis"]
108108
selenium = ["selenium"]

0 commit comments

Comments
 (0)