Merge latest main: combine client request properties with correlation ID tracing and global timeout

Author: Copilot

fabric_rti_mcp/__init__.py

@@ -1,6 +1,6 @@
 try:
     from importlib.metadata import version
 
-    __version__ = version("fabric-rti-mcp")
+    __version__ = version("microsoft-fabric-rti-mcp")
 except Exception:
     __version__ = "0.0.0.dev0"

fabric_rti_mcp/common.py

@@ -21,10 +21,14 @@ class GlobalFabricRTIConfig:
     @staticmethod
     def from_env() -> GlobalFabricRTIConfig:
         return GlobalFabricRTIConfig(
-            fabric_api_base=os.getenv(GlobalFabricRTIEnvVarNames.default_fabric_api_base, DEFAULT_FABRIC_API_BASE)
+            fabric_api_base=os.getenv(GlobalFabricRTIEnvVarNames.default_fabric_api_base, DEFAULT_FABRIC_API_BASE),
         )
 
     @staticmethod
     def existing_env_vars() -> list[str]:
         """Return a list of environment variable names that are currently set."""
         return [GlobalFabricRTIEnvVarNames.default_fabric_api_base]
+
+
+# Global configuration instance
+config = GlobalFabricRTIConfig.from_env()

fabric_rti_mcp/kusto/kusto_config.py

@@ -24,6 +24,7 @@ class KustoEnvVarNames:
     known_services = "KUSTO_KNOWN_SERVICES"
     eager_connect = "KUSTO_EAGER_CONNECT"
     allow_unknown_services = "KUSTO_ALLOW_UNKNOWN_SERVICES"
+    timeout = "FABRIC_RTI_KUSTO_TIMEOUT"
 
     @staticmethod
     def all() -> List[str]:
@@ -35,6 +36,7 @@ def all() -> List[str]:
             KustoEnvVarNames.known_services,
             KustoEnvVarNames.eager_connect,
             KustoEnvVarNames.allow_unknown_services,
+            KustoEnvVarNames.timeout,
         ]
 
 
@@ -52,6 +54,8 @@ class KustoConfig:
     # Security setting to allow unknown services. If this is set to False,
     # only services in known_services will be allowed.
     allow_unknown_services: bool = True
+    # Global timeout for all Kusto operations in seconds
+    timeout_seconds: Optional[int] = None
 
     @staticmethod
     def from_env() -> KustoConfig:
@@ -72,6 +76,16 @@ def from_env() -> KustoConfig:
         eager_connect = os.getenv(KustoEnvVarNames.eager_connect, "false").lower() in ("true", "1")
         allow_unknown_services = os.getenv(KustoEnvVarNames.allow_unknown_services, "true").lower() in ("true", "1")
 
+        # Parse timeout configuration
+        timeout_seconds = None
+        timeout_env = os.getenv(KustoEnvVarNames.timeout)
+        if timeout_env:
+            try:
+                timeout_seconds = int(timeout_env)
+            except ValueError:
+                # Ignore invalid timeout values
+                pass
+
         if known_services_string:
             try:
                 known_services_json = json.loads(known_services_string)
@@ -80,7 +94,12 @@ def from_env() -> KustoConfig:
                 logger.error(f"Failed to parse {KustoEnvVarNames.known_services}: {e}. Skipping known services.")
 
         return KustoConfig(
-            default_service, open_ai_embedding_endpoint, known_services, eager_connect, allow_unknown_services
+            default_service,
+            open_ai_embedding_endpoint,
+            known_services,
+            eager_connect,
+            allow_unknown_services,
+            timeout_seconds,
         )
 
     @staticmethod

fabric_rti_mcp/kusto/kusto_service.py

@@ -12,6 +12,7 @@
 )
 
 from fabric_rti_mcp import __version__  # type: ignore
+from fabric_rti_mcp.common import logger
 from fabric_rti_mcp.kusto.kusto_config import KustoConfig
 from fabric_rti_mcp.kusto.kusto_connection import KustoConnection, sanitize_uri
 from fabric_rti_mcp.kusto.kusto_response_formatter import format_results
