Refactor the instrument decorator to use inspect to find arguments

By using the inspect module to identify the bound arguments passed
to the function, we don't need to separate arg and kwarg attributes,
and we can handle cases different types of function call, where
keywords aren't enforced. We can also now handle keywords captured in
a **kwargs parameter.

airlock/file_browser_api.py

@@ -202,7 +202,7 @@ def build_string(node, indent):
         return "\n".join(build_string(self, ""))
 
 
-@instrument(arg_attributes={"workspace": 0})
+@instrument(func_attributes={"workspace": "workspace"})
 def get_workspace_tree(workspace, selected_path=ROOT_PATH, selected_only=False):
     """Recursively build workspace tree from the root dir.
 
@@ -244,7 +244,7 @@ def get_workspace_tree(workspace, selected_path=ROOT_PATH, selected_only=False):
     return root_node
 
 
-@instrument(arg_attributes={"release_request": 0})
+@instrument(func_attributes={"release_request": "release_request"})
 def get_request_tree(release_request, selected_path=ROOT_PATH, selected_only=False):
     """Build a tree recursively for a ReleaseRequest
 

airlock/views/request.py

@@ -45,7 +45,7 @@ def request_index(request):
 
 # we return different content if it is a HTMX request.
 @vary_on_headers("HX-Request")
-@instrument(kwarg_attributes={"release_request": "request_id"})
+@instrument(func_attributes={"release_request": "request_id"})
 def request_view(request, request_id: str, path: str = ""):
     release_request = get_release_request_or_raise(request.user, request_id)
 
@@ -99,7 +99,7 @@ def request_view(request, request_id: str, path: str = ""):
     return TemplateResponse(request, template, context)
 
 
-@instrument(kwarg_attributes={"release_request": "request_id"})
+@instrument(func_attributes={"release_request": "request_id"})
 @require_http_methods(["GET"])
 def request_contents(request, request_id: str, path: str):
     release_request = get_release_request_or_raise(request.user, request_id)
@@ -124,7 +124,7 @@ def request_contents(request, request_id: str, path: str):
     return serve_file(request, abspath, release_request.get_request_file(path))
 
 
-@instrument(kwarg_attributes={"release_request": "request_id"})
+@instrument(func_attributes={"release_request": "request_id"})
 @require_http_methods(["POST"])
 def request_submit(request, request_id):
     release_request = get_release_request_or_raise(request.user, request_id)
@@ -138,7 +138,7 @@ def request_submit(request, request_id):
     return redirect(release_request.get_url())
 
 
-@instrument(kwarg_attributes={"release_request": "request_id"})
+@instrument(func_attributes={"release_request": "request_id"})
 @require_http_methods(["POST"])
 def request_reject(request, request_id):
     release_request = get_release_request_or_raise(request.user, request_id)
@@ -152,7 +152,7 @@ def request_reject(request, request_id):
     return redirect(release_request.get_url())
 
 
-@instrument(kwarg_attributes={"release_request": "request_id"})
+@instrument(func_attributes={"release_request": "request_id"})
 @require_http_methods(["POST"])
 def request_release_files(request, request_id):
     release_request = get_release_request_or_raise(request.user, request_id)

airlock/views/workspace.py

@@ -39,7 +39,7 @@ def workspace_index(request):
 
 # we return different content if it is a HTMX request.
 @vary_on_headers("HX-Request")
-@instrument(kwarg_attributes={"workspace": "workspace_name"})
+@instrument(func_attributes={"workspace": "workspace_name"})
 def workspace_view(request, workspace_name: str, path: str = ""):
     workspace = get_workspace_or_raise(request.user, workspace_name)
     template = "file_browser/index.html"
@@ -99,7 +99,7 @@ def workspace_view(request, workspace_name: str, path: str = ""):
     )
 
 
-@instrument(kwarg_attributes={"workspace": "workspace_name"})
+@instrument(func_attributes={"workspace": "workspace_name"})
 @require_http_methods(["GET"])
 def workspace_contents(request, workspace_name: str, path: str):
     workspace = get_workspace_or_raise(request.user, workspace_name)
@@ -115,7 +115,7 @@ def workspace_contents(request, workspace_name: str, path: str):
     return serve_file(request, abspath)
 
 
-@instrument(kwarg_attributes={"workspace": "workspace_name"})
+@instrument(func_attributes={"workspace": "workspace_name"})
 @require_http_methods(["POST"])
 def workspace_add_file_to_request(request, workspace_name):
     workspace = get_workspace_or_raise(request.user, workspace_name)

services/tracing.py

@@ -1,5 +1,6 @@
 import os
 from functools import wraps
+from inspect import Parameter, signature
 from typing import Dict
 
 from opentelemetry import trace
