Repository: tedivm/paracelsus

Issue:
Option to Preserve Column Order
Thanks for the great tool!

Right now the table fields are sorted with PK fields first, then FK fields, then the rest of the fields.  In my case I'd like to preserve the order columns are defined in the class in the ouput diagrams.  I made a few mods and built the project locally to implement this for my immediate needs.  

Would you be open to a PR that leaves the current sorting behavior as the default and adds an option to preserve the field order in the tables?

For instance:
```sh
paracelsus graph example_app.models.base:Base \
  --import-module "example_app.models.users" \
  --import-module "example_app.models.posts" \
  --import-module "example_app.models.comments" \
  --preserve-order
```
or 
```toml
[tool.paracelsus]
base = "example.base:Base"
imports = [
  "example.models"
]
preserve-order = true
```
I have this working for `graph` and mermaid output but would need to verify and add some tests for `inject` and dot output.  I don't really have a time frame for this but thought it might be useful to others and would allow me to incorporate paracelsus into an automated build process.

Hints from issue comments:
@tedivm: I would absolutely be open to a PR allowing custom sort options. Rather than a `--preserve-order` flag I'd prefer a sort flag that takes options though: `--sort=default`, `--sort=preserve`. That way we can add other options in the future.@tedivm: Actually can you make it `--column-sort` instead? I'm thinking we'll have table sorts as well in the future (the difference being mermaids `TB`, `BT`, `RL`, and `LR`). I've already run into a few cases where changing the table sorting would have been helpful.

