Add omit_columns option to mermaid (#77)

* Add omit_columns option to mermaid * fmt * Fix help message * Add test cases for mermaid * rename option name * Add docs for omit-columns option
datnguye · Feb 8, 2024 · a6cf5f4 · a6cf5f4
1 parent 0d13cba
commit a6cf5f4
Show file tree

Hide file tree

Showing 4 changed files with 72 additions and 11 deletions.
diff --git a/dbterd/adapters/targets/mermaid/mermaid_test_relationship.py b/dbterd/adapters/targets/mermaid/mermaid_test_relationship.py
@@ -32,15 +32,19 @@ def parse(manifest, catalog, **kwargs):
     # https://mermaid.js.org/syntax/entityRelationshipDiagram.html
     mermaid = "erDiagram\n"
     for table in tables:
-        mermaid += '  "{table}" {{\n{columns}\n  }}\n'.format(
-            table=table.name.upper(),
-            columns="\n".join(
-                [
-                    f'    {x.data_type.replace(" ","-")} {x.name.replace(" ","-")}'
-                    for x in table.columns
-                ]
-            ),
+        table_name = table.name.upper()
+        columns = "\n".join(
+            [
+                f'    {x.data_type.replace(" ","-")} {x.name.replace(" ","-")}'
+                for x in table.columns
+            ]
         )
+        if kwargs.get("omit_columns", False):
+            mermaid += '  "{table_name}" {{\n  }}\n'.format(table_name=table_name)
+        else:
+            mermaid += '  "{table_name}" {{\n{columns}\n  }}\n'.format(
+                table_name=table_name, columns=columns
+            )
 
     for rel in relationships:
         key_from = f'"{rel.table_map[1]}"'

diff --git a/dbterd/cli/params.py b/dbterd/cli/params.py
@@ -63,6 +63,13 @@ def common_params(func):
         show_default=True,
         type=click.STRING,
     )
+    @click.option(
+        "--omit-columns",
+        help="Flag to omit columns in diagram. Currently only mermaid is supported",
+        is_flag=True,
+        default=False,
+        show_default=True,
+    )
     @functools.wraps(func)
     def wrapper(*args, **kwargs):
         return func(*args, **kwargs)  # pragma: no cover

diff --git a/docs/nav/guide/cli-references.md b/docs/nav/guide/cli-references.md
@@ -55,6 +55,8 @@ Command to generate diagram-as-a-code file
         -rt, --resource-type TEXT     Specified dbt resource type(seed, model,
                                         source, snapshot),default:model, use examples,
                                         -rt model -rt source
+        --omit-columns                Flag to omit columns in diagram. Currently
+                                        only mermaid is supported
         -h, --help                    Show this message and exit.
     ```
 
@@ -173,7 +175,6 @@ Configure the path to directory containing the output diagram file.
 
 ### dbterd run --target (-t)
 
-Target to the diagram-as-code platform
 > Default to `dbml`
 
 Supported target, please visit [Generate the Targets](https://dbterd.datnguyen.de/latest/nav/guide/targets/generate-dbml.html)
@@ -232,6 +233,19 @@ In the above:
     dbterd run --algo "test_relationship:(name:foreign_key|c_from:fk_column_name|c_to:pk_column_name)"
     ```
 
+### dbterd run --omit-columns
+
+Flag to omit columns in diagram. Currently only mermaid is supported
+
+> Default to `False`
+
+**Examples:**
+=== "CLI"
+
+    ```bash
+    dbterd run --target mermaid --omit-columns
+    ```
+
 ### dbterd run --manifest-version (-mv)
 
 Specified dbt manifest.json version

diff --git a/tests/unit/adapters/targets/mermaid/test_mermaid_test_relationship.py b/tests/unit/adapters/targets/mermaid/test_mermaid_test_relationship.py
@@ -8,7 +8,7 @@
 
 class TestMermaidTestRelationship:
     @pytest.mark.parametrize(
-        "tables, relationships, select, exclude, resource_type, expected",
+        "tables, relationships, select, exclude, resource_type, omit_columns, expected",
         [
             (
                 [
@@ -25,6 +25,7 @@ class TestMermaidTestRelationship:
                 [],
                 [],
                 ["model"],
+                False,
                 """erDiagram
                   "MODEL.DBT_RESTO.TABLE1" {
                       name1-type name1
@@ -73,6 +74,7 @@ class TestMermaidTestRelationship:
                 [],
                 [],
                 ["model", "source"],
+                False,
                 """erDiagram
                   "MODEL.DBT_RESTO.TABLE1" {
                       name1-type name1
@@ -118,6 +120,7 @@ class TestMermaidTestRelationship:
                 ["schema:--schema--"],
                 [],
                 ["model", "source"],
+                False,
                 """erDiagram
                     "MODEL.DBT_RESTO.TABLE1" {
                         name1-type name1
@@ -139,6 +142,7 @@ class TestMermaidTestRelationship:
                 [],
                 ["model.dbt_resto.table1"],
                 ["model"],
+                False,
                 """erDiagram
                 """,
             ),
@@ -165,6 +169,7 @@ class TestMermaidTestRelationship:
                 ["model.dbt_resto"],
                 ["model.dbt_resto.table2"],
                 ["model"],
+                False,
                 """erDiagram
                     "MODEL.DBT_RESTO.TABLE1" {
                         name1-type name1
@@ -186,6 +191,7 @@ class TestMermaidTestRelationship:
                 ["schema:", "wildcard:", ""],
                 [],
                 ["model"],
+                False,
                 """erDiagram
                     "MODEL.DBT_RESTO.TABLE1" {
                         name1-type name1
@@ -215,16 +221,45 @@ class TestMermaidTestRelationship:
                 ["schema:--schema--,wildcard:*dbt_resto.table*"],
                 ["wildcard:*table2"],
                 ["model"],
+                False,
                 """erDiagram
                     "MODEL.DBT_RESTO.TABLE1" {
                         name1-type name1
                     }
                 """,
             ),
+            (
+                [
+                    Table(
+                        name="model.dbt_resto.table1",
+                        node_name="model.dbt_resto.table1",
+                        database="--database--",
+                        schema="--schema--",
+                        columns=[Column(name="name1", data_type="name1-type")],
+                        raw_sql="--irrelevant--",
+                    )
+                ],
+                [],
+                [],
+                [],
+                ["model"],
+                True,
+                """erDiagram
+                  "MODEL.DBT_RESTO.TABLE1" {
+                  }
+                """,
+            ),
         ],
     )
     def test_parse(
-        self, tables, relationships, select, exclude, resource_type, expected
+        self,
+        tables,
+        relationships,
+        select,
+        exclude,
+        resource_type,
+        omit_columns,
+        expected,
     ):
         with mock.patch(
             "dbterd.adapters.algos.base.get_tables",
@@ -239,6 +274,7 @@ def test_parse(
                     catalog="--catalog--",
                     select=select,
                     exclude=exclude,
+                    omit_columns=omit_columns,
                     resource_type=resource_type,
                 )
                 print("mermaid ", mermaid.replace(" ", "").replace("\n", ""))