We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Expose `sfqid` attribute in snowflake exception messages
When contacting snowflake support they always start the conversation with `can you provide a query id`.  Exposing this id in all cases would be useful when contacting snowflake support.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `plugins/snowflake/dbt/adapters/snowflake/connections.py`
Content:
```
1 import base64
2 import datetime
3 import pytz
4 import re
5 from contextlib import contextmanager
6 from dataclasses import dataclass
7 from io import StringIO
8 from typing import Optional
9 
10 from cryptography.hazmat.backends import default_backend
11 from cryptography.hazmat.primitives import serialization
12 import requests
13 import snowflake.connector
14 import snowflake.connector.errors
15 
16 from dbt.exceptions import (
17     InternalException, RuntimeException, FailedToConnectException,
18     DatabaseException, warn_or_error
19 )
20 from dbt.adapters.base import Credentials
21 from dbt.adapters.sql import SQLConnectionManager
22 from dbt.logger import GLOBAL_LOGGER as logger
23 
24 
25 _TOKEN_REQUEST_URL = 'https://{}.snowflakecomputing.com/oauth/token-request'
26 
27 
28 @dataclass
29 class SnowflakeCredentials(Credentials):
30     account: str
31     user: str
32     warehouse: Optional[str]
33     role: Optional[str]
34     password: Optional[str]
35     authenticator: Optional[str]
36     private_key_path: Optional[str]
37     private_key_passphrase: Optional[str]
38     token: Optional[str]
39     oauth_client_id: Optional[str]
40     oauth_client_secret: Optional[str]
41     client_session_keep_alive: bool = False
42 
43     def __post_init__(self):
44         if (
45             self.authenticator != 'oauth' and
46             (self.oauth_client_secret or self.oauth_client_id or self.token)
47         ):
48             # the user probably forgot to set 'authenticator' like I keep doing
49             warn_or_error(
50                 'Authenticator is not set to oauth, but an oauth-only '
51                 'parameter is set! Did you mean to set authenticator: oauth?'
52             )
53 
54     @property
55     def type(self):
56         return 'snowflake'
57 
58     def _connection_keys(self):
59         return (
60             'account', 'user', 'database', 'schema', 'warehouse', 'role',
61             'client_session_keep_alive'
62         )
63 
64     def auth_args(self):
65         # Pull all of the optional authentication args for the connector,
66         # let connector handle the actual arg validation
67         result = {}
68         if self.password:
69             result['password'] = self.password
70         if self.authenticator:
71             result['authenticator'] = self.authenticator
72             if self.authenticator == 'oauth':
73                 token = self.token
74                 # if we have a client ID/client secret, the token is a refresh
75                 # token, not an access token
76                 if self.oauth_client_id and self.oauth_client_secret:
77                     token = self._get_access_token()
78                 elif self.oauth_client_id:
79                     warn_or_error(
80                         'Invalid profile: got an oauth_client_id, but not an '
81                         'oauth_client_secret!'
82                     )
83                 elif self.oauth_client_secret:
84                     warn_or_error(
85                         'Invalid profile: got an oauth_client_secret, but not '
86                         'an oauth_client_id!'
87                     )
88 
89                 result['token'] = token
90         result['private_key'] = self._get_private_key()
91         return result
92 
93     def _get_access_token(self) -> str:
94         if self.authenticator != 'oauth':
95             raise InternalException('Can only get access tokens for oauth')
96         missing = any(
97             x is None for x in
98             (self.oauth_client_id, self.oauth_client_secret, self.token)
99         )
100         if missing:
101             raise InternalException(
102                 'need a client ID a client secret, and a refresh token to get '
103                 'an access token'
104             )
105         # should the full url be a config item?
106         token_url = _TOKEN_REQUEST_URL.format(self.account)
107         # I think this is only used to redirect on success, which we ignore
108         # (it does not have to match the integration's settings in snowflake)
109         redirect_uri = 'http://localhost:9999'
110         data = {
111             'grant_type': 'refresh_token',
112             'refresh_token': self.token,
113             'redirect_uri': redirect_uri
114         }
115 
116         auth = base64.b64encode(
117             f'{self.oauth_client_id}:{self.oauth_client_secret}'
118             .encode('ascii')
119         ).decode('ascii')
120         headers = {
121             'Authorization': f'Basic {auth}',
122             'Content-type': 'application/x-www-form-urlencoded;charset=utf-8'
123         }
124         result = requests.post(token_url, headers=headers, data=data)
125         result_json = result.json()
126         if 'access_token' not in result_json:
127             raise DatabaseException(f'Did not get a token: {result_json}')
128         return result_json['access_token']
129 
130     def _get_private_key(self):
131         """Get Snowflake private key by path or None."""
132         if not self.private_key_path:
133             return None
134 
135         if self.private_key_passphrase:
136             encoded_passphrase = self.private_key_passphrase.encode()
137         else:
138             encoded_passphrase = None
139 
140         with open(self.private_key_path, 'rb') as key:
141             p_key = serialization.load_pem_private_key(
142                 key.read(),
143                 password=encoded_passphrase,
144                 backend=default_backend())
145 
146         return p_key.private_bytes(
147             encoding=serialization.Encoding.DER,
148             format=serialization.PrivateFormat.PKCS8,
149             encryption_algorithm=serialization.NoEncryption())
150 
151 
152 class SnowflakeConnectionManager(SQLConnectionManager):
153     TYPE = 'snowflake'
154 
155     @contextmanager
156     def exception_handler(self, sql):
157         try:
158             yield
159         except snowflake.connector.errors.ProgrammingError as e:
160             msg = str(e)
161 
162             logger.debug('Snowflake error: {}'.format(msg))
163 
164             if 'Empty SQL statement' in msg:
165                 logger.debug("got empty sql statement, moving on")
166             elif 'This session does not have a current database' in msg:
167                 self.release()
168                 raise FailedToConnectException(
169                     ('{}\n\nThis error sometimes occurs when invalid '
170                      'credentials are provided, or when your default role '
171                      'does not have access to use the specified database. '
172                      'Please double check your profile and try again.')
173                     .format(msg))
174             else:
175                 self.release()
176                 raise DatabaseException(msg)
177         except Exception as e:
178             logger.debug("Error running SQL: {}", sql)
179             logger.debug("Rolling back transaction.")
180             self.release()
181             if isinstance(e, RuntimeException):
182                 # during a sql query, an internal to dbt exception was raised.
183                 # this sounds a lot like a signal handler and probably has
184                 # useful information, so raise it without modification.
185                 raise
186             raise RuntimeException(str(e)) from e
187 
188     @classmethod
189     def open(cls, connection):
190         if connection.state == 'open':
191             logger.debug('Connection is already open, skipping open.')
192             return connection
193 
194         try:
195             creds = connection.credentials
196 
197             handle = snowflake.connector.connect(
198                 account=creds.account,
199                 user=creds.user,
200                 database=creds.database,
201                 schema=creds.schema,
202                 warehouse=creds.warehouse,
203                 role=creds.role,
204                 autocommit=False,
205                 client_session_keep_alive=creds.client_session_keep_alive,
206                 application='dbt',
207                 **creds.auth_args()
208             )
209 
210             connection.handle = handle
211             connection.state = 'open'
212         except snowflake.connector.errors.Error as e:
213             logger.debug("Got an error when attempting to open a snowflake "
214                          "connection: '{}'"
215                          .format(e))
216 
217             connection.handle = None
218             connection.state = 'fail'
219 
220             raise FailedToConnectException(str(e))
221 
222     def cancel(self, connection):
223         handle = connection.handle
224         sid = handle.session_id
225 
226         connection_name = connection.name
227 
228         sql = 'select system$abort_session({})'.format(sid)
229 
230         logger.debug("Cancelling query '{}' ({})".format(connection_name, sid))
231 
232         _, cursor = self.add_query(sql)
233         res = cursor.fetchone()
234 
235         logger.debug("Cancel query '{}': {}".format(connection_name, res))
236 
237     @classmethod
238     def get_status(cls, cursor):
239         state = cursor.sqlstate
240 
241         if state is None:
242             state = 'SUCCESS'
243 
244         return "{} {}".format(state, cursor.rowcount)
245 
246     @classmethod
247     def _split_queries(cls, sql):
248         "Splits sql statements at semicolons into discrete queries"
249 
250         sql_s = str(sql)
251         sql_buf = StringIO(sql_s)
252         split_query = snowflake.connector.util_text.split_statements(sql_buf)
253         return [part[0] for part in split_query]
254 
255     @classmethod
256     def process_results(cls, column_names, rows):
257         # Override for Snowflake. The datetime objects returned by
258         # snowflake-connector-python are not pickleable, so we need
259         # to replace them with sane timezones
260         fixed = []
261         for row in rows:
262             fixed_row = []
263             for col in row:
264                 if isinstance(col, datetime.datetime) and col.tzinfo:
265                     offset = col.utcoffset()
266                     offset_seconds = offset.total_seconds()
267                     new_timezone = pytz.FixedOffset(offset_seconds // 60)
268                     col = col.astimezone(tz=new_timezone)
269                 fixed_row.append(col)
270 
271             fixed.append(fixed_row)
272 
273         return super().process_results(column_names, fixed)
274 
275     def add_query(self, sql, auto_begin=True,
276                   bindings=None, abridge_sql_log=False):
277 
278         connection = None
279         cursor = None
280 
281         if bindings:
282             # The snowflake connector is more strict than, eg., psycopg2 -
283             # which allows any iterable thing to be passed as a binding.
284             bindings = tuple(bindings)
285 
286         queries = self._split_queries(sql)
287 
288         for individual_query in queries:
289             # hack -- after the last ';', remove comments and don't run
290             # empty queries. this avoids using exceptions as flow control,
291             # and also allows us to return the status of the last cursor
292             without_comments = re.sub(
293                 re.compile('^.*(--.*)$', re.MULTILINE),
294                 '', individual_query).strip()
295 
296             if without_comments == "":
297                 continue
298 
299             connection, cursor = super().add_query(
300                 individual_query, auto_begin,
301                 bindings=bindings,
302                 abridge_sql_log=abridge_sql_log
303             )
304 
305         if cursor is None:
306             conn = self.get_thread_connection()
307             if conn is None or conn.name is None:
308                 conn_name = '<None>'
309             else:
310                 conn_name = conn.name
311 
312             raise RuntimeException(
313                 "Tried to run an empty query on model '{}'. If you are "
314                 "conditionally running\nsql, eg. in a model hook, make "
315                 "sure your `else` clause contains valid sql!\n\n"
316                 "Provided SQL:\n{}"
317                 .format(conn_name, sql)
318             )
319 
320         return connection, cursor
321 
322     @classmethod
323     def _rollback_handle(cls, connection):
324         """On snowflake, rolling back the handle of an aborted session raises
325         an exception.
326         """
327         logger.debug('initiating rollback')
328         try:
329             connection.handle.rollback()
330         except snowflake.connector.errors.ProgrammingError as e:
331             msg = str(e)
332             if 'Session no longer exists' not in msg:
333                 raise
334
```
--- 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/plugins/snowflake/dbt/adapters/snowflake/connections.py b/plugins/snowflake/dbt/adapters/snowflake/connections.py
--- a/plugins/snowflake/dbt/adapters/snowflake/connections.py
+++ b/plugins/snowflake/dbt/adapters/snowflake/connections.py
@@ -159,6 +159,7 @@
         except snowflake.connector.errors.ProgrammingError as e:
             msg = str(e)
 
