Write the output of make-[tables|generators] to a file

.pre-commit-config.yaml

@@ -27,17 +27,17 @@ repos:
         language: system
         types: ['python']
         exclude: (?x)(
-          star|
-          tests/examples
+          tests/examples|
+          tests/workspace
           )
       - id: isort
         name: isort
         entry: poetry run isort
         language: system
         types: ['python']
         exclude: (?x)(
-          star|
-          tests/examples
+          tests/examples|
+          tests/workspace
           )
       - id: pylint
         name: Pylint
@@ -50,15 +50,15 @@ repos:
         language: system
         types: ['python']
         exclude: (?x)(
-          tests/|
-          docs/
+          docs/|
+          tests/
           )
       - id: mypy
         name: mypy
         entry: poetry run mypy --follow-imports=silent
         language: system
         exclude: (?x)(
-          star|
-          tests/examples
+          tests/examples|
+          tests/workspace
           )
         types: ['python']

.pylintrc

@@ -24,8 +24,8 @@ ignore=CVS
 
 # Add files or directories matching the regex patterns to the ignore-list. The
 # regex matches against paths.
-ignore-paths=sqlsynthgen/star.py,
-             tests/examples
+ignore-paths=tests/examples,
+             tests/workspace
 
 # Files or directories matching the regex patterns are skipped. The regex
 # matches against base names, not paths.

sqlsynthgen/main.py

@@ -1,51 +1,27 @@
 """Entrypoint for the SQLSynthGen package."""
 import sys
-from importlib import import_module
 from pathlib import Path
-from subprocess import CalledProcessError, run
 from sys import stderr
-from types import ModuleType
-from typing import Any, Optional
+from typing import Final, Optional
 
 import typer
-import yaml
 
 from sqlsynthgen.create import create_db_data, create_db_tables, create_db_vocab
-from sqlsynthgen.make import make_generators_from_tables
+from sqlsynthgen.make import make_generators_from_tables, make_tables_file
 from sqlsynthgen.settings import get_settings
+from sqlsynthgen.utils import import_file, read_yaml_file
 
-app = typer.Typer()
-
-
-def import_file(file_path: str) -> ModuleType:
-    """Import a file.
-
-    This utility function returns
-    the file at file_path as a module
-
-    Args:
-        file_path (str): Path to file to be imported
-
-    Returns:
-        ModuleType
-    """
-    file_path_path = Path(file_path)
-    module_path = ".".join(file_path_path.parts[:-1] + (file_path_path.stem,))
-    return import_module(module_path)
+ORM_FILENAME: Final[str] = "orm.py"
+SSG_FILENAME: Final[str] = "ssg.py"
 
-
-def read_yaml_file(path: str) -> Any:
-    """Read a yaml file in to dictionary, given a path."""
-    with open(path, "r", encoding="utf8") as f:
-        config = yaml.safe_load(f)
-    return config
+app = typer.Typer()
 
 
 @app.command()
 def create_data(
-    orm_file: str = typer.Argument(...),
-    ssg_file: str = typer.Argument(...),
-    num_passes: int = typer.Argument(...),
+    orm_file: str = typer.Option(ORM_FILENAME),
+    ssg_file: str = typer.Option(SSG_FILENAME),
+    num_passes: int = typer.Option(1),
 ) -> None:
     """Populate schema with synthetic data.
 
@@ -63,15 +39,14 @@ def create_data(
     Final input is the number of rows required.
 
     Example:
-        $ sqlsynthgen create-data example_orm.py expected_ssg.py 100
+        $ sqlsynthgen create-data
 
     Args:
-        orm_file (str): Path to object relational model.
-        ssg_file (str): Path to sqlsyngen output.
+        orm_file (str): Name of Python ORM file.
+          Must be in the current working directory.
+        ssg_file (str): Name of generators file.
+          Must be in the current working directory.
         num_passes (int): Number of passes to make.
-
-    Returns:
-        None
     """
     orm_module = import_file(orm_file)
     ssg_module = import_file(ssg_file)
