Skip to content

Commit ade5599

Browse files
committed
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
1 parent 35fc247 commit ade5599

File tree

3 files changed

+91
-26
lines changed

3 files changed

+91
-26
lines changed

core/testcontainers/core/generic.py

Lines changed: 54 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -20,27 +20,77 @@
2020
ADDITIONAL_TRANSIENT_ERRORS = []
2121
try:
2222
from sqlalchemy.exc import DBAPIError
23+
2324
ADDITIONAL_TRANSIENT_ERRORS.append(DBAPIError)
2425
except ImportError:
2526
pass
2627

2728

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+
2869
class DbContainer(DockerContainer):
2970
"""
3071
Generic database container.
3172
"""
73+
3274
@wait_container_is_ready(*ADDITIONAL_TRANSIENT_ERRORS)
3375
def _connect(self) -> None:
3476
import sqlalchemy
77+
3578
engine = sqlalchemy.create_engine(self.get_connection_url())
3679
engine.connect()
3780

3881
def get_connection_url(self) -> str:
3982
raise NotImplementedError
4083

41-
def _create_connection_url(self, dialect: str, username: str, password: str,
42-
host: Optional[str] = None, port: Optional[int] = None,
43-
dbname: Optional[str] = None, **kwargs) -> str:
84+
def _create_connection_url(
85+
self,
86+
dialect: str,
87+
username: str,
88+
password: str,
89+
host: Optional[str] = None,
90+
port: Optional[int] = None,
91+
dbname: Optional[str] = None,
92+
**kwargs,
93+
) -> str:
4494
if raise_for_deprecated_parameter(kwargs, "db_name", "dbname"):
4595
raise ValueError(f"Unexpected arguments: {','.join(kwargs)}")
4696
if self._container is None:
@@ -52,7 +102,7 @@ def _create_connection_url(self, dialect: str, username: str, password: str,
52102
url = f"{url}/{dbname}"
53103
return url
54104

55-
def start(self) -> 'DbContainer':
105+
def start(self) -> "DbContainer":
56106
self._configure()
57107
super().start()
58108
self._connect()

postgres/testcontainers/postgres/__init__.py

Lines changed: 28 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -11,12 +11,17 @@
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
15-
from testcontainers.core.generic import DbContainer
16+
17+
from testcontainers.core.config import MAX_TRIES, SLEEP_TIME
18+
from testcontainers.core.generic import DependencyFreeDbContainer
1619
from testcontainers.core.utils import raise_for_deprecated_parameter
20+
from testcontainers.core.waiting_utils import (wait_container_is_ready,
21+
wait_for_logs)
1722

1823

19-
class PostgresContainer(DbContainer):
24+
class PostgresContainer(DependencyFreeDbContainer):
2025
"""
2126
Postgres database container.
2227
@@ -41,14 +46,14 @@ class PostgresContainer(DbContainer):
4146
"""
4247
def __init__(self, image: str = "postgres:latest", port: int = 5432,
4348
username: Optional[str] = None, password: Optional[str] = None,
44-
dbname: Optional[str] = None, driver: str = "psycopg2", **kwargs) -> None:
49+
dbname: Optional[str] = None, driver: str | None = "psycopg2", **kwargs) -> None:
4550
raise_for_deprecated_parameter(kwargs, "user", "username")
4651
super(PostgresContainer, self).__init__(image=image, **kwargs)
47-
self.username = username or os.environ.get("POSTGRES_USER", "test")
48-
self.password = password or os.environ.get("POSTGRES_PASSWORD", "test")
49-
self.dbname = dbname or os.environ.get("POSTGRES_DB", "test")
52+
self.username: str = username or os.environ.get("POSTGRES_USER", "test")
53+
self.password: str = password or os.environ.get("POSTGRES_PASSWORD", "test")
54+
self.dbname: str = dbname or os.environ.get("POSTGRES_DB", "test")
5055
self.port = port
51-
self.driver = driver
56+
self.driver = f"+{driver}" if driver else ""
5257

5358
self.with_exposed_ports(self.port)
5459

@@ -59,7 +64,22 @@ def _configure(self) -> None:
5964

6065
def get_connection_url(self, host=None) -> str:
6166
return super()._create_connection_url(
62-
dialect=f"postgresql+{self.driver}", username=self.username,
67+
dialect=f"postgresql{self.driver}", username=self.username,
6368
password=self.password, dbname=self.dbname, host=host,
6469
port=self.port,
6570
)
71+
72+
@wait_container_is_ready()
73+
def _verify_status(self) -> None:
74+
wait_for_logs(self, ".*database system is ready to accept connections.*", MAX_TRIES, SLEEP_TIME)
75+
76+
count = 0
77+
while count < MAX_TRIES:
78+
status, _ = self.exec(f"pg_isready -hlocalhost -p{self.port} -U{self.username}")
79+
if status == 0:
80+
return
81+
82+
sleep(SLEEP_TIME)
83+
count += 1
84+
85+
raise RuntimeError("Postgres could not get into a ready state")

postgres/tests/test_postgres.py

Lines changed: 9 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,15 @@
1-
import sqlalchemy
21
from testcontainers.postgres import PostgresContainer
32

43

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

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}")
12+
13+
assert msg.decode("utf-8").endswith("accepting connections\n")
14+
assert status == 0
1415

15-
def test_docker_run_postgres_with_driver_pg8000():
16-
postgres_container = PostgresContainer("postgres:9.5", driver="pg8000")
17-
with postgres_container as postgres:
18-
engine = sqlalchemy.create_engine(postgres.get_connection_url())
19-
with engine.begin() as connection:
20-
connection.execute(sqlalchemy.text("select 1=1"))

0 commit comments

Comments
 (0)