diff --git a/snowflake_dbml/generator.py b/snowflake_dbml/generator.py index ae2d40b..4ab22e9 100644 --- a/snowflake_dbml/generator.py +++ b/snowflake_dbml/generator.py @@ -9,6 +9,7 @@ from pydbml.classes import Project, Table, Column, Reference, Note, TableGroup from snowflake.connector import DictCursor from snowflake_dbml.config import load_config, load_primary_key_hints, load_visualization_params +from snowflake_dbml.version import __version__ import logging @@ -344,13 +345,30 @@ def save_config_to_env(config): else: print("Configuration not saved.") +def mask_sensitive_info(conn_params): + """ + Masks sensitive information in connection parameters. + """ + masked_params = conn_params.copy() + if 'password' in masked_params: + masked_params['password'] = '****' + return masked_params + def main(): + # Configure logging + logging.basicConfig(filename='snowflake-dbml.log', level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') + # Setup Argument Parser with more options parser = argparse.ArgumentParser( - description="Reverese engineer ER diagrams files from Snowflake databases using DBML", + description=f"Reverese engineer ER diagrams files from Snowflake databases using DBML (version {__version__})", formatter_class=argparse.ArgumentDefaultsHelpFormatter # This will show default values in the help message ) + + # Version argument + parser.add_argument('-v', '--version', action='version', version=f'%(prog)s {__version__}', + help="Show version number and exit.") + parser.add_argument('--user', type=str, help='Snowflake user name') parser.add_argument('--password', type=str, help='Snowflake password') parser.add_argument('--account', type=str, help='Snowflake account identifier') @@ -400,18 +418,24 @@ def main(): if args.dynamic_table_color: visualization_params['dynamic_table_color'] = args.dynamic_table_color - logging.debug(f"Generating DBML") - logging.debug(f"Connection Params: {connection_params}") - logging.debug(f"Primary Key Hints: {primary_key_hints}") - logging.debug(f"Visualization Params: {visualization_params}") + logger.debug(f"Generating DBML") + masked_params = mask_sensitive_info(connection_params) + logger.debug(f"Connection Params: {masked_params}") + logger.debug(f"Primary Key Hints: {primary_key_hints}") + logger.debug(f"Visualization Params: {visualization_params}") + logger.info(f"Fetching data from Snowflake starting.") data = fetch_data(connection_params, config['included_schemas'], config['excluded_schemas']) + logger.info(f"Fetching data from Snowflake complete.") - logging.debug(f"Fetched data from Snowflake: {len(data)} rows") + # Logging the number of results for each key in the data dictionary + for key, value in data.items(): + logger.info(f"Number of results for {key}: {len(value)}") dbml_output = generate_dbml(data, connection_params=connection_params, primary_key_hints=primary_key_hints, visualization_params=visualization_params) print(dbml_output) + logger.info("Snowflake DBML generator complete.") if __name__ == "__main__":