Merge remote-tracking branch 'origin/main' into log-row-counts

Author: mhauru

sqlsynthgen/make.py

@@ -238,16 +238,24 @@ def _get_provider_for_column(column: Column) -> Tuple[list[str], str, list[str]]
 
     generator_function = mapping.get((column_type, column_size is not None), None)
 
+    # Try if we know how to generate for a superclass of this type.
     if not generator_function:
         for key, value in mapping.items():
             if issubclass(column_type, key[0]) and key[1] == (column_size is not None):
                 generator_function = value
                 break
 
+    # If we still don't have a generator, use null and warn.
     if not generator_function:
-        raise ValueError(f"Unsupported SQLAlchemy type: {column_type}")
-
-    if column_size:
+        generator_function = "generic.null_provider.null"
+        logger.warning(
+            "Unsupported SQLAlchemy type %s for column %s. "
+            "Setting this column to NULL always, "
+            "you may want to configure a row generator for it instead.",
+            column_type,
+            column.name,
+        )
+    elif column_size:
         generator_arguments.append(str(column_size))
 
     return variable_names, generator_function, generator_arguments

sqlsynthgen/providers.py

@@ -17,8 +17,9 @@ class Meta:
 
         name = "column_value_provider"
 
+    @staticmethod
     def column_value(
-        self, db_connection: Connection, orm_class: Any, column_name: str
+        db_connection: Connection, orm_class: Any, column_name: str
     ) -> Any:
         """Return a random value from the column specified."""
         query = select(orm_class).order_by(functions.random()).limit(1)
@@ -50,8 +51,8 @@ class Meta:
 
         name = "timedelta_provider"
 
