Merge pull request #138 from alan-turing-institute/sqla2-type-hints

Add more type hints

Author: mhauru

poetry.lock

@@ -1,21 +1,41 @@
 """Base table generator classes."""
 import logging
+from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
 
 import yaml
-from sqlalchemy import insert
+from sqlalchemy import Connection, insert
 from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.schema import Table
+
+
+class TableGenerator(ABC):
+    """Abstract base class for table generator classes."""
+
+    num_rows_per_pass: int = 1
+
+    @abstractmethod
+    def __call__(self, dst_db_conn: Connection) -> dict[str, Any]:
+        """Return, as a dictionary, a new row for the table that we are generating.
+
+        The only argument, `dst_db_conn`, should be a database connection to the
+        database to which the data is being written. Most generators won't use it, but
+        some do, and thus it's required by the interface.
+
+        The return value should be a dictionary with column names as strings for keys,
+        and the values being the values for the new row.
+        """
 
 
 @dataclass
 class FileUploader:
     """For uploading data files."""
 
-    table: Any
+    table: Table
 
-    def load(self, connection: Any) -> None:
+    def load(self, connection: Connection) -> None:
         """Load the data from file."""
         yaml_file = Path(self.table.fullname + ".yaml")
         if not yaml_file.exists():

sqlsynthgen/base.py

@@ -2,42 +2,46 @@
 import logging
 from typing import Any, Dict, Generator, List, Tuple
 
-from sqlalchemy import insert
+from sqlalchemy import Connection, insert
 from sqlalchemy.exc import IntegrityError
-from sqlalchemy.schema import CreateSchema
+from sqlalchemy.schema import CreateSchema, MetaData, Table
 
+from sqlsynthgen.base import FileUploader, TableGenerator
 from sqlsynthgen.settings import get_settings
-from sqlsynthgen.utils import create_db_engine
+from sqlsynthgen.utils import create_db_engine, get_sync_engine
 
 Story = Generator[Tuple[str, Dict[str, Any]], Dict[str, Any], None]
 
 
-def create_db_tables(metadata: Any) -> Any:
+def create_db_tables(metadata: MetaData) -> None:
     """Create tables described by the sqlalchemy metadata object."""
     settings = get_settings()
 
-    engine = create_db_engine(settings.dst_dsn)  # type: ignore
+    engine = get_sync_engine(create_db_engine(settings.dst_dsn))  # type: ignore
 
     # Create schema, if necessary.
     if settings.dst_schema:
         schema_name = settings.dst_schema
-        if not engine.dialect.has_schema(engine, schema=schema_name):
-            engine.execute(CreateSchema(schema_name, if_not_exists=True))
+        with engine.connect() as connection:
+            if not engine.dialect.has_schema(connection, schema_name=schema_name):
+                connection.execute(CreateSchema(schema_name, if_not_exists=True))
 
         # Recreate the engine, this time with a schema specified
-        engine = create_db_engine(
-            settings.dst_dsn, schema_name=schema_name  # type: ignore
+        engine = get_sync_engine(
+            create_db_engine(settings.dst_dsn, schema_name=schema_name)  # type: ignore
         )
 
     metadata.create_all(engine)
 
 
-def create_db_vocab(vocab_dict: Dict[str, Any]) -> None:
+def create_db_vocab(vocab_dict: Dict[str, FileUploader]) -> None:
     """Load vocabulary tables from files."""
     settings = get_settings()
 
-    dst_engine = create_db_engine(
-        settings.dst_dsn, schema_name=settings.dst_schema  # type: ignore
+    dst_engine = get_sync_engine(
+        create_db_engine(
+            settings.dst_dsn, schema_name=settings.dst_schema  # type: ignore
+        )
     )
 
     with dst_engine.connect() as dst_conn:
@@ -51,16 +55,18 @@ def create_db_vocab(vocab_dict: Dict[str, Any]) -> None:
 
 
 def create_db_data(
-    sorted_tables: list,
-    table_generator_dict: dict,
-    story_generator_list: list,
+    sorted_tables: list[Table],
+    table_generator_dict: dict[str, TableGenerator],
+    story_generator_list: list[dict[str, Any]],
     num_passes: int,
 ) -> None:
     """Connect to a database and populate it with data."""
     settings = get_settings()
 
-    dst_engine = create_db_engine(
-        settings.dst_dsn, schema_name=settings.dst_schema  # type: ignore
+    dst_engine = get_sync_engine(
+        create_db_engine(
+            settings.dst_dsn, schema_name=settings.dst_schema  # type: ignore
+        )
     )
 
     with dst_engine.connect() as dst_conn:
@@ -75,9 +81,9 @@ def create_db_data(
 
 def _populate_story(
     story: Story,
-    table_dict: Dict[str, Any],
-    table_generator_dict: Dict[str, Any],
-    dst_conn: Any,
+    table_dict: Dict[str, Table],
+    table_generator_dict: Dict[str, TableGenerator],
+    dst_conn: Connection,
 ) -> None:
     """Write to the database all the rows created by the given story."""
     # Loop over the rows generated by the story, insert them into their
@@ -100,7 +106,9 @@ def _populate_story(
             # because other parts of the story may refer to them.
             if cursor.returned_defaults:
                 # pylint: disable=protected-access
-                return_values = cursor.returned_defaults._mapping
+                return_values = {
+                    str(k): v for k, v in cursor.returned_defaults._mapping.items()
+                }
                 # pylint: enable=protected-access
             else:
                 return_values = {}
@@ -112,10 +120,10 @@ def _populate_story(
 
 
 def populate(
-    dst_conn: Any,
-    tables: list,
-    table_generator_dict: dict,
-    story_generator_list: list,
+    dst_conn: Connection,
+    tables: list[Table],
+    table_generator_dict: dict[str, TableGenerator],
+    story_generator_list: list[dict[str, Any]],
 ) -> None:
     """Populate a database schema with synthetic data."""
     table_dict = {table.name: table for table in tables}

sqlsynthgen/create.py

@@ -10,6 +10,7 @@
 import yaml
 from jsonschema.exceptions import ValidationError
 from jsonschema.validators import validate
+from sqlalchemy.schema import MetaData
 from typer import Option, Typer, echo
 
 from sqlsynthgen.create import create_db_data, create_db_tables, create_db_vocab
@@ -82,10 +83,13 @@ def create_data(
         echo("Creating data.")
     orm_module = import_file(orm_file)
     ssg_module = import_file(ssg_file)
+    orm_metadata: MetaData = orm_module.Base.metadata
+    table_generator_dict = ssg_module.table_generator_dict
+    story_generator_list = ssg_module.story_generator_list
     create_db_data(
-        orm_module.Base.metadata.sorted_tables,
-        ssg_module.table_generator_dict,
-        ssg_module.story_generator_list,
+        orm_metadata.sorted_tables,
+        table_generator_dict,
+        story_generator_list,
         num_passes,
     )
     if verbose:
@@ -136,7 +140,8 @@ def create_tables(
         echo("Creating tables.")
 
     orm_module = import_file(orm_file)
-    create_db_tables(orm_module.Base.metadata)
+    orm_metadata: MetaData = orm_module.Base.metadata
+    create_db_tables(orm_metadata)
 
     if verbose:
         echo("Tables created.")