From 12d5666ccdfd94424cc9de070562b8bd3af9cf34 Mon Sep 17 00:00:00 2001
From: Petr Reznikov <prez@yandex-team.ru>
Date: Tue, 7 Dec 2021 19:05:17 +0300
Subject: [PATCH] Support GRPC user-agent in SDK

* Prepend User-Agent set by grpc library with `yandex-cloud-python-sdk/$VERSION` entry. Version is taken from pkg_resources
* Allow SDK clients to add their own User-Agent entry, which will be prepended to previous result.
---
 examples/dataproc/main.py          |  6 ++++--
 examples/dataproc/using_wrapper.py |  6 ++++--
 yandexcloud/_channels.py           | 22 +++++++++++++++++++---
 yandexcloud/_sdk.py                | 17 +++++++++++++++--
 4 files changed, 42 insertions(+), 9 deletions(-)

diff --git a/examples/dataproc/main.py b/examples/dataproc/main.py
index 85bb0839..fb8e6aef 100644
--- a/examples/dataproc/main.py
+++ b/examples/dataproc/main.py
@@ -20,15 +20,17 @@
 
 import yandexcloud
 
+USER_AGENT = 'ycloud-python-sdk:dataproc.example.main'
+
 
 def main():
     logging.basicConfig(level=logging.INFO)
     arguments = parse_cmd()
     if arguments.token:
-        sdk = yandexcloud.SDK(token=arguments.token)
+        sdk = yandexcloud.SDK(token=arguments.token, user_agent=USER_AGENT)
     else:
         with open(arguments.sa_json_path) as infile:
-            sdk = yandexcloud.SDK(service_account_key=json.load(infile))
+            sdk = yandexcloud.SDK(service_account_key=json.load(infile), user_agent=USER_AGENT)
 
     fill_missing_arguments(sdk, arguments)
 
diff --git a/examples/dataproc/using_wrapper.py b/examples/dataproc/using_wrapper.py
index d7e5a955..bd9a9e78 100644
--- a/examples/dataproc/using_wrapper.py
+++ b/examples/dataproc/using_wrapper.py
@@ -8,15 +8,17 @@
 import yandexcloud
 from yandexcloud.operations import OperationError
 
+USER_AGENT = 'ycloud-python-sdk:dataproc.example.using_wrapper'
+
 
 def main():
     logging.basicConfig(level=logging.INFO)
     arguments = parse_cmd()
     if arguments.token:
-        sdk = yandexcloud.SDK(token=arguments.token)
+        sdk = yandexcloud.SDK(token=arguments.token, user_agent=USER_AGENT)
     else:
         with open(arguments.sa_json_path) as infile:
-            sdk = yandexcloud.SDK(service_account_key=json.load(infile))
+            sdk = yandexcloud.SDK(service_account_key=json.load(infile), user_agent=USER_AGENT)
     fill_missing_arguments(sdk, arguments)
 
     dataproc = sdk.wrappers.Dataproc(
diff --git a/yandexcloud/_channels.py b/yandexcloud/_channels.py
index 551ca415..8e2dca1f 100644
--- a/yandexcloud/_channels.py
+++ b/yandexcloud/_channels.py
@@ -6,9 +6,17 @@
 from yandexcloud import _auth_plugin
 from yandexcloud._auth_fabric import get_auth_token_requester
 
+import pkg_resources
+try:
+    VERSION = pkg_resources.get_distribution('yandexcloud').version
+except pkg_resources.DistributionNotFound:
+    VERSION = '0.0.0'
+
+SDK_USER_AGENT = 'yandex-cloud-python-sdk/{version}'.format(version=VERSION)
+
 
 class Channels(object):
-    def __init__(self, **kwargs):
+    def __init__(self, client_user_agent=None, **kwargs):
         self._channel_creds = grpc.ssl_channel_credentials(
             root_certificates=kwargs.get('root_certificates'),
             private_key=kwargs.get('private_key'),
@@ -20,11 +28,19 @@ def __init__(self, **kwargs):
 
         self._unauthenticated_channel = None
         self._channels = None
+        self._client_user_agent = client_user_agent
+
+    def channel_options(self):
+        return tuple(
+            ('grpc.primary_user_agent', user_agent)
+            for user_agent in [self._client_user_agent, SDK_USER_AGENT]
+            if user_agent is not None
+        )
 
     def channel(self, endpoint):
         if not self._channels:
             self._unauthenticated_channel = grpc.secure_channel(
-                self._endpoint, self._channel_creds)
+                self._endpoint, self._channel_creds, options=self.channel_options())
             endpoint_service = ApiEndpointServiceStub(
                 self._unauthenticated_channel)
             resp = endpoint_service.List(ListApiEndpointsRequest())
@@ -37,7 +53,7 @@ def channel(self, endpoint):
                 self._channel_creds, call_creds)
 
             self._channels = {
-                ep.id: grpc.secure_channel(ep.address, creds)
+                ep.id: grpc.secure_channel(ep.address, creds, options=self.channel_options())
                 for ep in endpoints
             }
 
diff --git a/yandexcloud/_sdk.py b/yandexcloud/_sdk.py
index 6ba6e99a..0204dcf4 100644
--- a/yandexcloud/_sdk.py
+++ b/yandexcloud/_sdk.py
@@ -10,8 +10,21 @@
 
 
 class SDK(object):
-    def __init__(self, interceptor=None, **kwargs):
-        self._channels = _channels.Channels(**kwargs)
+    def __init__(self, interceptor=None, user_agent=None, **kwargs):
+        """
+        API entry-point object.
+
+        :param interceptor: GRPC interceptor to be used instead of default RetryInterceptor
+        :type interceptor: Union[
+            UnaryUnaryClientInterceptor,
+            UnaryStreamClientInterceptor,
+            StreamUnaryClientInterceptor,
+            StreamStreamClientInterceptor
+        ]
+        :param user_agent: String to prepend User-Agent metadata header for all GRPC requests made via SDK object
+        :type user_agent: Optional[str]
+        """
+        self._channels = _channels.Channels(client_user_agent=user_agent, **kwargs)
         if interceptor is None:
             interceptor = RetryInterceptor(
                 max_retry_count=5,