Merge pull request #137 from alan-turing-institute/sqlalchemy2

mhauru · web-flow · commit 8259c0fe5a2f · 2023-09-06T18:09:54.000+02:00
Sqlalchemy2
diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml
@@ -57,7 +57,7 @@ jobs:
         uses: actions/cache@v3
         with:
           path: ${{ env.PRE_COMMIT_HOME }}
-          key: hooks-${{ runner.os }}-${{ hashFiles('.pre-commit-config.yaml') }}
+          key: hooks-${{ runner.os }}-${{ hashFiles('.pre-commit-config.yaml') }}-${{ env.PYTHON_VERSION }}-${{ hashFiles('poetry.lock') }}
       - name: Install Pre-Commit Hooks
         shell: bash
         if: steps.pre-commit-cache.outputs.cache-hit != 'true'
diff --git a/mypy.ini b/mypy.ini
@@ -10,3 +10,4 @@ warn_return_any = True
 show_error_codes = True
 warn_unused_ignores = True
 ignore_missing_imports = True
+plugins = pydantic.mypy
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -23,14 +23,14 @@ sqlalchemy-utils = "^0.38.3"
 mimesis = "^6.1.1"
 typer = "^0.7.0"
 pyyaml = "^5.0"
-sqlalchemy = {version = "^1.4", extras = ["asyncio"]}
+sqlalchemy = "^2"
 sphinx-rtd-theme = {version = "^1.2.0", optional = true}
 sphinxcontrib-napoleon = {version = "^0.7", optional = true}
-smartnoise-sql = "^0.2.11"
+smartnoise-sql = "^1"
 jinja2 = "^3.1.2"
 black = "^23.3.0"
 jsonschema = "^4.17.3"
-sqlacodegen = "3.0.0rc1"
+sqlacodegen = "^3.0.0rc3"
 asyncpg = "^0.27.0"
 greenlet = "^2.0.2"
 pymysql = "^1.1.0"
@@ -39,7 +39,7 @@ pandas = "^2"
 [tool.poetry.group.dev.dependencies]
 isort = "^5.10.1"
 pylint = "^2.15.8"
-mypy = "^0.991"
+mypy = "^1.5"
 types-pyyaml = "^6.0.12.4"
 pydocstyle = "^6.3.0"
 restructuredtext-lint = "^1.4.0"
diff --git a/sqlsynthgen/base.py b/sqlsynthgen/base.py
@@ -35,6 +35,7 @@ def load(self, connection: Any) -> None:
         try:
             stmt = insert(self.table).values(list(rows))
             connection.execute(stmt)
