support multiple files for demo data

Author: 3nids

pum/config_model.py

@@ -132,11 +132,22 @@ class DemoDataModel(PumCustomBaseModel):
     DemoDataModel represents a configuration for demo data.
 
     Attributes:
-        files: Optional list of named demo data files.
+        name: Name of the demo data.
+        file: Optional path to a single demo data file.
+        files: Optional list of paths to multiple demo data files.
     """
 
     name: str = Field(..., description="Name of the demo data.")
-    file: str = Field(..., description="Path to the demo data file.")
+
+    file: Optional[str] = None
+    files: Optional[List[str]] = None
+
+    @model_validator(mode="after")
+    def validate_args(self):
+        file, files = self.file, self.files
+        if (file and files) or (not file and not files):
+            raise PumConfigError("Exactly one of 'file' or 'files' must be set in a demo data set.")
+        return self
 
 
 class DependencyModel(PumCustomBaseModel):

pum/pum_config.py

@@ -256,9 +256,12 @@ def post_hook_handlers(self) -> list[HookHandler]:
             else []
         )
 
-    def demo_data(self) -> dict[str, str]:
+    def demo_data(self) -> dict[str, list[str]]:
         """Return a dictionary of demo data files defined in the configuration."""
-        return {dm.name: dm.file for dm in self.config.demo_data}
+        demo_data_files = {}
+        for dm in self.config.demo_data:
+            demo_data_files[dm.name] = dm.files or [dm.file]
+        return demo_data_files
 
     def __del__(self):
         # Cleanup temporary directories and sys.path modifications

test/data/demo_data_multi/.pum.yaml

@@ -0,0 +1,7 @@
+
+
+demo_data:
+  - name: some cool demo dataset
+    files:
+      - demo_data/demo_data_1.sql
+      - demo_data/demo_data_2.sql

test/data/demo_data_multi/changelogs/1.2.3/single_changelog.sql

@@ -0,0 +1,9 @@
+CREATE SCHEMA IF NOT EXISTS pum_test_data;
+
+CREATE TABLE pum_test_data.some_table (
+    id INT PRIMARY KEY,
+    name VARCHAR(100) NOT NULL,
+    created_date DATE DEFAULT CURRENT_DATE,
+    is_active BOOLEAN DEFAULT TRUE,
+    amount NUMERIC(10,2)
+);

test/data/demo_data_multi/demo_data/demo_data_1.sql

@@ -0,0 +1,6 @@
+
+INSERT INTO pum_test_data.some_table (id, name, created_date, is_active, amount)
+VALUES
+    (1, 'Test Item 1', '2023-01-01', TRUE, 100.00),
+    (2, 'Test Item 2', '2023-01-02', FALSE, 200.00)
+;

test/data/demo_data_multi/demo_data/demo_data_2.sql

@@ -0,0 +1,6 @@
+
+INSERT INTO pum_test_data.some_table (id, name, created_date, is_active, amount)
+VALUES
+    (3, 'Test Item 3', '2023-01-03', TRUE, 300.00),
+    (4, 'Test Item 4', '2023-01-04', TRUE, 400.00)
+;

test/test_upgrader.py

@@ -407,6 +407,24 @@ def test_demo_data(self) -> None:
             count = cursor.fetchone()[0]
             self.assertEqual(count, 4)
 
+    def test_demo_data_multi(self) -> None:
+        """Test the installation of demo data."""
+        test_dir = Path("test") / "data" / "demo_data_multi"
+        cfg = PumConfig.from_yaml(test_dir / ".pum.yaml")
+        sm = SchemaMigrations(cfg)
+        with psycopg.connect(f"service={self.pg_service}") as conn:
+            self.assertFalse(sm.exists(conn))
+            upgrader = Upgrader(config=cfg)
+            upgrader.install(connection=conn)
+            with self.assertRaises(PumException):
+                upgrader.install_demo_data(connection=conn, name="nope, nothing here fella")
+            upgrader.install_demo_data(connection=conn, name="some cool demo dataset")
+            self.assertTrue(sm.exists(conn))
+            cursor = conn.cursor()
+            cursor.execute("SELECT COUNT(*) FROM pum_test_data.some_table;")
+            count = cursor.fetchone()[0]
+            self.assertEqual(count, 4)
+
     def test_dependencies(self) -> None:
         """Test the installation of dependencies."""
         test_dir = Path("test") / "data" / "dependencies"