Fix from_expression when using Snowflake offline store

Signed-off-by: TomSteenbergen <[email protected]>
feast-dev · May 6, 2024 · a410da3 · a410da3
1 parent adc2939
commit a410da3
Showing 1 changed file with 6 additions and 2 deletions.
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 = (
             '"'