Log row counts for create-data

Author: mhauru

sqlsynthgen/create.py

@@ -1,5 +1,5 @@
 """Functions and classes to create and populate the target database."""
-from typing import Any, Generator, Mapping, Sequence, Tuple
+from typing import Any, Generator, Mapping, Optional, Sequence, Tuple
 
 from sqlalchemy import Connection, insert
 from sqlalchemy.exc import IntegrityError
@@ -10,6 +10,7 @@
 from sqlsynthgen.utils import create_db_engine, get_sync_engine, logger
 
 Story = Generator[Tuple[str, dict[str, Any]], dict[str, Any], None]
+RowCounts = dict[str, int]
 
 
 def create_db_tables(metadata: MetaData) -> None:
@@ -57,7 +58,7 @@ def create_db_data(
     table_generator_dict: Mapping[str, TableGenerator],
     story_generator_list: Sequence[Mapping[str, Any]],
     num_passes: int,
-) -> None:
+) -> RowCounts:
     """Connect to a database and populate it with data."""
     settings = get_settings()
     dst_dsn: str = settings.dst_dsn or ""
@@ -67,22 +68,26 @@ def create_db_data(
         create_db_engine(dst_dsn, schema_name=settings.dst_schema)
     )
 
+    row_counts: RowCounts = {}
     with dst_engine.connect() as dst_conn:
         for _ in range(num_passes):
-            populate(
+            row_counts = populate(
                 dst_conn,
                 sorted_tables,
                 table_generator_dict,
                 story_generator_list,
+                row_counts,
             )
+    return row_counts
 
 
 def _populate_story(
     story: Story,
     table_dict: Mapping[str, Table],
     table_generator_dict: Mapping[str, TableGenerator],
     dst_conn: Connection,
-) -> None:
+    row_counts: RowCounts,
+) -> RowCounts:
     """Write to the database all the rows created by the given story."""
     # Loop over the rows generated by the story, insert them into their
     # respective tables. Ideally this would say
@@ -111,19 +116,24 @@ def _populate_story(
             else:
                 return_values = {}
             final_values = {**insert_values, **return_values}
+            row_counts[table_name] = row_counts.get(table_name, 0) + 1
             table_name, provided_values = story.send(final_values)
     except StopIteration:
         # The story has finished, it has no more rows to generate
         pass
+    return row_counts
 
 
 def populate(
     dst_conn: Connection,
     tables: Sequence[Table],
     table_generator_dict: Mapping[str, TableGenerator],
     story_generator_list: Sequence[Mapping[str, Any]],
-) -> None:
+    row_counts: Optional[RowCounts] = None,
+) -> RowCounts:
     """Populate a database schema with synthetic data."""
+    if row_counts is None:
+        row_counts = {}
     table_dict = {table.name: table for table in tables}
     # Generate stories
     # Each story generator returns a python generator (an unfortunate naming clash with
@@ -143,7 +153,9 @@ def populate(
         # Run the inserts for each story within a transaction.
         logger.debug("Generating data for story %s", name)
         with dst_conn.begin():
-            _populate_story(story, table_dict, table_generator_dict, dst_conn)
+            row_counts = _populate_story(
+                story, table_dict, table_generator_dict, dst_conn, row_counts
+            )
 
     # Generate individual rows, table by table.
     for table in tables:
@@ -160,3 +172,5 @@ def populate(
             for _ in range(table_generator.num_rows_per_pass):
                 stmt = insert(table).values(table_generator(dst_conn))
                 dst_conn.execute(stmt)
+                row_counts[table.name] = row_counts.get(table.name, 0) + 1
+    return row_counts

sqlsynthgen/main.py

@@ -96,13 +96,15 @@ def create_data(
     orm_metadata = get_orm_metadata(orm_module, tables_config)
     table_generator_dict = ssg_module.table_generator_dict
     story_generator_list = ssg_module.story_generator_list
-    create_db_data(
+    row_counts = create_db_data(
         orm_metadata.sorted_tables,
         table_generator_dict,
         story_generator_list,
         num_passes,
     )
     logger.debug("Data created in %s passes.", num_passes)
+    for table_name, row_count in row_counts.items():
+        logger.debug("%s: %s rows created", table_name, row_count)
 
 
 @app.command()

tests/test_create.py

@@ -34,11 +34,13 @@ def test_create_db_data(
     ) -> None:
         """Test the generate function."""
         mock_get_settings.return_value = get_test_settings()
+        mock_populate.return_value = {}
 
         num_passes = 23
-        create_db_data([], {}, [], num_passes)
+        row_counts = create_db_data([], {}, [], num_passes)
 
         self.assertEqual(len(mock_populate.call_args_list), num_passes)
+        self.assertEqual(row_counts, {})
         mock_create_engine.assert_called()
 
     @patch("sqlsynthgen.create.get_settings")
@@ -62,13 +64,15 @@ def test_populate(self) -> None:
 
         def story() -> Generator[Tuple[str, dict], None, None]:
             """Mock story."""
-            yield "table_name", {}
+            yield table_name, {}
 
         def mock_story_gen(_: Any) -> Generator[Tuple[str, dict], None, None]:
             """A function that returns mock stories."""
             return story()
 
-        for num_stories_per_pass, num_rows_per_pass in itt.product([0, 2], [0, 3]):
+        for num_stories_per_pass, num_rows_per_pass, num_initial_rows in itt.product(
+            [0, 2], [0, 3], [0, 17]
+        ):
             with patch("sqlsynthgen.create.insert") as mock_insert:
                 mock_values = mock_insert.return_value.values
                 mock_dst_conn = MagicMock(spec=Connection)
@@ -78,9 +82,10 @@ def mock_story_gen(_: Any) -> Generator[Tuple[str, dict], None, None]:
                 mock_gen = MagicMock(spec=TableGenerator)
                 mock_gen.num_rows_per_pass = num_rows_per_pass
                 mock_gen.return_value = {}
+                row_counts = (
+                    {table_name: num_initial_rows} if num_initial_rows > 0 else {}
+                )
 
-                tables: list[Table] = [mock_table]
-                row_generators: dict[str, TableGenerator] = {table_name: mock_gen}
                 story_generators: list[dict[str, Any]] = (
                     [
                         {
@@ -92,13 +97,21 @@ def mock_story_gen(_: Any) -> Generator[Tuple[str, dict], None, None]:
                     if num_stories_per_pass > 0
                     else []
                 )
-                populate(
+                row_counts = populate(
                     mock_dst_conn,
-                    tables,
-                    row_generators,
+                    [mock_table],
+                    {table_name: mock_gen},
                     story_generators,
+                    row_counts,
                 )
 
+                expected_row_count = (
+                    num_stories_per_pass + num_rows_per_pass + num_initial_rows
+                )
+                self.assertEqual(
+                    row_counts,
+                    {table_name: expected_row_count} if expected_row_count > 0 else {},
+                )
                 self.assertListEqual(
                     [call(mock_dst_conn)] * (num_stories_per_pass + num_rows_per_pass),
                     mock_gen.call_args_list,
@@ -135,7 +148,8 @@ def test_populate_diff_length(self, mock_insert: MagicMock) -> None:
             "three": mock_gen_three,
         }
 
-        populate(mock_dst_conn, tables, row_generators, [])
+        row_counts = populate(mock_dst_conn, tables, row_generators, [], {})
+        self.assertEqual(row_counts, {"two": 1, "three": 1})
         self.assertListEqual(
             [call(mock_table_two), call(mock_table_three)], mock_insert.call_args_list
         )
@@ -207,4 +221,4 @@ def my_story() -> Story:
 
         with engine.connect() as conn:
             with conn.begin():
-                _populate_story(my_story(), dict(self.metadata.tables), {}, conn)
+                _populate_story(my_story(), dict(self.metadata.tables), {}, conn, {})

tests/test_functional.py

@@ -362,7 +362,14 @@ def test_workflow_maximal_args(self) -> None:
             "Generating data for table unique_constraint_test2\n"
             "Generating data for table test_entity\n"
             "Generating data for table hospital_visit\n"
-            "Data created in 2 passes.\n",
+            "Data created in 2 passes.\n"
+            f"person: {2*(3+1+2+2)} rows created\n"
+            f"hospital_visit: {2*(2*2+3)} rows created\n"
+            "data_type_test: 2 rows created\n"
+            "no_pk_test: 2 rows created\n"
+            "unique_constraint_test: 2 rows created\n"
+            "unique_constraint_test2: 2 rows created\n"
+            "test_entity: 2 rows created\n",
             completed_process.stdout.decode("utf-8"),
         )