diff --git a/src/fabric_cli/client/fab_api_client.py b/src/fabric_cli/client/fab_api_client.py index e0c9f0fc..0e896e27 100644 --- a/src/fabric_cli/client/fab_api_client.py +++ b/src/fabric_cli/client/fab_api_client.py @@ -113,7 +113,9 @@ def do_request( # Build headers from fabric_cli.core.fab_context import Context as FabContext - ctxt_cmd = FabContext().command + fab_context = FabContext() + ctxt_cmd = fab_context.command + fabric_skill = fab_context.fabric_skill headers = { "Authorization": "Bearer " + str(token), @@ -132,6 +134,10 @@ def do_request( fab_constant.ERROR_INVALID_OPERATION, ) + # The fabric skills header is only applicable for the Fabric audience (None or "fabric") + if fabric_skill and audience_value in (None, "fabric"): + headers[fab_constant.FABRIC_SKILL_HEADER] = fabric_skill + try: session = _get_session() retries_count = 3 diff --git a/src/fabric_cli/core/fab_constant.py b/src/fabric_cli/core/fab_constant.py index d883bb05..cf4b75e6 100644 --- a/src/fabric_cli/core/fab_constant.py +++ b/src/fabric_cli/core/fab_constant.py @@ -67,6 +67,7 @@ FAB_HOST_APP_ENV_VAR = "FAB_HOST_APP" FAB_HOST_APP_VERSION_ENV_VAR = "FAB_HOST_APP_VERSION" +FABRIC_SKILL_HEADER = "x-ms-fabric-skill" # Other constants FAB_CAPACITY_NAME_NONE = "none" diff --git a/src/fabric_cli/core/fab_context.py b/src/fabric_cli/core/fab_context.py index ed18b7c9..bc3c8fe6 100644 --- a/src/fabric_cli/core/fab_context.py +++ b/src/fabric_cli/core/fab_context.py @@ -6,6 +6,7 @@ import os import platform import sys +from typing import Optional import psutil @@ -23,6 +24,7 @@ class Context: def __init__(self): self._context: FabricElement = None self._command: str = None + self._fabric_skill: Optional[str] = None self._runtime_mode: str = fab_constant.FAB_MODE_COMMANDLINE session_id = self._get_context_session_id() self._context_file = os.path.join( @@ -72,6 +74,14 @@ def command(self) -> str: def command(self, command: str) -> None: self._command = command + @property + def fabric_skill(self) -> Optional[str]: + return self._fabric_skill + + @fabric_skill.setter + def fabric_skill(self, value: Optional[str]) -> None: + self._fabric_skill = value if isinstance(value, str) else None + def reset_context(self) -> None: self.cleanup_context_files(cleanup_all_stale=True, cleanup_current=True) self.context = self.context.tenant diff --git a/src/fabric_cli/core/fab_decorators.py b/src/fabric_cli/core/fab_decorators.py index fb245957..90acee6a 100644 --- a/src/fabric_cli/core/fab_decorators.py +++ b/src/fabric_cli/core/fab_decorators.py @@ -65,6 +65,7 @@ def wrapper(*args, **kwargs): # Import Context locally to avoid circular import from fabric_cli.core.fab_context import Context Context().command = args[0].command_path + Context().fabric_skill = getattr(args[0], "skill", None) return func(*args, **kwargs) return wrapper diff --git a/src/fabric_cli/core/fab_logger.py b/src/fabric_cli/core/fab_logger.py index 3c20420c..0995ed0f 100644 --- a/src/fabric_cli/core/fab_logger.py +++ b/src/fabric_cli/core/fab_logger.py @@ -73,8 +73,8 @@ def log_debug_http_request( for key, value in headers.items(): if key.lower() == "authorization": value = "*****" # Mask authorization token - elif key.lower() == "user-agent": - continue # Skip logging the User-Agent header + elif key.lower() in ("user-agent", fab_constant.FABRIC_SKILL_HEADER): + continue # Skip logging telemetry headers logger.debug(f" '{key}': '{value}'") # Body diff --git a/src/fabric_cli/parsers/fab_global_params.py b/src/fabric_cli/parsers/fab_global_params.py index edb72bd8..4288dfba 100644 --- a/src/fabric_cli/parsers/fab_global_params.py +++ b/src/fabric_cli/parsers/fab_global_params.py @@ -1,6 +1,8 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. +import argparse + def add_global_flags(parser) -> None: """ @@ -11,7 +13,7 @@ def add_global_flags(parser) -> None: """ # Add help flag parser.add_argument("-help", action="help") - + # Add format flag to override output format parser.add_argument( "--output_format", @@ -19,3 +21,10 @@ def add_global_flags(parser) -> None: choices=["json", "text"], help="Override output format type. Optional", ) + + parser.add_argument( + "--skill", + required=False, + default=argparse.SUPPRESS, + help=argparse.SUPPRESS, + ) diff --git a/tests/conftest.py b/tests/conftest.py index 176b8035..a8953166 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -85,6 +85,7 @@ def reset_context(): context_instance = Context() context_instance._context = None context_instance._command = None + context_instance._fabric_skill = None context_instance._loading_context = False context_instance._runtime_mode = fab_constant.FAB_MODE_COMMANDLINE diff --git a/tests/test_core/test_fab_api_client.py b/tests/test_core/test_fab_api_client.py index 2b16c07f..2d40c797 100644 --- a/tests/test_core/test_fab_api_client.py +++ b/tests/test_core/test_fab_api_client.py @@ -490,6 +490,7 @@ def test_do_request_user_agent_header( # Configure mocks mock_auth.return_value.get_access_token.return_value = "dummy-token" mock_context.return_value.command = "test-command" + mock_context.return_value.fabric_skill = None class DummyResponse: status_code = 200 diff --git a/tests/test_core/test_fab_skill_attribution.py b/tests/test_core/test_fab_skill_attribution.py new file mode 100644 index 00000000..3ce02da8 --- /dev/null +++ b/tests/test_core/test_fab_skill_attribution.py @@ -0,0 +1,181 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +import argparse +from argparse import Namespace +from unittest.mock import patch + +import pytest + +from fabric_cli.client.fab_api_client import do_request +from fabric_cli.core import fab_constant, fab_state_config +from fabric_cli.core import fab_logger as logger +from fabric_cli.core.fab_auth import FabAuth +from fabric_cli.core.fab_context import Context +from fabric_cli.core.fab_decorators import set_command_context +from fabric_cli.core.fab_parser_setup import create_parser_and_subparsers +from fabric_cli.parsers import fab_global_params + +pytestmark = pytest.mark.usefixtures("reset_context") + + +class DummyResponse: + status_code = 200 + text = "{}" + content = b"{}" + headers: dict[str, str] = {} + + +def test_skill_argument_is_hidden_success(): + parser = argparse.ArgumentParser() + fab_global_params.add_global_flags(parser) + + skill_flag = next( + action for action in parser._actions if "--skill" in action.option_strings + ) + + assert skill_flag.help == argparse.SUPPRESS + assert "--skill" not in parser.format_help() + + +@pytest.mark.parametrize( + "command", + [ + ["--skill", "semantic-model-authoring", "export", "ws.Workspace", "-o", "out"], + ["export", "--skill", "semantic-model-authoring", "ws.Workspace", "-o", "out"], + ["export", "ws.Workspace", "-o", "out", "--skill", "semantic-model-authoring"], + ["--skill", "semantic-model-authoring", "job", "run", "ws.Notebook"], + ["job", "run", "ws.Notebook", "--skill", "semantic-model-authoring"], + ], +) +def test_skill_argument_is_preserved_across_command_tree_success(command): + parser, _ = create_parser_and_subparsers() + + args = parser.parse_args(command) + + assert args.skill == "semantic-model-authoring" + + +def test_command_context_uses_skill_argument_success(): + @set_command_context() + def command(args: Namespace) -> None: + assert Context().fabric_skill == "argument-skill" + + command(Namespace(command_path="export", skill="argument-skill")) + + +def test_command_context_clears_previous_skill_success(): + Context().fabric_skill = "previous-skill" + + @set_command_context() + def command(args: Namespace) -> None: + assert Context().fabric_skill is None + + command(Namespace(command_path="export", skill=None)) + + +@pytest.mark.parametrize( + "value", + [ + "", + " ", + "bad\nvalue", + "bad\rvalue", + "-leading", + "a", + "semantic-model_authoring.v2", + "a" * 129, + ], + ids=[ + "empty", + "whitespace", + "newline", + "carriage-return", + "leading-hyphen", + "single-character", + "punctuation", + "long", + ], +) +def test_fabric_skill_setter_preserves_string_value_success(value): + Context().fabric_skill = value + + assert Context().fabric_skill == value + + +@pytest.mark.parametrize( + "value", + [123, True, [], {}, object()], + ids=["integer", "boolean", "list", "dictionary", "object"], +) +def test_fabric_skill_setter_normalizes_non_string_value_success(value): + Context().fabric_skill = value + + assert Context().fabric_skill is None + + +@patch.object(FabAuth(), "get_access_token", return_value="dummy-token") +@pytest.mark.parametrize( + "audience, expects_skill_header", + [ + (None, True), + ("fabric", True), + ("storage", False), + ("azure", False), + ("powerbi", False), + ], +) +def test_skill_header_is_scoped_to_fabric_api_success( + _mock_get_token, audience, expects_skill_header +): + Context().fabric_skill = "semantic-model-authoring" + args = Namespace(uri="items", method="get", audience=audience) + + with patch( + "requests.Session.request", return_value=DummyResponse() + ) as mock_request: + do_request(args) + + request_headers = mock_request.call_args.kwargs["headers"] + if expects_skill_header: + assert ( + request_headers[fab_constant.FABRIC_SKILL_HEADER] + == "semantic-model-authoring" + ) + else: + assert fab_constant.FABRIC_SKILL_HEADER not in request_headers + + +@patch.object(FabAuth(), "get_access_token", return_value="dummy-token") +def test_skill_header_is_omitted_when_skill_is_not_set_success(_mock_get_token): + args = Namespace(uri="items", method="get", audience="fabric") + + with patch( + "requests.Session.request", return_value=DummyResponse() + ) as mock_request: + do_request(args) + + request_headers = mock_request.call_args.kwargs["headers"] + assert fab_constant.FABRIC_SKILL_HEADER not in request_headers + + +@pytest.mark.parametrize( + "header_name", + [fab_constant.FABRIC_SKILL_HEADER, "X-MS-FABRIC-SKILL"], + ids=["canonical-case", "uppercase"], +) +def test_skill_header_is_not_logged_success(monkeypatch, header_name): + monkeypatch.setattr(fab_state_config, "get_config", lambda key: "true") + + with patch.object(logger, "get_logger") as mock_get_logger: + logger.log_debug_http_request( + "GET", + "http://example.com", + {header_name: "semantic-model-authoring"}, + 10, + ) + + logged_messages = [ + call.args[0] for call in mock_get_logger.return_value.debug.call_args_list + ] + assert not any("semantic-model-authoring" in message for message in logged_messages)