Merge branch 'main' into remove-generator-file

Author: Tim Band

datafaker/main.py

@@ -78,7 +78,7 @@ def _check_file_non_existence(file_path: Path) -> None:
     """Check that a given file does not exist. Exit with an error message if it does."""
     if file_path.exists():
         logger.error("%s should not already exist. Exiting...", file_path)
-        sys.exit(1)
+        raise Exit(1)
 
 
 def load_metadata_config(
@@ -202,6 +202,8 @@ def create_data(
         return
     except RuntimeError as e:
         logger.error(e.args[0])
+    except SettingsError as e:
+        logger.error(str(e))
     raise Exit(1)
 
 
@@ -554,7 +556,7 @@ def convert_table_names_to_tables(
                 "%s is not the name of a table in the destination database", name
             )
     if failed_count:
-        sys.exit(1)
+        raise Exit(1)
     return results
 
 
@@ -643,7 +645,7 @@ def dump_data(
                 "Must specify exactly one table if the output name is"
                 " specified, or specify an existing directory"
             )
-            sys.exit(1)
+            raise Exit(1)
     dst_dsn = get_destination_dsn()
     schema_name = get_destination_schema()
     config = read_config_file(config_file) if config_file is not None else {}
@@ -676,7 +678,7 @@ def validate_config(
         validate(config, schema_config)
     except ValidationError as e:
         logger.error(e)
-        sys.exit(1)
+        raise Exit(1) from e
     logger.debug("Config file is valid.")
 
 
@@ -772,7 +774,7 @@ def remove_tables(
             except InternalError as exc:
                 logger.error("Failed to drop tables: %s", exc)
                 logger.error("Please try again using the --all option.")
-                sys.exit(1)
+                raise Exit(1) from exc
         logger.debug("Tables dropped.")
     else:
         logger.info("Would remove tables if called with --yes.")

datafaker/make.py

@@ -11,6 +11,7 @@
 
 import pandas as pd
 import snsql
+import typer
 import yaml
 from black import FileMode, format_str
 from jinja2 import Environment, FileSystemLoader, Template
@@ -38,6 +39,7 @@
     create_db_engine,
     download_table,
     get_columns_assigned,
+    get_metadata,
     get_property,
     get_property_or_none,
     get_related_table_names,
@@ -605,9 +607,21 @@ def get_generation_info(
     row_generator_module_name = get_property_or_none(
         config, "row_generators_module", str
     )
+    if row_generator_module_name and "-" in row_generator_module_name:
+        logger.error(
+            "Row generator name %s should not contain a hyphen",
+            row_generator_module_name,
+        )
+        raise typer.Exit(1)
     story_generator_module_name = get_property_or_none(
         config, "story_generators_module", str
     )
+    if story_generator_module_name and "-" in story_generator_module_name:
+        logger.error(
+            "Story generator name %s should not contain a hyphen",
+            story_generator_module_name,
+        )
+        raise typer.Exit(1)
     object_instantiation: dict[str, Any] = get_property(
         config, "object_instantiation", {}
     )
@@ -703,8 +717,7 @@ def make_tables_file(
     """Construct the YAML file representing the schema."""
     engine = get_sync_engine(create_db_engine(db_dsn, schema_name=schema_name))
 
-    metadata = MetaData()
-    metadata.reflect(engine)
+    metadata = get_metadata(engine)
     meta_dict = metadata_to_dict(metadata, schema_name, engine, parquet_dir)
 
     if parquet_dir is not None:

datafaker/remove.py

@@ -6,6 +6,7 @@
 from datafaker.settings import get_destination_dsn, get_destination_schema
 from datafaker.utils import (
     create_db_engine,
+    get_metadata,
     get_sync_engine,
     get_vocabulary_table_names,
     logger,
@@ -67,6 +68,5 @@ def remove_db_tables(metadata: Optional[MetaData]) -> None:
         )
     )
     if metadata is None:
-        metadata = MetaData()
-        metadata.reflect(dst_engine)
+        metadata = get_metadata(dst_engine)
     metadata.drop_all(dst_engine)

datafaker/utils.py

@@ -32,7 +32,12 @@
 from jsonschema.validators import validate
 from sqlalchemy import Connection, Engine, ForeignKey, create_engine, event, select
 from sqlalchemy.engine.interfaces import DBAPIConnection
-from sqlalchemy.exc import IntegrityError, ProgrammingError
+from sqlalchemy.exc import (
+    IntegrityError,
+    NoSuchModuleError,
+    OperationalError,
+    ProgrammingError,
+)
 from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
 from sqlalchemy.orm import Session
 from sqlalchemy.schema import (
@@ -43,6 +48,9 @@
     MetaData,
     Table,
 )
+from typer import Exit
+
+from datafaker.settings import SettingsError
 
 # Define some types used repeatedly in the code base
 MaybeAsyncEngine = Union[Engine, AsyncEngine]
@@ -108,9 +116,15 @@ def import_file(file_path: str) -> ModuleType:
     """
     spec = importlib.util.spec_from_file_location("df", file_path)
     if spec is None or spec.loader is None:
-        raise ImportError(f"No loadable module at {file_path}")
+        raise SettingsError(f"No loadable module '{file_path}'")
     module = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(module)
+    try:
+        spec.loader.exec_module(module)
+    except ModuleNotFoundError as e:
+        logger.error("Failed to load module at %s with error:", file_path)
+        logger.error(e)
+    except FileNotFoundError as e:
+        raise SettingsError(f"No module found '{file_path}'") from e
     return module
 
 
@@ -193,11 +207,19 @@ def create_db_engine(
     **kwargs: Any,
 ) -> MaybeAsyncEngine:
     """Create a SQLAlchemy Engine."""
-    if use_asyncio:
-        async_dsn = db_dsn.replace("postgresql://", "postgresql+asyncpg://")
-        engine: MaybeAsyncEngine = create_async_engine(async_dsn, **kwargs)
-    else:
-        engine = create_engine(db_dsn, **kwargs)
+    try:
+        if use_asyncio:
+            async_dsn = db_dsn.replace("postgresql://", "postgresql+asyncpg://")
+            engine: MaybeAsyncEngine = create_async_engine(async_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 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:
@@ -248,6 +270,17 @@ def create_db_engine_dst(
     return create_db_engine(db_dsn, schema_name, use_asyncio)
 
 
+def get_metadata(engine: Engine) -> MetaData:
+    """Get the MetaData object associated with the engine passed."""
+    md = MetaData()
+    try:
+        md.reflect(engine)
+    except OperationalError as exc:
+        logger.error("Cannot connect to database: %s", exc)
+        raise Exit(1) from exc
+    return md
+
+
 def _find_parquet_directories(parquet_dir: Path) -> list[str]:
     """Find all the directories under ``parquet_dir`` that contain parquet files."""
     return [

tests/test_create.py

@@ -23,6 +23,7 @@
 from datafaker.make import FunctionCall, StoryGeneratorInfo
 from datafaker.populate import TableGenerator
 from datafaker.serialize_metadata import dict_to_metadata, metadata_to_dict
+from datafaker.settings import SettingsError
 from datafaker.utils import sorted_non_vocabulary_tables
 from tests.utils import DatafakerTestCase, GeneratesDBTestCase, RequiresDBTestCase
 
@@ -431,3 +432,34 @@ def test_unique_constraint_minimal(self) -> None:
             self.schema_name,
             metadata,
         )
+
+    def test_story_incorrect_name_minimal(self) -> None:
+        """Test we get a proper error message if the story generator module does not exist."""
+        config = {
+            "story_generators_module": "incorrect_module",
+        }
+        orm = {
+            "tables": {
+                "one": {
+                    "columns": {
+                        "id": {
+                            "primary": True,
+                            "type": "INTEGER",
+                        }
+                    }
+                }
+            }
+        }
+        metadata = dict_to_metadata(orm, config)
+        create_db_tables_into(metadata, self.dsn, self.schema_name)
+        self.assertRaises(
+            SettingsError,
+            create_db_data_into,
+            sorted_non_vocabulary_tables(metadata, config),
+            config,
+            None,
+            1,
+            self.dsn,
+            self.schema_name,
+            metadata,
+        )

tests/test_functional.py

@@ -5,6 +5,7 @@
 from pathlib import Path
 from typing import Any, Mapping
 
+import yaml
 from sqlalchemy import create_engine, inspect
 from typer.testing import CliRunner, Result
 
@@ -514,6 +515,71 @@ def test_create_schema(self) -> None:
         inspector = inspect(engine)
         self.assertTrue(inspector.has_schema(env["dst_schema"]))
 
+    def test_story_incorrect_name(self) -> None:
+        """Test we get a proper error message if the story generator module does not exist."""
+        config_file = "config_story_incorrect.yaml"
+        config = {
+            "story_generators_module": "incorrect_module",
+        }
+        with Path(config_file).open("w", encoding="utf-8") as fh:
+            fh.write(yaml.dump(config))
+        self.invoke(
+            "make-tables",
+            "--force",
+        )
+        self.invoke(
+            "create-tables",
+            "--config-file",
+            config_file,
+        )
+        completed_process = self.invoke(
+            "create-data",
+            "--config-file",
+            config_file,
+            expected_error="No module found 'incorrect_module",
+        )
+        self.assertReturnCode(completed_process, 1)
+
+    def test_story_hyphens_in_name(self) -> None:
+        """Test hyphens in story generator names cause an error to be emitted."""
+        config_file = "config_story_hyphens.yaml"
+        config = {
+            "story_generators_module": "story-generators",
+        }
+        with Path(config_file).open("w", encoding="utf-8") as fh:
+            fh.write(yaml.dump(config))
+        self.invoke(
+            "make-tables",
+            "--force",
+        )
+        completed_process = self.invoke(
+            "create-data",
+            "--config-file",
+            config_file,
+            expected_error="hyphen",
+        )
+        self.assertReturnCode(completed_process, 1)
+
+    def test_row_hyphens_in_name(self) -> None:
+        """Test hyphens in row generator names cause an error to be emitted."""
+        config_file = "config_row_hyphens.yaml"
+        config = {
+            "row_generators_module": "row-generators",
+        }
+        with Path(config_file).open("w", encoding="utf-8") as fh:
+            fh.write(yaml.dump(config))
+        self.invoke(
+            "make-tables",
+            "--force",
+        )
+        completed_process = self.invoke(
+            "create-data",
+            "--config-file",
+            config_file,
+            expected_error="hyphen",
+        )
+        self.assertReturnCode(completed_process, 1)
+
 
 class DuckDbFunctionalTestCase(DBFunctionalTestCaseBase):
     """End-to-end tests for the DuckDB workflow."""