add support async postgres driver

core/testcontainers/core/generic.py

@@ -28,11 +28,13 @@ 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()
+        conn = engine.connect()
+        conn.close()
 
     def get_connection_url(self) -> str:
         raise NotImplementedError

postgres/setup.py

@@ -14,6 +14,8 @@
         "testcontainers-core",
         "sqlalchemy",
         "psycopg2-binary",
+        "asyncpg",
+        "pytest-asyncio",
     ],
     python_requires=">=3.7",
 )

postgres/testcontainers/postgres/__init__.py

@@ -13,6 +13,14 @@
 import os
 from typing import Optional
 from testcontainers.core.generic import DbContainer
+from testcontainers.core.waiting_utils import wait_container_is_ready
+
+ADDITIONAL_TRANSIENT_ERRORS = []
+try:
+    from sqlalchemy.exc import DBAPIError
+    ADDITIONAL_TRANSIENT_ERRORS.append(DBAPIError)
+except ImportError:
+    pass
 
 
 class PostgresContainer(DbContainer):
@@ -41,10 +49,14 @@ class PostgresContainer(DbContainer):
     POSTGRES_USER = os.environ.get("POSTGRES_USER", "test")
     POSTGRES_PASSWORD = os.environ.get("POSTGRES_PASSWORD", "test")
     POSTGRES_DB = os.environ.get("POSTGRES_DB", "test")
+    DEFAULT_DRIVER = "psycopg2"
 
     def __init__(self, image: str = "postgres:latest", port: int = 5432, user: Optional[str] = None,
                  password: Optional[str] = None, dbname: Optional[str] = None,
-                 driver: str = "psycopg2", **kwargs) -> None:
+                 driver: Optional[str] = None, **kwargs) -> None:
+        if driver is None:
+            driver = self.DEFAULT_DRIVER
+
         super(PostgresContainer, self).__init__(image=image, **kwargs)
         self.POSTGRES_USER = user or self.POSTGRES_USER
         self.POSTGRES_PASSWORD = password or self.POSTGRES_PASSWORD
@@ -54,14 +66,24 @@ def __init__(self, image: str = "postgres:latest", port: int = 5432, user: Optio
 
         self.with_exposed_ports(self.port_to_expose)
 
+    @wait_container_is_ready(*ADDITIONAL_TRANSIENT_ERRORS)
+    def _connect(self) -> None:
+        import sqlalchemy
+        engine = sqlalchemy.create_engine(self.get_connection_url(driver=self.DEFAULT_DRIVER))
+        conn = engine.connect()
+        conn.close()
+
     def _configure(self) -> None:
         self.with_env("POSTGRES_USER", self.POSTGRES_USER)
         self.with_env("POSTGRES_PASSWORD", self.POSTGRES_PASSWORD)
         self.with_env("POSTGRES_DB", self.POSTGRES_DB)
 
-    def get_connection_url(self, host=None) -> str:
+    def get_connection_url(self, host: Optional[str] = None, driver: Optional[str] = None) -> str:
+        if driver is None:
+            driver = self.driver
+
         return super()._create_connection_url(
-            dialect="postgresql+{}".format(self.driver), username=self.POSTGRES_USER,
+            dialect="postgresql+{}".format(driver), username=self.POSTGRES_USER,
             password=self.POSTGRES_PASSWORD, db_name=self.POSTGRES_DB, host=host,
             port=self.port_to_expose,
         )

postgres/tests/test_postgres.py

@@ -1,4 +1,6 @@
+import pytest
 import sqlalchemy
+from sqlalchemy.ext.asyncio import create_async_engine
 from testcontainers.postgres import PostgresContainer
 
 
@@ -18,3 +20,13 @@ def test_docker_run_postgres_with_driver_pg8000():
         engine = sqlalchemy.create_engine(postgres.get_connection_url())
         with engine.begin() as connection:
             connection.execute(sqlalchemy.text("select 1=1"))
+
+
+@pytest.mark.asyncio
+async def test_docker_run_async_postgres():
+    with PostgresContainer("postgres:9.5", driver="asyncpg") as postgres:
+        engine = create_async_engine(postgres.get_connection_url())
+        async with engine.begin() as connection:
+            result = await connection.execute(sqlalchemy.text("select version()"))
+            for row in result:
+                assert row[0].lower().startswith("postgresql 9.5")