@@ -100,7 +101,15 @@ def _crp(
     if not is_destructive and not ignore_readonly:
         crp.set_option("request_readonly", True)
 
-    # Apply any additional properties provided by the user
+    # Set global timeout if configured
+    if CONFIG.timeout_seconds is not None:
+        # Convert seconds to timespan format (HH:MM:SS)
+        hours, remainder = divmod(CONFIG.timeout_seconds, 3600)
+        minutes, seconds = divmod(remainder, 60)
+        timeout_str = f"{hours:02d}:{minutes:02d}:{seconds:02d}"
+        crp.set_option("servertimeout", timeout_str)
+
+    # Apply any additional properties provided by the user (can override global settings)
     if additional_properties:
         for key, value in additional_properties.items():
             crp.set_option(key, value)
@@ -120,18 +129,27 @@ def _execute(
     caller_func = caller_frame.f_globals.get(action_name)  # type: ignore
     is_destructive = hasattr(caller_func, "_is_destructive")
 
-    connection = get_kusto_connection(cluster_uri)
-    client = connection.query_client
+    # Generate correlation ID for tracing
+    crp = _crp(action_name, is_destructive, readonly_override, client_request_properties)
+    correlation_id = crp.client_request_id  # type: ignore
 
-    # agents can send messy inputs
-    query = query.strip()
+    try:
+        connection = get_kusto_connection(cluster_uri)
+        client = connection.query_client
 
-    database = database or connection.default_database
-    database = database.strip()
+        # agents can send messy inputs
+        query = query.strip()
 
-    crp = _crp(action_name, is_destructive, readonly_override, client_request_properties)
-    result_set = client.execute(database, query, crp)
-    return format_results(result_set)
+        database = database or connection.default_database
+        database = database.strip()
+
+        result_set = client.execute(database, query, crp)
+        return format_results(result_set)
+
+    except Exception as e:
+        error_msg = f"Error executing Kusto operation '{action_name}' (correlation ID: {correlation_id}): {str(e)}"
+        logger.error(error_msg)
+        raise RuntimeError(error_msg) from e
 
 
 # NOTE: This is temporary. The intent is to not use environment variables for persistency.

tests/kusto/test_global_timeout.py

@@ -0,0 +1,90 @@
+"""Test global timeout configuration for Kusto tools."""
+
+import os
+from unittest.mock import Mock, patch
+
+from azure.kusto.data import ClientRequestProperties
+
+from fabric_rti_mcp.kusto.kusto_config import KustoConfig
+from fabric_rti_mcp.kusto.kusto_service import kusto_query
+
+
+def test_config_loads_timeout_from_env() -> None:
+    """Test that KustoConfig loads timeout from FABRIC_RTI_KUSTO_TIMEOUT environment variable."""
+    with patch.dict(os.environ, {"FABRIC_RTI_KUSTO_TIMEOUT": "300"}):
+        test_config = KustoConfig.from_env()
+        assert test_config.timeout_seconds == 300
+
+
+def test_config_handles_invalid_timeout() -> None:
+    """Test that KustoConfig handles invalid timeout values gracefully."""
+    with patch.dict(os.environ, {"FABRIC_RTI_KUSTO_TIMEOUT": "invalid"}):
+        test_config = KustoConfig.from_env()
+        assert test_config.timeout_seconds is None
+
+
+def test_config_no_timeout_env() -> None:
+    """Test that KustoConfig handles missing environment variable."""
+    with patch.dict(os.environ, {}, clear=True):
+        test_config = KustoConfig.from_env()
+        assert test_config.timeout_seconds is None
+
+
+@patch("fabric_rti_mcp.kusto.kusto_service.get_kusto_connection")
+@patch("fabric_rti_mcp.kusto.kusto_service.format_results")
+def test_global_timeout_applied_to_query(mock_format_results: Mock, mock_get_connection: Mock) -> None:
+    """Test that global timeout is applied to Kusto queries."""
+    # Mock connection
+    mock_connection = Mock()
+    mock_connection.default_database = "TestDB"
+    mock_client = Mock()
+    mock_connection.query_client = mock_client
+    mock_get_connection.return_value = mock_connection
+
+    # Mock format_results to return expected result
+    mock_format_results.return_value = [{"test": "result"}]
+
+    # Mock the kusto config with timeout
+    with patch("fabric_rti_mcp.kusto.kusto_service.CONFIG") as mock_config:
+        mock_config.timeout_seconds = 600
+        kusto_query("TestQuery", "https://test.kusto.windows.net")
+
+    # Verify that execute was called with ClientRequestProperties
+    mock_client.execute.assert_called_once()
+    call_args = mock_client.execute.call_args
+    crp = call_args[0][2]  # Third argument should be ClientRequestProperties
+
+    assert isinstance(crp, ClientRequestProperties)
+    # The timeout should be set as server timeout option in HH:MM:SS format
+    # 600 seconds = 10 minutes = 00:10:00
+    expected_timeout = "00:10:00"
+    assert crp._options.get("servertimeout") == expected_timeout
+
+
+@patch("fabric_rti_mcp.kusto.kusto_service.get_kusto_connection")
+@patch("fabric_rti_mcp.kusto.kusto_service.format_results")
+def test_no_timeout_when_not_configured(mock_format_results: Mock, mock_get_connection: Mock) -> None:
+    """Test that no timeout is set when not configured."""
+    # Mock connection
+    mock_connection = Mock()
+    mock_connection.default_database = "TestDB"
+    mock_client = Mock()
+    mock_connection.query_client = mock_client
+    mock_get_connection.return_value = mock_connection
+
+    # Mock format_results to return expected result
+    mock_format_results.return_value = [{"test": "result"}]
+
+    # Mock the kusto config without timeout
+    with patch("fabric_rti_mcp.kusto.kusto_service.CONFIG") as mock_config:
+        mock_config.timeout_seconds = None
+        kusto_query("TestQuery", "https://test.kusto.windows.net")
+
+    # Verify that execute was called with ClientRequestProperties
+    mock_client.execute.assert_called_once()
+    call_args = mock_client.execute.call_args
+    crp = call_args[0][2]  # Third argument should be ClientRequestProperties
+
+    assert isinstance(crp, ClientRequestProperties)
+    # No timeout should be set
+    assert "servertimeout" not in crp._options

tests/kusto/test_kusto_service.py

@@ -1,5 +1,7 @@
+import json
 from unittest.mock import MagicMock, Mock, patch
 
+import pytest
 from azure.kusto.data import ClientRequestProperties
 from azure.kusto.data.response import KustoResponseDataSet
 
@@ -21,6 +23,7 @@ def test_execute_basic_query(
     mock_connection = MagicMock()
     mock_connection.query_client = mock_client
     mock_connection.default_database = "default_db"
+    mock_connection.timeout_seconds = None  # Add timeout_seconds attribute
     mock_get_kusto_connection.return_value = mock_connection
 
     query = "  TestTable | take 10  "  # Added whitespace to test stripping
@@ -197,3 +200,92 @@ def test_destructive_operation_with_custom_client_request_properties(
     assert isinstance(result, list)
     assert len(result) == 1
     assert result[0]["TestColumn"] == "TestValue"
+
+
+@patch("fabric_rti_mcp.kusto.kusto_service.get_kusto_connection")
+def test_execute_error_includes_correlation_id(
+    mock_get_kusto_connection: Mock,
+    sample_cluster_uri: str,
+) -> None:
+    """Test that errors include correlation ID for easier debugging."""
+    # Arrange
+    mock_client = MagicMock()
+    mock_client.execute.side_effect = Exception("Kusto execution failed")
+
+    mock_connection = MagicMock()
+    mock_connection.query_client = mock_client
+    mock_connection.default_database = "default_db"
+    mock_get_kusto_connection.return_value = mock_connection
+
+    query = "TestTable | take 10"
+    database = "test_db"
+
+    # Act & Assert
+    with pytest.raises(RuntimeError) as exc_info:
+        kusto_query(query, sample_cluster_uri, database=database)
+
+    error_message = str(exc_info.value)
+
+    # Verify the error message includes correlation ID and operation name
+    assert "correlation ID:" in error_message
+    assert "KFRTI_MCP.kusto_query:" in error_message
+    assert "kusto_query" in error_message
+    assert "Kusto execution failed" in error_message
+
+
+@patch("fabric_rti_mcp.kusto.kusto_service.get_kusto_connection")
+def test_execute_json_parse_error_includes_correlation_id(
+    mock_get_kusto_connection: Mock,
+    sample_cluster_uri: str,
+) -> None:
+    """Test that JSON parsing errors include correlation ID for easier debugging."""
+    # Arrange - simulate the kind of error that was happening in the issue
+    mock_client = MagicMock()
+    mock_client.execute.side_effect = json.JSONDecodeError("Expecting value", "", 0)
+
+    mock_connection = MagicMock()
+    mock_connection.query_client = mock_client
+    mock_connection.default_database = "default_db"
+    mock_get_kusto_connection.return_value = mock_connection
+
+    command = ".show version"
+
+    # Act & Assert
+    with pytest.raises(RuntimeError) as exc_info:
+        kusto_command(command, sample_cluster_uri)
+
+    error_message = str(exc_info.value)
+
+    # Verify the error message includes correlation ID and operation name
+    assert "correlation ID:" in error_message
+    assert "KFRTI_MCP.kusto_command:" in error_message
+    assert "kusto_command" in error_message
+    assert "Expecting value" in error_message
+
+
+@patch("fabric_rti_mcp.kusto.kusto_service.logger")
+@patch("fabric_rti_mcp.kusto.kusto_service.get_kusto_connection")
+def test_successful_operations_do_not_log_correlation_id(
+    mock_get_kusto_connection: Mock,
+    mock_logger: Mock,
+    sample_cluster_uri: str,
+    mock_kusto_response: KustoResponseDataSet,
+) -> None:
+    """Test that successful operations do not log correlation IDs (only errors do)."""
+    # Arrange
+    mock_client = MagicMock()
+    mock_client.execute.return_value = mock_kusto_response
+
+    mock_connection = MagicMock()
+    mock_connection.query_client = mock_client
+    mock_connection.default_database = "default_db"
+    mock_get_kusto_connection.return_value = mock_connection
+
+    query = "TestTable | take 10"
+
+    # Act
+    kusto_query(query, sample_cluster_uri)
+
+    # Assert - verify no info or debug logging occurs for successful operations
+    assert not mock_logger.info.called
+    assert not mock_logger.debug.called