+            logger.debug('Snowflake query id: {}'.format(e.sfqid))
             logger.debug('Snowflake error: {}'.format(msg))
 
             if 'Empty SQL statement' in msg:
@@ -175,6 +176,9 @@
                 self.release()
                 raise DatabaseException(msg)
         except Exception as e:
+            if isinstance(e, snowflake.connector.errors.Error):
+                logger.debug('Snowflake query id: {}'.format(e.sfqid))
+
             logger.debug("Error running SQL: {}", sql)
             logger.debug("Rolling back transaction.")
             self.release()

{"golden_diff": "diff --git a/plugins/snowflake/dbt/adapters/snowflake/connections.py b/plugins/snowflake/dbt/adapters/snowflake/connections.py\n--- a/plugins/snowflake/dbt/adapters/snowflake/connections.py\n+++ b/plugins/snowflake/dbt/adapters/snowflake/connections.py\n@@ -159,6 +159,7 @@\n         except snowflake.connector.errors.ProgrammingError as e:\n             msg = str(e)\n \n+            logger.debug('Snowflake query id: {}'.format(e.sfqid))\n             logger.debug('Snowflake error: {}'.format(msg))\n \n             if 'Empty SQL statement' in msg:\n@@ -175,6 +176,9 @@\n                 self.release()\n                 raise DatabaseException(msg)\n         except Exception as e:\n+            if isinstance(e, snowflake.connector.errors.Error):\n+                logger.debug('Snowflake query id: {}'.format(e.sfqid))\n+\n             logger.debug(\"Error running SQL: {}\", sql)\n             logger.debug(\"Rolling back transaction.\")\n             self.release()\n", "issue": "Expose `sfqid` attribute in snowflake exception messages\nWhen contacting snowflake support they always start the conversation with `can you provide a query id`.  Exposing this id in all cases would be useful when contacting snowflake support.\n", "before_files": [{"content": "import base64\nimport datetime\nimport pytz\nimport re\nfrom contextlib import contextmanager\nfrom dataclasses import dataclass\nfrom io import StringIO\nfrom typing import Optional\n\nfrom cryptography.hazmat.backends import default_backend\nfrom cryptography.hazmat.primitives import serialization\nimport requests\nimport snowflake.connector\nimport snowflake.connector.errors\n\nfrom dbt.exceptions import (\n    InternalException, RuntimeException, FailedToConnectException,\n    DatabaseException, warn_or_error\n)\nfrom dbt.adapters.base import Credentials\nfrom dbt.adapters.sql import SQLConnectionManager\nfrom dbt.logger import GLOBAL_LOGGER as logger\n\n\n_TOKEN_REQUEST_URL = 'https://{}.snowflakecomputing.com/oauth/token-request'\n\n\n@dataclass\nclass SnowflakeCredentials(Credentials):\n    account: str\n    user: str\n    warehouse: Optional[str]\n    role: Optional[str]\n    password: Optional[str]\n    authenticator: Optional[str]\n    private_key_path: Optional[str]\n    private_key_passphrase: Optional[str]\n    token: Optional[str]\n    oauth_client_id: Optional[str]\n    oauth_client_secret: Optional[str]\n    client_session_keep_alive: bool = False\n\n    def __post_init__(self):\n        if (\n            self.authenticator != 'oauth' and\n            (self.oauth_client_secret or self.oauth_client_id or self.token)\n        ):\n            # the user probably forgot to set 'authenticator' like I keep doing\n            warn_or_error(\n                'Authenticator is not set to oauth, but an oauth-only '\n                'parameter is set! Did you mean to set authenticator: oauth?'\n            )\n\n    @property\n    def type(self):\n        return 'snowflake'\n\n    def _connection_keys(self):\n        return (\n            'account', 'user', 'database', 'schema', 'warehouse', 'role',\n            'client_session_keep_alive'\n        )\n\n    def auth_args(self):\n        # Pull all of the optional authentication args for the connector,\n        # let connector handle the actual arg validation\n        result = {}\n        if self.password:\n            result['password'] = self.password\n        if self.authenticator:\n            result['authenticator'] = self.authenticator\n            if self.authenticator == 'oauth':\n                token = self.token\n                # if we have a client ID/client secret, the token is a refresh\n                # token, not an access token\n                if self.oauth_client_id and self.oauth_client_secret:\n                    token = self._get_access_token()\n                elif self.oauth_client_id:\n                    warn_or_error(\n                        'Invalid profile: got an oauth_client_id, but not an '\n                        'oauth_client_secret!'\n                    )\n                elif self.oauth_client_secret:\n                    warn_or_error(\n                        'Invalid profile: got an oauth_client_secret, but not '\n                        'an oauth_client_id!'\n                    )\n\n                result['token'] = token\n        result['private_key'] = self._get_private_key()\n        return result\n\n    def _get_access_token(self) -> str:\n        if self.authenticator != 'oauth':\n            raise InternalException('Can only get access tokens for oauth')\n        missing = any(\n            x is None for x in\n            (self.oauth_client_id, self.oauth_client_secret, self.token)\n        )\n        if missing:\n            raise InternalException(\n                'need a client ID a client secret, and a refresh token to get '\n                'an access token'\n            )\n        # should the full url be a config item?\n        token_url = _TOKEN_REQUEST_URL.format(self.account)\n        # I think this is only used to redirect on success, which we ignore\n        # (it does not have to match the integration's settings in snowflake)\n        redirect_uri = 'http://localhost:9999'\n        data = {\n            'grant_type': 'refresh_token',\n            'refresh_token': self.token,\n            'redirect_uri': redirect_uri\n        }\n\n        auth = base64.b64encode(\n            f'{self.oauth_client_id}:{self.oauth_client_secret}'\n            .encode('ascii')\n        ).decode('ascii')\n        headers = {\n            'Authorization': f'Basic {auth}',\n            'Content-type': 'application/x-www-form-urlencoded;charset=utf-8'\n        }\n        result = requests.post(token_url, headers=headers, data=data)\n        result_json = result.json()\n        if 'access_token' not in result_json:\n            raise DatabaseException(f'Did not get a token: {result_json}')\n        return result_json['access_token']\n\n    def _get_private_key(self):\n        \"\"\"Get Snowflake private key by path or None.\"\"\"\n        if not self.private_key_path:\n            return None\n\n        if self.private_key_passphrase:\n            encoded_passphrase = self.private_key_passphrase.encode()\n        else:\n            encoded_passphrase = None\n\n        with open(self.private_key_path, 'rb') as key:\n            p_key = serialization.load_pem_private_key(\n                key.read(),\n                password=encoded_passphrase,\n                backend=default_backend())\n\n        return p_key.private_bytes(\n            encoding=serialization.Encoding.DER,\n            format=serialization.PrivateFormat.PKCS8,\n            encryption_algorithm=serialization.NoEncryption())\n\n\nclass SnowflakeConnectionManager(SQLConnectionManager):\n    TYPE = 'snowflake'\n\n    @contextmanager\n    def exception_handler(self, sql):\n        try:\n            yield\n        except snowflake.connector.errors.ProgrammingError as e:\n            msg = str(e)\n\n            logger.debug('Snowflake error: {}'.format(msg))\n\n            if 'Empty SQL statement' in msg:\n                logger.debug(\"got empty sql statement, moving on\")\n            elif 'This session does not have a current database' in msg:\n                self.release()\n                raise FailedToConnectException(\n                    ('{}\\n\\nThis error sometimes occurs when invalid '\n                     'credentials are provided, or when your default role '\n                     'does not have access to use the specified database. '\n                     'Please double check your profile and try again.')\n                    .format(msg))\n            else:\n                self.release()\n                raise DatabaseException(msg)\n        except Exception as e:\n            logger.debug(\"Error running SQL: {}\", sql)\n            logger.debug(\"Rolling back transaction.\")\n            self.release()\n            if isinstance(e, RuntimeException):\n                # during a sql query, an internal to dbt exception was raised.\n                # this sounds a lot like a signal handler and probably has\n                # useful information, so raise it without modification.\n                raise\n            raise RuntimeException(str(e)) from e\n\n    @classmethod\n    def open(cls, connection):\n        if connection.state == 'open':\n            logger.debug('Connection is already open, skipping open.')\n            return connection\n\n        try:\n            creds = connection.credentials\n\n            handle = snowflake.connector.connect(\n                account=creds.account,\n                user=creds.user,\n                database=creds.database,\n                schema=creds.schema,\n                warehouse=creds.warehouse,\n                role=creds.role,\n                autocommit=False,\n                client_session_keep_alive=creds.client_session_keep_alive,\n                application='dbt',\n                **creds.auth_args()\n            )\n\n            connection.handle = handle\n            connection.state = 'open'\n        except snowflake.connector.errors.Error as e:\n            logger.debug(\"Got an error when attempting to open a snowflake \"\n                         \"connection: '{}'\"\n                         .format(e))\n\n            connection.handle = None\n            connection.state = 'fail'\n\n            raise FailedToConnectException(str(e))\n\n    def cancel(self, connection):\n        handle = connection.handle\n        sid = handle.session_id\n\n        connection_name = connection.name\n\n        sql = 'select system$abort_session({})'.format(sid)\n\n        logger.debug(\"Cancelling query '{}' ({})\".format(connection_name, sid))\n\n        _, cursor = self.add_query(sql)\n        res = cursor.fetchone()\n\n        logger.debug(\"Cancel query '{}': {}\".format(connection_name, res))\n\n    @classmethod\n    def get_status(cls, cursor):\n        state = cursor.sqlstate\n\n        if state is None:\n            state = 'SUCCESS'\n\n        return \"{} {}\".format(state, cursor.rowcount)\n\n    @classmethod\n    def _split_queries(cls, sql):\n        \"Splits sql statements at semicolons into discrete queries\"\n\n        sql_s = str(sql)\n        sql_buf = StringIO(sql_s)\n        split_query = snowflake.connector.util_text.split_statements(sql_buf)\n        return [part[0] for part in split_query]\n\n    @classmethod\n    def process_results(cls, column_names, rows):\n        # Override for Snowflake. The datetime objects returned by\n        # snowflake-connector-python are not pickleable, so we need\n        # to replace them with sane timezones\n        fixed = []\n        for row in rows:\n            fixed_row = []\n            for col in row:\n                if isinstance(col, datetime.datetime) and col.tzinfo:\n                    offset = col.utcoffset()\n                    offset_seconds = offset.total_seconds()\n                    new_timezone = pytz.FixedOffset(offset_seconds // 60)\n                    col = col.astimezone(tz=new_timezone)\n                fixed_row.append(col)\n\n            fixed.append(fixed_row)\n\n        return super().process_results(column_names, fixed)\n\n    def add_query(self, sql, auto_begin=True,\n                  bindings=None, abridge_sql_log=False):\n\n        connection = None\n        cursor = None\n\n        if bindings:\n            # The snowflake connector is more strict than, eg., psycopg2 -\n            # which allows any iterable thing to be passed as a binding.\n            bindings = tuple(bindings)\n\n        queries = self._split_queries(sql)\n\n        for individual_query in queries:\n            # hack -- after the last ';', remove comments and don't run\n            # empty queries. this avoids using exceptions as flow control,\n            # and also allows us to return the status of the last cursor\n            without_comments = re.sub(\n                re.compile('^.*(--.*)$', re.MULTILINE),\n                '', individual_query).strip()\n\n            if without_comments == \"\":\n                continue\n\n            connection, cursor = super().add_query(\n                individual_query, auto_begin,\n                bindings=bindings,\n                abridge_sql_log=abridge_sql_log\n            )\n\n        if cursor is None:\n            conn = self.get_thread_connection()\n            if conn is None or conn.name is None:\n                conn_name = '<None>'\n            else:\n                conn_name = conn.name\n\n            raise RuntimeException(\n                \"Tried to run an empty query on model '{}'. If you are \"\n                \"conditionally running\\nsql, eg. in a model hook, make \"\n                \"sure your `else` clause contains valid sql!\\n\\n\"\n                \"Provided SQL:\\n{}\"\n                .format(conn_name, sql)\n            )\n\n        return connection, cursor\n\n    @classmethod\n    def _rollback_handle(cls, connection):\n        \"\"\"On snowflake, rolling back the handle of an aborted session raises\n        an exception.\n        \"\"\"\n        logger.debug('initiating rollback')\n        try:\n            connection.handle.rollback()\n        except snowflake.connector.errors.ProgrammingError as e:\n            msg = str(e)\n            if 'Session no longer exists' not in msg:\n                raise\n", "path": "plugins/snowflake/dbt/adapters/snowflake/connections.py"}], "after_files": [{"content": "import base64\nimport datetime\nimport pytz\nimport re\nfrom contextlib import contextmanager\nfrom dataclasses import dataclass\nfrom io import StringIO\nfrom typing import Optional\n\nfrom cryptography.hazmat.backends import default_backend\nfrom cryptography.hazmat.primitives import serialization\nimport requests\nimport snowflake.connector\nimport snowflake.connector.errors\n\nfrom dbt.exceptions import (\n    InternalException, RuntimeException, FailedToConnectException,\n    DatabaseException, warn_or_error\n)\nfrom dbt.adapters.base import Credentials\nfrom dbt.adapters.sql import SQLConnectionManager\nfrom dbt.logger import GLOBAL_LOGGER as logger\n\n\n_TOKEN_REQUEST_URL = 'https://{}.snowflakecomputing.com/oauth/token-request'\n\n\n@dataclass\nclass SnowflakeCredentials(Credentials):\n    account: str\n    user: str\n    warehouse: Optional[str]\n    role: Optional[str]\n    password: Optional[str]\n    authenticator: Optional[str]\n    private_key_path: Optional[str]\n    private_key_passphrase: Optional[str]\n    token: Optional[str]\n    oauth_client_id: Optional[str]\n    oauth_client_secret: Optional[str]\n    client_session_keep_alive: bool = False\n\n    def __post_init__(self):\n        if (\n            self.authenticator != 'oauth' and\n            (self.oauth_client_secret or self.oauth_client_id or self.token)\n        ):\n            # the user probably forgot to set 'authenticator' like I keep doing\n            warn_or_error(\n                'Authenticator is not set to oauth, but an oauth-only '\n                'parameter is set! Did you mean to set authenticator: oauth?'\n            )\n\n    @property\n    def type(self):\n        return 'snowflake'\n\n    def _connection_keys(self):\n        return (\n            'account', 'user', 'database', 'schema', 'warehouse', 'role',\n            'client_session_keep_alive'\n        )\n\n    def auth_args(self):\n        # Pull all of the optional authentication args for the connector,\n        # let connector handle the actual arg validation\n        result = {}\n        if self.password:\n            result['password'] = self.password\n        if self.authenticator:\n            result['authenticator'] = self.authenticator\n            if self.authenticator == 'oauth':\n                token = self.token\n                # if we have a client ID/client secret, the token is a refresh\n                # token, not an access token\n                if self.oauth_client_id and self.oauth_client_secret:\n                    token = self._get_access_token()\n                elif self.oauth_client_id:\n                    warn_or_error(\n                        'Invalid profile: got an oauth_client_id, but not an '\n                        'oauth_client_secret!'\n                    )\n                elif self.oauth_client_secret:\n                    warn_or_error(\n                        'Invalid profile: got an oauth_client_secret, but not '\n                        'an oauth_client_id!'\n                    )\n\n                result['token'] = token\n        result['private_key'] = self._get_private_key()\n        return result\n\n    def _get_access_token(self) -> str:\n        if self.authenticator != 'oauth':\n            raise InternalException('Can only get access tokens for oauth')\n        missing = any(\n            x is None for x in\n            (self.oauth_client_id, self.oauth_client_secret, self.token)\n        )\n        if missing:\n            raise InternalException(\n                'need a client ID a client secret, and a refresh token to get '\n                'an access token'\n            )\n        # should the full url be a config item?\n        token_url = _TOKEN_REQUEST_URL.format(self.account)\n        # I think this is only used to redirect on success, which we ignore\n        # (it does not have to match the integration's settings in snowflake)\n        redirect_uri = 'http://localhost:9999'\n        data = {\n            'grant_type': 'refresh_token',\n            'refresh_token': self.token,\n            'redirect_uri': redirect_uri\n        }\n\n        auth = base64.b64encode(\n            f'{self.oauth_client_id}:{self.oauth_client_secret}'\n            .encode('ascii')\n        ).decode('ascii')\n        headers = {\n            'Authorization': f'Basic {auth}',\n            'Content-type': 'application/x-www-form-urlencoded;charset=utf-8'\n        }\n        result = requests.post(token_url, headers=headers, data=data)\n        result_json = result.json()\n        if 'access_token' not in result_json:\n            raise DatabaseException(f'Did not get a token: {result_json}')\n        return result_json['access_token']\n\n    def _get_private_key(self):\n        \"\"\"Get Snowflake private key by path or None.\"\"\"\n        if not self.private_key_path:\n            return None\n\n        if self.private_key_passphrase:\n            encoded_passphrase = self.private_key_passphrase.encode()\n        else:\n            encoded_passphrase = None\n\n        with open(self.private_key_path, 'rb') as key:\n            p_key = serialization.load_pem_private_key(\n                key.read(),\n                password=encoded_passphrase,\n                backend=default_backend())\n\n        return p_key.private_bytes(\n            encoding=serialization.Encoding.DER,\n            format=serialization.PrivateFormat.PKCS8,\n            encryption_algorithm=serialization.NoEncryption())\n\n\nclass SnowflakeConnectionManager(SQLConnectionManager):\n    TYPE = 'snowflake'\n\n    @contextmanager\n    def exception_handler(self, sql):\n        try:\n            yield\n        except snowflake.connector.errors.ProgrammingError as e:\n            msg = str(e)\n\n            logger.debug('Snowflake query id: {}'.format(e.sfqid))\n            logger.debug('Snowflake error: {}'.format(msg))\n\n            if 'Empty SQL statement' in msg:\n                logger.debug(\"got empty sql statement, moving on\")\n            elif 'This session does not have a current database' in msg:\n                self.release()\n                raise FailedToConnectException(\n                    ('{}\\n\\nThis error sometimes occurs when invalid '\n                     'credentials are provided, or when your default role '\n                     'does not have access to use the specified database. '\n                     'Please double check your profile and try again.')\n                    .format(msg))\n            else:\n                self.release()\n                raise DatabaseException(msg)\n        except Exception as e:\n            if isinstance(e, snowflake.connector.errors.Error):\n                logger.debug('Snowflake query id: {}'.format(e.sfqid))\n\n            logger.debug(\"Error running SQL: {}\", sql)\n            logger.debug(\"Rolling back transaction.\")\n            self.release()\n            if isinstance(e, RuntimeException):\n                # during a sql query, an internal to dbt exception was raised.\n                # this sounds a lot like a signal handler and probably has\n                # useful information, so raise it without modification.\n                raise\n            raise RuntimeException(str(e)) from e\n\n    @classmethod\n    def open(cls, connection):\n        if connection.state == 'open':\n            logger.debug('Connection is already open, skipping open.')\n            return connection\n\n        try:\n            creds = connection.credentials\n\n            handle = snowflake.connector.connect(\n                account=creds.account,\n                user=creds.user,\n                database=creds.database,\n                schema=creds.schema,\n                warehouse=creds.warehouse,\n                role=creds.role,\n                autocommit=False,\n                client_session_keep_alive=creds.client_session_keep_alive,\n                application='dbt',\n                **creds.auth_args()\n            )\n\n            connection.handle = handle\n            connection.state = 'open'\n        except snowflake.connector.errors.Error as e:\n            logger.debug(\"Got an error when attempting to open a snowflake \"\n                         \"connection: '{}'\"\n                         .format(e))\n\n            connection.handle = None\n            connection.state = 'fail'\n\n            raise FailedToConnectException(str(e))\n\n    def cancel(self, connection):\n        handle = connection.handle\n        sid = handle.session_id\n\n        connection_name = connection.name\n\n        sql = 'select system$abort_session({})'.format(sid)\n\n        logger.debug(\"Cancelling query '{}' ({})\".format(connection_name, sid))\n\n        _, cursor = self.add_query(sql)\n        res = cursor.fetchone()\n\n        logger.debug(\"Cancel query '{}': {}\".format(connection_name, res))\n\n    @classmethod\n    def get_status(cls, cursor):\n        state = cursor.sqlstate\n\n        if state is None:\n            state = 'SUCCESS'\n\n        return \"{} {}\".format(state, cursor.rowcount)\n\n    @classmethod\n    def _split_queries(cls, sql):\n        \"Splits sql statements at semicolons into discrete queries\"\n\n        sql_s = str(sql)\n        sql_buf = StringIO(sql_s)\n        split_query = snowflake.connector.util_text.split_statements(sql_buf)\n        return [part[0] for part in split_query]\n\n    @classmethod\n    def process_results(cls, column_names, rows):\n        # Override for Snowflake. The datetime objects returned by\n        # snowflake-connector-python are not pickleable, so we need\n        # to replace them with sane timezones\n        fixed = []\n        for row in rows:\n            fixed_row = []\n            for col in row:\n                if isinstance(col, datetime.datetime) and col.tzinfo:\n                    offset = col.utcoffset()\n                    offset_seconds = offset.total_seconds()\n                    new_timezone = pytz.FixedOffset(offset_seconds // 60)\n                    col = col.astimezone(tz=new_timezone)\n                fixed_row.append(col)\n\n            fixed.append(fixed_row)\n\n        return super().process_results(column_names, fixed)\n\n    def add_query(self, sql, auto_begin=True,\n                  bindings=None, abridge_sql_log=False):\n\n        connection = None\n        cursor = None\n\n        if bindings:\n            # The snowflake connector is more strict than, eg., psycopg2 -\n            # which allows any iterable thing to be passed as a binding.\n            bindings = tuple(bindings)\n\n        queries = self._split_queries(sql)\n\n        for individual_query in queries:\n            # hack -- after the last ';', remove comments and don't run\n            # empty queries. this avoids using exceptions as flow control,\n            # and also allows us to return the status of the last cursor\n            without_comments = re.sub(\n                re.compile('^.*(--.*)$', re.MULTILINE),\n                '', individual_query).strip()\n\n            if without_comments == \"\":\n                continue\n\n            connection, cursor = super().add_query(\n                individual_query, auto_begin,\n                bindings=bindings,\n                abridge_sql_log=abridge_sql_log\n            )\n\n        if cursor is None:\n            conn = self.get_thread_connection()\n            if conn is None or conn.name is None:\n                conn_name = '<None>'\n            else:\n                conn_name = conn.name\n\n            raise RuntimeException(\n                \"Tried to run an empty query on model '{}'. If you are \"\n                \"conditionally running\\nsql, eg. in a model hook, make \"\n                \"sure your `else` clause contains valid sql!\\n\\n\"\n                \"Provided SQL:\\n{}\"\n                .format(conn_name, sql)\n            )\n\n        return connection, cursor\n\n    @classmethod\n    def _rollback_handle(cls, connection):\n        \"\"\"On snowflake, rolling back the handle of an aborted session raises\n        an exception.\n        \"\"\"\n        logger.debug('initiating rollback')\n        try:\n            connection.handle.rollback()\n        except snowflake.connector.errors.ProgrammingError as e:\n            msg = str(e)\n            if 'Session no longer exists' not in msg:\n                raise\n", "path": "plugins/snowflake/dbt/adapters/snowflake/connections.py"}]}