diff --git a/sdk/python/feast/infra/offline_stores/snowflake.py b/sdk/python/feast/infra/offline_stores/snowflake.py index 968055fcee..c88e1b1844 100644 --- a/sdk/python/feast/infra/offline_stores/snowflake.py +++ b/sdk/python/feast/infra/offline_stores/snowflake.py @@ -82,7 +82,7 @@ class SnowflakeOfflineStoreConfig(FeastConfigBaseModel): database: Optional[str] = None """ Snowflake database name """ - schema_: Optional[str] = Field("PUBLIC", alias="schema") + schema_: Optional[str] = Field(None, alias="schema") """ Snowflake schema name """ class Config: diff --git a/sdk/python/feast/infra/utils/snowflake_utils.py b/sdk/python/feast/infra/utils/snowflake_utils.py index 065a4bad46..3513daa878 100644 --- a/sdk/python/feast/infra/utils/snowflake_utils.py +++ b/sdk/python/feast/infra/utils/snowflake_utils.py @@ -53,13 +53,19 @@ def get_snowflake_conn(config, autocommit=True) -> SnowflakeConnection: else: kwargs = {} + if "schema" in kwargs: + kwargs["schema_"] = kwargs.pop("schema") + kwargs.update((k, v) for k, v in config_dict.items() if v is not None) - [ - kwargs.update({k: '"' + v + '"'}) - for k, v in kwargs.items() - if k in ["role", "warehouse", "database", "schema_"] - ] - kwargs["schema"] = kwargs.pop("schema_") + + for k, v in kwargs.items(): + if k in ["role", "warehouse", "database", "schema_"]: + kwargs[k] = f'"{v}"' + + if "schema_" in kwargs: + kwargs["schema"] = kwargs.pop("schema_") + else: + kwargs["schema"] = '"PUBLIC"' try: conn = snowflake.connector.connect(