@@ -68,50 +69,67 @@ def instrument(
     span_name: str = "",
     record_exception: bool = True,
     attributes: Dict[str, str] = None,
-    arg_attributes: Dict[str, int] = None,
-    kwarg_attributes: Dict[str, str] = None,
+    func_attributes: Dict[str, str] = None,
     existing_tracer: trace.Tracer = None,
 ):
     """
     A decorator to instrument a function with an OTEL tracing span.
+
+    span_name: custom name for the span, defaults to name of decorated functionvv
+    record_exception: passed to `start_as_current_span`; whether to record
+      exceptions when they happen.
     attributes: custom attributes to set on the span
-    arg_attributes: k, v pairs of attribute name to index of positional
-       arg. Sets the span attribute k to the str representation of the
-       the arg at index v
-    kwarg_attributes: k, v pairs of attribute name to function
-       kwarg. Sets the span attribute k to the str representation of
-       the function kwarg v
+    func_attributes: k, v pairs of attribute name to function parameter
+      name. Sets the span attribute k to the str representation of
+      the function argument v (can be either positional or keyword argument).
+      v must be either a string, or an object that can be passed to str().
+    existing_tracer: pass an optional existing tracer to use. Defaults to
+      a tracer named with the value of the environment variable
+      `OTEL_SERVICE_NAME` if available, or the name of the module containing
+      the decoraated function.
     """
 
     def span_decorator(func):
-        tracer = existing_tracer or trace.get_tracer("airlock")
-
-        def _set_attributes(span, attributes_dict):
-            for att in attributes_dict:
-                span.set_attribute(att, attributes_dict[att])
+        tracer = existing_tracer or trace.get_tracer(
+            os.environ.get("OTEL_SERVICE_NAME", "airlock")
+        )
+        name = span_name or func.__qualname__
+        attributes_dict = attributes or {}
+        func_signature = signature(func)
+        default_params = {
+            param_name: param.default
+            for param_name, param in func_signature.parameters.items()
+            if param and param.default is not Parameter.empty
+        }
 
         @wraps(func)
         def wrap_with_span(*args, **kwargs):
-            name = span_name or func.__qualname__
+            if func_attributes is not None:
+                bound_args = func_signature.bind(*args, **kwargs).arguments
+                for attribute, parameter_name in func_attributes.items():
+                    # Find the value of this parameter by(in order):
+                    # 1) the function kwargs directly; if a function signature takes a parameter
+                    # like `**kwargs`, we can retrieve a named parameter from the keyword arguments
+                    # there
+                    # 2) the bound args retrieved from the function signature; this will find any
+                    # explicity passed values when the function was called.
+                    # 3) the parameter default value, if there is one
+                    # 4) Finally, raises an exception if we can't find a value for the expected parameter
+                    if parameter_name in kwargs:
+                        func_arg = kwargs[parameter_name]
+                    elif parameter_name in bound_args:
+                        func_arg = bound_args[parameter_name]
+                    elif parameter_name in default_params:
+                        func_arg = default_params[parameter_name]
+                    else:
+                        raise AttributeError(
+                            f"Expected argument {parameter_name} not found in function signature"
+                        )
+                    attributes_dict[attribute] = str(func_arg)
+
             with tracer.start_as_current_span(
-                name, record_exception=record_exception
-            ) as span:
-                attributes_dict = attributes or {}
-                if kwarg_attributes is not None:
-                    for k, v in kwarg_attributes.items():
-                        assert (
-                            v in kwargs
-                        ), f"Expected kwarg {v} not found in function signature"
-                        attributes_dict[k] = str(kwargs[v])
-
-                if arg_attributes is not None:
-                    for k, v in arg_attributes.items():
-                        assert (
-                            len(args) > v
-                        ), f"Expected positional arg at index {v} not found in function signature"
-                        attributes_dict[k] = str(args[v])
-
-                _set_attributes(span, attributes_dict)
+                name, record_exception=record_exception, attributes=attributes_dict
+            ):
                 return func(*args, **kwargs)
 
         return wrap_with_span

tests/test_tracing.py

@@ -2,23 +2,24 @@
 
 import opentelemetry.exporter.otlp.proto.http.trace_exporter
 import pytest
+from opentelemetry import trace
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 from opentelemetry.sdk.trace.export import ConsoleSpanExporter
 
-import services.tracing as tracing
+from services.tracing import instrument, setup_default_tracing
 
 
 def test_setup_default_tracing_empty_env(monkeypatch):
     env = {"PYTHONPATH": ""}
     monkeypatch.setattr(os, "environ", env)
-    provider = tracing.setup_default_tracing(set_global=False)
+    provider = setup_default_tracing(set_global=False)
     assert provider._active_span_processor._span_processors == ()
 
 
 def test_setup_default_tracing_console(monkeypatch):
     env = {"PYTHONPATH": "", "OTEL_EXPORTER_CONSOLE": "1"}
     monkeypatch.setattr(os, "environ", env)
-    provider = tracing.setup_default_tracing(set_global=False)
+    provider = setup_default_tracing(set_global=False)
 
     processor = provider._active_span_processor._span_processors[0]
     assert isinstance(processor.span_exporter, ConsoleSpanExporter)
@@ -31,7 +32,7 @@ def test_setup_default_tracing_otlp_defaults(monkeypatch):
     monkeypatch.setattr(
         opentelemetry.exporter.otlp.proto.http.trace_exporter, "environ", env
     )
