Skip to content

Commit 396079a

Browse files
OverkillGuyJb DOYON
and
Jb DOYON
authored
fix(mysql): Add seed support in MySQL (#552)
Ref #541. New capability of "seeding" a db container using image's support for /docker-entrypoint-initdb.d/ folder. Using the "transferable" system, borrowed from Kafka. Updates DbContainer to have a new (NOOP-default) `_transfer_seed()` method, run after `_start()` and before `_connect()`, to allow the folder transfer. Currently implemented only in MySQL, but extensible to others that use the `/docker-entrypoint-initdb.d/` system. --------- Co-authored-by: Jb DOYON <[email protected]>
1 parent f761b98 commit 396079a

File tree

5 files changed

+58
-0
lines changed

5 files changed

+58
-0
lines changed

core/testcontainers/core/generic.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,8 +70,12 @@ def _create_connection_url(
7070
def start(self) -> "DbContainer":
7171
self._configure()
7272
super().start()
73+
self._transfer_seed()
7374
self._connect()
7475
return self
7576

7677
def _configure(self) -> None:
7778
raise NotImplementedError
79+
80+
def _transfer_seed(self) -> None:
81+
pass

modules/mysql/testcontainers/mysql/__init__.py

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,10 @@
1111
# License for the specific language governing permissions and limitations
1212
# under the License.
1313
import re
14+
import tarfile
15+
from io import BytesIO
1416
from os import environ
17+
from pathlib import Path
1518
from typing import Optional
1619

1720
from testcontainers.core.generic import DbContainer
@@ -40,6 +43,22 @@ class MySqlContainer(DbContainer):
4043
... with engine.begin() as connection:
4144
... result = connection.execute(sqlalchemy.text("select version()"))
4245
... version, = result.fetchone()
46+
47+
The optional :code:`seed` parameter enables arbitrary SQL files to be loaded.
48+
This is perfect for schema and sample data. This works by mounting the seed to
49+
`/docker-entrypoint-initdb./d`, which containerized MySQL are set up to load
50+
automatically.
51+
52+
.. doctest::
53+
>>> import sqlalchemy
54+
>>> from testcontainers.mysql import MySqlContainer
55+
>>> with MySqlContainer(seed="../../tests/seeds/") as mysql:
56+
... engine = sqlalchemy.create_engine(mysql.get_connection_url())
57+
... with engine.begin() as connection:
58+
... query = "select * from stuff" # Can now rely on schema/data
59+
... result = connection.execute(sqlalchemy.text(query))
60+
... first_stuff, = result.fetchone()
61+
4362
"""
4463

4564
def __init__(
@@ -50,6 +69,7 @@ def __init__(
5069
password: Optional[str] = None,
5170
dbname: Optional[str] = None,
5271
port: int = 3306,
72+
seed: Optional[str] = None,
5373
**kwargs,
5474
) -> None:
5575
raise_for_deprecated_parameter(kwargs, "MYSQL_USER", "username")
@@ -67,6 +87,7 @@ def __init__(
6787

6888
if self.username == "root":
6989
self.root_password = self.password
90+
self.seed = seed
7091

7192
def _configure(self) -> None:
7293
self.with_env("MYSQL_ROOT_PASSWORD", self.root_password)
@@ -86,3 +107,14 @@ def get_connection_url(self) -> str:
86107
return super()._create_connection_url(
87108
dialect="mysql+pymysql", username=self.username, password=self.password, dbname=self.dbname, port=self.port
88109
)
110+
111+
def _transfer_seed(self) -> None:
112+
if self.seed is None:
113+
return
114+
src_path = Path(self.seed)
115+
dest_path = "/docker-entrypoint-initdb.d/"
116+
with BytesIO() as archive, tarfile.TarFile(fileobj=archive, mode="w") as tar:
117+
for filename in src_path.iterdir():
118+
tar.add(filename.absolute(), arcname=filename.relative_to(src_path))
119+
archive.seek(0)
120+
self.get_wrapped_container().put_archive(dest_path, archive)
Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
-- Sample SQL schema, no data
2+
CREATE TABLE `stuff` (
3+
`id` mediumint NOT NULL AUTO_INCREMENT,
4+
`name` VARCHAR(63) NOT NULL,
5+
PRIMARY KEY (`id`)
6+
);
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
-- Sample data, to be loaded after the schema
2+
INSERT INTO stuff (name)
3+
VALUES ("foo"), ("bar"), ("qux"), ("frob");

modules/mysql/tests/test_mysql.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
from pathlib import Path
12
import re
23
from unittest import mock
34

@@ -29,6 +30,18 @@ def test_docker_run_legacy_mysql():
2930
assert row[0].startswith("5.7.44")
3031

3132

33+
@pytest.mark.skipif(is_arm(), reason="mysql container not available for ARM")
34+
def test_docker_run_mysql_8_seed():
35+
# Avoid pytest CWD path issues
36+
SEEDS_PATH = (Path(__file__).parent / "seeds").absolute()
37+
config = MySqlContainer("mysql:8", seed=SEEDS_PATH)
38+
with config as mysql:
39+
engine = sqlalchemy.create_engine(mysql.get_connection_url())
40+
with engine.begin() as connection:
41+
result = connection.execute(sqlalchemy.text("select * from stuff"))
42+
assert len(list(result)) == 4, "Should have gotten all the stuff"
43+
44+
3245
@pytest.mark.parametrize("version", ["11.3.2", "10.11.7"])
3346
def test_docker_run_mariadb(version: str):
3447
with MySqlContainer(f"mariadb:{version}") as mariadb:

0 commit comments

Comments
 (0)