Fix from_expression when using Snowflake offline store

TomSteenbergen · TomSteenbergen · commit a410da3fffa4 · 2024-05-06T13:48:48.000+02:00
Signed-off-by: TomSteenbergen &lt;tomsteenbergen1995@gmail.com&gt;
diff --git a/sdk/python/feast/infra/offline_stores/snowflake.py b/sdk/python/feast/infra/offline_stores/snowflake.py
@@ -139,8 +139,10 @@ def pull_latest_from_table_or_query(
         assert isinstance(data_source, SnowflakeSource)
 
         from_expression = data_source.get_table_query_string()
-        if not data_source.database and data_source.table:
+        if not data_source.database and not data_source.schema and data_source.table:
             from_expression = f'"{config.offline_store.database}"."{config.offline_store.schema_}".{from_expression}'
+        if not data_source.database and data_source.schema and data_source.table:
+            from_expression = f'"{config.offline_store.database}".{from_expression}'
 
         if join_key_columns:
             partition_by_join_key_string = '"' + '", "'.join(join_key_columns) + '"'
@@ -226,8 +228,10 @@ def pull_all_from_table_or_query(
         assert isinstance(data_source, SnowflakeSource)
 
         from_expression = data_source.get_table_query_string()
-        if not data_source.database and data_source.table:
+        if not data_source.database and not data_source.schema and data_source.table:
             from_expression = f'"{config.offline_store.database}"."{config.offline_store.schema_}".{from_expression}'
+        if not data_source.database and data_source.schema and data_source.table:
+            from_expression = f'"{config.offline_store.database}".{from_expression}'
 
         field_string = (
             '"'
