snowflake-cli 3.5.0__py3-none-any.whl → 3.7.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (64) hide show
  1. snowflake/cli/__about__.py +13 -1
  2. snowflake/cli/_app/commands_registration/builtin_plugins.py +4 -0
  3. snowflake/cli/_app/loggers.py +2 -2
  4. snowflake/cli/_app/snow_connector.py +7 -6
  5. snowflake/cli/_app/telemetry.py +3 -15
  6. snowflake/cli/_app/version_check.py +4 -4
  7. snowflake/cli/_plugins/auth/__init__.py +11 -0
  8. snowflake/cli/_plugins/auth/keypair/__init__.py +0 -0
  9. snowflake/cli/_plugins/auth/keypair/commands.py +151 -0
  10. snowflake/cli/_plugins/auth/keypair/manager.py +331 -0
  11. snowflake/cli/_plugins/auth/keypair/plugin_spec.py +30 -0
  12. snowflake/cli/_plugins/connection/commands.py +78 -1
  13. snowflake/cli/_plugins/helpers/commands.py +25 -1
  14. snowflake/cli/_plugins/helpers/snowsl_vars_reader.py +133 -0
  15. snowflake/cli/_plugins/init/commands.py +9 -6
  16. snowflake/cli/_plugins/logs/__init__.py +0 -0
  17. snowflake/cli/_plugins/logs/commands.py +105 -0
  18. snowflake/cli/_plugins/logs/manager.py +107 -0
  19. snowflake/cli/_plugins/logs/plugin_spec.py +16 -0
  20. snowflake/cli/_plugins/logs/utils.py +60 -0
  21. snowflake/cli/_plugins/nativeapp/entities/application.py +4 -1
  22. snowflake/cli/_plugins/nativeapp/sf_sql_facade.py +33 -6
  23. snowflake/cli/_plugins/notebook/commands.py +3 -0
  24. snowflake/cli/_plugins/notebook/notebook_entity.py +16 -27
  25. snowflake/cli/_plugins/object/command_aliases.py +3 -1
  26. snowflake/cli/_plugins/object/manager.py +4 -2
  27. snowflake/cli/_plugins/project/commands.py +89 -48
  28. snowflake/cli/_plugins/project/manager.py +57 -23
  29. snowflake/cli/_plugins/project/project_entity_model.py +22 -3
  30. snowflake/cli/_plugins/snowpark/commands.py +15 -2
  31. snowflake/cli/_plugins/spcs/compute_pool/commands.py +17 -5
  32. snowflake/cli/_plugins/sql/manager.py +43 -52
  33. snowflake/cli/_plugins/sql/source_reader.py +230 -0
  34. snowflake/cli/_plugins/stage/manager.py +25 -12
  35. snowflake/cli/_plugins/streamlit/commands.py +3 -0
  36. snowflake/cli/_plugins/streamlit/manager.py +19 -15
  37. snowflake/cli/api/artifacts/upload.py +30 -34
  38. snowflake/cli/api/artifacts/utils.py +8 -6
  39. snowflake/cli/api/cli_global_context.py +7 -2
  40. snowflake/cli/api/commands/decorators.py +11 -2
  41. snowflake/cli/api/commands/flags.py +35 -4
  42. snowflake/cli/api/commands/snow_typer.py +20 -2
  43. snowflake/cli/api/config.py +5 -3
  44. snowflake/cli/api/constants.py +2 -0
  45. snowflake/cli/api/entities/utils.py +29 -16
  46. snowflake/cli/api/errno.py +1 -0
  47. snowflake/cli/api/exceptions.py +75 -27
  48. snowflake/cli/api/feature_flags.py +1 -0
  49. snowflake/cli/api/identifiers.py +2 -0
  50. snowflake/cli/api/plugins/plugin_config.py +2 -2
  51. snowflake/cli/api/project/schemas/template.py +3 -3
  52. snowflake/cli/api/rendering/project_templates.py +3 -3
  53. snowflake/cli/api/rendering/sql_templates.py +2 -2
  54. snowflake/cli/api/rest_api.py +2 -3
  55. snowflake/cli/{_app → api}/secret.py +4 -1
  56. snowflake/cli/api/secure_path.py +16 -4
  57. snowflake/cli/api/sql_execution.py +8 -4
  58. snowflake/cli/api/utils/definition_rendering.py +14 -8
  59. snowflake/cli/api/utils/templating_functions.py +4 -4
  60. {snowflake_cli-3.5.0.dist-info → snowflake_cli-3.7.0.dist-info}/METADATA +11 -11
  61. {snowflake_cli-3.5.0.dist-info → snowflake_cli-3.7.0.dist-info}/RECORD +64 -52
  62. {snowflake_cli-3.5.0.dist-info → snowflake_cli-3.7.0.dist-info}/WHEEL +0 -0
  63. {snowflake_cli-3.5.0.dist-info → snowflake_cli-3.7.0.dist-info}/entry_points.txt +0 -0
  64. {snowflake_cli-3.5.0.dist-info → snowflake_cli-3.7.0.dist-info}/licenses/LICENSE +0 -0
