diff --git a/README.rst b/README.rst index 081456879..3ff92c3ff 100644 --- a/README.rst +++ b/README.rst @@ -44,13 +44,29 @@ Getting Started >>> import sqlalchemy >>> with PostgresContainer("postgres:9.5") as postgres: - ... engine = sqlalchemy.create_engine(postgres.get_connection_url()) + ... psql_url = postgres.get_connection_url() + ... engine = sqlalchemy.create_engine(psql_url) ... result = engine.execute("select version()") ... version, = result.fetchone() >>> version - 'PostgreSQL 9.5...' + 'PostgreSQL ......' + +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. + +.. doctest:: + + >>> import asyncpg + >>> from testcontainers.postgres import PostgresContainer + + >>> with PostgresContainer("postgres:9.5", driver=None) as postgres: + ... psql_url = container.get_connection_url() + ... with asyncpg.create_pool(dsn=psql_url,server_settings={"jit": "off"}) as pool: + ... conn = await pool.acquire() + ... ret = await conn.fetchval("SELECT 1") + ... assert ret == 1 + +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. -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. Installation ------------ diff --git a/core/testcontainers/core/generic.py b/core/testcontainers/core/generic.py index 7faac273a..b5323e5a1 100644 --- a/core/testcontainers/core/generic.py +++ b/core/testcontainers/core/generic.py @@ -20,27 +20,77 @@ ADDITIONAL_TRANSIENT_ERRORS = [] try: from sqlalchemy.exc import DBAPIError + ADDITIONAL_TRANSIENT_ERRORS.append(DBAPIError) except ImportError: pass +class DependencyFreeDbContainer(DockerContainer): + """ + A generic database without any package dependencies + """ + + def start(self) -> "DbContainer": + self._configure() + super().start() + self._verify_status() + return self + + def _verify_status(self) -> "DependencyFreeDbContainer": + """override this method to ensure the database is running and accepting connections""" + raise NotImplementedError + + def _configure(self) -> None: + raise NotImplementedError + + def _create_connection_url( + self, + dialect: str, + username: str, + password: str, + host: Optional[str] = None, + port: Optional[int] = None, + dbname: Optional[str] = None, + **kwargs, + ) -> str: + if raise_for_deprecated_parameter(kwargs, "db_name", "dbname"): + raise ValueError(f"Unexpected arguments: {','.join(kwargs)}") + if self._container is None: + raise ContainerStartException("container has not been started") + host = host or self.get_container_host_ip() + port = self.get_exposed_port(port) + url = f"{dialect}://{username}:{password}@{host}:{port}" + if dbname: + url = f"{url}/{dbname}" + return url + + class DbContainer(DockerContainer): """ Generic database container. """ + @wait_container_is_ready(*ADDITIONAL_TRANSIENT_ERRORS) def _connect(self) -> None: import sqlalchemy + engine = sqlalchemy.create_engine(self.get_connection_url()) engine.connect() def get_connection_url(self) -> str: raise NotImplementedError - def _create_connection_url(self, dialect: str, username: str, password: str, - host: Optional[str] = None, port: Optional[int] = None, - dbname: Optional[str] = None, **kwargs) -> str: + def _create_connection_url( + self, + dialect: str, + username: str, + password: str, + host: Optional[str] = None, + port: Optional[int] = None, + dbname: Optional[str] = None, + **kwargs, + ) -> str: if raise_for_deprecated_parameter(kwargs, "db_name", "dbname"): raise ValueError(f"Unexpected arguments: {','.join(kwargs)}") if self._container is None: @@ -52,7 +102,7 @@ def _create_connection_url(self, dialect: str, username: str, password: str, url = f"{url}/{dbname}" return url - def start(self) -> 'DbContainer': + def start(self) -> "DbContainer": self._configure() super().start() self._connect() diff --git a/postgres/testcontainers/postgres/__init__.py b/postgres/testcontainers/postgres/__init__.py index 85e0bac80..424aa5c1d 100644 --- a/postgres/testcontainers/postgres/__init__.py +++ b/postgres/testcontainers/postgres/__init__.py @@ -11,12 +11,16 @@ # License for the specific language governing permissions and limitations # under the License. import os +from time import sleep from typing import Optional -from testcontainers.core.generic import DbContainer + +from testcontainers.core.config import MAX_TRIES, SLEEP_TIME +from testcontainers.core.generic import DependencyFreeDbContainer from testcontainers.core.utils import raise_for_deprecated_parameter +from testcontainers.core.waiting_utils import wait_container_is_ready, wait_for_logs -class PostgresContainer(DbContainer): +class PostgresContainer(DependencyFreeDbContainer): """ Postgres database container. @@ -39,16 +43,24 @@ class PostgresContainer(DbContainer): >>> version 'PostgreSQL 9.5...' """ - def __init__(self, image: str = "postgres:latest", port: int = 5432, - username: Optional[str] = None, password: Optional[str] = None, - dbname: Optional[str] = None, driver: str = "psycopg2", **kwargs) -> None: + + def __init__( + self, + image: str = "postgres:latest", + port: int = 5432, + username: Optional[str] = None, + password: Optional[str] = None, + dbname: Optional[str] = None, + driver: str | None = "psycopg2", + **kwargs, + ) -> None: raise_for_deprecated_parameter(kwargs, "user", "username") super(PostgresContainer, self).__init__(image=image, **kwargs) - self.username = username or os.environ.get("POSTGRES_USER", "test") - self.password = password or os.environ.get("POSTGRES_PASSWORD", "test") - self.dbname = dbname or os.environ.get("POSTGRES_DB", "test") + self.username: str = username or os.environ.get("POSTGRES_USER", "test") + self.password: str = password or os.environ.get("POSTGRES_PASSWORD", "test") + self.dbname: str = dbname or os.environ.get("POSTGRES_DB", "test") self.port = port - self.driver = driver + self.driver = f"+{driver}" if driver else "" self.with_exposed_ports(self.port) @@ -59,7 +71,27 @@ def _configure(self) -> None: def get_connection_url(self, host=None) -> str: return super()._create_connection_url( - dialect=f"postgresql+{self.driver}", username=self.username, - password=self.password, dbname=self.dbname, host=host, + dialect=f"postgresql{self.driver}", + username=self.username, + password=self.password, + dbname=self.dbname, + host=host, port=self.port, ) + + @wait_container_is_ready() + def _verify_status(self) -> None: + wait_for_logs( + self, ".*database system is ready to accept connections.*", MAX_TRIES, SLEEP_TIME + ) + + count = 0 + while count < MAX_TRIES: + status, _ = self.exec(f"pg_isready -hlocalhost -p{self.port} -U{self.username}") + if status == 0: + return + + sleep(SLEEP_TIME) + count += 1 + + raise RuntimeError("Postgres could not get into a ready state") diff --git a/postgres/tests/test_postgres.py b/postgres/tests/test_postgres.py index c00c1b3fe..1d5d9adfb 100644 --- a/postgres/tests/test_postgres.py +++ b/postgres/tests/test_postgres.py @@ -1,20 +1,16 @@ -import sqlalchemy from testcontainers.postgres import PostgresContainer def test_docker_run_postgres(): - postgres_container = PostgresContainer("postgres:9.5") - with postgres_container as postgres: - engine = sqlalchemy.create_engine(postgres.get_connection_url()) - with engine.begin() as connection: - result = connection.execute(sqlalchemy.text("select version()")) - for row in result: - assert row[0].lower().startswith("postgresql 9.5") + # https://www.postgresql.org/support/versioning/ + supported_versions = ["11", "12", "13", "14", "latest"] + for version in supported_versions: + postgres_container = PostgresContainer(f"postgres:{version}") + with postgres_container as postgres: + status, msg = postgres.exec( + f"pg_isready -hlocalhost -p{postgres.port} -U{postgres.username}" + ) -def test_docker_run_postgres_with_driver_pg8000(): - postgres_container = PostgresContainer("postgres:9.5", driver="pg8000") - with postgres_container as postgres: - engine = sqlalchemy.create_engine(postgres.get_connection_url()) - with engine.begin() as connection: - connection.execute(sqlalchemy.text("select 1=1")) + assert msg.decode("utf-8").endswith("accepting connections\n") + assert status == 0