+            connection.commit()
         except SQLAlchemyError as e:
             logging.warning(
                 "Error inserting rows into table %s: %s", self.table.fullname, e
diff --git a/sqlsynthgen/create.py b/sqlsynthgen/create.py
@@ -94,11 +94,16 @@ def _populate_story(
             else:
                 default_values = {}
             insert_values = {**default_values, **provided_values}
-            stmt = insert(table).values(insert_values)
+            stmt = insert(table).values(insert_values).return_defaults()
             cursor = dst_conn.execute(stmt)
             # We need to return all the default values etc. to the generator,
             # because other parts of the story may refer to them.
-            return_values = dict(cursor.returned_defaults or {})
+            if cursor.returned_defaults:
+                # pylint: disable=protected-access
+                return_values = cursor.returned_defaults._mapping
+                # pylint: enable=protected-access
+            else:
+                return_values = {}
             final_values = {**insert_values, **return_values}
             table_name, provided_values = story.send(final_values)
     except StopIteration:
diff --git a/sqlsynthgen/make.py b/sqlsynthgen/make.py
@@ -222,6 +222,7 @@ def _get_provider_for_column(column: Any) -> Tuple[List[str], str, List[str]]:
         (sqltypes.DateTime, False): "generic.datetime.datetime",
         (sqltypes.Numeric, False): "generic.numeric.float_number",
         (sqltypes.LargeBinary, False): "generic.bytes_provider.bytes",
+        (sqltypes.Uuid, False): "generic.cryptographic.uuid",
         (postgresql.UUID, False): "generic.cryptographic.uuid",
         (sqltypes.String, False): "generic.text.color",
         (sqltypes.String, True): "generic.person.password",
diff --git a/sqlsynthgen/providers.py b/sqlsynthgen/providers.py
@@ -5,7 +5,7 @@
 
 from mimesis import Datetime, Text
 from mimesis.providers.base import BaseDataProvider, BaseProvider
-from sqlalchemy.sql import func, select
+from sqlalchemy.sql import functions, select
 
 
 class ColumnValueProvider(BaseProvider):
@@ -18,7 +18,7 @@ class Meta:
 
     def column_value(self, db_connection: Any, orm_class: Any, column_name: str) -> Any:
         """Return a random value from the column specified."""
-        query = select(orm_class).order_by(func.random()).limit(1)
+        query = select(orm_class).order_by(functions.random()).limit(1)
         random_row = db_connection.execute(query).first()
 
         if random_row:
diff --git a/sqlsynthgen/remove.py b/sqlsynthgen/remove.py
@@ -19,6 +19,7 @@ def remove_db_data(orm_module: ModuleType, ssg_module: ModuleType) -> None:
             # We presume that all tables that aren't vocab should be truncated
             if table.name not in ssg_module.vocab_dict:
                 dst_conn.execute(delete(table))
+                dst_conn.commit()
 
 
 def remove_db_vocab(orm_module: ModuleType, ssg_module: ModuleType) -> None:
@@ -33,6 +34,7 @@ def remove_db_vocab(orm_module: ModuleType, ssg_module: ModuleType) -> None:
             # We presume that all tables that are vocab should be truncated
             if table.name in ssg_module.vocab_dict:
                 dst_conn.execute(delete(table))
+                dst_conn.commit()
 
 
 def remove_db_tables(orm_module: ModuleType) -> None:
diff --git a/sqlsynthgen/utils.py b/sqlsynthgen/utils.py
@@ -5,7 +5,7 @@
 from importlib import import_module
 from pathlib import Path
 from types import ModuleType
-from typing import Any, Final, Optional
+from typing import Any, Final, Optional, Union
 
 import yaml
 from jsonschema.exceptions import ValidationError
@@ -64,11 +64,11 @@ def import_file(file_path: str) -> ModuleType:
     return module
 
 
-def download_table(table: Any, engine: Any, yaml_file_name: str) -> None:
+def download_table(table: Any, engine: Any, yaml_file_name: Union[str, Path]) -> None:
     """Download a Table and store it as a .yaml file."""
-    stmt = select([table])
+    stmt = select(table)
     with engine.connect() as conn:
-        result = [dict(row) for row in conn.execute(stmt)]
+        result = [dict(row) for row in conn.execute(stmt).mappings()]
 
     with Path(yaml_file_name).open("w", newline="", encoding="utf-8") as yamlfile:
         yamlfile.write(yaml.dump(result))
@@ -83,7 +83,7 @@ def create_db_engine(
     """Create a SQLAlchemy Engine."""
     if use_asyncio:
         async_dsn = db_dsn.replace("postgresql://", "postgresql+asyncpg://")
-        engine = create_async_engine(async_dsn, **kwargs)
+        engine: Any = create_async_engine(async_dsn, **kwargs)
         event_engine = engine.sync_engine
     else:
         engine = create_engine(db_dsn, **kwargs)
@@ -93,7 +93,7 @@ def create_db_engine(
 
         @event.listens_for(event_engine, "connect", insert=True)
         def connect(dbapi_connection: Any, _: Any) -> None:
-            set_search_path(dbapi_connection, schema_name)  # type: ignore
+            set_search_path(dbapi_connection, schema_name)
 
     return engine
 
@@ -105,7 +105,8 @@ def set_search_path(connection: Any, schema: str) -> None:
     connection.autocommit = True
 
     cursor = connection.cursor()
-    cursor.execute("SET search_path to %s;", (schema,))
+    # Parametrised queries don't work with asyncpg, hence the f-string.
+    cursor.execute(f"SET search_path TO {schema};")
     cursor.close()
 
     connection.autocommit = existing_autocommit
diff --git a/tests/examples/story_generators.py b/tests/examples/story_generators.py
@@ -44,7 +44,7 @@ def long_story(
     concept_query = sqla.text(
         "SELECT concept_id FROM concept WHERE concept_name = 'another concept'"
     )
-    concept_id = dst_db_conn.execute(concept_query).first()[0]
+    concept_id = dst_db_conn.execute(concept_query).mappings().first()["concept_id"]
 
     # Make a story where we first create a person, and then give that person two visits.
     person = yield ("person", {"research_opt_out": opt_out})
diff --git a/tests/test_base.py b/tests/test_base.py
@@ -50,6 +50,6 @@ def test_load(self) -> None:
 
         with self.engine.connect() as conn:
             vocab_gen.load(conn)
-            statement = select([BaseTable])
+            statement = select(BaseTable)
             rows = list(conn.execute(statement))
         self.assertEqual(3, len(rows))
diff --git a/tests/test_create.py b/tests/test_create.py
@@ -1,15 +1,21 @@
 """Tests for the create module."""
 import itertools as itt
+from pathlib import Path
 from typing import Any, Generator, Tuple
 from unittest.mock import MagicMock, call, patch
 
+from sqlalchemy import Column, Integer, create_engine
+from sqlalchemy.orm import declarative_base
+
 from sqlsynthgen.create import (
+    Story,
+    _populate_story,
     create_db_data,
     create_db_tables,
     create_db_vocab,
     populate,
 )
-from tests.utils import SSGTestCase, get_test_settings
+from tests.utils import RequiresDBTestCase, SSGTestCase, get_test_settings, run_psql
 
 
 class MyTestCase(SSGTestCase):
@@ -48,8 +54,7 @@ def test_create_db_tables(
         )
         mock_meta.create_all.assert_called_once_with(mock_create_engine.return_value)
 
-    @patch("sqlsynthgen.create.insert")
-    def test_populate(self, mock_insert: MagicMock) -> None:
+    def test_populate(self) -> None:
         """Test the populate function."""
         table_name = "table_name"
 
@@ -62,34 +67,47 @@ def mock_story_gen(_: Any) -> Generator[Tuple[str, dict], None, None]:
             return story()
 
         for num_stories_per_pass, num_rows_per_pass in itt.product([0, 2], [0, 3]):
-            mock_dst_conn = MagicMock()
-            mock_dst_conn.execute.return_value.returned_defaults = {}
-            mock_table = MagicMock()
-            mock_table.name = table_name
-            mock_gen = MagicMock()
-            mock_gen.num_rows_per_pass = num_rows_per_pass
-            mock_gen.return_value = {}
-
-            tables = [mock_table]
-            row_generators = {table_name: mock_gen}
-            story_generators = (
-                [{"name": mock_story_gen, "num_stories_per_pass": num_stories_per_pass}]
-                if num_stories_per_pass > 0
-                else []
-            )
-            populate(mock_dst_conn, tables, row_generators, story_generators)
-
-            mock_gen.assert_has_calls(
-                [call(mock_dst_conn)] * (num_stories_per_pass + num_rows_per_pass)
-            )
-            mock_insert.return_value.values.assert_has_calls(
-                [call(mock_gen.return_value)]
-                * (num_stories_per_pass + num_rows_per_pass)
-            )
-            mock_dst_conn.execute.assert_has_calls(
-                [call(mock_insert.return_value.values.return_value)]
-                * (num_stories_per_pass + num_rows_per_pass)
-            )
+            with patch("sqlsynthgen.create.insert") as mock_insert:
+                mock_values = mock_insert.return_value.values
+                mock_dst_conn = MagicMock()
+                mock_dst_conn.execute.return_value.returned_defaults = {}
+                mock_table = MagicMock()
+                mock_table.name = table_name
+                mock_gen = MagicMock()
+                mock_gen.num_rows_per_pass = num_rows_per_pass
+                mock_gen.return_value = {}
+
+                tables = [mock_table]
+                row_generators = {table_name: mock_gen}
+                story_generators = (
+                    [
+                        {
+                            "name": mock_story_gen,
+                            "num_stories_per_pass": num_stories_per_pass,
+                        }
+                    ]
+                    if num_stories_per_pass > 0
+                    else []
+                )
+                populate(mock_dst_conn, tables, row_generators, story_generators)
+
+                self.assertListEqual(
+                    [call(mock_dst_conn)] * (num_stories_per_pass + num_rows_per_pass),
+                    mock_gen.call_args_list,
+                )
+                self.assertListEqual(
+                    [call(mock_gen.return_value)]
+                    * (num_stories_per_pass + num_rows_per_pass),
+                    mock_values.call_args_list,
+                )
+                self.assertListEqual(
+                    (
+                        [call(mock_values.return_value.return_defaults.return_value)]
+                        * num_stories_per_pass
+                    )
+                    + ([call(mock_values.return_value)] * num_rows_per_pass),
+                    mock_dst_conn.execute.call_args_list,
+                )
 
     @patch("sqlsynthgen.create.insert")
     def test_populate_diff_length(self, mock_insert: MagicMock) -> None:
@@ -131,3 +149,42 @@ def test_create_db_vocab(
         )
         # Running the same insert twice should be fine.
         create_db_vocab(vocab_list)
+
+
+class TestStoryDefaults(RequiresDBTestCase):
+    """Test that we can handle column defaults in stories."""
+
+    # pylint: disable=invalid-name
+    Base = declarative_base()
+    # pylint: enable=invalid-name
+    metadata = Base.metadata
+
+    class ColumnDefaultsTable(Base):  # type: ignore
+        """A SQLAlchemy model."""
+
+        __tablename__ = "column_defaults"
+        someval = Column(Integer, primary_key=True)
+        otherval = Column(Integer, server_default="8")
+
+    def setUp(self) -> None:
+        """Ensure we have an empty DB to work with."""
+        dump_file_path = Path("dst.dump")
+        examples_dir = Path("tests/examples")
+        run_psql(examples_dir / dump_file_path)
+
+    def test_populate(self) -> None:
+        """Check that we can populate a table that has column defaults."""
+        engine = create_engine(
+            "postgresql://postgres:password@localhost:5432/dst",
+        )
+        self.metadata.create_all(engine)
+
+        def my_story() -> Story:
+            """A story generator."""
+            first_row = yield "column_defaults", {}
+            self.assertEqual(1, first_row["someval"])
+            self.assertEqual(8, first_row["otherval"])
+
+        with engine.connect() as conn:
+            with conn.begin():
+                _populate_story(my_story(), dict(self.metadata.tables), {}, conn)
diff --git a/tests/test_remove.py b/tests/test_remove.py
@@ -31,7 +31,11 @@ def test_remove_db_data(
     @patch("sqlsynthgen.remove.get_settings")
     def test_remove_db_data_raises(self, mock_get: MagicMock) -> None:
         """Check that remove_db_data raises if dst DSN is missing."""
-        mock_get.return_value = Settings(dst_dsn=None, _env_file=None)
+        mock_get.return_value = Settings(
+            dst_dsn=None,
+            # The mypy ignore can be removed once we upgrade to pydantic 2.
+            _env_file=None,  # type: ignore[call-arg]
+        )
         with self.assertRaises(AssertionError) as context_manager:
             remove_db_data(example_orm, remove_ssg)
         self.assertEqual(
@@ -56,7 +60,11 @@ def test_remove_db_vocab(
     @patch("sqlsynthgen.remove.get_settings")
     def test_remove_db_vocab_raises(self, mock_get: MagicMock) -> None:
         """Check that remove_db_vocab raises if dst DSN is missing."""
-        mock_get.return_value = Settings(dst_dsn=None, _env_file=None)
+        mock_get.return_value = Settings(
+            dst_dsn=None,
+            # The mypy ignore can be removed once we upgrade to pydantic 2.
+            _env_file=None,  # type: ignore[call-arg]
+        )
         with self.assertRaises(AssertionError) as context_manager:
             remove_db_vocab(example_orm, remove_ssg)
         self.assertEqual(
@@ -75,7 +83,11 @@ def test_remove_tables(self, mock_engine: MagicMock, _: MagicMock) -> None:
     @patch("sqlsynthgen.remove.get_settings")
     def test_remove_db_tables_raises(self, mock_get: MagicMock) -> None:
         """Check that remove_db_tables raises if dst DSN is missing."""
-        mock_get.return_value = Settings(dst_dsn=None, _env_file=None)
+        mock_get.return_value = Settings(
+            dst_dsn=None,
+            # The mypy ignore can be removed once we upgrade to pydantic 2.
+            _env_file=None,  # type: ignore[call-arg]
+        )
         with self.assertRaises(AssertionError) as context_manager:
             remove_db_tables(example_orm)
         self.assertEqual(
diff --git a/tests/test_settings.py b/tests/test_settings.py
@@ -12,7 +12,8 @@ def test_minimal_settings(self) -> None:
         """Test the minimal settings."""
         settings = Settings(
             # To stop any local .env files influencing the test
-            _env_file=None,
+            # The mypy ignore can be removed once we upgrade to pydantic 2.
+            _env_file=None,  # type: ignore[call-arg]
         )
         self.assertIsNone(settings.src_dsn)
         self.assertIsNone(settings.src_schema)
@@ -28,7 +29,8 @@ def test_maximal_settings(self) -> None:
             dst_dsn="postgresql://user:password@host:port/db_name?sslmode=require",
             dst_schema="src_schema",
             # To stop any local .env files influencing the test
-            _env_file=None,
+            # The mypy ignore can be removed once we upgrade to pydantic 2.
+            _env_file=None,  # type: ignore[call-arg]
         )
 
     def test_validation(self) -> None:
diff --git a/tests/test_utils.py b/tests/test_utils.py
diff --git a/tests/utils.py b/tests/utils.py
