Repository: yandex-cloud/yandex-cloud-ml-sdk

Issue:
Can't figure out how to pass gRPC metadata?
I'm aiming https://yandex.cloud/ru/docs/foundation-models/operations/disable-logging for sure, but can't find how to do it with `YCloudML`. 
Is should be something around interceptors?

Discussion:

Patch:
diff --git a/src/yandex_cloud_ml_sdk/_client.py b/src/yandex_cloud_ml_sdk/_client.py
index a146434..c4ff7d1 100644
--- a/src/yandex_cloud_ml_sdk/_client.py
+++ b/src/yandex_cloud_ml_sdk/_client.py
@@ -55,6 +55,7 @@ class AsyncCloudClient:
         interceptors: Sequence[grpc.aio.ClientInterceptor] | None,
         yc_profile: str | None,
         retry_policy: RetryPolicy,
+        enable_server_data_logging: bool | None,
     ):
         self._endpoint = endpoint
         self._auth = auth
@@ -75,6 +76,7 @@ class AsyncCloudClient:
         self._channels_lock = LazyLock()
 
         self._user_agent = _get_user_agent()
+        self._enable_server_data_logging = enable_server_data_logging
 
     async def _init_service_map(self, timeout: float):
         credentials = grpc.ssl_channel_credentials()
@@ -112,11 +114,17 @@ class AsyncCloudClient:
         timeout: float,
         retry_kind: RetryKind = RetryKind.NONE,
     ) -> tuple[tuple[str, str], ...]:
-        metadata = (
+        metadata: tuple[tuple[str, str], ...] = (
             (RETRY_KIND_METADATA_KEY, retry_kind.name),
             ("grpc.primary_user_agent", self._user_agent),
         )
 
+        if self._enable_server_data_logging is not None:
+            enable_server_data_logging = "true" if self._enable_server_data_logging else "false"
+            metadata += (
+                ("x-data-logging-enabled", enable_server_data_logging),
+            )
+
         if not auth_required:
             return metadata
 
diff --git a/src/yandex_cloud_ml_sdk/_sdk.py b/src/yandex_cloud_ml_sdk/_sdk.py
index 590371c..33b8393 100644
--- a/src/yandex_cloud_ml_sdk/_sdk.py
+++ b/src/yandex_cloud_ml_sdk/_sdk.py
@@ -45,6 +45,7 @@ class BaseSDK:
         yc_profile: UndefinedOr[str] = UNDEFINED,
         service_map: UndefinedOr[dict[str, str]] = UNDEFINED,
         interceptors: UndefinedOr[Sequence[aio.ClientInterceptor]] = UNDEFINED,
+        enable_server_data_logging: UndefinedOr[bool] = UNDEFINED,
     ):
         """
         Construct a new asynchronous sdk instance.
@@ -62,6 +63,11 @@ class BaseSDK:
         :param service_map: a way to redefine endpoints for one or more cloud subservices
             with a format of dict {service_name: service_address}.
         :type service_map: Dict[str, str]
+        :param enable_server_data_logging: when passed bool, we will add
+            `x-data-logging-enabled: <value>` to all of requests, which will
+            enable or disable logging of user data on server side.
+            It will do something only on those parts of backends which supports
+            this option.
 
         """
         endpoint = self._get_endpoint(endpoint)
@@ -74,6 +80,7 @@ class BaseSDK:
             retry_policy=retry_policy,
             interceptors=get_defined_value(interceptors, None),
             yc_profile=get_defined_value(yc_profile, None),
+            enable_server_data_logging=get_defined_value(enable_server_data_logging, None),
         )
         self._folder_id = folder_id
 


Test patch:
diff --git a/src/yandex_cloud_ml_sdk/_testing/client.py b/src/yandex_cloud_ml_sdk/_testing/client.py
index 12fc806..51098f0 100644
--- a/src/yandex_cloud_ml_sdk/_testing/client.py
+++ b/src/yandex_cloud_ml_sdk/_testing/client.py
@@ -31,6 +31,7 @@ class MockClient(AsyncCloudClient):
             interceptors=None,
             yc_profile=None,
             retry_policy=retry_policy,
+            enable_server_data_logging=None,
         )
         self.port = port
         self._sdk = sdk
diff --git a/tests/test_client.py b/tests/test_client.py
index ec2594b..83a52a0 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -18,7 +18,8 @@ from yandex.cloud.endpoint.api_endpoint_service_pb2 import ListApiEndpointsReque
 from yandex.cloud.endpoint.api_endpoint_service_pb2_grpc import ApiEndpointServiceStub
 
 from yandex_cloud_ml_sdk import AsyncYCloudML
-from yandex_cloud_ml_sdk._client import _get_user_agent, httpx_client
+from yandex_cloud_ml_sdk._client import AsyncCloudClient, _get_user_agent, httpx_client
+from yandex_cloud_ml_sdk.auth import NoAuth
 
 
 @pytest.fixture(name='servicers')
@@ -139,3 +140,48 @@ def test_multiple_threads(sdk_maker, caplog):
 async def test_httpx_client():
     async with httpx_client() as client:
         assert client.headers['User-Agent'] == _get_user_agent()
+
+
+# pylint: disable=protected-access
+@pytest.mark.asyncio
+async def test_x_data_logging(interceptors, retry_policy, user_agent_tuple):
+    base_result = (('yc-ml-sdk-retry', 'NONE'), user_agent_tuple)
+    client = AsyncCloudClient(
+        endpoint="foo",
+        auth=NoAuth(),
+        service_map={},
+        yc_profile=None,
+        retry_policy=retry_policy,
+        interceptors=interceptors,
+        enable_server_data_logging=None,
+    )
+
+    assert await client._get_metadata(auth_required=False, timeout=0) == base_result
+
+    client = AsyncCloudClient(
+        endpoint="foo",
+        auth=NoAuth(),
+        service_map={},
+        yc_profile=None,
+        retry_policy=retry_policy,
+        interceptors=interceptors,
+        enable_server_data_logging=True,
+    )
+
+    assert await client._get_metadata(auth_required=False, timeout=0) == base_result + (
+        ('x-data-logging-enabled', "true"),
+    )
+
+    client = AsyncCloudClient(
+        endpoint="foo",
+        auth=NoAuth(),
+        service_map={},
+        yc_profile=None,
+        retry_policy=retry_policy,
+        interceptors=interceptors,
+        enable_server_data_logging=False,
+    )
+
+    assert await client._get_metadata(auth_required=False, timeout=0) == base_result + (
+        ('x-data-logging-enabled', "false"),
+    )