fix: Check for snowflake functions when setting up materialization engine (feast-dev#4456)

benchoncy · web-flow · commit c365b4e71a16 · 2024-08-30T08:45:20.000+04:00
* Check for snowflake functions over stage

Signed-off-by: Ben Stuart &lt;ben@benstuart.ie&gt;

* Refactor stage_list to function_list

Signed-off-by: Ben Stuart &lt;ben@benstuart.ie&gt;

---------

Signed-off-by: Ben Stuart &lt;ben@benstuart.ie&gt;
diff --git a/sdk/python/feast/infra/materialization/snowflake_engine.py b/sdk/python/feast/infra/materialization/snowflake_engine.py
@@ -128,16 +128,16 @@ def update(
         stage_context = f'"{self.repo_config.batch_engine.database}"."{self.repo_config.batch_engine.schema_}"'
         stage_path = f'{stage_context}."feast_{project}"'
         with GetSnowflakeConnection(self.repo_config.batch_engine) as conn:
-            query = f"SHOW STAGES IN {stage_context}"
+            query = f"SHOW USER FUNCTIONS LIKE 'FEAST_{project.upper()}%' IN SCHEMA {stage_context}"
             cursor = execute_snowflake_statement(conn, query)
-            stage_list = pd.DataFrame(
+            function_list = pd.DataFrame(
                 cursor.fetchall(),
                 columns=[column.name for column in cursor.description],
             )
 
-            # if the stage already exists,
+            # if the SHOW FUNCTIONS query returns results,
             # assumes that the materialization functions have been deployed
-            if f"feast_{project}" in stage_list["name"].tolist():
+            if len(function_list.index) > 0:
                 click.echo(
                     f"Materialization functions for {Style.BRIGHT + Fore.GREEN}{project}{Style.RESET_ALL} already detected."
                 )
@@ -149,7 +149,7 @@ def update(
             )
             click.echo()
 
-            query = f"CREATE STAGE {stage_path}"
+            query = f"CREATE STAGE IF NOT EXISTS {stage_path}"
             execute_snowflake_statement(conn, query)
 
             copy_path, zip_path = package_snowpark_zip(project)
