Merge pull request #166 from alan-turing-institute/log-row-counts-iain

Log row counts (Iain's additions)

Author: mhauru

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sqlsynthgen"
-version = "0.3.3"
+version = "0.4.1"
 description = "Synthetic SQL data generator"
 authors = ["Iain <[email protected]>"]
 license = "MIT"

sqlsynthgen/create.py

@@ -1,5 +1,6 @@
 """Functions and classes to create and populate the target database."""
-from typing import Any, Generator, Mapping, Optional, Sequence, Tuple
+from collections import Counter
+from typing import Any, Generator, Mapping, Sequence, Tuple
 
 from sqlalchemy import Connection, insert
 from sqlalchemy.exc import IntegrityError
@@ -10,7 +11,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]
+RowCounts = Counter[str]
 
 
 def create_db_tables(metadata: MetaData) -> None:
@@ -68,15 +69,14 @@ def create_db_data(
         create_db_engine(dst_dsn, schema_name=settings.dst_schema)
     )
 
-    row_counts: RowCounts = {}
+    row_counts: Counter[str] = Counter()
     with dst_engine.connect() as dst_conn:
         for _ in range(num_passes):
-            row_counts = populate(
+            row_counts += populate(
                 dst_conn,
                 sorted_tables,
                 table_generator_dict,
                 story_generator_list,
-                row_counts,
             )
     return row_counts
 
@@ -86,13 +86,13 @@ def _populate_story(
     table_dict: Mapping[str, Table],
     table_generator_dict: Mapping[str, TableGenerator],
     dst_conn: Connection,
-    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
     # `for table_name, provided_values in story:`
     # but we have to loop more manually to be able to use the `send` function.
+    row_counts: Counter[str] = Counter()
     try:
         table_name, provided_values = next(story)
         while True:
@@ -129,11 +129,9 @@ def populate(
     tables: Sequence[Table],
     table_generator_dict: Mapping[str, TableGenerator],
     story_generator_list: Sequence[Mapping[str, Any]],
-    row_counts: Optional[RowCounts] = None,
 ) -> RowCounts:
     """Populate a database schema with synthetic data."""
-    if row_counts is None:
-        row_counts = {}
+    row_counts: Counter[str] = Counter()
     table_dict = {table.name: table for table in tables}
     # Generate stories
     # Each story generator returns a python generator (an unfortunate naming clash with
@@ -151,10 +149,10 @@ def populate(
     )
     for name, story in stories:
         # Run the inserts for each story within a transaction.
-        logger.debug("Generating data for story %s", name)
+        logger.debug('Generating data for story "%s".', name)
         with dst_conn.begin():
-            row_counts = _populate_story(
-                story, table_dict, table_generator_dict, dst_conn, row_counts
+            row_counts += _populate_story(
+                story, table_dict, table_generator_dict, dst_conn
             )
 
     # Generate individual rows, table by table.
@@ -166,7 +164,7 @@ def populate(
         table_generator = table_generator_dict[table.name]
         if table_generator.num_rows_per_pass == 0:
             continue
-        logger.debug("Generating data for table %s", table.name)
+        logger.debug('Generating data for table "%s".', table.name)
         # Run all the inserts for one table in a transaction
         with dst_conn.begin():
             for _ in range(table_generator.num_rows_per_pass):

sqlsynthgen/main.py

@@ -102,9 +102,16 @@ def create_data(
         story_generator_list,
         num_passes,
     )
-    logger.debug("Data created in %s passes.", num_passes)
+    logger.debug(
+        "Data created in %s %s.", num_passes, "pass" if num_passes == 1 else "passes"
+    )
     for table_name, row_count in row_counts.items():
-        logger.debug("%s: %s rows created", table_name, row_count)
+        logger.debug(
+            "%s: %s %s created.",
+            table_name,
+            row_count,
+            "row" if row_count == 1 else "rows",
+        )
 
 
 @app.command()

sqlsynthgen/remove.py

@@ -29,7 +29,7 @@ def remove_db_data(
         for table in reversed(metadata.sorted_tables):
             # We presume that all tables that aren't vocab should be truncated
             if table.name not in ssg_module.vocab_dict:
-                logger.debug("Truncating table %s", table.name)
+                logger.debug('Truncating table "%s".', table.name)
                 dst_conn.execute(delete(table))
                 dst_conn.commit()
 
@@ -50,7 +50,7 @@ def remove_db_vocab(
         for table in reversed(metadata.sorted_tables):
             # We presume that all tables that are vocab should be truncated
             if table.name in ssg_module.vocab_dict:
-                logger.debug("Truncating vocabulary table %s", table.name)
+                logger.debug('Truncating vocabulary table "%s".', table.name)
                 dst_conn.execute(delete(table))
                 dst_conn.commit()
 

tests/test_create.py

@@ -1,5 +1,6 @@
 """Tests for the create module."""
 import itertools as itt
+from collections import Counter
 from pathlib import Path
 from typing import Any, Generator, Tuple
 from unittest.mock import MagicMock, call, patch
@@ -82,7 +83,7 @@ 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 = (
+                row_counts = Counter(
                     {table_name: num_initial_rows} if num_initial_rows > 0 else {}
                 )
 
@@ -97,20 +98,23 @@ def mock_story_gen(_: Any) -> Generator[Tuple[str, dict], None, None]:
                     if num_stories_per_pass > 0
                     else []
                 )
-                row_counts = populate(
+                row_counts += populate(
                     mock_dst_conn,
                     [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(
+                    Counter(
+                        {table_name: expected_row_count}
+                        if expected_row_count > 0
+                        else {}
+                    ),
                     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),
@@ -148,7 +152,7 @@ def test_populate_diff_length(self, mock_insert: MagicMock) -> None:
             "three": mock_gen_three,
         }
 
-        row_counts = 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
@@ -221,4 +225,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

@@ -336,40 +336,40 @@ def test_workflow_maximal_args(self) -> None:
         self.assertSuccess(completed_process)
         self.assertEqual(
             "Creating data.\n"
-            "Generating data for story story_generators.short_story\n"
-            "Generating data for story story_generators.short_story\n"
-            "Generating data for story story_generators.short_story\n"
-            "Generating data for story story_generators.full_row_story\n"
-            "Generating data for story story_generators.long_story\n"
-            "Generating data for story story_generators.long_story\n"
-            "Generating data for table data_type_test\n"
-            "Generating data for table no_pk_test\n"
-            "Generating data for table person\n"
-            "Generating data for table unique_constraint_test\n"
-            "Generating data for table unique_constraint_test2\n"
-            "Generating data for table test_entity\n"
-            "Generating data for table hospital_visit\n"
-            "Generating data for story story_generators.short_story\n"
-            "Generating data for story story_generators.short_story\n"
-            "Generating data for story story_generators.short_story\n"
-            "Generating data for story story_generators.full_row_story\n"
-            "Generating data for story story_generators.long_story\n"
-            "Generating data for story story_generators.long_story\n"
-            "Generating data for table data_type_test\n"
-            "Generating data for table no_pk_test\n"
-            "Generating data for table person\n"
-            "Generating data for table unique_constraint_test\n"
-            "Generating data for table unique_constraint_test2\n"
-            "Generating data for table test_entity\n"
-            "Generating data for table hospital_visit\n"
+            'Generating data for story "story_generators.short_story".\n'
+            'Generating data for story "story_generators.short_story".\n'
+            'Generating data for story "story_generators.short_story".\n'
+            'Generating data for story "story_generators.full_row_story".\n'
+            'Generating data for story "story_generators.long_story".\n'
+            'Generating data for story "story_generators.long_story".\n'
+            'Generating data for table "data_type_test".\n'
+            'Generating data for table "no_pk_test".\n'
+            'Generating data for table "person".\n'
+            'Generating data for table "unique_constraint_test".\n'
+            'Generating data for table "unique_constraint_test2".\n'
+            'Generating data for table "test_entity".\n'
+            'Generating data for table "hospital_visit".\n'
+            'Generating data for story "story_generators.short_story".\n'
+            'Generating data for story "story_generators.short_story".\n'
+            'Generating data for story "story_generators.short_story".\n'
+            'Generating data for story "story_generators.full_row_story".\n'
+            'Generating data for story "story_generators.long_story".\n'
+            'Generating data for story "story_generators.long_story".\n'
+            'Generating data for table "data_type_test".\n'
+            'Generating data for table "no_pk_test".\n'
+            'Generating data for table "person".\n'
+            'Generating data for table "unique_constraint_test".\n'
+            '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"
-            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",
+            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"),
         )
 
@@ -390,13 +390,13 @@ def test_workflow_maximal_args(self) -> None:
         self.assertSuccess(completed_process)
         self.assertEqual(
             "Truncating non-vocabulary tables.\n"
-            "Truncating table hospital_visit\n"
-            "Truncating table test_entity\n"
-            "Truncating table unique_constraint_test2\n"
-            "Truncating table unique_constraint_test\n"
-            "Truncating table person\n"
-            "Truncating table no_pk_test\n"
-            "Truncating table data_type_test\n"
+            'Truncating table "hospital_visit".\n'
+            'Truncating table "test_entity".\n'
+            'Truncating table "unique_constraint_test2".\n'
+            'Truncating table "unique_constraint_test".\n'
+            'Truncating table "person".\n'
+            'Truncating table "no_pk_test".\n'
+            'Truncating table "data_type_test".\n'
             "Non-vocabulary tables truncated.\n",
             completed_process.stdout.decode("utf-8"),
         )
@@ -418,11 +418,11 @@ def test_workflow_maximal_args(self) -> None:
         self.assertSuccess(completed_process)
         self.assertEqual(
             "Truncating vocabulary tables.\n"
-            "Truncating vocabulary table concept\n"
-            "Truncating vocabulary table concept_type\n"
-            "Truncating vocabulary table ref_to_unignorable_table\n"
-            "Truncating vocabulary table mitigation_type\n"
-            "Truncating vocabulary table empty_vocabulary\n"
+            'Truncating vocabulary table "concept".\n'
+            'Truncating vocabulary table "concept_type".\n'
+            'Truncating vocabulary table "ref_to_unignorable_table".\n'
+            'Truncating vocabulary table "mitigation_type".\n'
+            'Truncating vocabulary table "empty_vocabulary".\n'
             "Vocabulary tables truncated.\n",
             completed_process.stdout.decode("utf-8"),
         )

tests/test_main.py

@@ -152,11 +152,15 @@ def test_create_tables(
         mock_create.assert_called_once_with(mock_import.return_value.Base.metadata)
         self.assertSuccess(result)
 
+    @patch("sqlsynthgen.main.logger")
     @patch("sqlsynthgen.main.import_file")
     @patch("sqlsynthgen.main.create_db_data")
-    def test_create_data(self, mock_create: MagicMock, mock_import: MagicMock) -> None:
+    def test_create_data(
+        self, mock_create: MagicMock, mock_import: MagicMock, mock_logger: MagicMock
+    ) -> None:
         """Test the create-data sub-command."""
 
+        mock_create.return_value = {"a": 1}
         result = runner.invoke(
             app,
             [
@@ -180,6 +184,14 @@ def test_create_data(self, mock_create: MagicMock, mock_import: MagicMock) -> No
         )
         self.assertSuccess(result)
 
+        mock_logger.debug.assert_has_calls(
+            [
+                call("Creating data."),
+                call("Data created in %s %s.", 1, "pass"),
+                call("%s: %s %s created.", "a", 1, "row"),
+            ]
+        )
+
     @patch("sqlsynthgen.main.Path")
     @patch("sqlsynthgen.main.make_tables_file")
     @patch("sqlsynthgen.main.get_settings")