@@ -81,36 +56,42 @@ def create_data(
 
 
 @app.command()
-def create_vocab(ssg_file: str = typer.Argument(...)) -> None:
-    """Create tables using the SQLAlchemy file."""
+def create_vocab(ssg_file: str = typer.Option(SSG_FILENAME)) -> None:
+    """Create tables using the SQLAlchemy file.
+
+    Example:
+        $ sqlsynthgen create-vocab
+
+    Args:
+        ssg_file (str): Name of generators file.
+          Must be in the current working directory.
+    """
     ssg_module = import_file(ssg_file)
     create_db_vocab(ssg_module.sorted_vocab)
 
 
 @app.command()
-def create_tables(orm_file: str = typer.Argument(...)) -> None:
+def create_tables(orm_file: str = typer.Option(ORM_FILENAME)) -> None:
     """Create schema from Python classes.
 
     This CLI command creates Postgresql schema using object relational model
     declared as Python tables. (eg.)
 
     Example:
-        $ sqlsynthgen create-tables example_orm.py
+        $ sqlsynthgen create-tables
 
     Args:
-        orm_file (str): Path to Python tables file.
-
-    Returns:
-        None
-
+        orm_file (str): Name of Python ORM file.
+          Must be in the current working directory.
     """
     orm_module = import_file(orm_file)
     create_db_tables(orm_module.Base.metadata)
 
 
 @app.command()
 def make_generators(
-    orm_file: str = typer.Argument(...),
+    orm_file: str = typer.Option(ORM_FILENAME),
+    ssg_file: str = typer.Option(SSG_FILENAME),
     config_file: Optional[str] = typer.Argument(None),
 ) -> None:
     """Make a SQLSynthGen file of generator classes.
@@ -119,55 +100,51 @@ def make_generators(
     returns a set of synthetic data generators for each attribute
 
     Example:
-        $ sqlsynthgen make-generators example_orm.py
+        $ sqlsynthgen make-generators
 
     Args:
-        orm_file (str): Path to Python tables file.
+        orm_file (str): Name of Python ORM file.
+          Must be in the current working directory.
+        ssg_file (str): Path to write the generators file to.
         config_file (str): Path to configuration file.
     """
+    ssg_file_path = Path(ssg_file)
+    if ssg_file_path.exists():
+        print(f"{ssg_file} should not already exist. Exiting...", file=stderr)
+        sys.exit(1)
+
     orm_module = import_file(orm_file)
     generator_config = read_yaml_file(config_file) if config_file is not None else {}
     result = make_generators_from_tables(orm_module, generator_config)
-    print(result)
+
+    ssg_file_path.write_text(result, encoding="utf-8")
 
 
 @app.command()
-def make_tables() -> None:
+def make_tables(
+    orm_file: str = typer.Option(ORM_FILENAME),
+) -> None:
     """Make a SQLAlchemy file of Table classes.
 
     This CLI command deploys sqlacodegen to discover a
-    schema structure, and generates a object relational model declared
+    schema structure, and generates an object relational model declared
     as Python classes.
 
     Example:
         $ sqlsynthgen make_tables
-    """
-    settings = get_settings()
-
-    command = ["sqlacodegen"]
-
-    if settings.src_schema:
-        command.append(f"--schema={settings.src_schema}")
-
-    command.append(str(get_settings().src_postgres_dsn))
 
-    try:
-        completed_process = run(
-            command, capture_output=True, encoding="utf-8", check=True
-        )
-    except CalledProcessError as e:
-        print(e.stderr, file=stderr)
-        sys.exit(e.returncode)
+    Args:
+        orm_file (str): Path to write the Python ORM file.
+    """
+    orm_file_path = Path(orm_file)
+    if orm_file_path.exists():
+        print(f"{orm_file} should not already exist. Exiting...", file=stderr)
+        sys.exit(1)
 
-    # sqlacodegen falls back on Tables() for tables without PKs,
-    # but we don't explicitly support Tables and behaviour is unpredictable.
-    if " = Table(" in completed_process.stdout:
-        print(
-            "WARNING: Table without PK detected. sqlsynthgen may not be able to continue.",
-            file=stderr,
-        )
+    settings = get_settings()
 
-    print(completed_process.stdout)
+    content = make_tables_file(str(settings.src_postgres_dsn), settings.src_schema)
+    orm_file_path.write_text(content, encoding="utf-8")
 
 
 if __name__ == "__main__":

sqlsynthgen/make.py

@@ -1,16 +1,18 @@
 """Functions to make a module of generator classes."""
-import csv
 import inspect
-from pathlib import Path
+import sys
+from subprocess import CalledProcessError, run
+from sys import stderr
 from types import ModuleType
 from typing import Any, Final, Optional
 
 from mimesis.providers.base import BaseProvider
-from sqlalchemy import create_engine, select
+from sqlalchemy import create_engine
 from sqlalchemy.sql import sqltypes
 
 from sqlsynthgen import providers
 from sqlsynthgen.settings import get_settings
+from sqlsynthgen.utils import download_table
 
 HEADER_TEXT: str = "\n".join(
     (
@@ -132,20 +134,6 @@ def _add_generator_for_table(
     return content, new_class_name
 
 
-def _download_table(table: Any, engine: Any) -> None:
-    """Download a table and store it as a .csv file."""
-    stmt = select([table])
-    with engine.connect() as conn:
-        result = list(conn.execute(stmt))
-    with Path(table.fullname + ".csv").open(
-        "w", newline="", encoding="utf-8"
-    ) as csvfile:
-        writer = csv.writer(csvfile, delimiter=",")
-        writer.writerow([x.name for x in table.columns])
-        for row in result:
-            writer.writerow(row)
-
-
 def make_generators_from_tables(
     tables_module: ModuleType, generator_config: dict
 ) -> str:
@@ -162,7 +150,7 @@ def make_generators_from_tables(
     new_content += f"\nimport {tables_module.__name__}"
     generator_module_name = generator_config.get("custom_generators_module", None)
     if generator_module_name is not None:
-        new_content += f"\nfrom . import {generator_module_name}"
+        new_content += f"\nimport {generator_module_name}"
 
     sorted_generators = "[\n"
     sorted_vocab = "[\n"
@@ -185,7 +173,7 @@ def make_generators_from_tables(
             )
             sorted_vocab += f"{INDENTATION}{class_name.lower()}_vocab,\n"
 
-            _download_table(table, engine)
+            download_table(table, engine)
 
         else:
             new_content, new_generator_name = _add_generator_for_table(
@@ -200,3 +188,34 @@ def make_generators_from_tables(
     new_content += "\n\n" + "sorted_vocab = " + sorted_vocab + "\n"
 
     return new_content
+
+
+def make_tables_file(db_dsn: str, schema_name: Optional[str]) -> str:
+    """Write a file with the SQLAlchemy ORM classes.
+
+    Exists with an error if sqlacodegen is unsuccessful.
+    """
+    command = ["sqlacodegen"]
+
+    if schema_name:
+        command.append(f"--schema={schema_name}")
+
+    command.append(db_dsn)
+
+    try:
+        completed_process = run(
+            command, capture_output=True, encoding="utf-8", check=True
+        )
+    except CalledProcessError as e:
+        print(e.stderr, file=stderr)
+        sys.exit(e.returncode)
+
+    # sqlacodegen falls back on Tables() for tables without PKs,
+    # but we don't explicitly support Tables and behaviour is unpredictable.
+    if " = Table(" in completed_process.stdout:
+        print(
+            "WARNING: Table without PK detected. sqlsynthgen may not be able to continue.",
+            file=stderr,
+        )
+
+    return completed_process.stdout