feat: Run seed scripts in DbContainer

core/testcontainers/core/generic.py

@@ -10,6 +10,7 @@
 #    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 #    License for the specific language governing permissions and limitations
 #    under the License.
+from pathlib import Path
 from typing import Optional
 
 from testcontainers.core.container import DockerContainer
@@ -69,7 +70,20 @@ def start(self) -> "DbContainer":
         self._configure()
         super().start()
         self._connect()
+        self._seed()
         return self
 
     def _configure(self) -> None:
         raise NotImplementedError
+
+    def _seed(self) -> None:
+        raise NotImplementedError
+
+    def _setup_seed(self, seed):
+        self.seed_scripts = None
+        self.seed_mount = "/seeds"  # TODO: Should it be configurable?
+        if seed is not None:
+            seed_localpath, seed_scripts = seed
+            seed_abspath = Path(seed_localpath).absolute()
+            self.with_volume_mapping(seed_abspath, self.seed_mount)
+            self.seed_scripts = seed_scripts

modules/mysql/testcontainers/mysql/__init__.py

@@ -40,6 +40,25 @@ class MySqlContainer(DbContainer):
             ...     with engine.begin() as connection:
             ...         result = connection.execute(sqlalchemy.text("select version()"))
             ...         version, = result.fetchone()
+
+        The optional :code:`seed` param enables arbitrary SQL files to be loaded. This
+        is perfect for schema and sample data. The format is a tuple, made up of the
+        path to local data, and list of script files to load. Each script will be loaded
+        by exec in the container, using the MySQL root user, before yielding the
+        container. Any errors loading the scripts, will cause a ValueError.
+
+        .. doctest::
+
+            >>> import sqlalchemy
+            >>> from testcontainers.mysql import MySqlContainer
+            >>> seed_data = ("../../tests/", ["schema.sql", "data.sql"])
+            >>> with MySqlContainer(seed=seed_data) as mysql:
+            ...     engine = sqlalchemy.create_engine(mysql.get_connection_url())
+            ...     with engine.begin() as connection:
+            ...         query = "select * from stuff"  # Can now rely on schema/data
+            ...         result = connection.execute(sqlalchemy.text(query))
+            ...         first_stuff, = result.fetchone()
+
     """
 
     def __init__(
@@ -50,6 +69,7 @@ def __init__(
         password: Optional[str] = None,
         dbname: Optional[str] = None,
         port: int = 3306,
+        seed: Optional[tuple[str, list[str]]] = None,
         **kwargs,
     ) -> None:
         raise_for_deprecated_parameter(kwargs, "MYSQL_USER", "username")
@@ -65,6 +85,7 @@ def __init__(
         self.password = password or environ.get("MYSQL_PASSWORD", "test")
         self.dbname = dbname or environ.get("MYSQL_DATABASE", "test")
 
+        self._setup_seed(seed)
         if self.username == "root":
             self.root_password = self.password
 
@@ -86,3 +107,15 @@ def get_connection_url(self) -> str:
         return super()._create_connection_url(
             dialect="mysql+pymysql", username=self.username, password=self.password, dbname=self.dbname, port=self.port
         )
+
+    def _seed(self) -> None:
+        """Apply the seed scripts given"""
+        if not self.seed_scripts:  # Defined in DbContainer._setup_seed(s)
+            return
+        container = self.get_wrapped_container()
+        for script in self.seed_scripts:
+            mysql_query = f"source {self.seed_mount}/{script}"
+            schema_cmd = ["mysql", f"-p{self.root_password}", self.dbname, "-e", mysql_query]
+            exit_code, _out = container.exec_run(schema_cmd)
+            if exit_code != 0:
+                raise ValueError(f"Error seeding the database with {script=}")

modules/mysql/tests/schema.sql

@@ -0,0 +1,6 @@
+-- Sample SQL schema, no data
+CREATE TABLE `stuff` (
+ `id` mediumint NOT NULL AUTO_INCREMENT,
+ `name` VARCHAR(63) NOT NULL,
+ PRIMARY KEY (`id`)
+);

modules/mysql/tests/seeds.sql

@@ -0,0 +1,3 @@
+-- Sample data, to be loaded after the schema
+INSERT INTO stuff (name)
+VALUES ("foo"), ("bar"), ("qux"), ("frob");

modules/mysql/tests/test_mysql.py

@@ -29,6 +29,26 @@ def test_docker_run_legacy_mysql():
                 assert row[0].startswith("5.7.44")
 
 
+@pytest.mark.skipif(is_arm(), reason="mysql container not available for ARM")
+def test_docker_run_mysql_8_seed():
+    seeds = ("modules/mysql/tests/", ["schema.sql", "seeds.sql"])
+    config = MySqlContainer("mysql:8", seed=seeds)
+    with config as mysql:
+        engine = sqlalchemy.create_engine(mysql.get_connection_url())
+        with engine.begin() as connection:
+            result = connection.execute(sqlalchemy.text("select * from stuff"))
+            assert len(list(result)) == 4, "Should have gotten all the stuff"
+
+
+@pytest.mark.skipif(is_arm(), reason="mysql container not available for ARM")
+def test_docker_run_mysql_8_seed_missing():
+    seeds = ("modules/mysql/tests/", ["schema.sql", "nosuchfile.sql"])
+    config = MySqlContainer("mysql:8", seed=seeds)
+    with pytest.raises(ValueError, match="Error seeding the database with script"):
+        with config as mysql:
+            pass
+
+
 @pytest.mark.parametrize("version", ["11.3.2", "10.11.7"])
 def test_docker_run_mariadb(version: str):
     with MySqlContainer(f"mariadb:{version}") as mariadb: