From 60ab9831986deebe46b06ace6833f50de2c19e30 Mon Sep 17 00:00:00 2001
From: Jonathan <jonathan.hilgart@snowflake.com>
Date: Mon, 15 Apr 2024 14:54:25 -0700
Subject: [PATCH] lint

---
 semantic_model_generator/generate_model.py | 25 +++++++++++++---------
 1 file changed, 15 insertions(+), 10 deletions(-)

diff --git a/semantic_model_generator/generate_model.py b/semantic_model_generator/generate_model.py
index c7d9af03..f3c3a88e 100644
--- a/semantic_model_generator/generate_model.py
+++ b/semantic_model_generator/generate_model.py
@@ -1,4 +1,5 @@
 from datetime import datetime
+from typing import Optional
 
 import jsonargparse
 from loguru import logger
@@ -114,7 +115,8 @@ def _raw_table_to_semantic_context_table(
     )
 
 
-def raw_schema_to_semantic_context( physical_tables : list[str], snowflake_account: str, semantic_model_name: str
+def raw_schema_to_semantic_context(
+    physical_tables: list[str], snowflake_account: str, semantic_model_name: str
 ) -> semantic_model_pb2.SemanticModel:
     """
     Converts a list of fully qualified Snowflake table names into a semantic model.
@@ -146,16 +148,18 @@ def raw_schema_to_semantic_context( physical_tables : list[str], snowflake_accou
         # Verify this is a valid FQN table. For now, we check that the table follows the following format.
         # {database}.{schema}.{table}
         fqn_table = create_fqn_table(table)
-        fqn_databse_schema = f"{fqn_table.database}_{fqn_table.schema}"
+        fqn_databse_schema = f"{fqn_table.database}_{fqn_table.schema_name}"
         if fqn_databse_schema not in unique_database_schema:
             unique_database_schema.append(fqn_databse_schema)
 
         with connector.connect(
-            db_name=fqn_table.database, schema_name=fqn_table.schema
+            db_name=fqn_table.database, schema_name=fqn_table.schema_name
         ) as conn:
             logger.info(f"Pulling column information from {fqn_table}")
             valid_schemas_tables_columns_df = get_valid_schemas_tables_columns_df(
-                conn=conn, table_schema=fqn_table.schema, table_names=[fqn_table.table]
+                conn=conn,
+                table_schema=fqn_table.schema_name,
+                table_names=[fqn_table.table],
             )
             assert not valid_schemas_tables_columns_df.empty
 
@@ -166,7 +170,7 @@ def raw_schema_to_semantic_context( physical_tables : list[str], snowflake_accou
 
             raw_table = get_table_representation(
                 conn=conn,
-                schema_name=fqn_table.schema,
+                schema_name=fqn_table.schema_name,
                 table_name=fqn_table.table,
                 table_index=0,
                 ndv_per_column=10,  # number of sample values to pull per column.
@@ -176,7 +180,7 @@ def raw_schema_to_semantic_context( physical_tables : list[str], snowflake_accou
 
             table_object = _raw_table_to_semantic_context_table(
                 database=fqn_table.database,
-                schema=fqn_table.schema,
+                schema=fqn_table.schema_name,
                 raw_table=raw_table,
             )
             table_objects.append(table_object)
@@ -237,10 +241,11 @@ def _to_snake_case(s: str) -> str:
     return snake_case_str
 
 
-def generate_base_semantic_context_from_snowflake(  physical_tables : list[str],
+def generate_base_semantic_context_from_snowflake(
+    physical_tables: list[str],
     snowflake_account: str,
     semantic_model_name: str,
-    output_yaml_path: str | None = None,
+    output_yaml_path: Optional[str] = None,
 ) -> None:
     """
     Generates a base semantic context from specified Snowflake tables and exports it to a YAML file.
@@ -255,7 +260,7 @@ def generate_base_semantic_context_from_snowflake(  physical_tables : list[str],
         None. Writes the semantic context to a YAML file.
     """
     context = raw_schema_to_semantic_context(
-    physical_tables ,
+        physical_tables,
         snowflake_account=snowflake_account,
         semantic_model_name=semantic_model_name,
     )
@@ -309,7 +314,7 @@ def generate_base_semantic_context_from_snowflake(  physical_tables : list[str],
     args = parser.parse_args()
 
     generate_base_semantic_context_from_snowflake(
-    physical_tables =args.physical_tables ,
+        physical_tables=args.physical_tables,
         snowflake_account=args.snowflake_account,
         semantic_model_name=args.semantic_model_name,
         output_yaml_path=args.output_yaml_path,