Add Db2 support

modules/db2/README.rst

@@ -0,0 +1,2 @@
+.. autoclass:: testcontainers.db2.Db2Container
+.. title:: testcontainers.db2.Db2Container

modules/db2/testcontainers/db2/__init__.py

@@ -0,0 +1,61 @@
+from os import environ
+from typing import Optional
+
+from testcontainers.core.generic import DbContainer
+from testcontainers.core.waiting_utils import wait_container_is_ready, wait_for_logs
+
+
+class Db2Container(DbContainer):
+    """
+    IBM Db2 database container.
+
+    Example:
+
+        .. doctest::
+
+            >>> import sqlalchemy
+            >>> from testcontainers.db2 import Db2Container
+
+            >>> with Db2Container("icr.io/db2_community/db2:latest") as db2:
+            ...    engine = sqlalchemy.create_engine(db2.get_connection_url())
+            ...    with engine.begin() as connection:
+            ...        result = connection.execute(sqlalchemy.text("select service_level from sysibmadm.env_inst_info"))
+    """
+
+    def __init__(
+        self,
+        image: str = "icr.io/db2_community/db2:latest",
+        username: str = "db2inst1",
+        password: Optional[str] = None,
+        port: int = 50000,
+        dbname: str = "testdb",
+        dialect: str = "db2+ibm_db",
+        **kwargs,
+    ) -> None:
+        super().__init__(image, **kwargs)
+
+        self.port = port
+        self.with_exposed_ports(self.port)
+
+        self.password = password or environ.get("DB2_PASSWORD", "password")
+        self.username = username
+        self.dbname = dbname
+        self.dialect = dialect
+
+    def _configure(self) -> None:
+        self.with_env("LICENSE", "accept")
+        self.with_env("DB2INSTANCE", self.username)
+        self.with_env("DB2INST1_PASSWORD", self.password)
+        self.with_env("DBNAME", self.dbname)
+        self.with_env("ARCHIVE_LOGS", "false")
+        self.with_env("AUTOCONFIG", "false")
+        self.with_kwargs(privileged=True)
+
+    @wait_container_is_ready()
+    def _connect(self) -> None:
+        wait_for_logs(self, predicate="Setup has completed")
+
+    def get_connection_url(self) -> str:
+        return super()._create_connection_url(
+            dialect=self.dialect, username=self.username, password=self.password, dbname=self.dbname, port=self.port
+        )

modules/db2/tests/test_db2.py

@@ -0,0 +1,43 @@
+from unittest import mock
+
+import pytest
+import sqlalchemy
+
+from testcontainers.core.utils import is_arm
+from testcontainers.db2 import Db2Container
+
+
+@pytest.mark.skipif(is_arm(), reason="db2 container not available for ARM")
+@pytest.mark.parametrize("version", ["11.5.9.0", "11.5.8.0"])
+def test_docker_run_db2(version: str):
+    with Db2Container(f"icr.io/db2_community/db2:{version}", password="password") as db2:
+        engine = sqlalchemy.create_engine(db2.get_connection_url())
+        with engine.begin() as connection:
+            result = connection.execute(sqlalchemy.text("select service_level from sysibmadm.env_inst_info"))
+            for row in result:
+                assert row[0] == f"DB2 v{version}"
+
+
+# This is a feature in the generic DbContainer class
+# but it can't be tested on its own
+# so is tested in various database modules:
+# - mysql / mariadb
+# - postgresql
+# - sqlserver
+# - mongodb
+# - db2
+def test_quoted_password():
+    user = "db2inst1"
+    dbname = "testdb"
+    password = "p@$%25+0&%rd :/!=?"
+    quoted_password = "p%40%24%2525+0%26%25rd %3A%2F%21%3D%3F"
+    kwargs = {
+        "username": user,
+        "password": password,
+        "dbname": dbname,
+    }
+    with Db2Container("icr.io/db2_community/db2:11.5.9.0", **kwargs) as container:
+        port = container.get_exposed_port(50000)
+        host = container.get_container_host_ip()
+        expected_url = f"db2+ibm_db://{user}:{quoted_password}@{host}:{port}/{dbname}"
+        assert expected_url == container.get_connection_url()

modules/mongodb/tests/test_mongodb.py

@@ -35,6 +35,7 @@ def test_docker_run_mongodb(version: str):
 # - postgresql
 # - sqlserver
 # - mongodb
+# - db2
 def test_quoted_password():
     user = "root"
     password = "p@$%25+0&%rd :/!=?"

modules/mysql/tests/test_mysql.py

@@ -69,6 +69,7 @@ def test_docker_env_variables():
 # - postgresql
 # - sqlserver
 # - mongodb
+# - db2
 def test_quoted_password():
     user = "root"
     password = "p@$%25+0&%rd :/!=?"

modules/postgres/tests/test_postgres.py

@@ -53,6 +53,7 @@ def test_docker_run_postgres_with_driver_pg8000():
 # - postgresql
 # - sqlserver
 # - mongodb
+# - db2
 def test_quoted_password():
     user = "root"
     password = "p@$%25+0&%rd :/!=?"

pyproject.toml

@@ -36,6 +36,7 @@ packages = [
     { include = "testcontainers", from = "modules/clickhouse" },
     { include = "testcontainers", from = "modules/cockroachdb" },
     { include = "testcontainers", from = "modules/cosmosdb" },
+    { include = "testcontainers", from = "modules/db2" },
     { include = "testcontainers", from = "modules/elasticsearch" },
     { include = "testcontainers", from = "modules/generic" },
     { include = "testcontainers", from = "modules/test_module_import"},
@@ -114,6 +115,7 @@ httpx = { version = "*", optional = true }
 azure-cosmos = { version = "*", optional = true }
 cryptography = { version = "*", optional = true }
 trino = { version = "*", optional = true }
+ibm_db_sa = { version = "*", optional = true }
 
 [tool.poetry.extras]
 arangodb = ["python-arango"]
@@ -123,6 +125,7 @@ cassandra = []
 clickhouse = ["clickhouse-driver"]
 cosmosdb = ["azure-cosmos"]
 cockroachdb = []
+db2 = ["sqlalchemy", "ibm_db_sa"]
 elasticsearch = []
 generic = ["httpx"]
 test_module_import = ["httpx"]