Snow 1733048 Add internal app data to cli

RELEASE-NOTES.md

@@ -31,7 +31,7 @@
   * Add ability to list release channels through `snow app release-channel list` command
   * Add ability to add and remove accounts from release channels through `snow app release-channel add-accounts` and snow app release-channel remove-accounts` commands.
   * Add ability to add/remove versions to/from release channels through `snow app release-channel add-version` and `snow app release-channel remove-version` commands.
-
+* Add support for restricting Snowflake user authentication policy to Snowflake CLI-only.
 ## Fixes and improvements
 * Fixed crashes with older x86_64 Intel CPUs.
 * Fixed inability to add patches to lowercase quoted versions

src/snowflake/cli/_app/constants.py

@@ -17,3 +17,4 @@
 from typing import Literal
 
 PARAM_APPLICATION_NAME: Literal["snowcli"] = "snowcli"
+INTERNAL_APPLICATION_NAME: Literal["SNOWFLAKE_CLI"] = "SNOWFLAKE_CLI"

src/snowflake/cli/_app/snow_connector.py

@@ -21,7 +21,9 @@
 
 import snowflake.connector
 from click.exceptions import ClickException
+from snowflake.cli.__about__ import VERSION
 from snowflake.cli._app.constants import (
+    INTERNAL_APPLICATION_NAME,
     PARAM_APPLICATION_NAME,
 )
 from snowflake.cli._app.secret import SecretType
@@ -35,6 +37,7 @@
     InvalidConnectionConfiguration,
     SnowflakeConnectionError,
 )
+from snowflake.cli.api.feature_flags import FeatureFlag
 from snowflake.cli.api.secure_path import SecurePath
 from snowflake.connector import SnowflakeConnection
 from snowflake.connector.errors import DatabaseError, ForbiddenError
@@ -150,6 +153,8 @@ def connect_to_snowflake(
 
     _update_connection_application_name(connection_parameters)
 
+    _update_internal_application_info(connection_parameters)
+
     try:
         # Whatever output is generated when creating connection,
         # we don't want it in our output. This is particularly important
@@ -238,6 +243,13 @@ def _update_connection_application_name(connection_parameters: Dict):
     connection_parameters.update(connection_application_params)
 
 
+def _update_internal_application_info(connection_parameters: Dict):
+    """Update internal application data if ENABLE_SEPARATE_AUTHENTICATION_POLICY_ID is enabled."""
+    if FeatureFlag.ENABLE_SEPARATE_AUTHENTICATION_POLICY_ID.is_enabled():
+        connection_parameters["internal_application_name"] = INTERNAL_APPLICATION_NAME
+        connection_parameters["internal_application_version"] = VERSION
+
+
 def _load_pem_from_file(private_key_file: str) -> SecretType:
     with SecurePath(private_key_file).open(
         "rb", read_file_limit_mb=DEFAULT_SIZE_LIMIT_MB

src/snowflake/cli/api/feature_flags.py

@@ -64,3 +64,6 @@ class FeatureFlag(FeatureFlagMixin):
         "ENABLE_STREAMLIT_VERSIONED_STAGE", False
     )
     ENABLE_SPCS_LOG_STREAMING = BooleanFlag("ENABLE_SPCS_LOG_STREAMING", False)
+    ENABLE_SEPARATE_AUTHENTICATION_POLICY_ID = BooleanFlag(
+        "ENABLE_SEPARATE_AUTHENTICATION_POLICY_ID", False
+    )

tests/test_snow_connector.py

@@ -222,3 +222,32 @@ def test_returns_nice_error_in_case_of_missing_master_token(runner):
         "When using a session token, you must provide the corresponding master token"
         in result.output
     )
+
+
+@mock.patch("snowflake.connector.connect")
+@pytest.mark.parametrize("feature_flag", [None, True, False])
+def test_internal_application_data_is_sent_if_feature_flag_is_set(
+    mock_connect, runner, feature_flag
+):
+    expected_kwargs = {
+        "application": "SNOWCLI.SQL",
+        "database": "db_for_test",
+        "schema": "test_public",
+        "role": "test_role",
+        "warehouse": "xs",
+        "password": "dummy_password",
+        "application_name": "snowcli",
+    }
+    env = {}
+    if feature_flag is not None:
+        env["SNOWFLAKE_CLI_FEATURES_ENABLE_SEPARATE_AUTHENTICATION_POLICY_ID"] = str(
+            feature_flag
+        )
+    if feature_flag:
+        # internal app data should be disabled by default
+        expected_kwargs["internal_application_name"] = "SNOWFLAKE_CLI"
+        expected_kwargs["internal_application_version"] = "0.0.0-test_patched"
+    with mock.patch.dict(os.environ, env):
+        result = runner.invoke(["sql", "-q", "select 1"])
+    assert result.exit_code == 0
+    mock_connect.assert_called_once_with(**expected_kwargs)