Skip to content

Commit

Permalink
feat: track invocation id for retry metrics (#741)
Browse files Browse the repository at this point in the history
* feat: track invocation id for retry metrics

* woops, fix the context

* slight adjustment to stringify

* old tests pass

* lint

* adjust based on new python-cloud-core changes

* updated cloud core, and all working

* test that invocation id changes between calls

* lint and fix test name

Co-authored-by: Anthonios Partheniou <partheniou@google.com>
  • Loading branch information
Aaron Gabriel Neyer and parthea committed Apr 12, 2022
1 parent 1491757 commit bd56931
Show file tree
Hide file tree
Showing 7 changed files with 267 additions and 176 deletions.
7 changes: 6 additions & 1 deletion google/cloud/storage/_helpers.py
Expand Up @@ -21,6 +21,7 @@
from hashlib import md5
import os
from urllib.parse import urlsplit
from uuid import uuid4

from google import resumable_media
from google.auth import environment_vars
Expand Down Expand Up @@ -584,6 +585,10 @@ def _api_core_retry_to_resumable_media_retry(retry, num_retries=None):
return resumable_media.RetryStrategy(max_retries=0)


def _get_invocation_id():
return "gccl-invocation-id/" + str(uuid4())


def _get_default_headers(
user_agent,
content_type="application/json; charset=UTF-8",
Expand All @@ -600,7 +605,7 @@ def _get_default_headers(
"Accept": "application/json",
"Accept-Encoding": "gzip, deflate",
"User-Agent": user_agent,
"x-goog-api-client": user_agent,
"X-Goog-API-Client": f"{user_agent} {_get_invocation_id()}",
"content-type": content_type,
"x-upload-content-type": x_upload_content_type or content_type,
}
2 changes: 2 additions & 0 deletions google/cloud/storage/_http.py
Expand Up @@ -18,6 +18,7 @@

from google.cloud import _http
from google.cloud.storage import __version__
from google.cloud.storage import _helpers


class Connection(_http.JSONConnection):
Expand Down Expand Up @@ -59,6 +60,7 @@ def __init__(self, client, client_info=None, api_endpoint=None):

def api_request(self, *args, **kwargs):
retry = kwargs.pop("retry", None)
kwargs["extra_api_info"] = _helpers._get_invocation_id()
call = functools.partial(super(Connection, self).api_request, *args, **kwargs)
if retry:
# If this is a ConditionalRetryPolicy, check conditions.
Expand Down
2 changes: 1 addition & 1 deletion setup.py
Expand Up @@ -30,7 +30,7 @@
dependencies = [
"google-auth >= 1.25.0, < 3.0dev",
"google-api-core >= 1.31.5, <3.0.0dev,!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.0",
"google-cloud-core >= 1.6.0, < 3.0dev",
"google-cloud-core >= 2.3.0, < 3.0dev",
"google-resumable-media >= 2.3.2",
"requests >= 2.18.0, < 3.0.0dev",
"protobuf",
Expand Down
2 changes: 2 additions & 0 deletions tests/unit/test__helpers.py
Expand Up @@ -19,6 +19,8 @@
from google.cloud.storage.retry import DEFAULT_RETRY
from google.cloud.storage.retry import DEFAULT_RETRY_IF_METAGENERATION_SPECIFIED

GCCL_INVOCATION_TEST_CONST = "gccl-invocation-id/test-invocation-123"


class Test__get_storage_host(unittest.TestCase):
@staticmethod
Expand Down
13 changes: 11 additions & 2 deletions tests/unit/test__http.py
Expand Up @@ -13,9 +13,13 @@
# limitations under the License.

import unittest
from unittest.mock import patch

import mock

from google.cloud.storage import _helpers
from tests.unit.test__helpers import GCCL_INVOCATION_TEST_CONST


class TestConnection(unittest.TestCase):
@staticmethod
Expand Down Expand Up @@ -44,12 +48,17 @@ def test_extra_headers(self):

conn = self._make_one(client)
req_data = "hey-yoooouuuuu-guuuuuyyssss"
result = conn.api_request("GET", "/rainbow", data=req_data, expect_json=False)
with patch.object(
_helpers, "_get_invocation_id", return_value=GCCL_INVOCATION_TEST_CONST
):
result = conn.api_request(
"GET", "/rainbow", data=req_data, expect_json=False
)
self.assertEqual(result, data)

expected_headers = {
"Accept-Encoding": "gzip",
base_http.CLIENT_INFO_HEADER: conn.user_agent,
base_http.CLIENT_INFO_HEADER: f"{conn.user_agent} {GCCL_INVOCATION_TEST_CONST}",
"User-Agent": conn.user_agent,
}
expected_uri = conn.build_api_url("/rainbow")
Expand Down

0 comments on commit bd56931

Please sign in to comment.