diff --git a/adbpyg_adapter/adapter.py b/adbpyg_adapter/adapter.py index c981cab..accd05a 100644 --- 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 index 069e67a..1004882 100644 --- 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