From 9ca93bdd93efa3b9d782a13ae1da1d4c4cb9fe7c Mon Sep 17 00:00:00 2001
From: Anthony Mahanna <anthony.mahanna@arangodb.com>
Date: Thu, 30 Nov 2023 08:55:23 -0500
Subject: [PATCH] new: `TracingManager.set_attributes()`

---
 adbpyg_adapter/adapter.py | 29 +++++++++--------------------
 adbpyg_adapter/tracing.py |  7 +++++++
 2 files changed, 16 insertions(+), 20 deletions(-)

diff --git a/adbpyg_adapter/adapter.py b/adbpyg_adapter/adapter.py
index 2bbf839..a392521 100644
--- a/adbpyg_adapter/adapter.py
+++ b/adbpyg_adapter/adapter.py
@@ -115,17 +115,6 @@ def set_tracer(self, tracer: Optional["trace.Tracer"]) -> None:
         if TRACING_ENABLED:
             TracingManager.set_tracer(tracer)
 
-    def __set_tracer_attributes(self, **attributes: Any) -> None:
-        """Set the OpenTelemetry tracer attributes for the adapter instance.
-
-        :param attributes: The OpenTelemetry tracer attributes.
-        :type attributes: Any
-        """
-        if TRACING_ENABLED:
-            current_span = trace.get_current_span()
-            for k, v in attributes.items():
-                current_span.set_attribute(k, v)
-
     ###########################
     # Public: ArangoDB -> PyG #
     ###########################
@@ -298,7 +287,7 @@ def udf_v1_x(v1_df):
         build a PyG-ready Tensor from a DataFrame equivalent to the
         associated ArangoDB collection.
         """
-        self.__set_tracer_attributes(name=name)
+        TracingManager.set_attributes(name=name)
         logger.debug(f"--arangodb_to_pyg('{name}')--")
 
         validate_adb_metagraph(metagraph)
@@ -573,7 +562,7 @@ def y_tensor_to_2_column_dataframe(pyg_tensor, adb_df):
         4) Dissasembles the 2-feature Tensor into two ArangoDB attributes,
             where each attribute holds one feature value.
         """
-        self.__set_tracer_attributes(name=name)
+        TracingManager.set_attributes(name=name)
         logger.debug(f"--pyg_to_arangodb('{name}')--")
 
         validate_pyg_metagraph(metagraph)
@@ -745,7 +734,7 @@ def __process_adb_v_col(
             node_data=node_data,
         )
 
-        self.__set_tracer_attributes(v_col=v_col, v_col_size=v_col_size)
+        TracingManager.set_attributes(v_col=v_col, v_col_size=v_col_size)
 
     @with_tracing
     def __process_adb_e_col(
@@ -804,7 +793,7 @@ def __process_adb_e_col(
             is_homogeneous=is_homogeneous,
         )
 
-        self.__set_tracer_attributes(e_col=e_col, e_col_size=e_col_size)
+        TracingManager.set_attributes(e_col=e_col, e_col_size=e_col_size)
 
     @with_tracing
     def __fetch_adb_docs(
@@ -1355,7 +1344,7 @@ def __process_pyg_node_batch(
         :return: The ArangoDB DataFrame representing the PyG Node batch.
         :rtype: pandas.DataFrame
         """
-        self.__set_tracer_attributes(start_index=start_index, end_index=end_index)
+        TracingManager.set_attributes(start_index=start_index, end_index=end_index)
 
         # 1. Set the ArangoDB Node Data
         df = self.__set_adb_data(
@@ -1415,7 +1404,7 @@ def __process_pyg_edge_batch(
         :return: The ArangoDB DataFrame representing the PyG Edge batch.
         :rtype: pandas.DataFrame
         """
-        self.__set_tracer_attributes(start_index=start_index, end_index=end_index)
+        TracingManager.set_attributes(start_index=start_index, end_index=end_index)
 
         src_n_type, _, dst_n_type = e_type
 
@@ -1495,7 +1484,7 @@ def __process_pyg_n_type(
             when inserting documents into the ArangoDB instance.
         :type adb_import_kwargs: Dict[str, Any]
         """
-        self.__set_tracer_attributes(n_type=n_type, n_type_size=node_data_total_size)
+        TracingManager.set_attributes(n_type=n_type, n_type_size=node_data_total_size)
 
         self.__process_batches(
             n_type,
@@ -1547,7 +1536,7 @@ def __process_pyg_e_type(
             when inserting documents into the ArangoDB instance.
         :type adb_import_kwargs: Dict[str, Any]
         """
-        self.__set_tracer_attributes(e_type=e_type, e_type_size=edge_data_total_size)
+        TracingManager.set_attributes(e_type=e_type, e_type_size=edge_data_total_size)
 
         self.__process_batches(
             e_type[1],
@@ -1796,7 +1785,7 @@ def __insert_adb_docs(
             https://docs.python-arango.com/en/main/specs.html#arango.collection.Collection.import_bulk
         :param adb_import_kwargs: Any
         """
-        self.__set_tracer_attributes(col=col, size=len(df))
+        TracingManager.set_attributes(col=col, size=len(df))
 
         action = f"ADB Import: '{col}' ({len(df)})"
         spinner_progress_task = spinner_progress.add_task("", action=action)
diff --git a/adbpyg_adapter/tracing.py b/adbpyg_adapter/tracing.py
index b00a1fb..51d2d34 100644
--- a/adbpyg_adapter/tracing.py
+++ b/adbpyg_adapter/tracing.py
@@ -28,6 +28,13 @@ def get_tracer(cls) -> Optional["Tracer"]:
     def set_tracer(cls, tracer: "Tracer") -> None:
         cls.__tracer = tracer
 
+    @classmethod
+    def set_attributes(self, **attributes: Any) -> None:
+        if TRACING_ENABLED and self.__tracer is not None:
+            current_span = trace.get_current_span()
+            for k, v in attributes.items():
+                current_span.set_attribute(k, v)
+
 
 T = TypeVar("T", bound=Callable[..., Any])