Added support for passing arguments to create_engine() (#391)

Author: LajosCseppento

CHANGES.rst

@@ -5,6 +5,8 @@ Version history
 
 - Type annotations for ARRAY column attributes now include the Python type of
   the array elements
+- Added support for specifying engine arguments via ``--engine-arg``
+  (PR by @LajosCseppento)
 
 **3.0.0**
 

README.rst

@@ -67,6 +67,9 @@ Examples::
     sqlacodegen postgresql:///some_local_db
     sqlacodegen --generator tables mysql+pymysql://user:password@localhost/dbname
     sqlacodegen --generator dataclasses sqlite:///database.db
+    # --engine-arg values are parsed with ast.literal_eval
+    sqlacodegen oracle+oracledb://user:[email protected]:1521/XE --engine-arg thick_mode=True
+    sqlacodegen oracle+oracledb://user:[email protected]:1521/XE --engine-arg thick_mode=True --engine-arg connect_args='{"user": "user", "dsn": "..."}'
 
 To see the list of generic options::
 

src/sqlacodegen/cli.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
 import argparse
+import ast
 import sys
 from contextlib import ExitStack
-from typing import TextIO
+from typing import Any, TextIO
 
 from sqlalchemy.engine import create_engine
 from sqlalchemy.schema import MetaData
@@ -29,6 +30,28 @@
     from importlib.metadata import entry_points, version
 
 
+def _parse_engine_arg(arg_str: str) -> tuple[str, Any]:
+    if "=" not in arg_str:
+        raise argparse.ArgumentTypeError("engine-arg must be in key=value format")
+
+    key, value = arg_str.split("=", 1)
+    try:
+        value = ast.literal_eval(value)
+    except Exception:
+        pass  # Leave as string if literal_eval fails
+
+    return key, value
+
+
+def _parse_engine_args(arg_list: list[str]) -> dict[str, Any]:
+    result = {}
+    for arg in arg_list or []:
+        key, value = _parse_engine_arg(arg)
+        result[key] = value
+
+    return result
+
+
 def main() -> None:
     generators = {ep.name: ep for ep in entry_points(group="sqlacodegen.generators")}
     parser = argparse.ArgumentParser(
@@ -58,6 +81,17 @@ def main() -> None:
         action="store_true",
         help="ignore views (always true for sqlmodels generator)",
     )
+    parser.add_argument(
+        "--engine-arg",
+        action="append",
+        help=(
+            "engine arguments in key=value format, e.g., "
+            '--engine-arg=connect_args=\'{"user": "scott"}\' '
+            "--engine-arg thick_mode=true or "
+            '--engine-arg thick_mode=\'{"lib_dir": "/path"}\' '
+            "(values are parsed with ast.literal_eval)"
+        ),
+    )
     parser.add_argument("--outfile", help="file to write output to (default: stdout)")
     args = parser.parse_args()
 
@@ -80,7 +114,8 @@ def main() -> None:
         print(f"Using pgvector {version('pgvector')}")
 
     # Use reflection to fill in the metadata
-    engine = create_engine(args.url)
+    engine_args = _parse_engine_args(args.engine_arg)
+    engine = create_engine(args.url, **engine_args)
     metadata = MetaData()
     tables = args.tables.split(",") if args.tables else None
     schemas = args.schemas.split(",") if args.schemas else [None]

tests/test_cli.py

@@ -150,6 +150,72 @@ class Foo(SQLModel, table=True):
     )
 
 
+def test_cli_engine_arg(db_path: Path, tmp_path: Path) -> None:
+    output_path = tmp_path / "outfile"
+    subprocess.run(
+        [
+            "sqlacodegen",
+            f"sqlite:///{db_path}",
+            "--generator",
+            "tables",
+            "--engine-arg",
+            'connect_args={"timeout": 10}',
+            "--outfile",
+            str(output_path),
+        ],
+        check=True,
+    )
+
+    assert (
+        output_path.read_text()
+        == """\
+from sqlalchemy import Column, Integer, MetaData, Table, Text
+
+metadata = MetaData()
+
+
+t_foo = Table(
+    'foo', metadata,
+    Column('id', Integer, primary_key=True),
+    Column('name', Text, nullable=False)
+)
+"""
+    )
+
+
+def test_cli_invalid_engine_arg(db_path: Path, tmp_path: Path) -> None:
+    output_path = tmp_path / "outfile"
+
+    # Expect exception:
+    # TypeError: 'this_arg_does_not_exist' is an invalid keyword argument for Connection()
+    with pytest.raises(subprocess.CalledProcessError) as exc_info:
+        subprocess.run(
+            [
+                "sqlacodegen",
+                f"sqlite:///{db_path}",
+                "--generator",
+                "tables",
+                "--engine-arg",
+                'connect_args={"this_arg_does_not_exist": 10}',
+                "--outfile",
+                str(output_path),
+            ],
+            check=True,
+            capture_output=True,
+        )
+
+    if sys.version_info < (3, 13):
+        assert (
+            "'this_arg_does_not_exist' is an invalid keyword argument"
+            in exc_info.value.stderr.decode()
+        )
+    else:
+        assert (
+            "got an unexpected keyword argument 'this_arg_does_not_exist'"
+            in exc_info.value.stderr.decode()
+        )
+
+
 def test_main() -> None:
     expected_version = version("sqlacodegen")
     completed = subprocess.run(