Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
27 commits
Select commit Hold shift + click to select a range
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions core/testcontainers/core/generic.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,8 @@
class DbContainer(DockerContainer):
"""
**DEPRECATED (for removal)**
Please use database-specific container classes or `SqlContainer` instead.
# from testcontainers.generic.sql import SqlContainer

Generic database container.
"""
Expand Down
37 changes: 37 additions & 0 deletions modules/generic/README.rst
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ FastAPI container that is using :code:`ServerContainer`

>>> from testcontainers.generic import ServerContainer
>>> from testcontainers.core.waiting_utils import wait_for_logs
>>> from testcontainers.core.image import DockerImage

>>> with DockerImage(path="./modules/generic/tests/samples/fastapi", tag="fastapi-test:latest") as image:
... with ServerContainer(port=80, image=image) as fastapi_server:
Expand Down Expand Up @@ -50,3 +51,39 @@ A more advance use-case, where we are using a FastAPI container that is using Re
... response = client.get(f"/get/{test_data['key']}")
... assert response.status_code == 200, "Failed to get data"
... assert response.json() == {"key": test_data["key"], "value": test_data["value"]}

.. autoclass:: testcontainers.generic.SqlContainer
.. title:: testcontainers.generic.SqlContainer

Postgres container that is using :code:`SqlContainer`

.. doctest::

>>> from testcontainers.generic import SqlContainer
>>> from testcontainers.generic.providers.sql_connection_wait_strategy import SqlAlchemyConnectWaitStrategy
>>> from sqlalchemy import text
>>> import sqlalchemy

>>> class CustomPostgresContainer(SqlContainer):
... def __init__(self, image="postgres:15-alpine",
... port=5432, username="test", password="test", dbname="test"):
... super().__init__(image=image, wait_strategy=SqlAlchemyConnectWaitStrategy())
... self.port_to_expose = port
... self.username = username
... self.password = password
... self.dbname = dbname
... def get_connection_url(self) -> str:
... host = self.get_container_host_ip()
... port = self.get_exposed_port(self.port_to_expose)
... return f"postgresql://{self.username}:{self.password}@{host}:{port}/{self.dbname}"
... def _configure(self) -> None:
... self.with_exposed_ports(self.port_to_expose)
... self.with_env("POSTGRES_USER", self.username)
... self.with_env("POSTGRES_PASSWORD", self.password)
... self.with_env("POSTGRES_DB", self.dbname)

>>> with CustomPostgresContainer() as postgres:
... engine = sqlalchemy.create_engine(postgres.get_connection_url())
... with engine.connect() as conn:
... result = conn.execute(text("SELECT 1"))
... assert result.scalar() == 1
1 change: 1 addition & 0 deletions modules/generic/testcontainers/generic/__init__.py
Original file line number Diff line number Diff line change
@@ -1 +1,2 @@
from .server import ServerContainer # noqa: F401
from .sql import SqlContainer # noqa: F401
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
from .sql_connection_wait_strategy import SqlAlchemyConnectWaitStrategy # noqa: F401
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
# This module provides a wait strategy for SQL database connectivity testing using SQLAlchemy.
# It includes handling for transient exceptions and connection retries.

import logging

from testcontainers.core.waiting_utils import WaitStrategy, WaitStrategyTarget

logger = logging.getLogger(__name__)

ADDITIONAL_TRANSIENT_ERRORS = []
try:
from sqlalchemy.exc import DBAPIError

ADDITIONAL_TRANSIENT_ERRORS.append(DBAPIError)
except ImportError:
logger.debug("SQLAlchemy not available, skipping DBAPIError handling")


class SqlAlchemyConnectWaitStrategy(WaitStrategy):
"""Wait strategy for database connectivity testing using SQLAlchemy."""

def __init__(self):
super().__init__()
self.with_transient_exceptions(TimeoutError, ConnectionError, *ADDITIONAL_TRANSIENT_ERRORS)

def wait_until_ready(self, container: WaitStrategyTarget) -> None:
"""Test database connectivity with retry logic until success or timeout."""
if not hasattr(container, "get_connection_url"):
raise AttributeError(f"Container {container} must have a get_connection_url method")

try:
import sqlalchemy
except ImportError as e:
raise ImportError("SQLAlchemy is required for database containers") from e

def _test_connection() -> bool:
"""Test database connection, returning True if successful."""
engine = sqlalchemy.create_engine(container.get_connection_url())
try:
with engine.connect():
logger.info("Database connection successful")
return True
finally:
engine.dispose()

result = self._poll(_test_connection)
if not result:
raise TimeoutError(f"Database connection failed after {self._startup_timeout}s timeout")
2 changes: 0 additions & 2 deletions modules/generic/testcontainers/generic/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,6 @@
from testcontainers.core.image import DockerImage
from testcontainers.core.waiting_utils import wait_container_is_ready

# This comment can be removed (Used for testing)


class ServerContainer(DockerContainer):
"""
Expand Down
139 changes: 139 additions & 0 deletions modules/generic/testcontainers/generic/sql.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,139 @@
import logging
from typing import Any, Optional
from urllib.parse import quote, urlencode

from testcontainers.core.container import DockerContainer
from testcontainers.core.exceptions import ContainerStartException
from testcontainers.core.waiting_utils import WaitStrategy

logger = logging.getLogger(__name__)


class SqlContainer(DockerContainer):
"""
Generic SQL database container providing common functionality.

This class can serve as a base for database-specific container implementations.
It provides connection management, URL construction, and basic lifecycle methods.
Database connection readiness is automatically handled by the provided wait strategy.

Note: `SqlAlchemyConnectWaitStrategy` from `sql_connection_wait_strategy` is a provided wait strategy for SQL databases.
"""

def __init__(self, image: str, wait_strategy: WaitStrategy, **kwargs):
"""
Initialize SqlContainer with optional wait strategy.

Args:
image: Docker image name
wait_strategy: Wait strategy for SQL database connectivity
**kwargs: Additional arguments passed to DockerContainer
"""
super().__init__(image, **kwargs)
self.wait_strategy = wait_strategy

def _create_connection_url(
self,
dialect: str,
username: str,
password: str,
host: Optional[str] = None,
port: Optional[int] = None,
dbname: Optional[str] = None,
query_params: Optional[dict[str, str]] = None,
**kwargs: Any,
) -> str:
"""
Create a database connection URL.

Args:
dialect: Database dialect (e.g., 'postgresql', 'mysql')
username: Database username
password: Database password
host: Database host (defaults to container host)
port: Database port
dbname: Database name
query_params: Additional query parameters for the URL
**kwargs: Additional parameters (checked for deprecated usage)

Returns:
str: Formatted database connection URL

Raises:
ValueError: If unexpected arguments are provided or required parameters are missing
ContainerStartException: If container is not started
"""

if self._container is None:
raise ContainerStartException("Container has not been started")

host = host or self.get_container_host_ip()
exposed_port = self.get_exposed_port(port)
quoted_password = quote(password, safe="")
quoted_username = quote(username, safe="")
url = f"{dialect}://{quoted_username}:{quoted_password}@{host}:{exposed_port}"

if dbname:
quoted_dbname = quote(dbname, safe="")
url = f"{url}/{quoted_dbname}"

if query_params:
query_string = urlencode(query_params)
url = f"{url}?{query_string}"

return url

def start(self) -> "SqlContainer":
"""
Start the database container and perform initialization.

Returns:
SqlContainer: Self for method chaining

Raises:
ContainerStartException: If container fails to start
Exception: If configuration, seed transfer, or connection fails
"""
logger.info(f"Starting database container: {self.image}")

try:
self._configure()
self.waiting_for(self.wait_strategy)
super().start()
self._transfer_seed()
logger.info("Database container started successfully")
except Exception as e:
logger.error(f"Failed to start database container: {e}")
raise

return self

def _configure(self) -> None:
"""
Configure the database container before starting.

Raises:
NotImplementedError: Must be implemented by subclasses
"""
raise NotImplementedError("Subclasses must implement _configure()")

def _transfer_seed(self) -> None:
"""
Transfer seed data to the database container.

This method can be overridden by subclasses to provide
database-specific seeding functionality.
"""
logger.debug("No seed data to transfer")

def get_connection_url(self) -> str:
"""
Get the database connection URL.

Returns:
str: Database connection URL

Raises:
NotImplementedError: Must be implemented by subclasses
"""
raise NotImplementedError("Subclasses must implement get_connection_url()")
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should this have a default implementation to just call _create_connection_url? The SimpleSqlContainer implements it in this way & this feels like a reasonable default. I suspect most users of this class would either duplicate the implementation, or call the private method anyway which I think we'd want to discourage.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

SimpleSqlContainer is just a test, are you sure this is the default for all SQL related implementations?
for example Postgres needs driver so its not the exact same, but an extended build using the baseline _create_connection_url which is great. I personally like the current design and believe/hope it allows for max flexibility in the future.

Loading
Loading