feat(postgres): Remove SqlAlchemy dependency from postgres container (#…

…445) Updates the pg testcontainer implementation to not use (and not install) SQLAlchemy nor psycopg2. Closes: #340 Closes: #336 Closes: #320 --------- Co-authored-by: Jason Turim <[email protected]>
testcontainers · Mar 10, 2024 · f30eb1d · f30eb1d
1 parent 2c4f171
commit f30eb1d
Show file tree

Hide file tree

Showing 6 changed files with 100 additions and 25 deletions.
diff --git 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
@@ -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
@@ -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,34 @@ 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.
+        """
+        driver_str = self.driver if driver is _UNSET else 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
@@ -1,9 +1,34 @@
-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
diff --git 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,20 +101,24 @@ neo4j = ["neo4j"]
 nginx = []
 opensearch = ["opensearch-py"]
 oracle = ["sqlalchemy", "cx_Oracle"]
-postgres = ["sqlalchemy", "psycopg2-binary"]
+postgres = []
 rabbitmq = ["pika"]
 redis = ["redis"]
 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 = "*"
+sqlalchemy = "*"
+
 
 [[tool.poetry.source]]
 name = "PyPI"