fix: custom span names for @with_tracing decorator

(private methods only)
arangoml · Jan 18, 2024 · ddbd331 · ddbd331
1 parent 3d0e3de
commit ddbd331
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 34 deletions.
diff --git a/adbpyg_adapter/adapter.py b/adbpyg_adapter/adapter.py
@@ -126,7 +126,7 @@ def set_tracer(self, tracer: Optional["Tracer"]) -> None:
     # Public: ArangoDB -> PyG #
     ###########################
 
-    @with_tracing
+    @with_tracing()
     def arangodb_to_pyg(
         self,
         name: str,
@@ -363,7 +363,7 @@ def udf_v1_x(v1_df):
         logger.info(f"Created PyG '{name}' Graph")
         return data
 
-    @with_tracing
+    @with_tracing()
     def arangodb_collections_to_pyg(
         self,
         name: str,
@@ -417,7 +417,7 @@ def arangodb_collections_to_pyg(
             name, metagraph, preserve_adb_keys, strict, **adb_export_kwargs
         )
 
-    @with_tracing
+    @with_tracing()
     def arangodb_graph_to_pyg(
         self,
         name: str,
@@ -469,7 +469,7 @@ def arangodb_graph_to_pyg(
     # Public: PyG -> ArangoDB #
     ###########################
 
-    @with_tracing
+    @with_tracing()
     def pyg_to_arangodb(
         self,
         name: str,
@@ -696,7 +696,7 @@ def y_tensor_to_2_column_dataframe(pyg_tensor, adb_df):
     # Private: ArangoDB -> PyG #
     ############################
 
-    @with_tracing
+    @with_tracing("__process_adb_v_col")
     def __process_adb_v_col(
         self,
         v_col: str,
@@ -742,7 +742,7 @@ def __process_adb_v_col(
             node_data=node_data,
         )
 
-    @with_tracing
+    @with_tracing("__process_adb_e_col")
     def __process_adb_e_col(
         self,
         e_col: str,
@@ -801,7 +801,7 @@ def __process_adb_e_col(
             is_homogeneous=is_homogeneous,
         )
 
-    @with_tracing
+    @with_tracing("__fetch_adb_docs")
     def __fetch_adb_docs(
         self,
         col: str,
@@ -865,7 +865,7 @@ def get_aql_return_value(
 
             return cursor, col_size
 
-    @with_tracing
+    @with_tracing("__process_adb_cursor")
     def __process_adb_cursor(
         self,
         progress_color: str,
@@ -916,7 +916,7 @@ def __process_adb_cursor(
                 if cursor.has_more():
                     cursor.fetch()
 
-    @with_tracing
+    @with_tracing("__process_adb_vertex_df")
     def __process_adb_vertex_df(
         self,
         i: int,
@@ -963,7 +963,7 @@ def __process_adb_vertex_df(
 
         return i
 
-    @with_tracing
+    @with_tracing("__process_adb_edge_df")
     def __process_adb_edge_df(
         self,
         _: int,
@@ -1067,7 +1067,7 @@ def __process_adb_edge_df(
 
         return 1  # Useless return value, but needed for type hinting
 
-    @with_tracing
+    @with_tracing("__split_adb_ids")
     def __split_adb_ids(self, s: Series) -> Series:
         """AranogDB -> PyG: Helper method to split the ArangoDB IDs
         within a Series into two columns
@@ -1080,7 +1080,7 @@ def __split_adb_ids(self, s: Series) -> Series:
         """
         return s.str.split(pat="/", n=1, expand=True)
 
-    @with_tracing
+    @with_tracing("__set_pyg_data")
     def __set_pyg_data(
         self,
         meta: Union[Set[str], Dict[str, ADBMetagraphValues]],
@@ -1116,7 +1116,7 @@ def __set_pyg_data(
                 m = f"'{k}' key in PyG Data must point to a Tensor"
                 raise TypeError(m)
 
-    @with_tracing
+    @with_tracing("__build_tensor_from_dataframe")
     def __build_tensor_from_dataframe(
         self,
         adb_df: DataFrame,
@@ -1173,7 +1173,7 @@ def __build_tensor_from_dataframe(
     # Private: PyG -> ArangoDB #
     ############################
 
-    @with_tracing
+    @with_tracing("__get_node_and_edge_types")
     def __get_node_and_edge_types(
         self,
         name: str,
@@ -1217,7 +1217,7 @@ def __get_node_and_edge_types(
 
         return node_types, edge_types
 
-    @with_tracing
+    @with_tracing("__etypes_to_edefinitions")
     def __etypes_to_edefinitions(self, edge_types: List[EdgeType]) -> List[Json]:
         """PyG -> ArangoDB: Converts PyG edge_types to ArangoDB edge_definitions
 
@@ -1262,7 +1262,7 @@ def __etypes_to_edefinitions(self, edge_types: List[EdgeType]) -> List[Json]:
 
         return edge_definitions
 
-    @with_tracing
+    @with_tracing("__ntypes_to_ocollections")
     def __ntypes_to_ocollections(
         self, node_types: List[str], edge_types: List[EdgeType]
     ) -> List[str]:
@@ -1286,7 +1286,7 @@ def __ntypes_to_ocollections(
         orphan_collections = set(node_types) ^ non_orphan_collections
         return list(orphan_collections)
 
-    @with_tracing
+    @with_tracing("__create_adb_graph")
     def __create_adb_graph(
         self,
         name: str,
@@ -1325,7 +1325,7 @@ def __create_adb_graph(
             orphan_collections,
         )
 
-    @with_tracing
+    @with_tracing("__process_pyg_node_batch")
     def __process_pyg_node_batch(
         self,
         n_type: str,
@@ -1385,7 +1385,7 @@ def __process_pyg_node_batch(
 
         return df
 
-    @with_tracing
+    @with_tracing("__process_pyg_edge_batch")
     def __process_pyg_edge_batch(
         self,
         e_type: EdgeType,
@@ -1461,7 +1461,7 @@ def __process_pyg_edge_batch(
 
         return df
 
-    @with_tracing
+    @with_tracing("__process_pyg_n_type")
     def __process_pyg_n_type(
         self,
         n_type: str,
@@ -1513,7 +1513,7 @@ def __process_pyg_n_type(
             adb_import_kwargs,
         )
 
-    @with_tracing
+    @with_tracing("__process_pyg_e_type")
     def __process_pyg_e_type(
         self,
         e_type: EdgeType,
@@ -1619,7 +1619,7 @@ def __process_batches(
             start_index = end_index
             end_index = min(end_index + batch_size, total_size)
 
-    @with_tracing
+    @with_tracing("__set_adb_data")
     def __set_adb_data(
         self,
         df: DataFrame,
@@ -1705,7 +1705,7 @@ def __set_adb_data(
 
         return df
 
-    @with_tracing
+    @with_tracing("__build_dataframe_from_tensor")
     def __build_dataframe_from_tensor(
         self,
         pyg_tensor: Tensor,
@@ -1782,7 +1782,7 @@ def __build_dataframe_from_tensor(
 
         raise PyGMetagraphError(f"Invalid {meta_val} type")  # pragma: no cover
 
-    @with_tracing
+    @with_tracing("__insert_adb_docs")
     def __insert_adb_docs(
         self,
         spinner_progress: Progress,

diff --git a/adbpyg_adapter/tracing.py b/adbpyg_adapter/tracing.py
@@ -50,19 +50,22 @@ def set_attributes(self, **attributes: Any) -> None:  # pragma: no cover
 T = TypeVar("T", bound=Callable[..., Any])
 
 
-def with_tracing(method: T) -> T:
-    if not TRACING_ENABLED:
-        return method  # pragma: no cover
+def with_tracing(span_name: Optional[str] = None) -> Callable[[T], T]:
+    def decorator(method: T) -> T:
+        if not TRACING_ENABLED:
+            return method  # pragma: no cover
 
-    @wraps(method)
-    def decorator(*args: Any, **kwargs: Any) -> Any:
-        if tracer := TracingManager.get_tracer():
-            with tracer.start_as_current_span(method.__name__):
-                return method(*args, **kwargs)
+        @wraps(method)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            if tracer := TracingManager.get_tracer():
+                with tracer.start_as_current_span(span_name or method.__name__):
+                    return method(*args, **kwargs)
 
-        return method(*args, **kwargs)
+            return method(*args, **kwargs)
 
-    return cast(T, decorator)
+        return cast(T, wrapper)
+
+    return decorator
 
 
 @contextmanager