Refactor and Improve CLI Functionality (#13)

* remove --show-ui

* Update test_cli.py

* Exclude source models from expanded models in DbtColumnLineageExtractor

* lint

* logger

* Update utils.py

* remove verbose

* Enhance model matching logic to prioritize exact name matches and improve logging for multiple potential matches

* lint

* Improve logging for lineage extraction by adding warnings for skipped output generation when no matches or lineage information is found.

* self ref

* seff ref

* bump to 0.1.7b2

Author: wenwu35

py_package/dbt_column_lineage_extractor/__init__.py

@@ -4,7 +4,8 @@
     read_json,
     pretty_print_dict,
     write_dict_to_file,
-    read_dict_from_file
+    read_dict_from_file,
+    setup_logging
 )
 
 __all__ = [
@@ -15,4 +16,5 @@
     "pretty_print_dict",
     "write_dict_to_file",
     "read_dict_from_file",
+    "setup_logging",
 ]

py_package/dbt_column_lineage_extractor/cli_direct.py

@@ -2,16 +2,29 @@
 import dbt_column_lineage_extractor.utils as utils
 from dbt_column_lineage_extractor import DbtColumnLineageExtractor
 
+
 def main():
     parser = argparse.ArgumentParser(description="DBT Column Lineage Extractor CLI")
-    parser.add_argument('--manifest', default='./inputs/manifest.json', help='Path to the manifest.json file, default to ./inputs/manifest.json')
-    parser.add_argument('--catalog', default='./inputs/catalog.json', help='Path to the catalog.json file, default to ./inputs/catalog.json')
-    parser.add_argument('--dialect', default='snowflake', help='SQL dialect to use, default is snowflake, more dialects at https://github.com/tobymao/sqlglot/tree/v25.24.5/sqlglot/dialects')
     parser.add_argument(
-        '--model', 
-        nargs='*', 
-        default=[], 
-        help='''List of models to extract lineage for using dbt-style selectors:
+        "--manifest",
+        default="./inputs/manifest.json",
+        help="Path to the manifest.json file, default to ./inputs/manifest.json",
+    )
+    parser.add_argument(
+        "--catalog",
+        default="./inputs/catalog.json",
+        help="Path to the catalog.json file, default to ./inputs/catalog.json",
+    )
+    parser.add_argument(
+        "--dialect",
+        default="snowflake",
+        help="SQL dialect to use, default is snowflake, more dialects at https://github.com/tobymao/sqlglot/tree/v25.24.5/sqlglot/dialects",
+    )
+    parser.add_argument(
+        "--model",
+        nargs="*",
+        default=[],
+        help="""List of models to extract lineage for using dbt-style selectors:
             - Simple model names: model_name
             - Include ancestors: +model_name (include upstream/parent models)
             - Include descendants: model_name+ (include downstream/child models)
@@ -20,15 +33,28 @@ def main():
             - Tag filtering: tag:my_tag (models with specific tag)
             - Path filtering: path:models/finance (models in specific path)
             - Package filtering: package:my_package (models in specific package)
-            Default behavior extracts lineage for all models.'''
+            Default behavior extracts lineage for all models.""",
+    )
+    parser.add_argument(
+        "--model-list-json",
+        help="Path to a JSON file containing a list of models to extract lineage for. If specified, this takes precedence over --model",
+    )
+    parser.add_argument(
+        "--output-dir",
+        default="./outputs",
+        help="Directory to write output json files, default to ./outputs",
+    )
+    parser.add_argument(
+        "--continue-on-error",
+        action="store_true",
+        help="Continue processing even if some models fail",
     )
-    parser.add_argument('--model-list-json', help='Path to a JSON file containing a list of models to extract lineage for. If specified, this takes precedence over --model')
-    parser.add_argument('--output-dir', default='./outputs', help='Directory to write output json files, default to ./outputs')
-    parser.add_argument('--show-ui', action='store_true', help='Flag to show lineage outputs in the console')
-    parser.add_argument('--continue-on-error', action='store_true', help='Continue processing even if some models fail')
 
     args = parser.parse_args()
 
+    # Set up logging
+    logger = utils.setup_logging()
+
     try:
         selected_models = args.model
         if args.model_list_json:
@@ -37,7 +63,7 @@ def main():
                 if not isinstance(selected_models, list):
                     raise ValueError("The JSON file must contain a list of model names")
             except Exception as e:
-                print(f"Error reading model list from JSON file: {e}")
+                logger.error(f"Error reading model list from JSON file: {e}")
                 return 1
 
         extractor = DbtColumnLineageExtractor(
@@ -47,17 +73,19 @@ def main():
             dialect=args.dialect,
         )
 
-        print(f"Processing {len(extractor.selected_models)} models after selector expansion")
-        
+        logger.info(f"Processing {len(extractor.selected_models)} models after selector expansion")
+
         try:
             lineage_map = extractor.build_lineage_map()
-            
+
             if not lineage_map:
-                print("Warning: No valid lineage was generated. Check for errors above.")
+                logger.warning("Warning: No valid lineage was generated. Check for errors above.")
                 if not args.continue_on_error:
                     return 1
-                    
-            lineage_to_direct_parents = extractor.get_columns_lineage_from_sqlglot_lineage_map(lineage_map)
+
+            lineage_to_direct_parents = extractor.get_columns_lineage_from_sqlglot_lineage_map(
+                lineage_map
+            )
             lineage_to_direct_children = (
                 extractor.get_lineage_to_direct_children_from_lineage_to_direct_parents(
                     lineage_to_direct_parents
@@ -72,27 +100,24 @@ def main():
                 lineage_to_direct_children, f"{args.output_dir}/lineage_to_direct_children.json"
             )
 
-            if args.show_ui:
-                print("===== Lineage to Direct Parents =====")
-                utils.pretty_print_dict(lineage_to_direct_parents)
-                print("===== Lineage to Direct Children =====")
-                utils.pretty_print_dict(lineage_to_direct_children)
-
-            print("Lineage extraction complete. Output files written to output directory.")
+            logger.info("Lineage extraction complete. Output files written to output directory.")
             return 0
-            
+
         except Exception as e:
-            print(f"Error during lineage extraction: {str(e)}")
+            logger.error(f"Error during lineage extraction: {str(e)}")
             if not args.continue_on_error:
                 raise
             return 1
-            
+
     except Exception as e:
-        print(f"Error: {str(e)}")
+        logger.error(f"Error: {str(e)}")
         import traceback
+
         traceback.print_exc()
         return 1
 
-if __name__ == '__main__':
+
+if __name__ == "__main__":
     import sys
+
     sys.exit(main())