Mssql

.github/workflows/tests.yml

@@ -15,24 +15,34 @@ jobs:
     steps:
       - name: Checkout Code
         uses: actions/checkout@v6
-      - name: Start PostgreSQL
+      - name: Install prerequisites
+        id: install-prereqs
         shell: bash
         run: |
-          sudo systemctl start postgresql.service
-      - name: Install poetry
+          sudo apt install python3-poetry mssql-tools18
+        background: true
+      - name: Start PostgreSQL
         shell: bash
         run: |
-          sudo apt install python3-poetry
+          sudo systemctl start postgresql.service
+      - name: Wait for poetry to be installed
+        wait: install-prereqs
       - name: Configure poetry
         shell: bash
-        run: |
-          python -m poetry config virtualenvs.in-project true
+        run: python -m poetry config virtualenvs.in-project true
       - name: Install dependencies
+        id: install-deps
         shell: bash
         if: steps.poetry-cache.outputs.cache-hit != 'true'
-        run: |
-          python -m poetry install --all-extras
+        run: python -m poetry install --all-extras
+        background: true
+      - name: Set up MS SQL Server
+        uses: hoverkraft-tech/compose-action@11beaa1c2dae4e8ed7b1665aa074723b6cecb0e4 # v3.0.0
+        with:
+          compose-file: ./docker-compose.yml
+      - name: Wait for the dependencies
+        wait:
+          - install-deps
       - name: Run tests
         shell: bash
-        run: |
-          poetry run python -m unittest
+        run: poetry run python -m pytest -rs

CONTRIBUTING.md

@@ -41,10 +41,16 @@ Please install the following software on your workstation:
 
 ## Running unit tests
 
-Executing unit tests is straightforward:
+Executing unit tests is straightforward. For example inside a Poetry shell you can run:
 
 ```bash
-python -m unittest discover --verbose tests/
+pytest -rs
+```
+
+or outside:
+
+```bash
+poetry run pytest -rs
 ```
 
 ## Building documentation locally

datafaker/create.py

@@ -1,5 +1,4 @@
 """Functions and classes to create and populate the target database."""
-import re
 from collections import Counter
 from pathlib import Path
 from typing import Any, Generator, Iterable, Iterator, Mapping, Sequence, Tuple
@@ -8,9 +7,8 @@
 import yaml
 from sqlalchemy import Connection, insert, inspect
 from sqlalchemy.exc import IntegrityError
-from sqlalchemy.ext.compiler import compiles
 from sqlalchemy.orm import Session
