PY-171 Enrich and enable sending the query context to log_query (#158)

* PY-171 Enrich and enable sending the query context to log_query

Add the following to the `X-Neptune-Client-Metadata` header sent in all HTTP requests sent to Neptune:

* unique `nq_query_id` to the sent data. (each `nq.*` function called gets a new `nq_query_id`)
* custom metadata read from `NEPTUNE_QUERY_METADATA` env (if set)

Enable the mechanism:

* set `ADD_QUERY_METADATA = True` in `neptune_query.internal.query_metadata_context`

* Fix unit test

* Code review notes addresses

1. Update docs in `env.py`
2. Update the `log.debug` message for when `NEPTUNE_QUERY_METADATA` is too long
3. Switch to generating the random nq_query_id using `secrets` rather than `random`
4. Add the test files that I missed in the previous commit

* Remove test_query_metadata_query_id

* Stop sending `x-neptune-client-metadata` to `neptune_api.api.storage.signed_url_generic`

Author: gabrys

src/neptune_query/internal/env.py

@@ -89,6 +89,7 @@ def _map_logging_level(value: str) -> str:
 NEPTUNE_API_TOKEN = EnvVariable[str]("NEPTUNE_API_TOKEN", _map_str)
 NEPTUNE_ENABLE_COLORS = EnvVariable[bool]("NEPTUNE_ENABLE_COLORS", _map_bool, True)
 NEPTUNE_HTTP_REQUEST_TIMEOUT_SECONDS = EnvVariable[int]("NEPTUNE_HTTP_REQUEST_TIMEOUT_SECONDS", int, 60)
+NEPTUNE_LOGGER_LEVEL = EnvVariable[str]("NEPTUNE_LOGGER_LEVEL", _map_logging_level, "WARN")
 NEPTUNE_PROJECT = EnvVariable[str]("NEPTUNE_PROJECT", _map_str)
 NEPTUNE_QUERY_ATTRIBUTE_DEFINITIONS_BATCH_SIZE = EnvVariable[int](
     "NEPTUNE_QUERY_ATTRIBUTE_DEFINITIONS_BATCH_SIZE", int, 10_000
@@ -97,11 +98,15 @@ def _map_logging_level(value: str) -> str:
 NEPTUNE_QUERY_FILES_BATCH_SIZE = EnvVariable[int]("NEPTUNE_QUERY_FILES_BATCH_SIZE", int, 200)
 NEPTUNE_QUERY_FILES_MAX_CONCURRENCY = EnvVariable[int]("NEPTUNE_QUERY_FILES_MAX_CONCURRENCY", int, 1)
 NEPTUNE_QUERY_FILES_TIMEOUT = EnvVariable[Optional[int]]("NEPTUNE_QUERY_FILES_TIMEOUT", _lift_optional(int), None)
-NEPTUNE_QUERY_MAX_REQUEST_SIZE = EnvVariable[int]("NEPTUNE_QUERY_MAX_REQUEST_SIZE", int, 220_000)
 NEPTUNE_QUERY_MAX_ATTRIBUTE_FILTER_SIZE = EnvVariable[int](
     "NEPTUNE_QUERY_MAX_ATTRIBUTE_FILTER_SIZE", int, 60_000
 )  # 120_000 bytes seems to be the hard limit, but let's be safe
+NEPTUNE_QUERY_MAX_REQUEST_SIZE = EnvVariable[int]("NEPTUNE_QUERY_MAX_REQUEST_SIZE", int, 220_000)
 NEPTUNE_QUERY_MAX_WORKERS = EnvVariable[int]("NEPTUNE_QUERY_MAX_WORKERS", int, 32)
+# Metadata to be attached to each query for tracking purposes (e.g. user info, experiment context, etc.).
+# Accepts a plain string or JSON-encoded string. After we JSON-encode the value (using Python's json.dumps),
+# it must fit in 82 bytes. Otherwise, the client sends a placeholder string instead.
+NEPTUNE_QUERY_METADATA = EnvVariable[str]("NEPTUNE_QUERY_METADATA", _map_str, "")
 NEPTUNE_QUERY_RETRY_HARD_TIMEOUT = EnvVariable[Optional[int]](
     "NEPTUNE_QUERY_RETRY_HARD_TIMEOUT", _lift_optional(int), 3600
 )
@@ -111,4 +116,3 @@ def _map_logging_level(value: str) -> str:
 NEPTUNE_QUERY_SERIES_BATCH_SIZE = EnvVariable[int]("NEPTUNE_QUERY_SERIES_BATCH_SIZE", int, 10_000)
 NEPTUNE_QUERY_SYS_ATTRS_BATCH_SIZE = EnvVariable[int]("NEPTUNE_QUERY_EXPERIMENT_SYS_ATTRS_BATCH_SIZE", int, 10_000)
 NEPTUNE_VERIFY_SSL = EnvVariable[bool]("NEPTUNE_VERIFY_SSL", _map_bool, True)
-NEPTUNE_LOGGER_LEVEL = EnvVariable[str]("NEPTUNE_LOGGER_LEVEL", _map_logging_level, "WARN")

src/neptune_query/internal/query_metadata_context.py

@@ -15,11 +15,12 @@
 #
 from __future__ import annotations
 
+import base64
 import contextlib
-import dataclasses
 import functools
 import json
-from dataclasses import dataclass
+import logging
+import secrets
 from importlib.metadata import (
     PackageNotFoundError,
     version,
@@ -34,38 +35,56 @@
 
 from neptune_api.types import Response
 
+from neptune_query.internal import env
 from neptune_query.internal.composition import concurrency
 
 # This flag is used to control whether query metadata should be added to the request headers
-# Can remove this after backend supports this properly (under ticket PY-171)
-ADD_QUERY_METADATA = False
+ADD_QUERY_METADATA = True
 
-
-@functools.cache
-def get_client_version() -> str:
-    package_name = "neptune-query"
-    try:
-        package_version = version(package_name)
-    except PackageNotFoundError:
-        package_version = "unknown"
-    return f"{package_name}/{package_version}"
+logger = logging.getLogger(__name__)
 
 
-@dataclass
 class QueryMetadata:
-    api_function: str
-    client_version: str
-
-    def __post_init__(self) -> None:
-        self.api_function = self.api_function[:50]
-        self.client_version = self.client_version[:50]
+    def __init__(
+        self,
+        api_function: str,
+        client_version: str,
+        nq_query_id: str,
+        user_data: str | None = None,
+    ) -> None:
+
+        # longest observed: "fetch_experiments_table_global" - 30 characters
+        self.api_function = api_function[:32]
+
+        # longest observed: "nq/1.8.0b1.post8+cea1d73" - 24 characters
+        self.client_version = client_version[:24]
+
+        # nq_query_id should be exactly 8 characters
+        self.nq_query_id = nq_query_id[:8]
+
+        # user_data may contribute at most 80 characters when JSON-encoded
+        self.user_data: str | dict | None = _process_user_data(user_data)
+
+    def to_json(self) -> str:
+        return json.dumps(
+            {
+                "fn": self.api_function,
+                "v": self.client_version,
+                "qid": self.nq_query_id,
+                "ud": self.user_data,
+            }
+        )
 
 
 @contextlib.contextmanager
 def use_query_metadata(api_function: str) -> Generator[None, None, None]:
+    user_data = env.NEPTUNE_QUERY_METADATA.get()
+
     query_metadata = QueryMetadata(
         api_function=api_function,
-        client_version=get_client_version(),
+        client_version=_get_client_version(),
+        nq_query_id=_generate_nq_query_id(),
+        user_data=user_data if user_data else None,
     )
     with concurrency.use_thread_local({"query_metadata": query_metadata}):
         yield
@@ -82,7 +101,47 @@ def wrapper(*args: T.args, **kwargs: T.kwargs) -> Response[R]:
             "query_metadata", expected_type=QueryMetadata
         )
         if ADD_QUERY_METADATA and query_metadata:
-            kwargs["x_neptune_client_metadata"] = json.dumps(dataclasses.asdict(query_metadata))
+            kwargs["x_neptune_client_metadata"] = query_metadata.to_json()
         return func(*args, **kwargs)
 
     return wrapper
+
+
+@functools.cache
+def _get_client_version() -> str:
+    try:
+        package_version = version("neptune-query")
+    except PackageNotFoundError:
+        package_version = "unknown"
+    return f"nq/{package_version}"
+
+
+@functools.cache
+def _process_user_data(user_data: str | None) -> str | dict | None:
+    if not user_data:
+        return None
+
+    try:
+        user_data = json.loads(user_data)
+        user_data_is_proper_json = True
+    except json.JSONDecodeError:
+        user_data_is_proper_json = False
+
+    if len(json.dumps(user_data)) > 82:
+        user_data = "NEPTUNE_QUERY_METADATA too long"
+        if user_data_is_proper_json:
+            logger.debug("User data in NEPTUNE_QUERY_METADATA env too long. Skipping user data in query metadata.")
+        else:
+            logger.debug(
+                "User data in NEPTUNE_QUERY_METADATA env too long (and not JSON-encoded). "
+                "Skipping user data in query metadata."
+            )
+
+    return user_data
+
+
+def _generate_nq_query_id() -> str:
+    # Generate a random string of exactly 8 URL-safe base64 characters (a-z, A-Z, 0-9, -, _):
+    nq_query_id_bytes = secrets.token_bytes(6)
+    nq_query_id = base64.urlsafe_b64encode(nq_query_id_bytes).decode("ascii")
+    return nq_query_id

src/neptune_query/internal/retrieval/files.py

@@ -35,8 +35,6 @@
     Provider,
 )
 
-from neptune_query.internal.query_metadata_context import with_neptune_client_metadata
-
 from ...exceptions import NeptuneFileDownloadError
 from ...types import File
 from .. import (
@@ -72,7 +70,8 @@ def fetch_signed_urls(
 
     logger.debug(f"Calling signed_url_generic with body: {body}")
 
-    call_api = retry.handle_errors_default(with_neptune_client_metadata(signed_url_generic.sync_detailed))
+    # skip with_neptune_client_metadata - storage bridge API doesn't support it
+    call_api = retry.handle_errors_default(signed_url_generic.sync_detailed)
     response = call_api(client=client, body=body)
 
     logger.debug(

tests/unit/internal/test_query_metadata_context.py

@@ -16,6 +16,7 @@
 from __future__ import annotations
 
 import json
+from base64 import b64decode
 from unittest.mock import Mock
 
 import mock
@@ -30,14 +31,18 @@
 
 def test_query_metadata_truncation() -> None:
     # given
-    long_string = "a" * 100
-    metadata = QueryMetadata(api_function=long_string, client_version=long_string)
+    metadata = QueryMetadata(
+        api_function="a" * 100,
+        client_version="b" * 100,
+        nq_query_id="c" * 100,
+        user_data=None,
+    )
 
     # then
-    assert len(metadata.api_function) == 50
-    assert len(metadata.client_version) == 50
-    assert metadata.api_function == "a" * 50
-    assert metadata.client_version == "a" * 50
+    assert len(metadata.to_json()) < 200
+    assert metadata.api_function == "a" * 32
+    assert metadata.client_version == "b" * 24
+    assert metadata.nq_query_id == "c" * 8
 
 
 def test_use_query_metadata() -> None:
@@ -74,7 +79,15 @@ def test_with_neptune_client_metadata_with_context() -> None:
     # when
     with (
         mock.patch("neptune_query.internal.query_metadata_context.ADD_QUERY_METADATA", True),
-        mock.patch("neptune_query.internal.query_metadata_context.get_client_version", return_value="1.2.3"),
+        mock.patch("neptune_query.internal.query_metadata_context._get_client_version", return_value="1.2.3"),
+        mock.patch(
+            "neptune_query.internal.query_metadata_context.secrets.token_bytes",
+            return_value=b64decode("abcd1234"),
+        ),
+        mock.patch(
+            "neptune_query.internal.query_metadata_context.env.NEPTUNE_QUERY_METADATA.get",
+            return_value='{"magic_number": 42, "names": ["John", "Larry"]}',
+        ),
     ):
         decorated_call = with_neptune_client_metadata(mock_api_call)
         with use_query_metadata(api_function="test_api"):
@@ -84,5 +97,7 @@ def test_with_neptune_client_metadata_with_context() -> None:
     mock_api_call.assert_called_once()
     _, kwargs = mock_api_call.call_args
     assert "x_neptune_client_metadata" in kwargs
-    expected_json = json.dumps({"api_function": "test_api", "client_version": "1.2.3"})
+    expected_json = json.dumps(
+        {"fn": "test_api", "v": "1.2.3", "qid": "abcd1234", "ud": {"magic_number": 42, "names": ["John", "Larry"]}}
+    )
     assert kwargs["x_neptune_client_metadata"] == expected_json