@@ -0,0 +1,331 @@
1
+ from enum import Enum
2
+ from typing import Dict, List, Optional, Tuple
3
+
4
+ from click import ClickException
5
+ from cryptography.hazmat.primitives import serialization
6
+ from cryptography.hazmat.primitives.asymmetric import rsa
7
+ from snowflake.cli._plugins.object.manager import ObjectManager
8
+ from snowflake.cli.api import exceptions
9
+ from snowflake.cli.api.cli_global_context import (
10
+ _CliGlobalContextAccess,
11
+ get_cli_context,
12
+ )
13
+ from snowflake.cli.api.config import (
14
+ connection_exists,
15
+ get_connection_dict,
16
+ set_config_value,
17
+ )
18
+ from snowflake.cli.api.console import cli_console
19
+ from snowflake.cli.api.constants import ObjectType
20
+ from snowflake.cli.api.identifiers import FQN
21
+ from snowflake.cli.api.secret import SecretType
22
+ from snowflake.cli.api.secure_path import SecurePath
23
+ from snowflake.cli.api.sql_execution import SqlExecutionMixin
24
+ from snowflake.connector import DictCursor
25
+ from snowflake.connector.cursor import SnowflakeCursor
26
+
27
+
28
+ class PublicKeyProperty(Enum):
29
+ RSA_PUBLIC_KEY = "RSA_PUBLIC_KEY"
30
+ RSA_PUBLIC_KEY_2 = "RSA_PUBLIC_KEY_2"
31
+
32
+
33
+ class AuthManager(SqlExecutionMixin):
34
+ def setup(
35
+ self,
36
+ connection_name: Optional[str],
37
+ key_length: int,
38
+ output_path: SecurePath,
39
+ private_key_passphrase: SecretType,
40
+ ):
41
+ # When the user provide new connection name
42
+ if connection_name and connection_exists(connection_name):
43
+ raise ClickException(
44
+ f"Connection with name {connection_name} already exists."
45
+ )
46
+
47
+ cli_context = get_cli_context()
48
+ # When the use not provide connection name, so we overwrite the current connection
49
+ if not connection_name:
50
+ connection_name = cli_context.connection_context.connection_name
51
+
52
+ key_name = AuthManager._get_free_key_name(output_path, connection_name) # type: ignore[arg-type]
53
+ self._generate_key_pair_and_set_public_key(
54
+ user=cli_context.connection.user,
55
+ key_length=key_length,
56
+ output_path=output_path,
57
+ key_name=key_name, # type: ignore[arg-type]
58
+ private_key_passphrase=private_key_passphrase,
59
+ )
60
+
61
+ self._create_or_update_connection(
62
+ current_connection=cli_context.connection_context.connection_name,
63
+ connection_name=connection_name, # type: ignore[arg-type]
64
+ private_key_path=self._get_private_key_path(
65
+ output_path=output_path, key_name=key_name # type: ignore[arg-type]
66
+ ),
67
+ )
68
+
69
+ def rotate(
70
+ self,
71
+ key_length: int,
72
+ output_path: SecurePath,
73
+ private_key_passphrase: SecretType,
74
+ ):
75
+ cli_context = get_cli_context()
76
+ connection_name = cli_context.connection_context.connection_name
77
+
78
+ self._ensure_connection_has_private_key(
79
+ cli_context.connection_context.connection_name
80
+ )
81
+
82
+ public_key, public_key_2 = self._get_public_keys()
83
+
84
+ if not public_key and not public_key_2:
85
+ raise ClickException("No public key found. Use the setup command first.")
86
+
87
+ if public_key_2:
88
+ self.set_public_key(
89
+ cli_context.connection.user,
90
+ PublicKeyProperty.RSA_PUBLIC_KEY,
91
+ public_key_2,
92
+ )
93
+
94
+ key_name = AuthManager._get_free_key_name(output_path, connection_name)
95
+ public_key = self._generate_keys_and_return_public_key(
96
+ key_length=key_length,
97
+ output_path=output_path,
98
+ key_name=key_name,
99
+ private_key_passphrase=private_key_passphrase,
100
+ )
101
+ self.set_public_key(
102
+ cli_context.connection.user, PublicKeyProperty.RSA_PUBLIC_KEY_2, public_key
103
+ )
104
+ self._create_or_update_connection(
105
+ current_connection=cli_context.connection_context.connection_name,
106
+ connection_name=connection_name,
107
+ private_key_path=self._get_private_key_path(
108
+ output_path=output_path, key_name=key_name
109
+ ),
110
+ )
111
+
112
+ def _generate_key_pair_and_set_public_key(
113
+ self,
114
+ user: str,
115
+ key_length: int,
116
+ output_path: SecurePath,
117
+ key_name: str,
118
+ private_key_passphrase: SecretType,
119
+ ):
120
+ public_key_exists, public_key_2_exists = self._get_public_keys()
121
+
122
+ if public_key_exists or public_key_2_exists:
123
+ raise exceptions.CouldNotSetKeyPairError()
124
+
125
+ if not output_path.exists():
126
+ output_path.mkdir(parents=True)
127
+
128
+ public_key = self._generate_keys_and_return_public_key(
129
+ key_length=key_length,
130
+ output_path=output_path,
131
+ key_name=key_name, # type: ignore[arg-type]
132
+ private_key_passphrase=private_key_passphrase,
133
+ )
134
+ self.set_public_key(user, PublicKeyProperty.RSA_PUBLIC_KEY, public_key)
135
+
136
+ def list_keys(self) -> List[Dict]:
137
+ key_properties = [
138
+ "RSA_PUBLIC_KEY",
139
+ "RSA_PUBLIC_KEY_FP",
140
+ "RSA_PUBLIC_KEY_LAST_SET_TIME",
141
+ "RSA_PUBLIC_KEY_2",
142
+ "RSA_PUBLIC_KEY_2_FP",
143
+ "RSA_PUBLIC_KEY_2_LAST_SET_TIME",
144
+ ]
145
+ cursor = ObjectManager(connection=self._conn).describe(
146
+ object_type=ObjectType.USER.value.sf_name,
147
+ fqn=FQN.from_string(self._conn.user),
148
+ cursor_class=DictCursor,
149
+ )
150
+ only_public_key_properties = [
151
+ p for p in cursor.fetchall() if p.get("property") in key_properties
152
+ ]
153
+ return only_public_key_properties
154
+
155
+ def set_public_key(
156
+ self, user: str, public_key_property: PublicKeyProperty, public_key: str
157
+ ) -> SnowflakeCursor:
158
+ return self.execute_query(
159
+ f"ALTER USER {user} SET {public_key_property.value}='{public_key}'"
160
+ )
161
+
162
+ def remove_public_key(
163
+ self, public_key_property: PublicKeyProperty
164
+ ) -> SnowflakeCursor:
165
+ cli_context = get_cli_context()
166
+ return self.execute_query(
167
+ f"ALTER USER {cli_context.connection.user} UNSET {public_key_property.value}"
168
+ )
169
+
170
+ def status(self):
171
+ cli_context = get_cli_context()
172
+ self._ensure_connection_has_private_key(
173
+ cli_context.connection_context.connection_name
174
+ )
175
+ cli_console.step("Private key set for connection - OK")
176
+ self._check_connection(cli_context)
177
+ cli_console.step("Test connection - OK")
178
+
179
+ def extend_connection_add(
180
+ self,
181
+ connection_name: str,
182
+ connection_options: Dict,
183
+ key_length: int,
184
+ output_path: SecurePath,
185
+ private_key_passphrase: SecretType,
186
+ ) -> Dict:
187
+ key_name = AuthManager._get_free_key_name(output_path, connection_name)
188
+
189
+ self._generate_key_pair_and_set_public_key(
190
+ user=connection_options["user"],
191
+ key_length=key_length,
192
+ output_path=output_path,
193
+ key_name=key_name,
194
+ private_key_passphrase=private_key_passphrase,
195
+ )
196
+
197
+ connection_options["authenticator"] = "SNOWFLAKE_JWT"
198
+ connection_options["private_key_file"] = str(
199
+ self._get_private_key_path(output_path=output_path, key_name=key_name).path
200
+ )
201
+ if connection_options.get("password"):
202
+ del connection_options["password"]
203
+
204
+ return connection_options
205
+
206
+ @staticmethod
207
+ def _ensure_connection_has_private_key(connection_name: str) -> None:
208
+ connection = get_connection_dict(connection_name)
209
+ if not connection.get("private_key_file") and not connection.get(
210
+ "private_key_path"
211
+ ):
212
+ raise ClickException(
213
+ f"The private key is not set in {connection_name} connection."
214
+ )
215
+
216
+ @staticmethod
217
+ def _check_connection(cli_context: _CliGlobalContextAccess) -> None:
218
+ cli_context.connection
219
+
220
+ def _get_public_keys(self) -> Tuple[str, str]:
221
+ keys = self.list_keys()
222
+ public_key = ""
223
+ public_key_2 = ""
224
+ for p in keys:
225
+ if (
226
+ p.get("property") == PublicKeyProperty.RSA_PUBLIC_KEY.value
227
+ and p.get("value") != "null"
228
+ ):
229
+ public_key = p.get("value") # type: ignore
230
+ if (
231
+ p.get("property") == PublicKeyProperty.RSA_PUBLIC_KEY_2.value
232
+ and p.get("value") != "null"
233
+ ):
234
+ public_key_2 = p.get("value") # type: ignore
235
+ return public_key, public_key_2
236
+
237
+ @staticmethod
238
+ def _generate_keys_and_return_public_key(
239
+ key_length: int,
240
+ output_path: SecurePath,
241
+ key_name: str,
242
+ private_key_passphrase: SecretType,
243
+ ) -> str:
244
+ private_key = rsa.generate_private_key(
245
+ public_exponent=65537,
246
+ key_size=key_length,
247
+ )
248
+
249
+ if private_key_passphrase:
250
+ pem = SecretType(
251
+ private_key.private_bytes(
252
+ encoding=serialization.Encoding.PEM,
253
+ format=serialization.PrivateFormat.PKCS8,
254
+ encryption_algorithm=serialization.BestAvailableEncryption(
255
+ private_key_passphrase.value.encode("utf-8")
256
+ ),
257
+ )
258
+ )
259
+ cli_console.message(
260
+ "Set the `PRIVATE_KEY_PASSPHRASE` environment variable before using the connection."
261
+ )
262
+ else:
263
+ pem = SecretType(
264
+ private_key.private_bytes(
265
+ encoding=serialization.Encoding.PEM,
266
+ format=serialization.PrivateFormat.PKCS8,
267
+ encryption_algorithm=serialization.NoEncryption(),
268
+ )
269
+ )
270
+
271
+ with AuthManager._get_private_key_path(output_path, key_name).open(
272
+ mode="wb"
273
+ ) as file:
274
+ file.write(pem.value)
275
+
276
+ public_key = private_key.public_key()
277
+ public_pem = public_key.public_bytes(
278
+ encoding=serialization.Encoding.PEM,
279
+ format=serialization.PublicFormat.SubjectPublicKeyInfo,
280
+ )
281
+ with AuthManager._get_public_key_path(output_path, key_name).open(
282
+ mode="wb"
283
+ ) as file:
284
+ file.write(public_pem)
285
+
286
+ return public_pem.decode("utf-8")
287
+
288
+ @staticmethod
289
+ def _get_free_key_name(output_path: SecurePath, key_name: str) -> str:
290
+ new_private_key = f"{key_name}.p8"
291
+ new_public_key = f"{key_name}.pub"
292
+ new_key_name = key_name
293
+ counter = 1
294
+
295
+ while (
296
+ (output_path / new_private_key).exists()
297
+ and (output_path / new_public_key).exists()
298
+ and counter <= 100
299
+ ):
300
+ new_key_name = f"{key_name}_{counter}"
301
+ new_private_key = f"{new_key_name}.p8"
302
+ new_public_key = f"{new_key_name}.pub"
303
+ counter += 1
304
+
305
+ if counter == 100:
306
+ raise ClickException(
307
+ "Too many key pairs with the same name in the output directory."
308
+ )
309
+
310
+ return new_key_name
311
+
312
+ @staticmethod
313
+ def _get_private_key_path(output_path: SecurePath, key_name: str) -> SecurePath:
314
+ return (output_path / f"{key_name}.p8").resolve()
315
+
316
+ @staticmethod
317
+ def _get_public_key_path(output_path: SecurePath, key_name: str) -> SecurePath:
318
+ return (output_path / f"{key_name}.pub").resolve()
319
+
320
+ @staticmethod
321
+ def _create_or_update_connection(
322
+ current_connection: Optional[str],
323
+ connection_name: str,
324
+ private_key_path: SecurePath,
325
+ ):
326
+ connection = get_connection_dict(current_connection)
327
+ connection.pop("password", None)
328
+ connection["authenticator"] = "SNOWFLAKE_JWT"
329
+ connection["private_key_file"] = str(private_key_path.path)
330
+
331
+ set_config_value(["connections", connection_name], value=connection)
@@ -0,0 +1,30 @@
1
+ # Copyright (c) 2024 Snowflake Inc.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ from snowflake.cli._plugins.auth import app
16
+ from snowflake.cli.api.plugins.command import (
17
+ SNOWCLI_ROOT_COMMAND_PATH,
18
+ CommandSpec,
19
+ CommandType,
20
+ plugin_hook_impl,
21
+ )
22
+
23
+
24
+ @plugin_hook_impl
25
+ def command_spec():
26
+ return CommandSpec(
27
+ parent_command_path=SNOWCLI_ROOT_COMMAND_PATH,
28
+ command_type=CommandType.COMMAND_GROUP,
29
+ typer_instance=app.create_instance(),
30
+ )
@@ -16,8 +16,9 @@ from __future__ import annotations
16
16
 
