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
snowflakedb · Dec 12, 2024 · b22b5b4 · b22b5b4
1 parent 5a3e36e
commit b22b5b4
Show file tree

Hide file tree

Showing 5 changed files with 70 additions and 3 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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"