-    provider = tracing.setup_default_tracing(set_global=False)
+    provider = setup_default_tracing(set_global=False)
     assert provider.resource.attributes["service.name"] == "airlock"
 
     exporter = provider._active_span_processor._span_processors[0].span_exporter
@@ -52,7 +53,7 @@ def test_setup_default_tracing_otlp_with_env(monkeypatch):
     monkeypatch.setattr(
         opentelemetry.exporter.otlp.proto.http.trace_exporter, "environ", env
     )
-    provider = tracing.setup_default_tracing(set_global=False)
+    provider = setup_default_tracing(set_global=False)
     assert provider.resource.attributes["service.name"] == "service"
 
     exporter = provider._active_span_processor._span_processors[0].span_exporter
@@ -63,60 +64,84 @@ def test_setup_default_tracing_otlp_with_env(monkeypatch):
 
 
 def test_not_instrument_decorator():
-    assert tracing.trace.get_current_span().is_recording() is False
+    assert trace.get_current_span().is_recording() is False
 
 
-@tracing.instrument
+@instrument
 def test_instrument_decorator():
-    current_span = tracing.trace.get_current_span()
+    current_span = trace.get_current_span()
     assert current_span.is_recording() is True
     assert current_span.name == "test_instrument_decorator"
 
 
-@tracing.instrument(span_name="testing", attributes={"foo": "bar"})
+@instrument(span_name="testing", attributes={"foo": "bar"})
 def test_instrument_decorator_with_name_and_attributes():
-    current_span = tracing.trace.get_current_span()
+    current_span = trace.get_current_span()
     assert current_span.is_recording() is True
     assert current_span.name == "testing"
     assert current_span.attributes == {"foo": "bar"}
 
 
-def test_instrument_decorator_with_function_kwarg_attributes():
-    @tracing.instrument(kwarg_attributes={"number": "num"})
-    def assert_function_kwarg_attributes(*, num):
-        current_span = tracing.trace.get_current_span()
-        assert current_span.attributes == {"number": str(num)}
-
-    assert_function_kwarg_attributes(num=1)
-
-
-def test_instrument_decorator_with_function_arg_attributes():
-    @tracing.instrument(arg_attributes={"number": 0})
-    def assert_function_arg_attributes(num):
-        current_span = tracing.trace.get_current_span()
-        assert current_span.attributes == {"number": str(num)}
+@pytest.mark.parametrize(
+    "func_attributes,func_args,func_kwargs,expected_attributes",
+    [
+        # positional arg
+        ({"func_attributes": {"number": "num"}}, (1,), {}, {"number": "1"}),
+        # keyword arg
+        (
+            {"func_attributes": {"text": "string"}},
+            (1,),
+            {"string": "bar"},
+            {"text": "bar"},
+        ),
+        # default keyword arg
+        ({"func_attributes": {"text": "string"}}, (1,), {}, {"text": "Foo"}),
+        # all args passed as keywords
+        (
+            {"func_attributes": {"number": "num"}},
+            (),
+            {"num": 1, "string": "bar"},
+            {"number": "1"},
+        ),
+        # all args passed as positional
+        ({"func_attributes": {"number": "num"}}, (1, "bar"), {}, {"number": "1"}),
+        # multiple func attributes
+        (
+            {"func_attributes": {"number": "num", "text": "string"}},
+            (1,),
+            {},
+            {"number": "1", "text": "Foo"},
+        ),
+    ],
+)
+def test_instrument_decorator_with_function_attributes(
+    func_attributes, func_args, func_kwargs, expected_attributes
+):
+    @instrument(**func_attributes)
+    def assert_function_kwarg_attributes(num, string="Foo"):
+        current_span = trace.get_current_span()
+        assert current_span.attributes == expected_attributes
+        return num, string
 
-    assert_function_arg_attributes(1)
+    assert_function_kwarg_attributes(*func_args, **func_kwargs)
 
 
 @pytest.mark.parametrize(
-    "instrument_params,expect_ok",
+    "func_kwargs,expect_ok",
     [
-        ({"kwarg_attributes": {"number": "foo"}}, False),
-        ({"arg_attributes": {"number": 1}}, False),
-        ({"arg_attributes": {"number": 0}}, True),
+        ({}, False),
+        ({"foo": 1}, False),
+        ({"bar": 1}, True),
     ],
 )
-def test_instrument_decorator_with_invalid_function_attributes(
-    instrument_params, expect_ok
-):
-    @tracing.instrument(**instrument_params)
-    def decorated_function(num):
-        current_span = tracing.trace.get_current_span()
-        assert current_span.attributes == {"number": str(num)}
+def test_instrument_decorator_with_unnamed_kwargs(func_kwargs, expect_ok):
+    @instrument(func_attributes={"foo": "bar"})
+    def decorated_function(**kwargs):
+        current_span = trace.get_current_span()
+        assert current_span.attributes == {"foo": str(kwargs["bar"])}
 
     if expect_ok:
-        decorated_function(1)
+        decorated_function(**func_kwargs)
     else:
-        with pytest.raises(AssertionError, match="not found in function signature"):
-            decorated_function(1)
+        with pytest.raises(AttributeError, match="not found in function signature"):
+            decorated_function(**func_kwargs)