propogate user id to all children of SynapseStorage::__init__

schematic/__init__.py

@@ -27,7 +27,6 @@
 from schematic.configuration.configuration import CONFIG
 from schematic.loader import LOADER
 from schematic.version import __version__
-from schematic_api.api.security_controller import info_from_bearer_auth
 from dotenv import load_dotenv
 
 Synapse.allow_client_caching(False)
@@ -43,11 +42,9 @@ def __init__(self, attributes_to_propagate) -> None:
 
     def on_start(self, span: Span, parent_context: SpanContext) -> None:
         """Propagates attributes from the parent span to the child span.
-
         Arguments:
             span: The child span to which the attributes should be propagated.
             parent_context: The context of the parent span.
-
         Returns:
             None
         """
@@ -136,10 +133,7 @@ def set_up_tracing(session: requests.Session) -> None:
         exporter = OTLPSpanExporter(session=session)
         trace.get_tracer_provider().add_span_processor(BatchSpanProcessor(exporter))
         # Add the custom AttributePropagatingSpanProcessor to propagate attributes
-        attributes_to_propagate = ["user.id"]
-        attribute_propagator = AttributePropagatingSpanProcessor(
-            attributes_to_propagate
-        )
+        attribute_propagator = AttributePropagatingSpanProcessor(["user.id"])
         trace.get_tracer_provider().add_span_processor(attribute_propagator)
     else:
         trace.set_tracer_provider(TracerProvider(sampler=ALWAYS_OFF))
@@ -182,32 +176,6 @@ def request_hook(span: Span, environ: Dict) -> None:
     """
     if not span or not span.is_recording():
         return
-    try:
-        auth_header = environ.get("HTTP_AUTHORIZATION", None)
-        access_token = os.getenv("SYNAPSE_ACCESS_TOKEN", None)
-
-        if auth_header and len(auth_header.split(" ")) > 1:
-            token = auth_header.split(" ")[1]
-        else:
-            token = access_token
-
-        if token:
-            user_info = info_from_bearer_auth(token)
-
-        if user_info:
-            span.set_attribute("user.id", user_info.get("sub"))
-
-    except Exception:
-        logger.exception("Failed to set user info in span")
-
-    try:
-        if (request := environ.get("werkzeug.request", None)) and isinstance(
-            request, Request
-        ):
-            for arg in request.args:
-                span.set_attribute(key=f"schematic.{arg}", value=request.args[arg])
-    except Exception:
-        logger.exception("Failed to set request info in span")
 
 
 def response_hook(span: Span, status: str, response_headers: List) -> None:

schematic/store/synapse.py

@@ -75,6 +75,7 @@
 from schematic.utils.io_utils import cleanup_temporary_storage
 from schematic.utils.schema_utils import get_class_label_from_display_name
 from schematic.utils.validate_utils import comma_separated_list_regex, rule_in_rule_list
+from schematic_api.api.security_controller import info_from_bearer_auth
 
 logger = logging.getLogger("Synapse storage")
 
@@ -321,6 +322,13 @@ def __init__(
             Consider necessity of adding "columns" and "where_clauses" params to the constructor. Currently with how `query_fileview` is implemented, these params are not needed at this step but could be useful in the future if the need for more scoped querys expands.
         """
         self.syn = self.login(synapse_cache_path, access_token)
+
+        current_span = trace.get_current_span()
+        if current_span.is_recording() and access_token:
+            user_info = info_from_bearer_auth(access_token)
+            if user_info:
+                current_span.set_attribute("user.id", user_info.get("sub"))
+
         self.project_scope = project_scope
         self.storageFileview = CONFIG.synapse_master_fileview_id
         self.manifest = CONFIG.synapse_manifest_basename
@@ -496,21 +504,24 @@ def login(
         Returns:
             synapseclient.Synapse: A Synapse object that is logged in
         """
-        # If no token is provided, try retrieving access token from environment
         if not access_token:
             access_token = os.getenv("SYNAPSE_ACCESS_TOKEN")
 
         # login using a token
         if access_token:
             try:
+                current_span = trace.get_current_span()
+                if current_span.is_recording() and access_token:
+                    user_info = info_from_bearer_auth(access_token)
+                    if user_info:
+                        current_span.set_attribute("user.id", user_info.get("sub"))
                 syn = synapseclient.Synapse(
                     cache_root_dir=synapse_cache_path,
                     debug=False,
                     skip_checks=True,
                     cache_client=False,
                 )
                 syn.login(authToken=access_token, silent=True)
-                current_span = trace.get_current_span()
             except SynapseHTTPError as exc:
                 raise ValueError(
                     "No access to resources. Please make sure that your token is correct"
@@ -528,7 +539,6 @@ def login(
             current_span = trace.get_current_span()
             if current_span.is_recording():
                 current_span.set_attribute("user.id", syn.credentials.owner_id)
-
         return syn
 
     def missing_entity_handler(method):