From 9180ea50e2422f98ba20e6db1afaee87875bc702 Mon Sep 17 00:00:00 2001 From: Jason Turim Date: Sat, 22 Apr 2023 18:50:48 -0400 Subject: [PATCH] feat(postgres): Remove SqlAlchemy dependency from postgres container Remove test that was testing sqlalchemy support for driver types Add tests for supported versions of Postgres Modify the `get_connection_url` convenience method to support a driverless url Co-authored-by: Jason Turim Co-authored-by: Jan Katins --- INDEX.rst | 26 ++++++++-- README.md | 4 +- .../testcontainers/postgres/__init__.py | 48 +++++++++++++++---- modules/postgres/tests/test_postgres.py | 23 ++++----- pyproject.toml | 2 +- 5 files changed, 73 insertions(+), 30 deletions(-) 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..87a1a6601 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: str | None = "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..7ae3b4ede 100644 --- a/modules/postgres/tests/test_postgres.py +++ b/modules/postgres/tests/test_postgres.py @@ -1,21 +1,14 @@ -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 = ["12", "13", "14", "15", "16", "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 diff --git a/pyproject.toml b/pyproject.toml index 7afb4cd96..b537a52df 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -102,7 +102,7 @@ neo4j = ["neo4j"] nginx = [] opensearch = ["opensearch-py"] oracle = ["sqlalchemy", "cx_Oracle"] -postgres = ["sqlalchemy", "psycopg2-binary"] +postgres = [] rabbitmq = ["pika"] redis = ["redis"] selenium = ["selenium"]