+    @staticmethod
     def timedelta(
-        self,
         min_dt: dt.timedelta = dt.timedelta(seconds=0),
         # ints bigger than this cause trouble
         max_dt: dt.timedelta = dt.timedelta(seconds=2**32),
@@ -75,8 +76,8 @@ class Meta:
 
         name = "timespan_provider"
 
+    @staticmethod
     def timespan(
-        self,
         earliest_start_year: int,
         last_start_year: int,
         min_dt: dt.timedelta = dt.timedelta(seconds=0),
@@ -194,6 +195,7 @@ class Meta:
 
         name = "null_provider"
 
-    def null(self) -> None:
+    @staticmethod
+    def null() -> None:
         """Return `None`."""
         return None

tests/examples/example_orm.py

@@ -1,4 +1,4 @@
-from typing import List, Optional
+from typing import Any, List, Optional
 
 from sqlalchemy import (
     BigInteger,
@@ -18,6 +18,7 @@
     UniqueConstraint,
     Uuid,
 )
+from sqlalchemy.dialects.postgresql import BIT, CIDR
 from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
 import datetime
 
@@ -71,6 +72,15 @@ class Person(Base):
     )
 
 
+class StrangeTypeTable(Base):
+    __tablename__ = "strange_type_table"
+    __table_args__ = (PrimaryKeyConstraint("id", name="strange_type_table_pkey"),)
+
+    id: Mapped[int] = mapped_column(Integer, primary_key=True)
+    column_with_unusual_type: Mapped[Optional[Any]] = mapped_column(CIDR)
+    column_with_unusual_type_and_length: Mapped[Optional[Any]] = mapped_column(BIT(3))
+
+
 class UnignorableTable(Base):
     __tablename__ = "unignorable_table"
     __table_args__ = (PrimaryKeyConstraint("id", name="unignorable_table_pkey"),)

tests/examples/expected_ssg.py

@@ -90,6 +90,19 @@ def __call__(self, dst_db_conn):
         return result
 
 
+class strange_type_tableGenerator(TableGenerator):
+    num_rows_per_pass = 1
+
+    def __init__(self):
+        pass
+
+    def __call__(self, dst_db_conn):
+        result = {}
+        result["column_with_unusual_type"] = generic.null_provider.null()
+        result["column_with_unusual_type_and_length"] = generic.null_provider.null()
+        return result
+
+
 class unique_constraint_testGenerator(TableGenerator):
     num_rows_per_pass = 1
 
@@ -188,6 +201,7 @@ def __call__(self, dst_db_conn):
     "data_type_test": data_type_testGenerator(),
     "no_pk_test": no_pk_testGenerator(),
     "person": personGenerator(),
+    "strange_type_table": strange_type_tableGenerator(),
     "unique_constraint_test": unique_constraint_testGenerator(),
     "unique_constraint_test2": unique_constraint_test2Generator(),
     "test_entity": test_entityGenerator(),

tests/examples/src.dump

@@ -218,6 +218,18 @@ CREATE TABLE public.table_to_be_ignored (
 
 ALTER TABLE public.table_to_be_ignored OWNER TO postgres;
 
+--
+-- Name: strange_type_table; Type: TABLE; Schema: public; Owner: postgres
+--
+
+CREATE TABLE public.strange_type_table (
+    id integer NOT NULL,
+    column_with_unusual_type cidr,
+    column_with_unusual_type_and_length bit(3)
+);
+
+ALTER TABLE public.strange_type_table OWNER TO postgres;
+
 --
 -- Data for Name: concept; Type: TABLE DATA; Schema: public; Owner: postgres
 --
@@ -1390,6 +1402,13 @@ ALTER TABLE ONLY public.ref_to_unignorable_table
 ALTER TABLE ONLY public.unignorable_table
     ADD CONSTRAINT unignorable_table_pkey PRIMARY KEY (id);
 
+--
+-- Name: strange_type_table strange_type_table_pkey; Type: CONSTRAINT; Schema: public; Owner: postgres
+--
+
+ALTER TABLE ONLY public.strange_type_table
+    ADD CONSTRAINT strange_type_table_pkey PRIMARY KEY (id);
+
 --
 -- Name: fki_concept_concept_type_id_fkey; Type: INDEX; Schema: public; Owner: postgres
 --

tests/test_functional.py

@@ -14,7 +14,6 @@ class FunctionalTestCase(RequiresDBTestCase):
 
     def test_version_command(self) -> None:
         """Check that the version command works."""
-
         completed_process = run(
             ["sqlsynthgen", "version"],
             capture_output=True,
@@ -40,10 +39,10 @@ class DBFunctionalTestCase(RequiresDBTestCase):
     alt_ssg_file_path = Path("my_ssg.py")
 
     vocabulary_file_paths = tuple(
-        map(Path, ("concept.yaml", "concept_type.yaml", "mitigation_type.yaml"))
+        map(Path, ("concept.yaml", "concept_type.yaml", "mitigation_type.yaml")),
     )
     generator_file_paths = tuple(
-        map(Path, ("story_generators.py", "row_generators.py"))
+        map(Path, ("story_generators.py", "row_generators.py")),
     )
     dump_file_path = Path("dst.dump")
     config_file_path = Path("example_config.yaml")
@@ -60,7 +59,6 @@ class DBFunctionalTestCase(RequiresDBTestCase):
 
     def setUp(self) -> None:
         """Pre-test setup."""
-
         # Create a mostly-blank destination database
         run_psql(self.examples_dir / self.dump_file_path)
 
@@ -103,6 +101,16 @@ def test_workflow_minimal_args(self) -> None:
         # this could mean that we might accidentally violate the constraints. In
         # practice this won't happen because we only write one row to an empty table.
         self.assertEqual(
+            "Unsupported SQLAlchemy type "
+            "<class 'sqlalchemy.dialects.postgresql.types.CIDR'> "
+            "for column column_with_unusual_type. "
+            "Setting this column to NULL always, "
+            "you may want to configure a row generator for it instead.\n"
+            "Unsupported SQLAlchemy type "
+            "<class 'sqlalchemy.dialects.postgresql.types.BIT'> "
+            "for column column_with_unusual_type_and_length. "
+            "Setting this column to NULL always, "
+            "you may want to configure a row generator for it instead.\n"
             "A unique constraint (ab_uniq) isn't fully covered by one "
             "row generator (['a']). Enforcement of the constraint may not work.\n"
             "A unique constraint (ab_uniq) isn't fully covered by one "
@@ -188,7 +196,6 @@ def test_workflow_minimal_args(self) -> None:
 
     def test_workflow_maximal_args(self) -> None:
         """Test the CLI workflow runs with optional arguments."""
-
         completed_process = run(
             [
                 "sqlsynthgen",
@@ -257,7 +264,19 @@ def test_workflow_maximal_args(self) -> None:
             capture_output=True,
             env=self.env,
         )
-        self.assertEqual("", completed_process.stderr.decode("utf-8"))
+        self.assertEqual(
+            "Unsupported SQLAlchemy type "
+            "<class 'sqlalchemy.dialects.postgresql.types.CIDR'> "
+            "for column column_with_unusual_type. "
+            "Setting this column to NULL always, "
+            "you may want to configure a row generator for it instead.\n"
+            "Unsupported SQLAlchemy type "
+            "<class 'sqlalchemy.dialects.postgresql.types.BIT'> "
+            "for column column_with_unusual_type_and_length. "
+            "Setting this column to NULL always, "
+            "you may want to configure a row generator for it instead.\n",
+            completed_process.stderr.decode("utf-8"),
+        )
         self.assertSuccess(completed_process)
         self.assertEqual(
             f"Making {self.alt_ssg_file_path}.\n"
@@ -345,6 +364,7 @@ def test_workflow_maximal_args(self) -> None:
             '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 "strange_type_table".\n'
             'Generating data for table "unique_constraint_test".\n'
             'Generating data for table "unique_constraint_test2".\n'
             'Generating data for table "test_entity".\n'
@@ -358,6 +378,7 @@ def test_workflow_maximal_args(self) -> None:
             '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 "strange_type_table".\n'
             'Generating data for table "unique_constraint_test".\n'
             'Generating data for table "unique_constraint_test2".\n'
             'Generating data for table "test_entity".\n'
@@ -367,6 +388,7 @@ def test_workflow_maximal_args(self) -> None:
             f"hospital_visit: {2*(2*2+3)} rows created.\n"
             "data_type_test: 2 rows created.\n"
             "no_pk_test: 2 rows created.\n"
+            "strange_type_table: 2 rows created.\n"
             "unique_constraint_test: 2 rows created.\n"
             "unique_constraint_test2: 2 rows created.\n"
             "test_entity: 2 rows created.\n",
@@ -394,6 +416,7 @@ def test_workflow_maximal_args(self) -> None:
             'Truncating table "test_entity".\n'
             'Truncating table "unique_constraint_test2".\n'
             'Truncating table "unique_constraint_test".\n'
+            'Truncating table "strange_type_table".\n'
             'Truncating table "person".\n'
             'Truncating table "no_pk_test".\n'
             'Truncating table "data_type_test".\n'
@@ -459,7 +482,6 @@ def test_unique_constraint_fail(self) -> None:
         We also deliberately call create-data multiple times to make sure that the
         loading of existing keys from the database at start up works as expected.
         """
-
         # This is all exactly the same stuff we run in test_workflow_maximal_args.
         run(
             [

tests/test_remove.py

@@ -12,14 +12,14 @@ class RemoveTestCase(SSGTestCase):
 
     @patch("sqlsynthgen.remove.get_settings", side_effect=get_test_settings)
     @patch("sqlsynthgen.remove.create_db_engine")
-    @patch("sqlsynthgen.remove.delete", side_effect=range(1, 8))
+    @patch("sqlsynthgen.remove.delete", side_effect=range(1, 9))
     def test_remove_db_data(
         self, mock_delete: MagicMock, mock_engine: MagicMock, _: MagicMock
     ) -> None:
         """Test the remove_db_data function."""
         config = {"tables": {"unignorable_table": {"ignore": True}}}
         remove_db_data(example_orm, remove_ssg, config)
-        self.assertEqual(mock_delete.call_count, 7)
+        self.assertEqual(mock_delete.call_count, 8)
         mock_delete.assert_has_calls(
             [
                 call(example_orm.Base.metadata.tables[t])
@@ -29,6 +29,7 @@ def test_remove_db_data(
                     "unique_constraint_test2",
                     "unique_constraint_test",
                     "person",
+                    "strange_type_table",
                     "no_pk_test",
                     "data_type_test",
                 )