From b22b5b4c3f53f32ba6ccf540bfd67cc8910a02cf Mon Sep 17 00:00:00 2001 From: Abby Shen Date: Thu, 12 Dec 2024 15:56:12 -0800 Subject: [PATCH] SNOW-1846319: Add Feature Flag for log streaming (#1925) * Add Feature Flag for log streaming * New error message * update release notes * fix comment and use new error message --- RELEASE-NOTES.md | 1 + .../cli/_plugins/spcs/services/commands.py | 11 +++- src/snowflake/cli/api/exceptions.py | 10 ++++ src/snowflake/cli/api/feature_flags.py | 1 + tests/spcs/test_services.py | 50 ++++++++++++++++++- 5 files changed, 70 insertions(+), 3 deletions(-) diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md index d2b0da52b7..1f002a8c12 100644 --- a/RELEASE-NOTES.md +++ b/RELEASE-NOTES.md @@ -30,6 +30,7 @@ * Fixed crashes with older x86_64 Intel CPUs. * Fixed inability to add patches to lowercase quoted versions * Fixes label being set to blank instead of None when not provided. +* Added a feature flag `ENABLE_SPCS_LOG_STREAMING` to control the rollout of the log streaming feature # v3.2.0 diff --git a/src/snowflake/cli/_plugins/spcs/services/commands.py b/src/snowflake/cli/_plugins/spcs/services/commands.py index 063a2e3e11..a5d9fdcba0 100644 --- a/src/snowflake/cli/_plugins/spcs/services/commands.py +++ b/src/snowflake/cli/_plugins/spcs/services/commands.py @@ -37,7 +37,11 @@ ) from snowflake.cli.api.commands.snow_typer import SnowTyperFactory from snowflake.cli.api.constants import ObjectType -from snowflake.cli.api.exceptions import IncompatibleParametersError +from snowflake.cli.api.exceptions import ( + FeatureNotEnabledError, + IncompatibleParametersError, +) +from snowflake.cli.api.feature_flags import FeatureFlag from snowflake.cli.api.identifiers import FQN from snowflake.cli.api.output.types import ( CommandResult, @@ -250,6 +254,11 @@ def logs( Retrieves local logs from a service container. """ if follow: + if FeatureFlag.ENABLE_SPCS_LOG_STREAMING.is_disabled(): + raise FeatureNotEnabledError( + "ENABLE_SPCS_LOG_STREAMING", + "Streaming logs from spcs containers is disabled.", + ) if num_lines != DEFAULT_NUM_LINES: raise IncompatibleParametersError(["--follow", "--num-lines"]) if previous_logs: diff --git a/src/snowflake/cli/api/exceptions.py b/src/snowflake/cli/api/exceptions.py index fad1e97c10..2aac4e9608 100644 --- a/src/snowflake/cli/api/exceptions.py +++ b/src/snowflake/cli/api/exceptions.py @@ -229,3 +229,13 @@ def __init__(self, show_obj_query: str): super().__init__( f"Received multiple rows from result of SQL statement: {show_obj_query}. Usage of 'show_specific_object' may not be properly scoped." ) + + +class FeatureNotEnabledError(ClickException): + def __init__(self, feature_name: str, custom_message: Optional[str] = None): + base_message = f"To enable it, add '{feature_name} = true' to '[cli.features]' section of your configuration file." + if custom_message: + message = f"{custom_message} {base_message}" + else: + message = base_message + super().__init__(message) diff --git a/src/snowflake/cli/api/feature_flags.py b/src/snowflake/cli/api/feature_flags.py index d504056e02..2a56458083 100644 --- a/src/snowflake/cli/api/feature_flags.py +++ b/src/snowflake/cli/api/feature_flags.py @@ -63,3 +63,4 @@ class FeatureFlag(FeatureFlagMixin): ENABLE_STREAMLIT_VERSIONED_STAGE = BooleanFlag( "ENABLE_STREAMLIT_VERSIONED_STAGE", False ) + ENABLE_SPCS_LOG_STREAMING = BooleanFlag("ENABLE_SPCS_LOG_STREAMING", False) diff --git a/tests/spcs/test_services.py b/tests/spcs/test_services.py index 3c62648020..b48dedd1ce 100644 --- a/tests/spcs/test_services.py +++ b/tests/spcs/test_services.py @@ -605,7 +605,11 @@ def test_stream_logs_with_include_timestamps_true(mock_sleep, mock_logs): @patch("snowflake.cli._plugins.spcs.services.manager.ServiceManager.execute_query") -def test_logs_incompatible_flags(mock_execute_query, runner): +@patch( + "snowflake.cli.api.feature_flags.FeatureFlag.ENABLE_SPCS_LOG_STREAMING.is_disabled" +) +def test_logs_incompatible_flags(mock_is_disabled, mock_execute_query, runner): + mock_is_disabled.return_value = False result = runner.invoke( [ "spcs", @@ -628,7 +632,13 @@ def test_logs_incompatible_flags(mock_execute_query, runner): @patch("snowflake.cli._plugins.spcs.services.manager.ServiceManager.execute_query") -def test_logs_incompatible_flags_follow_previous_logs(mock_execute_query, runner): +@patch( + "snowflake.cli.api.feature_flags.FeatureFlag.ENABLE_SPCS_LOG_STREAMING.is_disabled" +) +def test_logs_incompatible_flags_follow_previous_logs( + mock_is_disabled, mock_execute_query, runner +): + mock_is_disabled.return_value = False result = runner.invoke( [ "spcs", @@ -653,6 +663,42 @@ def test_logs_incompatible_flags_follow_previous_logs(mock_execute_query, runner ) +@patch( + "snowflake.cli.api.feature_flags.FeatureFlag.ENABLE_SPCS_LOG_STREAMING.is_disabled" +) +def test_logs_streaming_disabled(mock_is_disabled, runner): + mock_is_disabled.return_value = True + result = runner.invoke( + [ + "spcs", + "service", + "logs", + "test_service", + "--container-name", + "test_container", + "--instance-id", + "0", + "--follow", + "--num-lines", + "100", + ] + ) + assert ( + result.exit_code != 0 + ), "Expected a non-zero exit code due to feature flag disabled" + + expected_output = ( + "+- Error ----------------------------------------------------------------------+\n" + "| Streaming logs from spcs containers is disabled. To enable it, add |\n" + "| 'ENABLE_SPCS_LOG_STREAMING = true' to '[cli.features]' section of your |\n" + "| configuration file. |\n" + "+------------------------------------------------------------------------------+\n" + ) + assert ( + result.output == expected_output + ), f"Expected formatted output not found: {result.output}" + + def test_read_yaml(other_directory): tmp_dir = Path(other_directory) spec_path = tmp_dir / "spec.yml"