Solution patch:
diff --git a/README.md b/README.md
index 454da39..6b16aaf 100644
--- a/README.md
+++ b/README.md
@@ -91,6 +91,32 @@ paracelsus graph example_app.models.base:Base \
   --include-tables "posts"
 ```
 
+### Specify Column Sort Order
+
+By default Paracelsus will sort the columns in all models such as primary keys are first, foreign keys are next and all other
+columns are sorted alphabetically by name. 
+
+```bash
+paracelsus graph example_app.models.base:Base \
+  --import-module "example_app.models.users" \
+```
+
+produces the same results as:
+
+```bash
+paracelsus graph example_app.models.base:Base \
+  --import-module "example_app.models.users" \
+  --column-sort key-based
+```
+
+Pass the --column-sort option to change this behavior. To preserve the order of fields present in the models use "preserve-order":
+
+```bash
+paracelsus graph example_app.models.base:Base \
+  --import-module "example_app.models.users" \
+  --column-sort preserve-order
+```
+
 ### Generate Mermaid Diagrams
 
 
diff --git a/paracelsus/cli.py b/paracelsus/cli.py
index 88e891f..22e0186 100644
--- a/paracelsus/cli.py
+++ b/paracelsus/cli.py
@@ -20,7 +20,12 @@ class Formats(str, Enum):
     gv = "gv"
 
 
-def get_base_class(base_class_path: str | None, settings=Dict[str, Any] | None) -> str:
+class ColumnSorts(str, Enum):
+    key_based = "key-based"
+    preserve = "preserve-order"
+
+
+def get_base_class(base_class_path: str | None, settings: Dict[str, Any] | None) -> str:
     if base_class_path:
         return base_class_path
     if not settings:
@@ -63,6 +68,12 @@ def graph(
     format: Annotated[
         Formats, typer.Option(help="The file format to output the generated graph to.")
     ] = Formats.mermaid,
+    column_sort: Annotated[
+        ColumnSorts,
+        typer.Option(
+            help="Specifies the method of sorting columns in diagrams.",
+        ),
+    ] = ColumnSorts.key_based,
 ):
     settings = get_pyproject_settings()
     base_class = get_base_class(base_class_path, settings)
@@ -78,6 +89,7 @@ def graph(
             exclude_tables=set(exclude_tables + settings.get("exclude_tables", [])),
             python_dir=python_dir,
             format=format.value,
+            column_sort=column_sort,
         )
     )
 
@@ -140,6 +152,12 @@ def inject(
             help="Perform a dry run and return a success code of 0 if there are no changes or 1 otherwise.",
         ),
     ] = False,
+    column_sort: Annotated[
+        ColumnSorts,
+        typer.Option(
+            help="Specifies the method of sorting columns in diagrams.",
+        ),
+    ] = ColumnSorts.key_based,
 ):
     settings = get_pyproject_settings()
     if "imports" in settings:
@@ -153,6 +171,7 @@ def inject(
         exclude_tables=set(exclude_tables + settings.get("exclude_tables", [])),
         python_dir=python_dir,
         format=format.value,
+        column_sort=column_sort,
     )
 
     comment_format = transformers[format].comment_format  # type: ignore
diff --git a/paracelsus/graph.py b/paracelsus/graph.py
index ed2d048..5f98048 100644
--- a/paracelsus/graph.py
+++ b/paracelsus/graph.py
@@ -25,6 +25,7 @@ def get_graph_string(
     exclude_tables: Set[str],
     python_dir: List[Path],
     format: str,
+    column_sort: str,
 ) -> str:
     # Update the PYTHON_PATH to allow more module imports.
     sys.path.append(str(os.getcwd()))
@@ -60,7 +61,7 @@ def get_graph_string(
     filtered_metadata = filter_metadata(metadata=metadata, include_tables=include_tables)
 
     # Save the graph structure to string.
-    return str(transformer(filtered_metadata))
+    return str(transformer(filtered_metadata, column_sort))
 
 
 def resolve_included_tables(
diff --git a/paracelsus/transformers/dot.py b/paracelsus/transformers/dot.py
index 19aac66..4b90a89 100644
--- a/paracelsus/transformers/dot.py
+++ b/paracelsus/transformers/dot.py
@@ -2,7 +2,7 @@ import pydot  # type: ignore
 import logging
 from sqlalchemy.sql.schema import MetaData, Table
 
-from . import utils
+from .utils import sort_columns
 
 
 logger = logging.getLogger(__name__)
@@ -12,10 +12,12 @@ class Dot:
     comment_format: str = "dot"
     metadata: MetaData
     graph: pydot.Dot
+    column_sort: str
 
-    def __init__(self, metaclass: MetaData) -> None:
+    def __init__(self, metaclass: MetaData, column_sort: str) -> None:
         self.metadata = metaclass
         self.graph = pydot.Dot("database", graph_type="graph")
+        self.column_sort = column_sort
 
         for table in self.metadata.tables.values():
             node = pydot.Node(name=table.name)
@@ -55,7 +57,7 @@ class Dot:
 
     def _table_label(self, table: Table) -> str:
         column_output = ""
-        columns = sorted(table.columns, key=utils.column_sort_key)
+        columns = sort_columns(table_columns=table.columns, column_sort=self.column_sort)
         for column in columns:
             attributes = set([])
             if column.primary_key:
diff --git a/paracelsus/transformers/mermaid.py b/paracelsus/transformers/mermaid.py
index 76783d4..df92674 100644
--- a/paracelsus/transformers/mermaid.py
+++ b/paracelsus/transformers/mermaid.py
@@ -1,7 +1,7 @@
 import logging
 from sqlalchemy.sql.schema import Column, MetaData, Table
 
-from . import utils
+from .utils import sort_columns
 
 
 logger = logging.getLogger(__name__)
@@ -10,14 +10,16 @@ logger = logging.getLogger(__name__)
 class Mermaid:
     comment_format: str = "mermaid"
     metadata: MetaData
+    column_sort: str
 
-    def __init__(self, metaclass: MetaData) -> None:
+    def __init__(self, metaclass: MetaData, column_sort: str) -> None:
         self.metadata = metaclass
+        self.column_sort = column_sort
 
     def _table(self, table: Table) -> str:
         output = f"  {table.name}"
         output += " {\n"
-        columns = sorted(table.columns, key=utils.column_sort_key)
+        columns = sort_columns(table_columns=table.columns, column_sort=self.column_sort)
         for column in columns:
             output += self._column(column)
         output += "  }\n\n"
diff --git a/paracelsus/transformers/utils.py b/paracelsus/transformers/utils.py
index aacdd95..2e7d576 100644
--- a/paracelsus/transformers/utils.py
+++ b/paracelsus/transformers/utils.py
@@ -1,7 +1,8 @@
 from sqlalchemy.sql.schema import Column
+from sqlalchemy.sql import ColumnCollection
 
 
-def column_sort_key(column: Column):
+def key_based_column_sort(column: Column) -> str:
     if column.primary_key:
         prefix = "01"
     elif len(column.foreign_keys):
@@ -9,3 +10,13 @@ def column_sort_key(column: Column):
     else:
         prefix = "03"
     return f"{prefix}_{column.name}"
+
+
+def sort_columns(table_columns: ColumnCollection, column_sort: str) -> list:
+    match column_sort:
+        case "preserve-order":
+            columns = [column for column in table_columns]
+        case _:
+            columns = sorted(table_columns, key=key_based_column_sort)
+
+    return columns


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 4b326b8..1ac4664 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -27,8 +27,8 @@ def metaclass():
         __tablename__ = "posts"
 
         id = mapped_column(Uuid, primary_key=True, default=uuid4())
-        author = mapped_column(ForeignKey(User.id), nullable=False)
         created = mapped_column(DateTime, nullable=False, default=datetime.now(UTC))
+        author = mapped_column(ForeignKey(User.id), nullable=False)
         live = mapped_column(Boolean, default=False, comment="True if post is published")
         content = mapped_column(Text, default="")
 
@@ -36,9 +36,9 @@ def metaclass():
         __tablename__ = "comments"
 
         id = mapped_column(Uuid, primary_key=True, default=uuid4())
+        created = mapped_column(DateTime, nullable=False, default=datetime.now(UTC))
         post = mapped_column(Uuid, ForeignKey(Post.id), default=uuid4())
         author = mapped_column(ForeignKey(User.id), nullable=False)
-        created = mapped_column(DateTime, nullable=False, default=datetime.now(UTC))
         live = mapped_column(Boolean, default=False)
         content = mapped_column(Text, default="")
 
@@ -51,3 +51,74 @@ def package_path():
     with tempfile.TemporaryDirectory() as package_path:
         shutil.copytree(template_path, package_path, dirs_exist_ok=True)
         yield Path(package_path)
+
+
+@pytest.fixture()
+def mermaid_full_string_preseve_column_sort() -> str:
+    return """erDiagram
+  users {
+    CHAR(32) id PK
+    VARCHAR(100) display_name "nullable"
+    DATETIME created
+  }
+
+  posts {
+    CHAR(32) id PK
+    DATETIME created
+    CHAR(32) author FK
+    BOOLEAN live "True if post is published,nullable"
+    TEXT content "nullable"
+  }
+
+  comments {
+    CHAR(32) id PK
+    DATETIME created
+    CHAR(32) post FK "nullable"
+    CHAR(32) author FK
+    BOOLEAN live "nullable"
+    TEXT content "nullable"
+  }
+
+  users ||--o{ posts : author
+  posts ||--o{ comments : post
+  users ||--o{ comments : author
+"""
+
+
+@pytest.fixture()
+def dot_full_string_preseve_column_sort() -> str:
+    return """graph database {
+users [label=<
+    <table border="0" cellborder="1" cellspacing="0" cellpadding="4">
+        <tr><td colspan="3" bgcolor="lightblue"><b>users</b></td></tr>
+        <tr><td align="left">CHAR(32)</td><td align="left">id</td><td>Primary Key</td></tr>
+        <tr><td align="left">VARCHAR(100)</td><td align="left">display_name</td><td></td></tr>
+        <tr><td align="left">DATETIME</td><td align="left">created</td><td></td></tr>
+    </table>
+>, margin=0, shape=none];
+posts [label=<
+    <table border="0" cellborder="1" cellspacing="0" cellpadding="4">
+        <tr><td colspan="3" bgcolor="lightblue"><b>posts</b></td></tr>
+        <tr><td align="left">CHAR(32)</td><td align="left">id</td><td>Primary Key</td></tr>
+        <tr><td align="left">DATETIME</td><td align="left">created</td><td></td></tr>
+        <tr><td align="left">CHAR(32)</td><td align="left">author</td><td>Foreign Key</td></tr>
+        <tr><td align="left">BOOLEAN</td><td align="left">live</td><td></td></tr>
+        <tr><td align="left">TEXT</td><td align="left">content</td><td></td></tr>
+    </table>
+>, margin=0, shape=none];
+users -- posts  [arrowhead=crow, arrowtail=none, dir=both, label=author];
+comments [label=<
+    <table border="0" cellborder="1" cellspacing="0" cellpadding="4">
+        <tr><td colspan="3" bgcolor="lightblue"><b>comments</b></td></tr>
+        <tr><td align="left">CHAR(32)</td><td align="left">id</td><td>Primary Key</td></tr>
+        <tr><td align="left">DATETIME</td><td align="left">created</td><td></td></tr>
+        <tr><td align="left">CHAR(32)</td><td align="left">post</td><td>Foreign Key</td></tr>
+        <tr><td align="left">CHAR(32)</td><td align="left">author</td><td>Foreign Key</td></tr>
+        <tr><td align="left">BOOLEAN</td><td align="left">live</td><td></td></tr>
+        <tr><td align="left">TEXT</td><td align="left">content</td><td></td></tr>
+    </table>
+>, margin=0, shape=none];
+posts -- comments  [arrowhead=crow, arrowtail=none, dir=both, label=post];
+users -- comments  [arrowhead=crow, arrowtail=none, dir=both, label=author];
+}
+"""
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 20bcd24..667a253 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -1,3 +1,5 @@
+import pytest
+
 from typer.testing import CliRunner
 
 from paracelsus.cli import app
@@ -8,6 +10,17 @@ runner = CliRunner()
 
 
 def test_graph(package_path):
+    result = runner.invoke(
+        app,
+        ["graph", "example.base:Base", "--import-module", "example.models", "--python-dir", str(package_path)],
+    )
+
+    assert result.exit_code == 0
+    mermaid_assert(result.stdout)
+
+
+@pytest.mark.parametrize("column_sort_arg", ["key-based", "preserve-order"])
+def test_graph_column_sort(package_path, column_sort_arg):
     result = runner.invoke(
         app,
         [
@@ -17,6 +30,8 @@ def test_graph(package_path):
             "example.models",
             "--python-dir",
             str(package_path),
+            "--column-sort",
+            column_sort_arg,
         ],
     )
 
@@ -99,6 +114,29 @@ def test_inject(package_path):
         mermaid_assert(readme)
 
 
+@pytest.mark.parametrize("column_sort_arg", ["key-based", "preserve-order"])
+def test_inject_column_sort(package_path, column_sort_arg):
+    result = runner.invoke(
+        app,
+        [
+            "inject",
+            str(package_path / "README.md"),
+            "example.base:Base",
+            "--import-module",
+            "example.models",
+            "--python-dir",
+            str(package_path),
+            "--column-sort",
+            column_sort_arg,
+        ],
+    )
+    assert result.exit_code == 0
+
+    with open(package_path / "README.md") as fp:
+        readme = fp.read()
+        mermaid_assert(readme)
+
+
 def test_version():
     result = runner.invoke(app, ["version"])
     assert result.exit_code == 0
diff --git a/tests/test_graph.py b/tests/test_graph.py
index cdc4fdf..e08ca53 100644
--- a/tests/test_graph.py
+++ b/tests/test_graph.py
@@ -1,9 +1,12 @@
+import pytest
+
 from paracelsus.graph import get_graph_string
 
 from .utils import mermaid_assert
 
 
-def test_get_graph_string(package_path):
+@pytest.mark.parametrize("column_sort_arg", ["key-based", "preserve-order"])
+def test_get_graph_string(column_sort_arg, package_path):
     graph_string = get_graph_string(
         base_class_path="example.base:Base",
         import_module=["example.models"],
@@ -11,6 +14,7 @@ def test_get_graph_string(package_path):
         exclude_tables=set(),
         python_dir=[package_path],
         format="mermaid",
+        column_sort=column_sort_arg,
     )
     mermaid_assert(graph_string)
 
@@ -23,6 +27,7 @@ def test_get_graph_string_with_exclude(package_path):
         include_tables=set(),
         exclude_tables={"comments"},
         python_dir=[package_path],
+        column_sort="key-based",
         format="mermaid",
     )
     assert "comments {" not in graph_string
@@ -38,6 +43,7 @@ def test_get_graph_string_with_exclude(package_path):
         exclude_tables={"users", "comments"},
         python_dir=[package_path],
         format="mermaid",
+        column_sort="key-based",
     )
     assert "posts {" in graph_string
     assert "users ||--o{ posts" not in graph_string
@@ -51,6 +57,7 @@ def test_get_graph_string_with_include(package_path):
         include_tables={"users", "posts"},
         exclude_tables=set(),
         python_dir=[package_path],
+        column_sort="key-based",
         format="mermaid",
     )
     assert "comments {" not in graph_string
@@ -66,6 +73,7 @@ def test_get_graph_string_with_include(package_path):
         include_tables={"posts"},
         exclude_tables=set(),
         python_dir=[package_path],
+        column_sort="key-based",
         format="mermaid",
     )
     assert "posts {" in graph_string
diff --git a/tests/transformers/test_dot.py b/tests/transformers/test_dot.py
index d66b63c..c241dc9 100644
--- a/tests/transformers/test_dot.py
+++ b/tests/transformers/test_dot.py
@@ -1,18 +1,14 @@
 from paracelsus.transformers.dot import Dot
 
+from ..utils import dot_assert
 
-def test_mermaid(metaclass):
-    mermaid = Dot(metaclass=metaclass)
-    graph_string = str(mermaid)
 
-    assert '<tr><td colspan="3" bgcolor="lightblue"><b>users</b></td></tr>' in graph_string
-    assert '<tr><td colspan="3" bgcolor="lightblue"><b>posts</b></td></tr>' in graph_string
-    assert '<tr><td colspan="3" bgcolor="lightblue"><b>comments</b></td></tr>' in graph_string
+def test_dot(metaclass):
+    dot = Dot(metaclass=metaclass, column_sort="key-based")
+    graph_string = str(dot)
+    dot_assert(graph_string)
 
-    assert "users -- posts  [arrowhead=crow, arrowtail=none, dir=both, label=author];" in graph_string
-    assert "posts -- comments  [arrowhead=crow, arrowtail=none, dir=both, label=post];" in graph_string
-    assert "users -- comments  [arrowhead=crow, arrowtail=none, dir=both, label=author];" in graph_string
 
-    assert '<tr><td align="left">CHAR(32)</td><td align="left">author</td><td>Foreign Key</td></tr>' in graph_string
-    assert '<tr><td align="left">CHAR(32)</td><td align="left">post</td><td>Foreign Key</td></tr>' in graph_string
-    assert '<tr><td align="left">DATETIME</td><td align="left">created</td><td></td></tr>' in graph_string
+def test_dot_column_sort_preserve_order(metaclass, dot_full_string_preseve_column_sort):
+    dot = Dot(metaclass=metaclass, column_sort="preserve-order")
+    assert str(dot) == dot_full_string_preseve_column_sort
diff --git a/tests/transformers/test_mermaid.py b/tests/transformers/test_mermaid.py
index 4a5be2a..db09935 100644
--- a/tests/transformers/test_mermaid.py
+++ b/tests/transformers/test_mermaid.py
@@ -4,6 +4,11 @@ from ..utils import mermaid_assert
 
 
 def test_mermaid(metaclass):
-    mermaid = Mermaid(metaclass=metaclass)
+    mermaid = Mermaid(metaclass=metaclass, column_sort="key-based")
     graph_string = str(mermaid)
     mermaid_assert(graph_string)
+
+
+def test_mermaid_column_sort_preserve_order(metaclass, mermaid_full_string_preseve_column_sort):
+    mermaid = Mermaid(metaclass=metaclass, column_sort="preserve-order")
+    assert str(mermaid) == mermaid_full_string_preseve_column_sort
diff --git a/tests/utils.py b/tests/utils.py
index 4aa3782..403fdaf 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -11,3 +11,17 @@ def mermaid_assert(output: str) -> None:
     assert 'CHAR(32) post FK "nullable"' in output
     assert 'BOOLEAN live "True if post is published,nullable"' in output
     assert "DATETIME created" in output
+
+
+def dot_assert(output: str) -> None:
+    assert '<tr><td colspan="3" bgcolor="lightblue"><b>users</b></td></tr>' in output
+    assert '<tr><td colspan="3" bgcolor="lightblue"><b>posts</b></td></tr>' in output
+    assert '<tr><td colspan="3" bgcolor="lightblue"><b>comments</b></td></tr>' in output
+
+    assert "users -- posts  [arrowhead=crow, arrowtail=none, dir=both, label=author];" in output
+    assert "posts -- comments  [arrowhead=crow, arrowtail=none, dir=both, label=post];" in output
+    assert "users -- comments  [arrowhead=crow, arrowtail=none, dir=both, label=author];" in output
+
+    assert '<tr><td align="left">CHAR(32)</td><td align="left">author</td><td>Foreign Key</td></tr>' in output
+    assert '<tr><td align="left">CHAR(32)</td><td align="left">post</td><td>Foreign Key</td></tr>' in output
+    assert '<tr><td align="left">DATETIME</td><td align="left">created</td><td></td></tr>' in output