-from sqlalchemy.schema import CreateSchema, CreateTable, MetaData, Table
+from sqlalchemy.schema import CreateSchema, MetaData, Table
 
 from datafaker.base import FileUploader
 from datafaker.db_utils import (
@@ -33,31 +31,6 @@
 Story = Generator[Tuple[str, dict[str, Any]], dict[str, Any], None]
 RowCounts = Counter[str]
 
-serial_re = re.compile(r"\bSERIAL\b")
-
-
-@compiles(CreateTable, "duckdb")
-def remove_on_delete_cascade(element: CreateTable, compiler: Any, **kw: Any) -> str:
-    """
-    Intercede in compilation for column creation.
-
-    DuckDB does not understand cascades, and we don't care about
-    that in datafaker so we remove ``ON DELETE CASCASE``.
-
-    DuckDB does not understand ``SERIAL`` and we don't care
-    about autoincrementing, so we will replace it simply with
-    ``INTEGER``.
-
-    Ideally ``duckdb_engine`` would remove these for us.
-    :param element: The CreateTable being executed.
-    :param compiler: Actually a DDLCompiler, but that type is not exported.
-    :param kw: Further arguments.
-    :return: Corrected SQL.
-    """
-    text: str = compiler.visit_create_table(element, **kw)
-    t2 = serial_re.sub("INTEGER", text)
-    return t2.replace(" ON DELETE CASCADE", "")
-
 
 def create_db_tables(metadata: MetaData) -> None:
     """Create tables described by the sqlalchemy metadata object."""
@@ -320,7 +293,7 @@ def insert(self) -> None:
         table = self._table_dict[self._table_name]
         if table.name in self._table_generator_dict:
             table_generator = self._table_generator_dict[table.name]
-            default_values = table_generator(self._dst_conn)
+            default_values = table_generator.generate_row(self._dst_conn)
         else:
             default_values = {}
         insert_values = {**default_values, **self._provided_values}
@@ -403,7 +376,7 @@ def populate(
         with dst_conn.begin():
             try:
                 for _ in range(table_generator.num_rows_per_pass):
-                    stmt = insert(table).values(table_generator(dst_conn))
+                    stmt = insert(table).values(table_generator.generate_row(dst_conn))
                     dst_conn.execute(stmt)
                     row_counts[table.name] = row_counts.get(table.name, 0) + 1
                 dst_conn.commit()

datafaker/db_utils.py

@@ -10,8 +10,16 @@
 import sqlalchemy.dialects
 import yaml
 
-# pylint: disable=no-name-in-module
-from psycopg2.errors import UndefinedObject  # ty: ignore[unresolved-import]
+try:
+    # pylint: disable=no-name-in-module
+    from psycopg2.errors import UndefinedObject  # ty: ignore[unresolved-import]
+except ImportError:
+    # psycopg2 is only installed with the "postgres" extra; this error can
+    # only be raised by the psycopg2 driver, so it never matches otherwise.
+    class UndefinedObject:  # type: ignore[no-redef]
+        """Placeholder when psycopg2 is not installed."""
+
+
 from sqlalchemy import Connection, Engine, ForeignKey, create_engine, event, select
 from sqlalchemy.engine.interfaces import DBAPIConnection
 from sqlalchemy.exc import (
@@ -38,7 +46,9 @@
     get_ignored_table_names,
     get_vocabulary_table_names,
     logger,
+    make_async_dsn,
     make_foreign_key_name,
+    unqualify_fk_target,
 )
 
 # Define some types used repeatedly in the code base
@@ -139,23 +149,31 @@ def create_db_engine(
     **kwargs: Any,
 ) -> MaybeAsyncEngine:
     """Create a SQLAlchemy Engine."""
+    kwargs.setdefault("pool_pre_ping", True)
     try:
         if use_asyncio:
-            async_dsn = db_dsn.replace("postgresql://", "postgresql+asyncpg://")
-            engine: MaybeAsyncEngine = create_async_engine(async_dsn, **kwargs)
+            engine: MaybeAsyncEngine = create_async_engine(
+                make_async_dsn(db_dsn), **kwargs
+            )
         else:
             engine = create_engine(db_dsn, **kwargs)
     except NoSuchModuleError as exc:
         logger.error("Failed to connect to the database: %s", exc)
         logger.error("Perhaps the dialect '%s' is invalid.", db_dsn.split(":")[0])
         raise Exit(1) from exc
+    except ModuleNotFoundError as exc:
+        logger.error("Failed to connect to the database: %s", exc)
+        raise Exit(1) from exc
     except ValueError as exc:
         logger.error("DSN %s is malformed: %s", db_dsn, exc)
         raise Exit(1) from exc
 
     settings = {}
     if schema_name is not None:
-        settings["search_path"] = schema_name
+        if get_sync_engine(engine).dialect.name == "mssql":
+            engine = engine.execution_options(schema_translate_map={None: schema_name})
+        else:
+            settings["search_path"] = schema_name
     if parquet_dir is not None:
         joined = ",".join(_find_parquet_directories(parquet_dir))
         # double up single quotes
@@ -202,11 +220,11 @@ def create_db_engine_dst(
     return create_db_engine(db_dsn, schema_name, use_asyncio)
 
 
-def get_metadata(engine: Engine) -> MetaData:
+def get_metadata(engine: Engine, schema_name: Optional[str] = None) -> MetaData:
     """Get the MetaData object associated with the engine passed."""
     md = MetaData()
     try:
-        md.reflect(engine)
+        md.reflect(engine, schema=schema_name)
     except OperationalError as exc:
         logger.error("Cannot connect to database: %s", exc)
         raise Exit(1) from exc
@@ -417,10 +435,13 @@ def reinstate_vocab_foreign_key_constraints(
             ].items():
                 fk_targets = column_dict.get("foreign_keys", [])
                 if fk_targets:
+                    table_names = frozenset(meta_dict.get("tables", {}).keys())
                     fk = ForeignKeyConstraint(
                         columns=[column_name],
                         name=make_foreign_key_name(vocab_table_name, column_name),
-                        refcolumns=fk_targets,
+                        refcolumns=[
+                            unqualify_fk_target(t, table_names) for t in fk_targets
+                        ],
                     )
                     logger.debug("Restoring foreign key constraint %s", fk.name)
                     with Session(dst_engine) as session: