Skip to content

Commit 2d9eee3

Browse files
feat: Add ExecWaitStrategy and migrate Postgres from deprecated decorator (#935)
## Problem The `PostgresContainer` currently uses the deprecated `@wait_container_is_ready()` decorator, which is slated for removal. The container was previously migrated away from log-based waiting due to locale-dependent issues (#703, #695), but still relies on the deprecated decorator pattern. ## Solution This PR introduces a new `ExecWaitStrategy` wait strategy and migrates `PostgresContainer` to use it. ### Changes 1. **New `ExecWaitStrategy` in `core/testcontainers/core/wait_strategies.py`** - Executes commands inside containers and waits for expected exit codes - Follows the modern structured wait strategy pattern - Reusable for any `DockerContainer` with CLI-based health checks - Composable with other wait strategies via `CompositeWaitStrategy` - Includes runtime check to ensure container supports `exec()` 2. **Updated `modules/postgres/testcontainers/postgres/__init__.py`** - Migrated `_connect()` method from `@wait_container_is_ready` decorator to `ExecWaitStrategy` - Uses `psql` command execution to verify database readiness - Maintains existing behavior and password escaping logic ### Design Consideration: Protocol Limitations **Issue**: The `WaitStrategyTarget` protocol is designed to support both `DockerContainer` and `ComposeContainer`, but `ExecWaitStrategy` only works with `DockerContainer` (which has an `exec()` method). **Current Solution**: Runtime check with `hasattr(container, "exec")` that raises a clear error if used with incompatible containers. **Discussion Points**: - Should we have separate protocol types for different container capabilities? - Should `ExecWaitStrategy` accept a more specific type (e.g., `DockerContainer` directly)? - Is the runtime check acceptable, or should this be enforced at the type level? I've opted for the runtime check to maintain consistency with the existing `WaitStrategyTarget` protocol pattern, but I'm open to alternatives if the maintainers prefer a different approach. # PR Checklist - [x] Your PR title follows the [Conventional Commits](https://www.conventionalcommits.org/en/v1.0.0/) syntax as we make use of this for detecting Semantic Versioning changes. - Additions to the community modules do not contribute to SemVer scheme: all community features will be tagged [community-feat](https://github.com/testcontainers/testcontainers-python/issues?q=label%3Acommunity-feat+), but we do not want to release minor or major versions due to features or breaking changes outside of core. So please use `fix(postgres):` or `fix(my_new_vector_db):` if you want to add or modify community modules. This may change in the future if we have a separate package released with community modules. - [x] Your PR allows maintainers to edit your branch, this will speed up resolving minor issues! - [x] The new container is implemented under `modules/*` - Your module follows [PEP 420](https://peps.python.org/pep-0420/) with implicit namespace packages (if unsure, look at other existing community modules) - Your package namespacing follows `testcontainers.<modulename>.*` and you DO NOT have an `__init__.py` above your module's level. - Your module has its own tests under `modules/*/tests` - Your module has a `README.rst` and hooks in the `.. auto-class` and `.. title` of your container - Implement the new feature (typically in `__init__.py`) and corresponding tests. - [x] Your module is added in `pyproject.toml` - it is declared under `tool.poetry.packages` - see other community modules - it is declared under `tool.poetry.extras` with the same name as your module name, we still prefer adding _NO EXTRA DEPENDENCIES_, meaning `mymodule = []` is the preferred addition (see the notes at the bottom) - [x] Your branch is up-to-date (or your branch will be rebased with `git rebase`)
1 parent febccb7 commit 2d9eee3

File tree

2 files changed

+102
-5
lines changed

2 files changed

+102
-5
lines changed

core/testcontainers/core/wait_strategies.py

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
- HealthcheckWaitStrategy: Wait for Docker health checks to pass
77
- PortWaitStrategy: Wait for TCP ports to be available
88
- FileExistsWaitStrategy: Wait for files to exist on the filesystem
9+
- ExecWaitStrategy: Wait for command execution inside container to succeed
910
- CompositeWaitStrategy: Combine multiple wait strategies
1011
1112
Example:
@@ -19,6 +20,9 @@
1920
# Wait for log message
2021
container.waiting_for(LogMessageWaitStrategy("Server started"))
2122
23+
# Wait for command execution
24+
container.waiting_for(ExecWaitStrategy(["pg_isready", "-U", "postgres"]))
25+
2226
# Combine multiple strategies
2327
container.waiting_for(CompositeWaitStrategy(
2428
LogMessageWaitStrategy("Database ready"),
@@ -779,9 +783,103 @@ def wait_until_ready(self, container: WaitStrategyTarget) -> None:
779783
logger.debug("CompositeWaitStrategy: All strategies completed successfully")
780784

781785

786+
class ExecWaitStrategy(WaitStrategy):
787+
"""
788+
Wait for a command execution inside the container to succeed.
789+
790+
This strategy executes a command inside the container and waits for it to
791+
return a successful exit code. It's useful for databases and services
792+
that provide CLI tools to check readiness.
793+
794+
Args:
795+
command: Command to execute (list of strings or single string)
796+
expected_exit_code: Expected exit code for success (default: 0)
797+
798+
Example:
799+
# Wait for Postgres readiness
800+
strategy = ExecWaitStrategy(
801+
["sh", "-c",
802+
"PGPASSWORD='password' psql -U user -d db -h 127.0.0.1 -c 'select 1;'"]
803+
)
804+
805+
# Wait for Redis readiness
806+
strategy = ExecWaitStrategy(["redis-cli", "ping"])
807+
808+
# Check for specific exit code
809+
strategy = ExecWaitStrategy(["custom-healthcheck.sh"], expected_exit_code=0)
810+
"""
811+
812+
def __init__(
813+
self,
814+
command: Union[str, list[str]],
815+
expected_exit_code: int = 0,
816+
) -> None:
817+
super().__init__()
818+
self._command = command if isinstance(command, list) else [command]
819+
self._expected_exit_code = expected_exit_code
820+
821+
def wait_until_ready(self, container: WaitStrategyTarget) -> None:
822+
"""
823+
Wait until command execution succeeds with the expected exit code.
824+
825+
Args:
826+
container: The container to execute commands in
827+
828+
Raises:
829+
TimeoutError: If the command doesn't succeed within the timeout period
830+
RuntimeError: If the container doesn't support exec
831+
"""
832+
# Check if container supports exec (DockerContainer does, ComposeContainer doesn't)
833+
if not hasattr(container, "exec"):
834+
raise RuntimeError(
835+
f"ExecWaitStrategy requires a container with exec support. "
836+
f"Container type {type(container).__name__} does not support exec."
837+
)
838+
839+
start_time = time.time()
840+
last_exit_code = None
841+
last_output = None
842+
843+
while True:
844+
duration = time.time() - start_time
845+
if duration > self._startup_timeout:
846+
command_str = " ".join(self._command)
847+
raise TimeoutError(
848+
f"Command execution did not succeed within {self._startup_timeout:.3f} seconds. "
849+
f"Command: {command_str}. "
850+
f"Expected exit code: {self._expected_exit_code}, "
851+
f"last exit code: {last_exit_code}. "
852+
f"Last output: {last_output}. "
853+
f"Hint: Check if the service is starting correctly, the command is valid, "
854+
f"and all required environment variables or credentials are properly configured."
855+
)
856+
857+
try:
858+
result = container.exec(self._command)
859+
last_exit_code = result.exit_code
860+
last_output = result.output.decode() if hasattr(result.output, "decode") else str(result.output)
861+
862+
if result.exit_code == self._expected_exit_code:
863+
logger.debug(
864+
f"ExecWaitStrategy: Command succeeded with exit code {result.exit_code} after {duration:.2f}s"
865+
)
866+
return
867+
868+
logger.debug(
869+
f"ExecWaitStrategy: Command failed with exit code {result.exit_code}, "
870+
f"expected {self._expected_exit_code}. Retrying..."
871+
)
872+
except Exception as e:
873+
logger.debug(f"ExecWaitStrategy: Command execution failed with exception: {e}. Retrying...")
874+
last_output = str(e)
875+
876+
time.sleep(self._poll_interval)
877+
878+
782879
__all__ = [
783880
"CompositeWaitStrategy",
784881
"ContainerStatusWaitStrategy",
882+
"ExecWaitStrategy",
785883
"FileExistsWaitStrategy",
786884
"HealthcheckWaitStrategy",
787885
"HttpWaitStrategy",

modules/postgres/testcontainers/postgres/__init__.py

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515

1616
from testcontainers.core.generic import DbContainer
1717
from testcontainers.core.utils import raise_for_deprecated_parameter
18-
from testcontainers.core.waiting_utils import wait_container_is_ready
18+
from testcontainers.core.wait_strategies import ExecWaitStrategy
1919

2020
_UNSET = object()
2121

@@ -87,15 +87,14 @@ def get_connection_url(self, host: Optional[str] = None, driver: Optional[str] =
8787
port=self.port,
8888
)
8989

90-
@wait_container_is_ready()
9190
def _connect(self) -> None:
91+
"""Wait for Postgres to be ready by executing a query via psql."""
9292
escaped_single_password = self.password.replace("'", "'\"'\"'")
93-
result = self.exec(
93+
strategy = ExecWaitStrategy(
9494
[
9595
"sh",
9696
"-c",
9797
f"PGPASSWORD='{escaped_single_password}' psql --username {self.username} --dbname {self.dbname} --host 127.0.0.1 -c 'select version();'",
9898
]
9999
)
100-
if result.exit_code:
101-
raise ConnectionError("pg_isready is not ready yet")
100+
strategy.wait_until_ready(self)

0 commit comments

Comments
 (0)