diff --git a/INDEX.rst b/INDEX.rst index be5e3d1cd..9612bb86d 100644 --- a/INDEX.rst +++ b/INDEX.rst @@ -45,15 +45,33 @@ Getting Started >>> from testcontainers.postgres import PostgresContainer >>> import sqlalchemy - >>> with PostgresContainer("postgres:9.5") as postgres: - ... engine = sqlalchemy.create_engine(postgres.get_connection_url()) + >>> with PostgresContainer("postgres:latest") as postgres: + ... psql_url = postgres.get_connection_url() + ... engine = sqlalchemy.create_engine(psql_url) ... with engine.begin() as connection: ... result = connection.execute(sqlalchemy.text("select version()")) ... version, = result.fetchone() >>> version - 'PostgreSQL 9.5...' + 'PostgreSQL ...' + +The snippet above will spin up the current latest version of a postgres database in a container. The :code:`get_connection_url()` convenience method returns a :code:`sqlalchemy` compatible url (using the :code:`psycopg2` driver per default) to connect to the database and retrieve the database version. + +.. doctest:: + + >>> import asyncpg + >>> from testcontainers.postgres import PostgresContainer + + >>> with PostgresContainer("postgres:16", 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 using a specific version and the driver is set to None, to influence the :code:`get_connection_url()` convenience method to not include a driver in the URL (e.g. for compatibility with :code:`psycopg` v3). + +Note, that the :code:`sqlalchemy` and :code:`psycopg2` packages are no longer a dependency of :code:`testcontainers[postgres]` and not needed to launch the Postgres container. Your project therefore needs to declare a dependency on the used driver and db access methods you use in your code. -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/README.md b/README.md index 58f5eca52..84f40b61c 100644 --- a/README.md +++ b/README.md @@ -12,13 +12,13 @@ For more information, see [the docs][readthedocs]. >>> from testcontainers.postgres import PostgresContainer >>> import sqlalchemy ->>> with PostgresContainer("postgres:9.5") as postgres: +>>> with PostgresContainer("postgres:16") as postgres: ... engine = sqlalchemy.create_engine(postgres.get_connection_url()) ... with engine.begin() as connection: ... result = connection.execute(sqlalchemy.text("select version()")) ... version, = result.fetchone() >>> version -'PostgreSQL 9.5...' +'PostgreSQL 16...' ``` The snippet above will spin up a postgres database in a container. The `get_connection_url()` convenience method returns a `sqlalchemy` compatible url we use to connect to the database and retrieve the database version. diff --git a/modules/postgres/testcontainers/postgres/__init__.py b/modules/postgres/testcontainers/postgres/__init__.py index a61ad2cf8..2a2371c26 100644 --- a/modules/postgres/testcontainers/postgres/__init__.py +++ b/modules/postgres/testcontainers/postgres/__init__.py @@ -11,16 +11,23 @@ # 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.config import MAX_TRIES, SLEEP_TIME from testcontainers.core.generic import DbContainer from testcontainers.core.utils import raise_for_deprecated_parameter +from testcontainers.core.waiting_utils import wait_container_is_ready, wait_for_logs + +_UNSET = object() class PostgresContainer(DbContainer): """ Postgres database container. + To get a URL without a driver, pass in :code:`driver=None`. + Example: The example spins up a Postgres database and connects to it using the :code:`psycopg` @@ -31,7 +38,7 @@ class PostgresContainer(DbContainer): >>> from testcontainers.postgres import PostgresContainer >>> import sqlalchemy - >>> postgres_container = PostgresContainer("postgres:9.5") + >>> postgres_container = PostgresContainer("postgres:16") >>> with postgres_container as postgres: ... engine = sqlalchemy.create_engine(postgres.get_connection_url()) ... with engine.begin() as connection: @@ -48,16 +55,16 @@ def __init__( username: Optional[str] = None, password: Optional[str] = None, dbname: Optional[str] = None, - driver: str = "psycopg2", + driver: Optional[str] = "psycopg2", **kwargs, ) -> None: raise_for_deprecated_parameter(kwargs, "user", "username") super().__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) @@ -66,12 +73,37 @@ def _configure(self) -> None: self.with_env("POSTGRES_PASSWORD", self.password) self.with_env("POSTGRES_DB", self.dbname) - def get_connection_url(self, host=None) -> str: + def get_connection_url(self, host: Optional[str] = None, driver: Optional[str] = _UNSET) -> str: + """Get a DB connection URL to connect to the PG DB. + + If a driver is set in the constructor (defaults to psycopg2!), the URL will contain the + driver. The optional driver argument to :code:`get_connection_url` overwrites the constructor + set value. Pass :code:`driver=None` to get URLs without a driver. + """ + if driver is _UNSET: + driver_str = self.driver + else: + driver_str = f"+{driver}" return super()._create_connection_url( - dialect=f"postgresql+{self.driver}", + dialect=f"postgresql{driver_str}", username=self.username, password=self.password, dbname=self.dbname, host=host, port=self.port, ) + + @wait_container_is_ready() + def _connect(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/modules/postgres/tests/test_postgres.py b/modules/postgres/tests/test_postgres.py index c1963531c..611c3485a 100644 --- a/modules/postgres/tests/test_postgres.py +++ b/modules/postgres/tests/test_postgres.py @@ -1,9 +1,30 @@ -import sqlalchemy +import sys + +import pytest from testcontainers.postgres import PostgresContainer +import sqlalchemy + +# https://www.postgresql.org/support/versioning/ +@pytest.mark.parametrize("version", ["12", "13", "14", "15", "16", "latest"]) +def test_docker_run_postgres(version: str, monkeypatch): + def fail(*args, **kwargs): + raise AssertionError("SQLA was called during PG container setup") + monkeypatch.setattr(sqlalchemy, "create_engine", fail) + 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}") + + assert msg.decode("utf-8").endswith("accepting connections\n") + assert status == 0 + status, msg = postgres.exec( + f"psql -hlocalhost -p{postgres.port} -U{postgres.username} -c 'select 2*3*5*7*11*13*17 as a;' ") + assert "510510" in msg.decode("utf-8") + assert "(1 row)" in msg.decode("utf-8") + assert status == 0 -def test_docker_run_postgres(): +def test_docker_run_postgres_with_sqlalchemy(): postgres_container = PostgresContainer("postgres:9.5") with postgres_container as postgres: engine = sqlalchemy.create_engine(postgres.get_connection_url()) diff --git a/poetry.lock b/poetry.lock index f97690266..d6a5453be 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1747,7 +1747,7 @@ files = [ name = "psycopg2-binary" version = "2.9.9" description = "psycopg2 - Python-PostgreSQL Database Adapter" -optional = true +optional = false python-versions = ">=3.7" files = [ {file = "psycopg2-binary-2.9.9.tar.gz", hash = "sha256:7f01846810177d829c7692f1f5ada8096762d9172af1b1a28d4ab5b77c923c1c"}, @@ -2331,7 +2331,6 @@ files = [ {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"}, {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"}, {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"}, - {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"}, {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"}, {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"}, {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"}, @@ -3151,7 +3150,7 @@ neo4j = ["neo4j"] nginx = [] opensearch = ["opensearch-py"] oracle = ["cx_Oracle", "sqlalchemy"] -postgres = ["psycopg2-binary", "sqlalchemy"] +postgres = [] rabbitmq = ["pika"] redis = ["redis"] selenium = ["selenium"] @@ -3159,4 +3158,4 @@ selenium = ["selenium"] [metadata] lock-version = "2.0" python-versions = ">=3.9,<4.0" -content-hash = "f4cb027301e265217ccb581b0ddd06fe6d91319fbcfbc3d20504a1fdbc45d7b1" +content-hash = "1aaa137a54688cdbad8ea0c162b2eb5185323074facb7cfad1f5bb88ca73c4e5" diff --git a/pyproject.toml b/pyproject.toml index 7afb4cd96..39c777614 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -79,7 +79,6 @@ pymysql = { version = "*", extras = ["rsa"], optional = true } neo4j = { version = "*", optional = true } opensearch-py = { version = "*", optional = true } cx_Oracle = { version = "*", optional = true } -psycopg2-binary = { version = "*", optional = true } pika = { version = "*", optional = true } redis = { version = "*", optional = true } selenium = { version = "*", optional = true } @@ -102,7 +101,7 @@ neo4j = ["neo4j"] nginx = [] opensearch = ["opensearch-py"] oracle = ["sqlalchemy", "cx_Oracle"] -postgres = ["sqlalchemy", "psycopg2-binary"] +postgres = [] rabbitmq = ["pika"] redis = ["redis"] selenium = ["selenium"] @@ -110,12 +109,15 @@ selenium = ["selenium"] [tool.poetry.group.dev.dependencies] mypy = "1.7.1" pre-commit = "^3.6" -pg8000 = "*" pytest = "7.4.3" pytest-cov = "4.1.0" sphinx = "^7.2.6" twine = "^4.0.2" anyio = "^4.3.0" +# for tests only +psycopg2-binary = "*" +pg8000 = "*" + [[tool.poetry.source]] name = "PyPI"