Skip to content

Commit

Permalink
feat(postgres): Remove SqlAlchemy dependency from postgres container
Browse files Browse the repository at this point in the history
Create a new generic database subclass - DependencyFreeDbContainer
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 <jason@opscanvas.com>
Co-authored-by: Jan Katins <jasc@gmx.net>
  • Loading branch information
Jason Turim and jankatins committed Mar 6, 2024
1 parent ed3b9fa commit 0132a27
Show file tree
Hide file tree
Showing 5 changed files with 95 additions and 28 deletions.
20 changes: 18 additions & 2 deletions INDEX.rst
Original file line number Diff line number Diff line change
Expand Up @@ -46,15 +46,31 @@ 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)
... 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 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.


Installation
------------

Expand Down
40 changes: 40 additions & 0 deletions core/testcontainers/core/generic.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,46 @@
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.
Expand Down
38 changes: 28 additions & 10 deletions modules/postgres/testcontainers/postgres/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,13 +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.
Expand All @@ -31,7 +34,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:
Expand All @@ -48,16 +51,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")
super(PostgresContainer, self).__init__(image=image, **kwargs)
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)

Expand All @@ -68,10 +71,25 @@ def _configure(self) -> None:

def get_connection_url(self, host=None) -> str:
return super()._create_connection_url(
dialect=f"postgresql+{self.driver}",
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")
23 changes: 8 additions & 15 deletions modules/postgres/tests/test_postgres.py
Original file line number Diff line number Diff line change
@@ -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
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand Down

0 comments on commit 0132a27

Please sign in to comment.