From 5661991475a564d6028a8ee555e2eb20fbd9e788 Mon Sep 17 00:00:00 2001
From: Abby Shen <abby.shen@snowflake.com>
Date: Thu, 5 Dec 2024 14:46:10 -0800
Subject: [PATCH] Add Feature Flag for log streaming

---
 .../cli/_plugins/spcs/services/commands.py    |  8 +++-
 src/snowflake/cli/api/exceptions.py           |  5 +++
 src/snowflake/cli/api/feature_flags.py        |  1 +
 tests/spcs/test_services.py                   | 43 +++++++++++++++++--
 4 files changed, 52 insertions(+), 5 deletions(-)

diff --git a/src/snowflake/cli/_plugins/spcs/services/commands.py b/src/snowflake/cli/_plugins/spcs/services/commands.py
index 063a2e3e11..e1bdbbcf89 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,8 @@ def logs(
     Retrieves local logs from a service container.
     """
     if follow:
+        if not FeatureFlag.ENABLE_SPCS_LOG_STREAMING.is_enabled():
+            raise FeatureNotEnabledError("Log Streaming")
         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..d0c2703b18 100644
--- a/src/snowflake/cli/api/exceptions.py
+++ b/src/snowflake/cli/api/exceptions.py
@@ -229,3 +229,8 @@ 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(RuntimeError):
+    def __init__(self, feature_name: str):
+        super().__init__(f"The feature '{feature_name}' is not enabled.")
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..4b664ab9cb 100644
--- a/tests/spcs/test_services.py
+++ b/tests/spcs/test_services.py
@@ -25,6 +25,7 @@
 from snowflake.cli._plugins.spcs.services.commands import _service_name_callback
 from snowflake.cli._plugins.spcs.services.manager import ServiceManager
 from snowflake.cli.api.constants import ObjectType
+from snowflake.cli.api.exceptions import FeatureNotEnabledError
 from snowflake.cli.api.identifiers import FQN
 from snowflake.cli.api.project.util import to_string_literal
 from snowflake.connector.cursor import SnowflakeCursor
@@ -605,7 +606,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_enabled"
+)
+def test_logs_incompatible_flags(mock_is_enabled, mock_execute_query, runner):
+    mock_is_enabled.return_value = True
     result = runner.invoke(
         [
             "spcs",
@@ -618,17 +623,23 @@ def test_logs_incompatible_flags(mock_execute_query, runner):
             "0",
             "--follow",
             "--num-lines",
-            "100",
+            "200",
         ]
     )
     assert (
         result.exit_code != 0
-    ), "Expected a non-zero exit code due to incompatible flags"
+    ), "Expected a non-zero exit code due to incompatible parameters"
     assert "Parameters '--follow' and '--num-lines' are incompatible" in result.output
 
 
 @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_enabled"
+)
+def test_logs_incompatible_flags_follow_previous_logs(
+    mock_is_enabled, mock_execute_query, runner
+):
+    mock_is_enabled.return_value = True
     result = runner.invoke(
         [
             "spcs",
@@ -653,6 +664,30 @@ def test_logs_incompatible_flags_follow_previous_logs(mock_execute_query, runner
     )
 
 
+@patch(
+    "snowflake.cli.api.feature_flags.FeatureFlag.ENABLE_SPCS_LOG_STREAMING.is_enabled"
+)
+def test_logs_streaming_disabled(mock_is_enabled, runner):
+    mock_is_enabled.return_value = False
+    with pytest.raises(FeatureNotEnabledError) as exc_info:
+        runner.invoke(
+            [
+                "spcs",
+                "service",
+                "logs",
+                "test_service",
+                "--container-name",
+                "test_container",
+                "--instance-id",
+                "0",
+                "--follow",
+                "--num-lines",
+                "100",
+            ]
+        )
+    assert "The feature 'Log Streaming' is not enabled" in str(exc_info.value)
+
+
 def test_read_yaml(other_directory):
     tmp_dir = Path(other_directory)
     spec_path = tmp_dir / "spec.yml"