17
17
  import logging
18
18
  import os.path
19
+ from copy import deepcopy
19
20
  from pathlib import Path
20
- from typing import Optional
21
+ from typing import Dict, Optional, Tuple
21
22
 
22
23
  import typer
23
24
  from click import ( # type: ignore
@@ -28,10 +29,14 @@ from click import ( # type: ignore
28
29
  )
29
30
  from click.core import ParameterSource # type: ignore
30
31
  from snowflake import connector
32
+ from snowflake.cli._app.snow_connector import connect_to_snowflake
33
+ from snowflake.cli._plugins.auth.keypair.commands import KEY_PAIR_DEFAULT_PATH
34
+ from snowflake.cli._plugins.auth.keypair.manager import AuthManager
31
35
  from snowflake.cli._plugins.connection.util import (
32
36
  strip_if_value_present,
33
37
  )
34
38
  from snowflake.cli._plugins.object.manager import ObjectManager
39
+ from snowflake.cli.api import exceptions
35
40
  from snowflake.cli.api.cli_global_context import get_cli_context
36
41
  from snowflake.cli.api.commands.flags import (
37
42
  PLAIN_PASSWORD_MSG,
@@ -61,12 +66,15 @@ from snowflake.cli.api.config import (
61
66
  )
62
67
  from snowflake.cli.api.console import cli_console
63
68
  from snowflake.cli.api.constants import ObjectType
69
+ from snowflake.cli.api.feature_flags import FeatureFlag
64
70
  from snowflake.cli.api.output.types import (
65
71
  CollectionResult,
66
72
  CommandResult,
67
73
  MessageResult,
68
74
  ObjectResult,
69
75
  )
76
+ from snowflake.cli.api.secret import SecretType
77
+ from snowflake.cli.api.secure_path import SecurePath
70
78
  from snowflake.connector import ProgrammingError
71
79
  from snowflake.connector.constants import CONNECTIONS_FILE
72
80
 
@@ -282,6 +290,13 @@ def add(
282
290
  if connection_exists(connection_name):
283
291
  raise UsageError(f"Connection {connection_name} already exists")
284
292
 
293
+ if FeatureFlag.ENABLE_AUTH_KEYPAIR.is_enabled() and not no_interactive:
294
+ connection_options, keypair_error = _extend_add_with_key_pair(
295
+ connection_name, connection_options
296
+ )
297
+ else:
298
+ keypair_error = ""
299
+
285
300
  connections_file = add_connection_to_proper_file(
286
301
  connection_name,
287
302
  ConnectionConfig(**connection_options),
@@ -289,6 +304,12 @@ def add(
289
304
  if set_as_default:
290
305
  set_config_value(path=["default_connection_name"], value=connection_name)
291
306
 
307
+ if keypair_error:
308
+ return MessageResult(
309
+ f"Wrote new password-based connection {connection_name} to {connections_file}, "
310
+ f"however there were some issues during key pair setup. Review the following error and check 'snow auth keypair' "
311
+ f"commands to setup key pair authentication:\n * {keypair_error}"
312
+ )
292
313
  return MessageResult(
293
314
  f"Wrote new connection {connection_name} to {connections_file}"
294
315
  )
@@ -402,3 +423,59 @@ def generate_jwt(
402
423
  return MessageResult(token)
403
424
  except (ValueError, TypeError) as err:
404
425
  raise ClickException(str(err))
426
+
427
+
428
+ def _extend_add_with_key_pair(
429
+ connection_name: str, connection_options: Dict
430
+ ) -> Tuple[Dict, str]:
431
+ if not _should_extend_with_key_pair(connection_options):
432
+ return connection_options, ""
433
+
434
+ configure_key_pair = typer.confirm(
435
+ "Do you want to configure key pair authentication?",
436
+ default=False,
437
+ )
438
+ if not configure_key_pair:
439
+ return connection_options, ""
440
+
441
+ key_length = typer.prompt(
442
+ "Key length",
443
+ default=2048,
444
+ show_default=True,
445
+ )
446
+
447
+ output_path = typer.prompt(
448
+ "Output path",
449
+ default=KEY_PAIR_DEFAULT_PATH,
450
+ show_default=True,
451
+ value_proc=lambda value: SecurePath(value),
452
+ )
453
+ private_key_passphrase = typer.prompt(
454
+ "Private key passphrase",
455
+ default="",
456
+ hide_input=True,
457
+ show_default=False,
458
+ value_proc=lambda value: SecretType(value),
459
+ )
460
+ connection = connect_to_snowflake(temporary_connection=True, **connection_options)
461
+ try:
462
+ connection_options = AuthManager(connection=connection).extend_connection_add(
463
+ connection_name=connection_name,
464
+ connection_options=deepcopy(connection_options),
465
+ key_length=key_length,
466
+ output_path=output_path,
467
+ private_key_passphrase=private_key_passphrase,
468
+ )
469
+ except exceptions.CouldNotSetKeyPairError:
470
+ return connection_options, "The public key is set already."
471
+ except Exception as e:
472
+ return connection_options, str(e)
473
+ return connection_options, ""
474
+
475
+
476
+ def _should_extend_with_key_pair(connection_options: Dict) -> bool:
477
+ return (
478
+ connection_options.get("password") is not None
479
+ and connection_options.get("private_key_file") is None
480
+ and connection_options.get("private_key_path") is None
481
+ )
@@ -15,12 +15,14 @@
15
15
  from __future__ import annotations
16
16
 
17
17
  import logging
18
+ import os
18
19
  from pathlib import Path
19
20
  from typing import Any, List, Optional
20
21
 
21
22
  import typer
22
23
  import yaml
23
24
  from click import ClickException
25
+ from snowflake.cli._plugins.helpers.snowsl_vars_reader import check_env_vars
24
26
  from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
25
27
  from snowflake.cli.api.config import (
26
28
  ConnectionConfig,
@@ -29,7 +31,12 @@ from snowflake.cli.api.config import (
29
31
  set_config_value,
30
32
  )
31
33
  from snowflake.cli.api.console import cli_console
32
- from snowflake.cli.api.output.types import CommandResult, MessageResult
34
+ from snowflake.cli.api.output.types import (
35
+ CollectionResult,
36
+ CommandResult,
37
+ MessageResult,
38
+ MultipleResults,
39
+ )
33
40
  from snowflake.cli.api.project.definition_conversion import (
34
41
  convert_project_definition_to_v2,
35
42
  )
@@ -293,3 +300,20 @@ def _validate_and_save_connections_imported_from_snowsql(
293
300
  path=["default_connection_name"],
294
301
  value=default_cli_connection_name,
295
302
  )
303
+
304
+
305
+ @app.command(name="check-snowsql-env-vars", requires_connection=False)
306
+ def check_snowsql_env_vars(**options):
307
+ """Check if there are any SnowSQL environment variables set."""
308
+
309
+ env_vars = os.environ.copy()
310
+ discovered, unused, summary = check_env_vars(env_vars)
311
+
312
+ results = []
313
+ if discovered:
314
+ results.append(CollectionResult(discovered))
315
+ if unused:
316
+ results.append(CollectionResult(unused))
317
+
318
+ results.append(MessageResult(summary))
319
+ return MultipleResults(results)
@@ -0,0 +1,133 @@
1
+ from __future__ import annotations
2
+
3
+ EnvVars = dict[str, str]
4
+
5
+ DicoveredVars = list[dict[str, str]]
6
+ UnusedVars = list[dict[str, str]]
7
+ Summary = str
8
+
9
+ CheckResult = tuple[DicoveredVars, UnusedVars, Summary]
10
+
11
+ KNOWN_SNOWSQL_ENV_VARS = {
12
+ "SNOWSQL_ACCOUNT": {
13
+ "Found": "SNOWSQL_ACCOUNT",
14
+ "Suggested": "SNOWFLAKE_ACCOUNT",
15
+ "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-connections#use-environment-variables-for-snowflake-credentials",
16
+ },
17
+ "SNOWSQL_PWD": {
18
+ "Found": "SNOWSQL_PASSWORD",
19
+ "Suggested": "SNOWFLAKE_PASSWORD",
20
+ "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-connections#use-environment-variables-for-snowflake-credentials",
21
+ },
22
+ "SNOWSQL_USER": {
23
+ "Found": "SNOWSQL_USER",
24
+ "Suggested": "SNOWFLAKE_USER",
25
+ "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-connections#use-environment-variables-for-snowflake-credentials",
26
+ },
27
+ "SNOWSQL_REGION": {
28
+ "Found": "SNOWSQL_REGION",
29
+ "Suggested": "SNOWFLAKE_REGION",
30
+ "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-connections#use-environment-variables-for-snowflake-credentials",
31
+ },
32
+ "SNOWSQL_ROLE": {
33
+ "Found": "SNOWSQL_ROLE",
34
+ "Suggested": "SNOWFLAKE_ROLE",
35
+ "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-connections#use-environment-variables-for-snowflake-credentials",
36
+ },
37
+ "SNOWSQL_WAREHOUSE": {
38
+ "Found": "SNOWSQL_WAREHOUSE",
39
+ "Suggested": "SNOWFLAKE_WAREHOUSE",
40
+ "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-connections#use-environment-variables-for-snowflake-credentials",
41
+ },
42
+ "SNOWSQL_DATABASE": {
43
+ "Found": "SNOWSQL_DATABASE",
44
+ "Suggested": "SNOWFLAKE_DATABASE",
45
+ "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-connections#use-environment-variables-for-snowflake-credentials",
46
+ },
47
+ "SNOWSQL_SCHEMA": {
48
+ "found": "SNOWSQL_SCHEMA",
49
+ "Suggested": "SNOWFLAKE_SCHEMA",
50
+ "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-connections#use-environment-variables-for-snowflake-credentials",
51
+ },
52
+ "SNOWSQL_HOST": {
53
+ "Found": "SNOWSQL_HOST",
54
+ "Suggested": "SNOWFLAKE_HOST",
55
+ "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-connections#use-environment-variables-for-snowflake-credentials",
56
+ },
57
+ "SNOWSQL_PORT": {
58
+ "Found": "SNOWSQL_PORT",
59
+ "Suggested": "SNOWFLAKE_PORT",
60
+ "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-connections#use-environment-variables-for-snowflake-credentials",
61
+ },
62
+ "SNOWSQL_PROTOCOL": {
63
+ "Found": "SNOWSQL_PROTOCOL",
64
+ "Suggested": "SNOWFLAKE_PROTOCOL",
65
+ "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-connections#use-environment-variables-for-snowflake-credentials",
66
+ },
67
+ "SNOWSQL_PROXY_HOST": {
68
+ "Found": "SNOWSQL_PROXY_HOST",
69
+ "Suggested": "PROXY_HOST",
70
+ "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-cli#use-a-proxy-server",
71
+ },
72
+ "SNOWSQL_PROXY_PORT": {
73
+ "Found": "SNOWSQL_PROXY_HOST",
74
+ "Suggested": "PROXY_HOST",
75
+ "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-cli#use-a-proxy-server",
76
+ },
77
+ "SNOWSQL_PROXY_USER": {
78
+ "Found": "SNOWSQL_PROXY_PORT",
79
+ "Suggested": "PROXY_PORT",
80
+ "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-cli#use-a-proxy-server",
81
+ },
82
+ "SNOWSQL_PROXY_PWD": {
83
+ "Found": "SNOWSQL_PROXY_PWD",
84
+ "Suggested": "PROXY_PWD",
85
+ "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-cli#use-a-proxy-server",
86
+ },
87
+ "SNOWSQL_PRIVATE_KEY_PASSPHRASE": {
88
+ "Found": "SNOWSQL_PRIVATE_KEY_PASSPHRASE",
89
+ "Suggested": "PRIVATE_KEY_PASSPHRASE",
90
+ "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-connections",
91
+ },
92
+ "EXIT_ON_ERROR": {
93
+ "Found": "EXIT_ON_ERROR",
94
+ "Suggested": "SNOWFLAKE_ENHANCED_EXIT_CODES",
95
+ "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/sql/execute-sql",
96
+ },
97
+ }
98
+
99
+
100
+ def check_env_vars(variables: EnvVars) -> CheckResult:
101
+ """Checks passed dict objects for possible SnowSQL variables.
102
+
103
+ Returns tuple of
104
+ - sequence of variables that can be adjusted
105
+ - sequence of variables that have no corresponding variables
106
+ - a summary messages
107
+ """
108
+ discovered: DicoveredVars = []
109
+ unused: UnusedVars = []
110
+
111
+ prefix_matched = (e for e in variables if e.lower().startswith("snowsql"))
112
+
113
+ for var in prefix_matched:
114
+ if suggestion := KNOWN_SNOWSQL_ENV_VARS.get(var, None):
115
+ discovered.append(suggestion)
116
+ else:
117
+ unused.append(
118
+ {
119
+ "Found": var,
120
+ "Suggested": "n/a",
121
+ "Additional info": "Unused variable",
122
+ }
123
+ )
124
+
125
+ discovered_cnt = len(discovered)
126
+ unused_cnt = len(unused)
127
+
128
+ summary: Summary = (
129
+ f"Found {discovered_cnt + unused_cnt} SnowSQL environment variables,"
130
+ f" {discovered_cnt} with replacements, {unused_cnt} unused."
131
+ )
132
+
133
+ return discovered, unused, summary