Add naming constraint as cli args

Author: leonarduschen

src/sqlacodegen/cli.py

@@ -3,7 +3,7 @@
 import argparse
 import sys
 from contextlib import ExitStack
-from typing import TextIO
+from typing import Sequence, TextIO
 
 from sqlalchemy.engine import create_engine
 from sqlalchemy.schema import MetaData
@@ -14,6 +14,18 @@
     from importlib.metadata import entry_points, version
 
 
+def parse_naming_convs(naming_convs: Sequence[str]) -> dict[str, str]:
+    d = {}
+    for naming_conv in naming_convs:
+        try:
+            key, value = naming_conv.split("=", 1)
+        except ValueError:
+            raise ValueError('Naming convention must be in "key=template" format')
+
+        d[key] = value
+    return d
+
+
 def main() -> None:
     generators = {ep.name: ep for ep in entry_points(group="sqlacodegen.generators")}
     parser = argparse.ArgumentParser(
@@ -40,6 +52,13 @@ def main() -> None:
     )
     parser.add_argument("--noviews", action="store_true", help="ignore views")
     parser.add_argument("--outfile", help="file to write output to (default: stdout)")
+    parser.add_argument(
+        "--conv",
+        nargs="*",
+        help='constraint naming conventions in "key=template" format \
+e.g., --conv "pk=pk_%%(table_name)s" "uq=uq_%%(table_name)s_%%(column_0_name)s"',
+    )
+
     args = parser.parse_args()
 
     if args.version:
@@ -58,6 +77,11 @@ def main() -> None:
     for schema in schemas:
         metadata.reflect(engine, schema, not args.noviews, tables)
 
+    # Naming convention must be added after reflection to
+    # avoid the token %(constraint_name)s duplicating the name
+    if args.conv:
+        metadata.naming_convention = parse_naming_convs(args.conv)
+
     # Instantiate the generator
     generator_class = generators[args.generator].load()
     generator = generator_class(metadata, engine, set(args.option or ()))

tests/test_cli.py

@@ -172,3 +172,36 @@ def test_main() -> None:
         check=True,
     )
     assert completed.stdout.decode().strip() == expected_version
+
+
+@pytest.fixture
+def empty_db_path(tmp_path: Path) -> Path:
+    path = tmp_path / "test.db"
+
+    return path
+
+
+def test_naming_convention(empty_db_path: Path, tmp_path: Path) -> None:
+    output_path = tmp_path / "outfile"
+    subprocess.run(
+        [
+            "sqlacodegen",
+            f"sqlite:///{empty_db_path}",
+            "--outfile",
+            str(output_path),
+            "--conv",
+            "pk=pk_%(table_name)s",
+        ],
+        check=True,
+    )
+
+    assert (
+        output_path.read_text()
+        == """\
+from sqlalchemy import MetaData
+
+metadata = MetaData()
+metadata.naming_convention = {'pk': 'pk_%(table_name)s'}
+
+"""
+    )