Skip to content

Commit 2ca2321

Browse files
authored
feat(generic): Reintroducing the generic SQL module (#892)
Related to #884 Trying to replace the old generic.py in core to a nicer version under the generic module. `SqlContainer` its not as good (in being generic) as `ServerContainer` but it should allow us to deprecate `core/testcontainers/core/generic.py` with minimal effort for users, it could lead to a more Generic version like `DBContainer` in the future. Update 1: Refactor to use `SqlConnectWaitStrategy` Update 2: Now `SqlConnectWaitStrategy` is required and the users can provide `SqlContainer` with any wait strategy. Update 3: Now utilizes all the latest improvements from `WaitStrategy` Note: I think the added tests + documentation (provided in this PR) are by themselves a great improvement over the current generic.py
1 parent af382f7 commit 2ca2321

File tree

10 files changed

+467
-2
lines changed

10 files changed

+467
-2
lines changed

core/testcontainers/core/generic.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@
2929
class DbContainer(DockerContainer):
3030
"""
3131
**DEPRECATED (for removal)**
32+
Please use database-specific container classes or `SqlContainer` instead.
33+
# from testcontainers.generic.sql import SqlContainer
3234
3335
Generic database container.
3436
"""

modules/generic/README.rst

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ FastAPI container that is using :code:`ServerContainer`
99

1010
>>> from testcontainers.generic import ServerContainer
1111
>>> from testcontainers.core.waiting_utils import wait_for_logs
12+
>>> from testcontainers.core.image import DockerImage
1213

1314
>>> with DockerImage(path="./modules/generic/tests/samples/fastapi", tag="fastapi-test:latest") as image:
1415
... with ServerContainer(port=80, image=image) as fastapi_server:
@@ -50,3 +51,39 @@ A more advance use-case, where we are using a FastAPI container that is using Re
5051
... response = client.get(f"/get/{test_data['key']}")
5152
... assert response.status_code == 200, "Failed to get data"
5253
... assert response.json() == {"key": test_data["key"], "value": test_data["value"]}
54+
55+
.. autoclass:: testcontainers.generic.SqlContainer
56+
.. title:: testcontainers.generic.SqlContainer
57+
58+
Postgres container that is using :code:`SqlContainer`
59+
60+
.. doctest::
61+
62+
>>> from testcontainers.generic import SqlContainer
63+
>>> from testcontainers.generic.providers.sql_connection_wait_strategy import SqlAlchemyConnectWaitStrategy
64+
>>> from sqlalchemy import text
65+
>>> import sqlalchemy
66+
67+
>>> class CustomPostgresContainer(SqlContainer):
68+
... def __init__(self, image="postgres:15-alpine",
69+
... port=5432, username="test", password="test", dbname="test"):
70+
... super().__init__(image=image, wait_strategy=SqlAlchemyConnectWaitStrategy())
71+
... self.port_to_expose = port
72+
... self.username = username
73+
... self.password = password
74+
... self.dbname = dbname
75+
... def get_connection_url(self) -> str:
76+
... host = self.get_container_host_ip()
77+
... port = self.get_exposed_port(self.port_to_expose)
78+
... return f"postgresql://{self.username}:{self.password}@{host}:{port}/{self.dbname}"
79+
... def _configure(self) -> None:
80+
... self.with_exposed_ports(self.port_to_expose)
81+
... self.with_env("POSTGRES_USER", self.username)
82+
... self.with_env("POSTGRES_PASSWORD", self.password)
83+
... self.with_env("POSTGRES_DB", self.dbname)
84+
85+
>>> with CustomPostgresContainer() as postgres:
86+
... engine = sqlalchemy.create_engine(postgres.get_connection_url())
87+
... with engine.connect() as conn:
88+
... result = conn.execute(text("SELECT 1"))
89+
... assert result.scalar() == 1
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,2 @@
11
from .server import ServerContainer # noqa: F401
2+
from .sql import SqlContainer # noqa: F401
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
from .sql_connection_wait_strategy import SqlAlchemyConnectWaitStrategy # noqa: F401
Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
# This module provides a wait strategy for SQL database connectivity testing using SQLAlchemy.
2+
# It includes handling for transient exceptions and connection retries.
3+
4+
import logging
5+
6+
from testcontainers.core.waiting_utils import WaitStrategy, WaitStrategyTarget
7+
8+
logger = logging.getLogger(__name__)
9+
10+
ADDITIONAL_TRANSIENT_ERRORS = []
11+
try:
12+
from sqlalchemy.exc import DBAPIError
13+
14+
ADDITIONAL_TRANSIENT_ERRORS.append(DBAPIError)
15+
except ImportError:
16+
logger.debug("SQLAlchemy not available, skipping DBAPIError handling")
17+
18+
19+
class SqlAlchemyConnectWaitStrategy(WaitStrategy):
20+
"""Wait strategy for database connectivity testing using SQLAlchemy."""
21+
22+
def __init__(self):
23+
super().__init__()
24+
self.with_transient_exceptions(TimeoutError, ConnectionError, *ADDITIONAL_TRANSIENT_ERRORS)
25+
26+
def wait_until_ready(self, container: WaitStrategyTarget) -> None:
27+
"""Test database connectivity with retry logic until success or timeout."""
28+
if not hasattr(container, "get_connection_url"):
29+
raise AttributeError(f"Container {container} must have a get_connection_url method")
30+
31+
try:
32+
import sqlalchemy
33+
except ImportError as e:
34+
raise ImportError("SQLAlchemy is required for database containers") from e
35+
36+
def _test_connection() -> bool:
37+
"""Test database connection, returning True if successful."""
38+
engine = sqlalchemy.create_engine(container.get_connection_url())
39+
try:
40+
with engine.connect():
41+
logger.info("Database connection successful")
42+
return True
43+
finally:
44+
engine.dispose()
45+
46+
result = self._poll(_test_connection)
47+
if not result:
48+
raise TimeoutError(f"Database connection failed after {self._startup_timeout}s timeout")

modules/generic/testcontainers/generic/server.py

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,6 @@
99
from testcontainers.core.image import DockerImage
1010
from testcontainers.core.waiting_utils import wait_container_is_ready
1111

12-
# This comment can be removed (Used for testing)
13-
1412

1513
class ServerContainer(DockerContainer):
1614
"""
Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
import logging
2+
from typing import Any, Optional
3+
from urllib.parse import quote, urlencode
4+
5+
from testcontainers.core.container import DockerContainer
6+
from testcontainers.core.exceptions import ContainerStartException
7+
from testcontainers.core.waiting_utils import WaitStrategy
8+
9+
logger = logging.getLogger(__name__)
10+
11+
12+
class SqlContainer(DockerContainer):
13+
"""
14+
Generic SQL database container providing common functionality.
15+
16+
This class can serve as a base for database-specific container implementations.
17+
It provides connection management, URL construction, and basic lifecycle methods.
18+
Database connection readiness is automatically handled by the provided wait strategy.
19+
20+
Note: `SqlAlchemyConnectWaitStrategy` from `sql_connection_wait_strategy` is a provided wait strategy for SQL databases.
21+
"""
22+
23+
def __init__(self, image: str, wait_strategy: WaitStrategy, **kwargs):
24+
"""
25+
Initialize SqlContainer with optional wait strategy.
26+
27+
Args:
28+
image: Docker image name
29+
wait_strategy: Wait strategy for SQL database connectivity
30+
**kwargs: Additional arguments passed to DockerContainer
31+
"""
32+
super().__init__(image, **kwargs)
33+
self.wait_strategy = wait_strategy
34+
35+
def _create_connection_url(
36+
self,
37+
dialect: str,
38+
username: str,
39+
password: str,
40+
host: Optional[str] = None,
41+
port: Optional[int] = None,
42+
dbname: Optional[str] = None,
43+
query_params: Optional[dict[str, str]] = None,
44+
**kwargs: Any,
45+
) -> str:
46+
"""
47+
Create a database connection URL.
48+
49+
Args:
50+
dialect: Database dialect (e.g., 'postgresql', 'mysql')
51+
username: Database username
52+
password: Database password
53+
host: Database host (defaults to container host)
54+
port: Database port
55+
dbname: Database name
56+
query_params: Additional query parameters for the URL
57+
**kwargs: Additional parameters (checked for deprecated usage)
58+
59+
Returns:
60+
str: Formatted database connection URL
61+
62+
Raises:
63+
ValueError: If unexpected arguments are provided or required parameters are missing
64+
ContainerStartException: If container is not started
65+
"""
66+
67+
if self._container is None:
68+
raise ContainerStartException("Container has not been started")
69+
70+
host = host or self.get_container_host_ip()
71+
exposed_port = self.get_exposed_port(port)
72+
quoted_password = quote(password, safe="")
73+
quoted_username = quote(username, safe="")
74+
url = f"{dialect}://{quoted_username}:{quoted_password}@{host}:{exposed_port}"
75+
76+
if dbname:
77+
quoted_dbname = quote(dbname, safe="")
78+
url = f"{url}/{quoted_dbname}"
79+
80+
if query_params:
81+
query_string = urlencode(query_params)
82+
url = f"{url}?{query_string}"
83+
84+
return url
85+
86+
def start(self) -> "SqlContainer":
87+
"""
88+
Start the database container and perform initialization.
89+
90+
Returns:
91+
SqlContainer: Self for method chaining
92+
93+
Raises:
94+
ContainerStartException: If container fails to start
95+
Exception: If configuration, seed transfer, or connection fails
96+
"""
97+
logger.info(f"Starting database container: {self.image}")
98+
99+
try:
100+
self._configure()
101+
self.waiting_for(self.wait_strategy)
102+
super().start()
103+
self._transfer_seed()
104+
logger.info("Database container started successfully")
105+
except Exception as e:
106+
logger.error(f"Failed to start database container: {e}")
107+
raise
108+
109+
return self
110+
111+
def _configure(self) -> None:
112+
"""
113+
Configure the database container before starting.
114+
115+
Raises:
116+
NotImplementedError: Must be implemented by subclasses
117+
"""
118+
raise NotImplementedError("Subclasses must implement _configure()")
119+
120+
def _transfer_seed(self) -> None:
121+
"""
122+
Transfer seed data to the database container.
123+
124+
This method can be overridden by subclasses to provide
125+
database-specific seeding functionality.
126+
"""
127+
logger.debug("No seed data to transfer")
128+
129+
def get_connection_url(self) -> str:
130+
"""
131+
Get the database connection URL.
132+
133+
Returns:
134+
str: Database connection URL
135+
136+
Raises:
137+
NotImplementedError: Must be implemented by subclasses
138+
"""
139+
raise NotImplementedError("Subclasses must implement get_connection_url()")

0 commit comments

Comments
 (0)