We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Snowflake login support private key or web browser authentication
The current snowflake support seems only allow password authentication but we are using Azure AD login without password for the account.
Can we add functionality to allow different mechanism of authentication?

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `sdk/python/feast/infra/utils/snowflake_utils.py`
Content:
```
1 import configparser
2 import os
3 import random
4 import string
5 from logging import getLogger
6 from tempfile import TemporaryDirectory
7 from typing import Dict, Iterator, List, Optional, Tuple, cast
8 
9 import pandas as pd
10 from tenacity import (
11     retry,
12     retry_if_exception_type,
13     stop_after_attempt,
14     wait_exponential,
15 )
16 
17 from feast.errors import SnowflakeIncompleteConfig, SnowflakeQueryUnknownError
18 
19 try:
20     import snowflake.connector
21     from snowflake.connector import ProgrammingError, SnowflakeConnection
22     from snowflake.connector.cursor import SnowflakeCursor
23 except ImportError as e:
24     from feast.errors import FeastExtrasDependencyImportError
25 
26     raise FeastExtrasDependencyImportError("snowflake", str(e))
27 
28 
29 getLogger("snowflake.connector.cursor").disabled = True
30 getLogger("snowflake.connector.connection").disabled = True
31 getLogger("snowflake.connector.network").disabled = True
32 logger = getLogger(__name__)
33 
34 
35 def execute_snowflake_statement(conn: SnowflakeConnection, query) -> SnowflakeCursor:
36     cursor = conn.cursor().execute(query)
37     if cursor is None:
38         raise SnowflakeQueryUnknownError(query)
39     return cursor
40 
41 
42 def get_snowflake_conn(config, autocommit=True) -> SnowflakeConnection:
43     if config.type == "snowflake.offline":
44         config_header = "connections.feast_offline_store"
45 
46     config = dict(config)
47 
48     # read config file
49     config_reader = configparser.ConfigParser()
50     config_reader.read([config["config_path"]])
51     if config_reader.has_section(config_header):
52         kwargs = dict(config_reader[config_header])
53     else:
54         kwargs = {}
55 
56     kwargs.update((k, v) for k, v in config.items() if v is not None)
57 
58     try:
59         conn = snowflake.connector.connect(
60             account=kwargs["account"],
61             user=kwargs["user"],
62             password=kwargs["password"],
63             role=f'''"{kwargs['role']}"''',
64             warehouse=f'''"{kwargs['warehouse']}"''',
65             database=f'''"{kwargs['database']}"''',
66             schema=f'''"{kwargs['schema_']}"''',
67             application="feast",
68             autocommit=autocommit,
69         )
70 
71         return conn
72     except KeyError as e:
73         raise SnowflakeIncompleteConfig(e)
74 
75 
76 # TO DO -- sfc-gh-madkins
77 # Remove dependency on write_pandas function by falling back to native snowflake python connector
78 # Current issue is datetime[ns] types are read incorrectly in Snowflake, need to coerce to datetime[ns, UTC]
79 def write_pandas(
80     conn: SnowflakeConnection,
81     df: pd.DataFrame,
82     table_name: str,
83     database: Optional[str] = None,
84     schema: Optional[str] = None,
85     chunk_size: Optional[int] = None,
86     compression: str = "gzip",
87     on_error: str = "abort_statement",
88     parallel: int = 4,
89     quote_identifiers: bool = True,
90     auto_create_table: bool = False,
91     create_temp_table: bool = False,
92 ):
93     """Allows users to most efficiently write back a pandas DataFrame to Snowflake.
94 
95     It works by dumping the DataFrame into Parquet files, uploading them and finally copying their data into the table.
96 
97     Returns whether all files were ingested correctly, number of chunks uploaded, and number of rows ingested
98     with all of the COPY INTO command's output for debugging purposes.
99 
100         Example usage:
101             import pandas
102             from snowflake.connector.pandas_tools import write_pandas
103 
104             df = pandas.DataFrame([('Mark', 10), ('Luke', 20)], columns=['name', 'balance'])
105             success, nchunks, nrows, _ = write_pandas(cnx, df, 'customers')
106 
107     Args:
108         conn: Connection to be used to communicate with Snowflake.
109         df: Dataframe we'd like to write back.
110         table_name: Table name where we want to insert into.
111         database: Database schema and table is in, if not provided the default one will be used (Default value = None).
112         schema: Schema table is in, if not provided the default one will be used (Default value = None).
113         chunk_size: Number of elements to be inserted once, if not provided all elements will be dumped once
114             (Default value = None).
115         compression: The compression used on the Parquet files, can only be gzip, or snappy. Gzip gives supposedly a
116             better compression, while snappy is faster. Use whichever is more appropriate (Default value = 'gzip').
117         on_error: Action to take when COPY INTO statements fail, default follows documentation at:
118             https://docs.snowflake.com/en/sql-reference/sql/copy-into-table.html#copy-options-copyoptions
119             (Default value = 'abort_statement').
120         parallel: Number of threads to be used when uploading chunks, default follows documentation at:
121             https://docs.snowflake.com/en/sql-reference/sql/put.html#optional-parameters (Default value = 4).
122         quote_identifiers: By default, identifiers, specifically database, schema, table and column names
123             (from df.columns) will be quoted. If set to False, identifiers are passed on to Snowflake without quoting.
124             I.e. identifiers will be coerced to uppercase by Snowflake.  (Default value = True)
125         auto_create_table: When true, will automatically create a table with corresponding columns for each column in
126             the passed in DataFrame. The table will not be created if it already exists
127         create_temp_table: Will make the auto-created table as a temporary table
128     """
129     if database is not None and schema is None:
130         raise ProgrammingError(
131             "Schema has to be provided to write_pandas when a database is provided"
132         )
133     # This dictionary maps the compression algorithm to Snowflake put copy into command type
134     # https://docs.snowflake.com/en/sql-reference/sql/copy-into-table.html#type-parquet
135     compression_map = {"gzip": "auto", "snappy": "snappy"}
136     if compression not in compression_map.keys():
137         raise ProgrammingError(
138             "Invalid compression '{}', only acceptable values are: {}".format(
139                 compression, compression_map.keys()
140             )
141         )
142     if quote_identifiers:
143         location = (
144             (('"' + database + '".') if database else "")
145             + (('"' + schema + '".') if schema else "")
146             + ('"' + table_name + '"')
147         )
148     else:
149         location = (
150             (database + "." if database else "")
151             + (schema + "." if schema else "")
152             + (table_name)
153         )
154     if chunk_size is None:
155         chunk_size = len(df)
156     cursor: SnowflakeCursor = conn.cursor()
157     stage_name = create_temporary_sfc_stage(cursor)
158 
159     with TemporaryDirectory() as tmp_folder:
160         for i, chunk in chunk_helper(df, chunk_size):
161             chunk_path = os.path.join(tmp_folder, "file{}.txt".format(i))
162             # Dump chunk into parquet file
163             chunk.to_parquet(
164                 chunk_path,
165                 compression=compression,
166                 use_deprecated_int96_timestamps=True,
167             )
168             # Upload parquet file
169             upload_sql = (
170                 "PUT /* Python:snowflake.connector.pandas_tools.write_pandas() */ "
171                 "'file://{path}' @\"{stage_name}\" PARALLEL={parallel}"
172             ).format(
173                 path=chunk_path.replace("\\", "\\\\").replace("'", "\\'"),
174                 stage_name=stage_name,
175                 parallel=parallel,
176             )
177             logger.debug(f"uploading files with '{upload_sql}'")
178             cursor.execute(upload_sql, _is_internal=True)
179             # Remove chunk file
180             os.remove(chunk_path)
181     if quote_identifiers:
182         columns = '"' + '","'.join(list(df.columns)) + '"'
183     else:
184         columns = ",".join(list(df.columns))
185 
186     if auto_create_table:
187         file_format_name = create_file_format(compression, compression_map, cursor)
188         infer_schema_sql = f"SELECT COLUMN_NAME, TYPE FROM table(infer_schema(location=>'@\"{stage_name}\"', file_format=>'{file_format_name}'))"
189         logger.debug(f"inferring schema with '{infer_schema_sql}'")
190         result_cursor = cursor.execute(infer_schema_sql, _is_internal=True)
191         if result_cursor is None:
192             raise SnowflakeQueryUnknownError(infer_schema_sql)
193         result = cast(List[Tuple[str, str]], result_cursor.fetchall())
194         column_type_mapping: Dict[str, str] = dict(result)
195         # Infer schema can return the columns out of order depending on the chunking we do when uploading
196         # so we have to iterate through the dataframe columns to make sure we create the table with its
197         # columns in order
198         quote = '"' if quote_identifiers else ""
199         create_table_columns = ", ".join(
200             [f"{quote}{c}{quote} {column_type_mapping[c]}" for c in df.columns]
201         )
202         create_table_sql = (
203             f"CREATE {'TEMP ' if create_temp_table else ''}TABLE IF NOT EXISTS {location} "
204             f"({create_table_columns})"
205             f" /* Python:snowflake.connector.pandas_tools.write_pandas() */ "
206         )
207         logger.debug(f"auto creating table with '{create_table_sql}'")
208         cursor.execute(create_table_sql, _is_internal=True)
209         drop_file_format_sql = f"DROP FILE FORMAT IF EXISTS {file_format_name}"
210         logger.debug(f"dropping file format with '{drop_file_format_sql}'")
211         cursor.execute(drop_file_format_sql, _is_internal=True)
212 
213     # in Snowflake, all parquet data is stored in a single column, $1, so we must select columns explicitly
214     # see (https://docs.snowflake.com/en/user-guide/script-data-load-transform-parquet.html)
215     if quote_identifiers:
216         parquet_columns = "$1:" + ",$1:".join(f'"{c}"' for c in df.columns)
217     else:
218         parquet_columns = "$1:" + ",$1:".join(df.columns)
219     copy_into_sql = (
220         "COPY INTO {location} /* Python:snowflake.connector.pandas_tools.write_pandas() */ "
221         "({columns}) "
222         'FROM (SELECT {parquet_columns} FROM @"{stage_name}") '
223         "FILE_FORMAT=(TYPE=PARQUET COMPRESSION={compression}) "
224         "PURGE=TRUE ON_ERROR={on_error}"
225     ).format(
226         location=location,
227         columns=columns,
228         parquet_columns=parquet_columns,
229         stage_name=stage_name,
230         compression=compression_map[compression],
231         on_error=on_error,
232     )
233     logger.debug("copying into with '{}'".format(copy_into_sql))
234     # Snowflake returns the original cursor if the query execution succeeded.
235     result_cursor = cursor.execute(copy_into_sql, _is_internal=True)
236     if result_cursor is None:
237         raise SnowflakeQueryUnknownError(copy_into_sql)
238     result_cursor.close()
239 
240 
241 @retry(
242     wait=wait_exponential(multiplier=1, max=4),
243     retry=retry_if_exception_type(ProgrammingError),
244     stop=stop_after_attempt(5),
245     reraise=True,
246 )
247 def create_file_format(
248     compression: str, compression_map: Dict[str, str], cursor: SnowflakeCursor
249 ) -> str:
250     file_format_name = (
251         '"' + "".join(random.choice(string.ascii_lowercase) for _ in range(5)) + '"'
252     )
253     file_format_sql = (
254         f"CREATE FILE FORMAT {file_format_name} "
255         f"/* Python:snowflake.connector.pandas_tools.write_pandas() */ "
256         f"TYPE=PARQUET COMPRESSION={compression_map[compression]}"
257     )
258     logger.debug(f"creating file format with '{file_format_sql}'")
259     cursor.execute(file_format_sql, _is_internal=True)
260     return file_format_name
261 
262 
263 @retry(
264     wait=wait_exponential(multiplier=1, max=4),
265     retry=retry_if_exception_type(ProgrammingError),
266     stop=stop_after_attempt(5),
267     reraise=True,
268 )
269 def create_temporary_sfc_stage(cursor: SnowflakeCursor) -> str:
270     stage_name = "".join(random.choice(string.ascii_lowercase) for _ in range(5))
271     create_stage_sql = (
272         "create temporary stage /* Python:snowflake.connector.pandas_tools.write_pandas() */ "
273         '"{stage_name}"'
274     ).format(stage_name=stage_name)
275     logger.debug(f"creating stage with '{create_stage_sql}'")
276     result_cursor = cursor.execute(create_stage_sql, _is_internal=True)
277     if result_cursor is None:
278         raise SnowflakeQueryUnknownError(create_stage_sql)
279     result_cursor.fetchall()
280     return stage_name
281 
282 
283 def chunk_helper(lst: pd.DataFrame, n: int) -> Iterator[Tuple[int, pd.DataFrame]]:
284     """Helper generator to chunk a sequence efficiently with current index like if enumerate was called on sequence."""
285     for i in range(0, len(lst), n):
286         yield int(i / n), lst[i : i + n]
287
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/sdk/python/feast/infra/utils/snowflake_utils.py b/sdk/python/feast/infra/utils/snowflake_utils.py
--- a/sdk/python/feast/infra/utils/snowflake_utils.py
+++ b/sdk/python/feast/infra/utils/snowflake_utils.py
@@ -43,29 +43,27 @@
     if config.type == "snowflake.offline":
         config_header = "connections.feast_offline_store"
 
-    config = dict(config)
+    config_dict = dict(config)
 
     # read config file
     config_reader = configparser.ConfigParser()
-    config_reader.read([config["config_path"]])
+    config_reader.read([config_dict["config_path"]])
     if config_reader.has_section(config_header):
         kwargs = dict(config_reader[config_header])
     else:
         kwargs = {}
 
-    kwargs.update((k, v) for k, v in config.items() if v is not None)
+    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_")
 
     try:
         conn = snowflake.connector.connect(
-            account=kwargs["account"],
-            user=kwargs["user"],
-            password=kwargs["password"],
-            role=f'''"{kwargs['role']}"''',
-            warehouse=f'''"{kwargs['warehouse']}"''',
-            database=f'''"{kwargs['database']}"''',
-            schema=f'''"{kwargs['schema_']}"''',
-            application="feast",
-            autocommit=autocommit,
+            application="feast", autocommit=autocommit, **kwargs
         )
 
         return conn

{"golden_diff": "diff --git a/sdk/python/feast/infra/utils/snowflake_utils.py b/sdk/python/feast/infra/utils/snowflake_utils.py\n--- a/sdk/python/feast/infra/utils/snowflake_utils.py\n+++ b/sdk/python/feast/infra/utils/snowflake_utils.py\n@@ -43,29 +43,27 @@\n     if config.type == \"snowflake.offline\":\n         config_header = \"connections.feast_offline_store\"\n \n-    config = dict(config)\n+    config_dict = dict(config)\n \n     # read config file\n     config_reader = configparser.ConfigParser()\n-    config_reader.read([config[\"config_path\"]])\n+    config_reader.read([config_dict[\"config_path\"]])\n     if config_reader.has_section(config_header):\n         kwargs = dict(config_reader[config_header])\n     else:\n         kwargs = {}\n \n-    kwargs.update((k, v) for k, v in config.items() if v is not None)\n+    kwargs.update((k, v) for k, v in config_dict.items() if v is not None)\n+    [\n+        kwargs.update({k: '\"' + v + '\"'})\n+        for k, v in kwargs.items()\n+        if k in [\"role\", \"warehouse\", \"database\", \"schema_\"]\n+    ]\n+    kwargs[\"schema\"] = kwargs.pop(\"schema_\")\n \n     try:\n         conn = snowflake.connector.connect(\n-            account=kwargs[\"account\"],\n-            user=kwargs[\"user\"],\n-            password=kwargs[\"password\"],\n-            role=f'''\"{kwargs['role']}\"''',\n-            warehouse=f'''\"{kwargs['warehouse']}\"''',\n-            database=f'''\"{kwargs['database']}\"''',\n-            schema=f'''\"{kwargs['schema_']}\"''',\n-            application=\"feast\",\n-            autocommit=autocommit,\n+            application=\"feast\", autocommit=autocommit, **kwargs\n         )\n \n         return conn\n", "issue": "Snowflake login support private key or web browser authentication\nThe current snowflake support seems only allow password authentication but we are using Azure AD login without password for the account.\r\nCan we add functionality to allow different mechanism of authentication?\n", "before_files": [{"content": "import configparser\nimport os\nimport random\nimport string\nfrom logging import getLogger\nfrom tempfile import TemporaryDirectory\nfrom typing import Dict, Iterator, List, Optional, Tuple, cast\n\nimport pandas as pd\nfrom tenacity import (\n    retry,\n    retry_if_exception_type,\n    stop_after_attempt,\n    wait_exponential,\n)\n\nfrom feast.errors import SnowflakeIncompleteConfig, SnowflakeQueryUnknownError\n\ntry:\n    import snowflake.connector\n    from snowflake.connector import ProgrammingError, SnowflakeConnection\n    from snowflake.connector.cursor import SnowflakeCursor\nexcept ImportError as e:\n    from feast.errors import FeastExtrasDependencyImportError\n\n    raise FeastExtrasDependencyImportError(\"snowflake\", str(e))\n\n\ngetLogger(\"snowflake.connector.cursor\").disabled = True\ngetLogger(\"snowflake.connector.connection\").disabled = True\ngetLogger(\"snowflake.connector.network\").disabled = True\nlogger = getLogger(__name__)\n\n\ndef execute_snowflake_statement(conn: SnowflakeConnection, query) -> SnowflakeCursor:\n    cursor = conn.cursor().execute(query)\n    if cursor is None:\n        raise SnowflakeQueryUnknownError(query)\n    return cursor\n\n\ndef get_snowflake_conn(config, autocommit=True) -> SnowflakeConnection:\n    if config.type == \"snowflake.offline\":\n        config_header = \"connections.feast_offline_store\"\n\n    config = dict(config)\n\n    # read config file\n    config_reader = configparser.ConfigParser()\n    config_reader.read([config[\"config_path\"]])\n    if config_reader.has_section(config_header):\n        kwargs = dict(config_reader[config_header])\n    else:\n        kwargs = {}\n\n    kwargs.update((k, v) for k, v in config.items() if v is not None)\n\n    try:\n        conn = snowflake.connector.connect(\n            account=kwargs[\"account\"],\n            user=kwargs[\"user\"],\n            password=kwargs[\"password\"],\n            role=f'''\"{kwargs['role']}\"''',\n            warehouse=f'''\"{kwargs['warehouse']}\"''',\n            database=f'''\"{kwargs['database']}\"''',\n            schema=f'''\"{kwargs['schema_']}\"''',\n            application=\"feast\",\n            autocommit=autocommit,\n        )\n\n        return conn\n    except KeyError as e:\n        raise SnowflakeIncompleteConfig(e)\n\n\n# TO DO -- sfc-gh-madkins\n# Remove dependency on write_pandas function by falling back to native snowflake python connector\n# Current issue is datetime[ns] types are read incorrectly in Snowflake, need to coerce to datetime[ns, UTC]\ndef write_pandas(\n    conn: SnowflakeConnection,\n    df: pd.DataFrame,\n    table_name: str,\n    database: Optional[str] = None,\n    schema: Optional[str] = None,\n    chunk_size: Optional[int] = None,\n    compression: str = \"gzip\",\n    on_error: str = \"abort_statement\",\n    parallel: int = 4,\n    quote_identifiers: bool = True,\n    auto_create_table: bool = False,\n    create_temp_table: bool = False,\n):\n    \"\"\"Allows users to most efficiently write back a pandas DataFrame to Snowflake.\n\n    It works by dumping the DataFrame into Parquet files, uploading them and finally copying their data into the table.\n\n    Returns whether all files were ingested correctly, number of chunks uploaded, and number of rows ingested\n    with all of the COPY INTO command's output for debugging purposes.\n\n        Example usage:\n            import pandas\n            from snowflake.connector.pandas_tools import write_pandas\n\n            df = pandas.DataFrame([('Mark', 10), ('Luke', 20)], columns=['name', 'balance'])\n            success, nchunks, nrows, _ = write_pandas(cnx, df, 'customers')\n\n    Args:\n        conn: Connection to be used to communicate with Snowflake.\n        df: Dataframe we'd like to write back.\n        table_name: Table name where we want to insert into.\n        database: Database schema and table is in, if not provided the default one will be used (Default value = None).\n        schema: Schema table is in, if not provided the default one will be used (Default value = None).\n        chunk_size: Number of elements to be inserted once, if not provided all elements will be dumped once\n            (Default value = None).\n        compression: The compression used on the Parquet files, can only be gzip, or snappy. Gzip gives supposedly a\n            better compression, while snappy is faster. Use whichever is more appropriate (Default value = 'gzip').\n        on_error: Action to take when COPY INTO statements fail, default follows documentation at:\n            https://docs.snowflake.com/en/sql-reference/sql/copy-into-table.html#copy-options-copyoptions\n            (Default value = 'abort_statement').\n        parallel: Number of threads to be used when uploading chunks, default follows documentation at:\n            https://docs.snowflake.com/en/sql-reference/sql/put.html#optional-parameters (Default value = 4).\n        quote_identifiers: By default, identifiers, specifically database, schema, table and column names\n            (from df.columns) will be quoted. If set to False, identifiers are passed on to Snowflake without quoting.\n            I.e. identifiers will be coerced to uppercase by Snowflake.  (Default value = True)\n        auto_create_table: When true, will automatically create a table with corresponding columns for each column in\n            the passed in DataFrame. The table will not be created if it already exists\n        create_temp_table: Will make the auto-created table as a temporary table\n    \"\"\"\n    if database is not None and schema is None:\n        raise ProgrammingError(\n            \"Schema has to be provided to write_pandas when a database is provided\"\n        )\n    # This dictionary maps the compression algorithm to Snowflake put copy into command type\n    # https://docs.snowflake.com/en/sql-reference/sql/copy-into-table.html#type-parquet\n    compression_map = {\"gzip\": \"auto\", \"snappy\": \"snappy\"}\n    if compression not in compression_map.keys():\n        raise ProgrammingError(\n            \"Invalid compression '{}', only acceptable values are: {}\".format(\n                compression, compression_map.keys()\n            )\n        )\n    if quote_identifiers:\n        location = (\n            (('\"' + database + '\".') if database else \"\")\n            + (('\"' + schema + '\".') if schema else \"\")\n            + ('\"' + table_name + '\"')\n        )\n    else:\n        location = (\n            (database + \".\" if database else \"\")\n            + (schema + \".\" if schema else \"\")\n            + (table_name)\n        )\n    if chunk_size is None:\n        chunk_size = len(df)\n    cursor: SnowflakeCursor = conn.cursor()\n    stage_name = create_temporary_sfc_stage(cursor)\n\n    with TemporaryDirectory() as tmp_folder:\n        for i, chunk in chunk_helper(df, chunk_size):\n            chunk_path = os.path.join(tmp_folder, \"file{}.txt\".format(i))\n            # Dump chunk into parquet file\n            chunk.to_parquet(\n                chunk_path,\n                compression=compression,\n                use_deprecated_int96_timestamps=True,\n            )\n            # Upload parquet file\n            upload_sql = (\n                \"PUT /* Python:snowflake.connector.pandas_tools.write_pandas() */ \"\n                \"'file://{path}' @\\\"{stage_name}\\\" PARALLEL={parallel}\"\n            ).format(\n                path=chunk_path.replace(\"\\\\\", \"\\\\\\\\\").replace(\"'\", \"\\\\'\"),\n                stage_name=stage_name,\n                parallel=parallel,\n            )\n            logger.debug(f\"uploading files with '{upload_sql}'\")\n            cursor.execute(upload_sql, _is_internal=True)\n            # Remove chunk file\n            os.remove(chunk_path)\n    if quote_identifiers:\n        columns = '\"' + '\",\"'.join(list(df.columns)) + '\"'\n    else:\n        columns = \",\".join(list(df.columns))\n\n    if auto_create_table:\n        file_format_name = create_file_format(compression, compression_map, cursor)\n        infer_schema_sql = f\"SELECT COLUMN_NAME, TYPE FROM table(infer_schema(location=>'@\\\"{stage_name}\\\"', file_format=>'{file_format_name}'))\"\n        logger.debug(f\"inferring schema with '{infer_schema_sql}'\")\n        result_cursor = cursor.execute(infer_schema_sql, _is_internal=True)\n        if result_cursor is None:\n            raise SnowflakeQueryUnknownError(infer_schema_sql)\n        result = cast(List[Tuple[str, str]], result_cursor.fetchall())\n        column_type_mapping: Dict[str, str] = dict(result)\n        # Infer schema can return the columns out of order depending on the chunking we do when uploading\n        # so we have to iterate through the dataframe columns to make sure we create the table with its\n        # columns in order\n        quote = '\"' if quote_identifiers else \"\"\n        create_table_columns = \", \".join(\n            [f\"{quote}{c}{quote} {column_type_mapping[c]}\" for c in df.columns]\n        )\n        create_table_sql = (\n            f\"CREATE {'TEMP ' if create_temp_table else ''}TABLE IF NOT EXISTS {location} \"\n            f\"({create_table_columns})\"\n            f\" /* Python:snowflake.connector.pandas_tools.write_pandas() */ \"\n        )\n        logger.debug(f\"auto creating table with '{create_table_sql}'\")\n        cursor.execute(create_table_sql, _is_internal=True)\n        drop_file_format_sql = f\"DROP FILE FORMAT IF EXISTS {file_format_name}\"\n        logger.debug(f\"dropping file format with '{drop_file_format_sql}'\")\n        cursor.execute(drop_file_format_sql, _is_internal=True)\n\n    # in Snowflake, all parquet data is stored in a single column, $1, so we must select columns explicitly\n    # see (https://docs.snowflake.com/en/user-guide/script-data-load-transform-parquet.html)\n    if quote_identifiers:\n        parquet_columns = \"$1:\" + \",$1:\".join(f'\"{c}\"' for c in df.columns)\n    else:\n        parquet_columns = \"$1:\" + \",$1:\".join(df.columns)\n    copy_into_sql = (\n        \"COPY INTO {location} /* Python:snowflake.connector.pandas_tools.write_pandas() */ \"\n        \"({columns}) \"\n        'FROM (SELECT {parquet_columns} FROM @\"{stage_name}\") '\n        \"FILE_FORMAT=(TYPE=PARQUET COMPRESSION={compression}) \"\n        \"PURGE=TRUE ON_ERROR={on_error}\"\n    ).format(\n        location=location,\n        columns=columns,\n        parquet_columns=parquet_columns,\n        stage_name=stage_name,\n        compression=compression_map[compression],\n        on_error=on_error,\n    )\n    logger.debug(\"copying into with '{}'\".format(copy_into_sql))\n    # Snowflake returns the original cursor if the query execution succeeded.\n    result_cursor = cursor.execute(copy_into_sql, _is_internal=True)\n    if result_cursor is None:\n        raise SnowflakeQueryUnknownError(copy_into_sql)\n    result_cursor.close()\n\n\n@retry(\n    wait=wait_exponential(multiplier=1, max=4),\n    retry=retry_if_exception_type(ProgrammingError),\n    stop=stop_after_attempt(5),\n    reraise=True,\n)\ndef create_file_format(\n    compression: str, compression_map: Dict[str, str], cursor: SnowflakeCursor\n) -> str:\n    file_format_name = (\n        '\"' + \"\".join(random.choice(string.ascii_lowercase) for _ in range(5)) + '\"'\n    )\n    file_format_sql = (\n        f\"CREATE FILE FORMAT {file_format_name} \"\n        f\"/* Python:snowflake.connector.pandas_tools.write_pandas() */ \"\n        f\"TYPE=PARQUET COMPRESSION={compression_map[compression]}\"\n    )\n    logger.debug(f\"creating file format with '{file_format_sql}'\")\n    cursor.execute(file_format_sql, _is_internal=True)\n    return file_format_name\n\n\n@retry(\n    wait=wait_exponential(multiplier=1, max=4),\n    retry=retry_if_exception_type(ProgrammingError),\n    stop=stop_after_attempt(5),\n    reraise=True,\n)\ndef create_temporary_sfc_stage(cursor: SnowflakeCursor) -> str:\n    stage_name = \"\".join(random.choice(string.ascii_lowercase) for _ in range(5))\n    create_stage_sql = (\n        \"create temporary stage /* Python:snowflake.connector.pandas_tools.write_pandas() */ \"\n        '\"{stage_name}\"'\n    ).format(stage_name=stage_name)\n    logger.debug(f\"creating stage with '{create_stage_sql}'\")\n    result_cursor = cursor.execute(create_stage_sql, _is_internal=True)\n    if result_cursor is None:\n        raise SnowflakeQueryUnknownError(create_stage_sql)\n    result_cursor.fetchall()\n    return stage_name\n\n\ndef chunk_helper(lst: pd.DataFrame, n: int) -> Iterator[Tuple[int, pd.DataFrame]]:\n    \"\"\"Helper generator to chunk a sequence efficiently with current index like if enumerate was called on sequence.\"\"\"\n    for i in range(0, len(lst), n):\n        yield int(i / n), lst[i : i + n]\n", "path": "sdk/python/feast/infra/utils/snowflake_utils.py"}], "after_files": [{"content": "import configparser\nimport os\nimport random\nimport string\nfrom logging import getLogger\nfrom tempfile import TemporaryDirectory\nfrom typing import Dict, Iterator, List, Optional, Tuple, cast\n\nimport pandas as pd\nfrom tenacity import (\n    retry,\n    retry_if_exception_type,\n    stop_after_attempt,\n    wait_exponential,\n)\n\nfrom feast.errors import SnowflakeIncompleteConfig, SnowflakeQueryUnknownError\n\ntry:\n    import snowflake.connector\n    from snowflake.connector import ProgrammingError, SnowflakeConnection\n    from snowflake.connector.cursor import SnowflakeCursor\nexcept ImportError as e:\n    from feast.errors import FeastExtrasDependencyImportError\n\n    raise FeastExtrasDependencyImportError(\"snowflake\", str(e))\n\n\ngetLogger(\"snowflake.connector.cursor\").disabled = True\ngetLogger(\"snowflake.connector.connection\").disabled = True\ngetLogger(\"snowflake.connector.network\").disabled = True\nlogger = getLogger(__name__)\n\n\ndef execute_snowflake_statement(conn: SnowflakeConnection, query) -> SnowflakeCursor:\n    cursor = conn.cursor().execute(query)\n    if cursor is None:\n        raise SnowflakeQueryUnknownError(query)\n    return cursor\n\n\ndef get_snowflake_conn(config, autocommit=True) -> SnowflakeConnection:\n    if config.type == \"snowflake.offline\":\n        config_header = \"connections.feast_offline_store\"\n\n    config_dict = dict(config)\n\n    # read config file\n    config_reader = configparser.ConfigParser()\n    config_reader.read([config_dict[\"config_path\"]])\n    if config_reader.has_section(config_header):\n        kwargs = dict(config_reader[config_header])\n    else:\n        kwargs = {}\n\n    kwargs.update((k, v) for k, v in config_dict.items() if v is not None)\n    [\n        kwargs.update({k: '\"' + v + '\"'})\n        for k, v in kwargs.items()\n        if k in [\"role\", \"warehouse\", \"database\", \"schema_\"]\n    ]\n    kwargs[\"schema\"] = kwargs.pop(\"schema_\")\n\n    try:\n        conn = snowflake.connector.connect(\n            application=\"feast\", autocommit=autocommit, **kwargs\n        )\n\n        return conn\n    except KeyError as e:\n        raise SnowflakeIncompleteConfig(e)\n\n\n# TO DO -- sfc-gh-madkins\n# Remove dependency on write_pandas function by falling back to native snowflake python connector\n# Current issue is datetime[ns] types are read incorrectly in Snowflake, need to coerce to datetime[ns, UTC]\ndef write_pandas(\n    conn: SnowflakeConnection,\n    df: pd.DataFrame,\n    table_name: str,\n    database: Optional[str] = None,\n    schema: Optional[str] = None,\n    chunk_size: Optional[int] = None,\n    compression: str = \"gzip\",\n    on_error: str = \"abort_statement\",\n    parallel: int = 4,\n    quote_identifiers: bool = True,\n    auto_create_table: bool = False,\n    create_temp_table: bool = False,\n):\n    \"\"\"Allows users to most efficiently write back a pandas DataFrame to Snowflake.\n\n    It works by dumping the DataFrame into Parquet files, uploading them and finally copying their data into the table.\n\n    Returns whether all files were ingested correctly, number of chunks uploaded, and number of rows ingested\n    with all of the COPY INTO command's output for debugging purposes.\n\n        Example usage:\n            import pandas\n            from snowflake.connector.pandas_tools import write_pandas\n\n            df = pandas.DataFrame([('Mark', 10), ('Luke', 20)], columns=['name', 'balance'])\n            success, nchunks, nrows, _ = write_pandas(cnx, df, 'customers')\n\n    Args:\n        conn: Connection to be used to communicate with Snowflake.\n        df: Dataframe we'd like to write back.\n        table_name: Table name where we want to insert into.\n        database: Database schema and table is in, if not provided the default one will be used (Default value = None).\n        schema: Schema table is in, if not provided the default one will be used (Default value = None).\n        chunk_size: Number of elements to be inserted once, if not provided all elements will be dumped once\n            (Default value = None).\n        compression: The compression used on the Parquet files, can only be gzip, or snappy. Gzip gives supposedly a\n            better compression, while snappy is faster. Use whichever is more appropriate (Default value = 'gzip').\n        on_error: Action to take when COPY INTO statements fail, default follows documentation at:\n            https://docs.snowflake.com/en/sql-reference/sql/copy-into-table.html#copy-options-copyoptions\n            (Default value = 'abort_statement').\n        parallel: Number of threads to be used when uploading chunks, default follows documentation at:\n            https://docs.snowflake.com/en/sql-reference/sql/put.html#optional-parameters (Default value = 4).\n        quote_identifiers: By default, identifiers, specifically database, schema, table and column names\n            (from df.columns) will be quoted. If set to False, identifiers are passed on to Snowflake without quoting.\n            I.e. identifiers will be coerced to uppercase by Snowflake.  (Default value = True)\n        auto_create_table: When true, will automatically create a table with corresponding columns for each column in\n            the passed in DataFrame. The table will not be created if it already exists\n        create_temp_table: Will make the auto-created table as a temporary table\n    \"\"\"\n    if database is not None and schema is None:\n        raise ProgrammingError(\n            \"Schema has to be provided to write_pandas when a database is provided\"\n        )\n    # This dictionary maps the compression algorithm to Snowflake put copy into command type\n    # https://docs.snowflake.com/en/sql-reference/sql/copy-into-table.html#type-parquet\n    compression_map = {\"gzip\": \"auto\", \"snappy\": \"snappy\"}\n    if compression not in compression_map.keys():\n        raise ProgrammingError(\n            \"Invalid compression '{}', only acceptable values are: {}\".format(\n                compression, compression_map.keys()\n            )\n        )\n    if quote_identifiers:\n        location = (\n            (('\"' + database + '\".') if database else \"\")\n            + (('\"' + schema + '\".') if schema else \"\")\n            + ('\"' + table_name + '\"')\n        )\n    else:\n        location = (\n            (database + \".\" if database else \"\")\n            + (schema + \".\" if schema else \"\")\n            + (table_name)\n        )\n    if chunk_size is None:\n        chunk_size = len(df)\n    cursor: SnowflakeCursor = conn.cursor()\n    stage_name = create_temporary_sfc_stage(cursor)\n\n    with TemporaryDirectory() as tmp_folder:\n        for i, chunk in chunk_helper(df, chunk_size):\n            chunk_path = os.path.join(tmp_folder, \"file{}.txt\".format(i))\n            # Dump chunk into parquet file\n            chunk.to_parquet(\n                chunk_path,\n                compression=compression,\n                use_deprecated_int96_timestamps=True,\n            )\n            # Upload parquet file\n            upload_sql = (\n                \"PUT /* Python:snowflake.connector.pandas_tools.write_pandas() */ \"\n                \"'file://{path}' @\\\"{stage_name}\\\" PARALLEL={parallel}\"\n            ).format(\n                path=chunk_path.replace(\"\\\\\", \"\\\\\\\\\").replace(\"'\", \"\\\\'\"),\n                stage_name=stage_name,\n                parallel=parallel,\n            )\n            logger.debug(f\"uploading files with '{upload_sql}'\")\n            cursor.execute(upload_sql, _is_internal=True)\n            # Remove chunk file\n            os.remove(chunk_path)\n    if quote_identifiers:\n        columns = '\"' + '\",\"'.join(list(df.columns)) + '\"'\n    else:\n        columns = \",\".join(list(df.columns))\n\n    if auto_create_table:\n        file_format_name = create_file_format(compression, compression_map, cursor)\n        infer_schema_sql = f\"SELECT COLUMN_NAME, TYPE FROM table(infer_schema(location=>'@\\\"{stage_name}\\\"', file_format=>'{file_format_name}'))\"\n        logger.debug(f\"inferring schema with '{infer_schema_sql}'\")\n        result_cursor = cursor.execute(infer_schema_sql, _is_internal=True)\n        if result_cursor is None:\n            raise SnowflakeQueryUnknownError(infer_schema_sql)\n        result = cast(List[Tuple[str, str]], result_cursor.fetchall())\n        column_type_mapping: Dict[str, str] = dict(result)\n        # Infer schema can return the columns out of order depending on the chunking we do when uploading\n        # so we have to iterate through the dataframe columns to make sure we create the table with its\n        # columns in order\n        quote = '\"' if quote_identifiers else \"\"\n        create_table_columns = \", \".join(\n            [f\"{quote}{c}{quote} {column_type_mapping[c]}\" for c in df.columns]\n        )\n        create_table_sql = (\n            f\"CREATE {'TEMP ' if create_temp_table else ''}TABLE IF NOT EXISTS {location} \"\n            f\"({create_table_columns})\"\n            f\" /* Python:snowflake.connector.pandas_tools.write_pandas() */ \"\n        )\n        logger.debug(f\"auto creating table with '{create_table_sql}'\")\n        cursor.execute(create_table_sql, _is_internal=True)\n        drop_file_format_sql = f\"DROP FILE FORMAT IF EXISTS {file_format_name}\"\n        logger.debug(f\"dropping file format with '{drop_file_format_sql}'\")\n        cursor.execute(drop_file_format_sql, _is_internal=True)\n\n    # in Snowflake, all parquet data is stored in a single column, $1, so we must select columns explicitly\n    # see (https://docs.snowflake.com/en/user-guide/script-data-load-transform-parquet.html)\n    if quote_identifiers:\n        parquet_columns = \"$1:\" + \",$1:\".join(f'\"{c}\"' for c in df.columns)\n    else:\n        parquet_columns = \"$1:\" + \",$1:\".join(df.columns)\n    copy_into_sql = (\n        \"COPY INTO {location} /* Python:snowflake.connector.pandas_tools.write_pandas() */ \"\n        \"({columns}) \"\n        'FROM (SELECT {parquet_columns} FROM @\"{stage_name}\") '\n        \"FILE_FORMAT=(TYPE=PARQUET COMPRESSION={compression}) \"\n        \"PURGE=TRUE ON_ERROR={on_error}\"\n    ).format(\n        location=location,\n        columns=columns,\n        parquet_columns=parquet_columns,\n        stage_name=stage_name,\n        compression=compression_map[compression],\n        on_error=on_error,\n    )\n    logger.debug(\"copying into with '{}'\".format(copy_into_sql))\n    # Snowflake returns the original cursor if the query execution succeeded.\n    result_cursor = cursor.execute(copy_into_sql, _is_internal=True)\n    if result_cursor is None:\n        raise SnowflakeQueryUnknownError(copy_into_sql)\n    result_cursor.close()\n\n\n@retry(\n    wait=wait_exponential(multiplier=1, max=4),\n    retry=retry_if_exception_type(ProgrammingError),\n    stop=stop_after_attempt(5),\n    reraise=True,\n)\ndef create_file_format(\n    compression: str, compression_map: Dict[str, str], cursor: SnowflakeCursor\n) -> str:\n    file_format_name = (\n        '\"' + \"\".join(random.choice(string.ascii_lowercase) for _ in range(5)) + '\"'\n    )\n    file_format_sql = (\n        f\"CREATE FILE FORMAT {file_format_name} \"\n        f\"/* Python:snowflake.connector.pandas_tools.write_pandas() */ \"\n        f\"TYPE=PARQUET COMPRESSION={compression_map[compression]}\"\n    )\n    logger.debug(f\"creating file format with '{file_format_sql}'\")\n    cursor.execute(file_format_sql, _is_internal=True)\n    return file_format_name\n\n\n@retry(\n    wait=wait_exponential(multiplier=1, max=4),\n    retry=retry_if_exception_type(ProgrammingError),\n    stop=stop_after_attempt(5),\n    reraise=True,\n)\ndef create_temporary_sfc_stage(cursor: SnowflakeCursor) -> str:\n    stage_name = \"\".join(random.choice(string.ascii_lowercase) for _ in range(5))\n    create_stage_sql = (\n        \"create temporary stage /* Python:snowflake.connector.pandas_tools.write_pandas() */ \"\n        '\"{stage_name}\"'\n    ).format(stage_name=stage_name)\n    logger.debug(f\"creating stage with '{create_stage_sql}'\")\n    result_cursor = cursor.execute(create_stage_sql, _is_internal=True)\n    if result_cursor is None:\n        raise SnowflakeQueryUnknownError(create_stage_sql)\n    result_cursor.fetchall()\n    return stage_name\n\n\ndef chunk_helper(lst: pd.DataFrame, n: int) -> Iterator[Tuple[int, pd.DataFrame]]:\n    \"\"\"Helper generator to chunk a sequence efficiently with current index like if enumerate was called on sequence.\"\"\"\n    for i in range(0, len(lst), n):\n        yield int(i / n), lst[i : i + n]\n", "path": "sdk/python/feast/infra/utils/snowflake_utils.py"}]}