Enable exporting spans to snowflake stage if a TruLensSnowflakeSpanExporter is provided

examples/experimental/otel_exporter.ipynb

@@ -118,10 +118,10 @@
     "    session=session,\n",
     ")\n",
     "\n",
-    "with custom_app() as recording:\n",
+    "with custom_app(run_name=\"test run\", input_id=\"456\") as recording:\n",
     "    test_app.respond_to_query(\"test\")\n",
     "\n",
-    "with custom_app(run_name=\"test run\", input_id=\"456\") as recording:\n",
+    "with custom_app() as recording:\n",
     "    test_app.respond_to_query(\"throw\")"
    ]
   },

src/core/trulens/experimental/otel_tracing/core/exporter/connector.py

@@ -0,0 +1,44 @@
+import logging
+from typing import Sequence
+
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter
+from opentelemetry.sdk.trace.export import SpanExportResult
+from trulens.core.database import connector as core_connector
+from trulens.experimental.otel_tracing.core.exporter.utils import (
+    check_if_trulens_span,
+)
+from trulens.experimental.otel_tracing.core.exporter.utils import (
+    construct_event,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class TruLensOTELSpanExporter(SpanExporter):
+    """
+    Implementation of `SpanExporter` that flushes the spans in the TruLens session to the connector.
+    """
+
+    connector: core_connector.DBConnector
+    """
+    The database connector used to export the spans.
+    """
+
+    def __init__(self, connector: core_connector.DBConnector):
+        self.connector = connector
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        trulens_spans = list(filter(check_if_trulens_span, spans))
+
+        try:
+            events = list(map(construct_event, trulens_spans))
+            self.connector.add_events(events)
+
+        except Exception as e:
+            logger.error(
+                f"Error exporting spans to the database: {e}",
+            )
+            return SpanExportResult.FAILURE
+
+        return SpanExportResult.SUCCESS

src/core/trulens/experimental/otel_tracing/core/exporter/snowflake.py

@@ -0,0 +1,107 @@
+import logging
+import os
+import tempfile
+from typing import Sequence
+
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter
+from opentelemetry.sdk.trace.export import SpanExportResult
+from trulens.connectors.snowflake import SnowflakeConnector
+from trulens.core.database import connector as core_connector
+from trulens.experimental.otel_tracing.core.exporter.utils import (
+    check_if_trulens_span,
+)
+from trulens.experimental.otel_tracing.core.exporter.utils import (
+    convert_readable_span_to_proto,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class TruLensSnowflakeSpanExporter(SpanExporter):
+    """
+    Implementation of `SpanExporter` that flushes the spans in the TruLens session to a Snowflake Stage.
+    """
+
+    connector: core_connector.DBConnector
+    """
+    The database connector used to export the spans.
+    """
+
+    def __init__(self, connector: core_connector.DBConnector):
+        self.connector = connector
+
+    def _export_to_snowflake_stage(
+        self, spans: Sequence[ReadableSpan]
+    ) -> SpanExportResult:
+        """
+        Exports a list of spans to a Snowflake stage as a protobuf file.
+        This function performs the following steps:
+        1. Writes the provided spans to a temporary protobuf file.
+        2. Creates a Snowflake stage if it does not already exist.
+        3. Uploads the temporary protobuf file to the Snowflake stage.
+        4. Removes the temporary protobuf file.
+        Args:
+            spans (Sequence[ReadableSpan]): A sequence of spans to be exported.
+        Returns:
+            SpanExportResult: The result of the export operation, either SUCCESS or FAILURE.
+        """
+        if not isinstance(self.connector, SnowflakeConnector):
+            return SpanExportResult.FAILURE
+
+        # Avoid uploading empty files to the stage
+        if not spans:
+            return SpanExportResult.SUCCESS
+
+        snowpark_session = self.connector.snowpark_session
+        tmp_file_path = ""
+
+        try:
+            with tempfile.NamedTemporaryFile(
+                delete=False, suffix=".pb", mode="wb"
+            ) as tmp_file:
+                tmp_file_path = tmp_file.name
+                logger.debug(
+                    f"Writing spans to the protobuf file: {tmp_file_path}"
+                )
+
+                for span in spans:
+                    span_proto = convert_readable_span_to_proto(span)
+                    tmp_file.write(span_proto.SerializeToString())
+                logger.debug(
+                    f"Spans written to the protobuf file: {tmp_file_path}"
+                )
+        except Exception as e:
+            logger.error(f"Error writing spans to the protobuf file: {e}")
+            return SpanExportResult.FAILURE
+
+        try:
+            logger.debug("Uploading file to Snowflake stage")
+
+            logger.debug("Creating Snowflake stage if it does not exist")
+            snowpark_session.sql(
+                "CREATE TEMP STAGE IF NOT EXISTS trulens_spans"
+            ).collect()
+
+            logger.debug("Uploading the protobuf file to the stage")
+            snowpark_session.sql(
+                f"PUT file://{tmp_file_path} @trulens_spans"
+            ).collect()
+
+        except Exception as e:
+            logger.error(f"Error uploading the protobuf file to the stage: {e}")
+            return SpanExportResult.FAILURE
+
+        try:
+            logger.debug("Removing the temporary protobuf file")
+            os.remove(tmp_file_path)
+        except Exception as e:
+            # Not returning failure here since the export was technically a success
+            logger.error(f"Error removing the temporary protobuf file: {e}")
+
+        return SpanExportResult.SUCCESS
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        trulens_spans = list(filter(check_if_trulens_span, spans))
+
+        return self._export_to_snowflake_stage(trulens_spans)

src/core/trulens/experimental/otel_tracing/core/exporter.py → src/core/trulens/experimental/otel_tracing/core/exporter/utils.py

@@ -1,20 +1,14 @@
 from datetime import datetime
 import logging
-import os
-import tempfile
-from typing import Any, Optional, Sequence
+from typing import Any, Optional
 
 from opentelemetry.proto.common.v1.common_pb2 import AnyValue
 from opentelemetry.proto.common.v1.common_pb2 import ArrayValue
 from opentelemetry.proto.common.v1.common_pb2 import KeyValue
 from opentelemetry.proto.common.v1.common_pb2 import KeyValueList
 from opentelemetry.proto.trace.v1.trace_pb2 import Span as SpanProto
 from opentelemetry.sdk.trace import ReadableSpan
-from opentelemetry.sdk.trace.export import SpanExporter
-from opentelemetry.sdk.trace.export import SpanExportResult
 from opentelemetry.trace import StatusCode
-from trulens.connectors.snowflake import SnowflakeConnector
-from trulens.core.database import connector as core_connector
 from trulens.core.schema import event as event_schema
 from trulens.otel.semconv.trace import SpanAttributes
 
@@ -147,121 +141,3 @@ def construct_event(span: ReadableSpan) -> event_schema.Event:
             "parent_id": str(parent.span_id if parent else ""),
         },
     )
-
-
-class TruLensSnowflakeSpanExporter(SpanExporter):
-    """
-    Implementation of `SpanExporter` that flushes the spans in the TruLens session to a Snowflake Stage.
-    """
-
-    connector: core_connector.DBConnector
-    """
-    The database connector used to export the spans.
-    """
-
-    def __init__(self, connector: core_connector.DBConnector):
-        self.connector = connector
-
-    def _export_to_snowflake_stage(
-        self, spans: Sequence[ReadableSpan]
-    ) -> SpanExportResult:
-        """
-        Exports a list of spans to a Snowflake stage as a protobuf file.
-        This function performs the following steps:
-        1. Writes the provided spans to a temporary protobuf file.
-        2. Creates a Snowflake stage if it does not already exist.
-        3. Uploads the temporary protobuf file to the Snowflake stage.
-        4. Removes the temporary protobuf file.
-        Args:
-            spans (Sequence[ReadableSpan]): A sequence of spans to be exported.
-        Returns:
-            SpanExportResult: The result of the export operation, either SUCCESS or FAILURE.
-        """
-        if not isinstance(self.connector, SnowflakeConnector):
-            return SpanExportResult.FAILURE
-
-        # Avoid uploading empty files to the stage
-        if not spans:
-            return SpanExportResult.SUCCESS
-
-        snowpark_session = self.connector.snowpark_session
-        tmp_file_path = ""
-
-        try:
-            with tempfile.NamedTemporaryFile(
-                delete=False, suffix=".pb", mode="wb"
-            ) as tmp_file:
-                tmp_file_path = tmp_file.name
-                logger.debug(
-                    f"Writing spans to the protobuf file: {tmp_file_path}"
-                )
-
-                for span in spans:
-                    span_proto = convert_readable_span_to_proto(span)
-                    tmp_file.write(span_proto.SerializeToString())
-                logger.debug(
-                    f"Spans written to the protobuf file: {tmp_file_path}"
-                )
-        except Exception as e:
-            logger.error(f"Error writing spans to the protobuf file: {e}")
-            return SpanExportResult.FAILURE
-
-        try:
-            logger.debug("Uploading file to Snowflake stage")
-
-            logger.debug("Creating Snowflake stage if it does not exist")
-            snowpark_session.sql(
-                "CREATE TEMP STAGE IF NOT EXISTS trulens_spans"
-            ).collect()
-
-            logger.debug("Uploading the protobuf file to the stage")
-            snowpark_session.sql(
-                f"PUT file://{tmp_file_path} @trulens_spans"
-            ).collect()
-
-        except Exception as e:
-            logger.error(f"Error uploading the protobuf file to the stage: {e}")
-            return SpanExportResult.FAILURE
-
-        try:
-            logger.debug("Removing the temporary protobuf file")
-            os.remove(tmp_file_path)
-        except Exception as e:
-            # Not returning failure here since the export was technically a success
-            logger.error(f"Error removing the temporary protobuf file: {e}")
-
-        return SpanExportResult.SUCCESS
-
-    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
-        trulens_spans = list(filter(check_if_trulens_span, spans))
-
-        return self._export_to_snowflake_stage(trulens_spans)
-
-
-class TruLensOTELSpanExporter(SpanExporter):
-    """
-    Implementation of `SpanExporter` that flushes the spans in the TruLens session to the connector.
-    """
-
-    connector: core_connector.DBConnector
-    """
-    The database connector used to export the spans.
-    """
-
-    def __init__(self, connector: core_connector.DBConnector):
-        self.connector = connector
-
-    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
-        trulens_spans = list(filter(check_if_trulens_span, spans))
-
-        try:
-            events = list(map(construct_event, trulens_spans))
-            self.connector.add_events(events)
-
-        except Exception as e:
-            logger.error(
-                f"Error exporting spans to the database: {e}",
-            )
-            return SpanExportResult.FAILURE
-
-        return SpanExportResult.SUCCESS

src/core/trulens/experimental/otel_tracing/core/instrument.py

@@ -4,6 +4,7 @@
 import uuid
 
 from opentelemetry import trace
+from opentelemetry.baggage import get_baggage
 from opentelemetry.baggage import remove_baggage
 from opentelemetry.baggage import set_baggage
 import opentelemetry.context as context_api
@@ -110,17 +111,24 @@ class OTELRecordingContext:
     context api keeps track of what is changed in the context, and used to undo the changes.
     """
 
+    context_keys_added: List[str] = []
+    """
+    Keys added to the OTEL context.
+    """
+
     def __init__(self, *, app: core_app.App, run_name: str, input_id: str):
         self.app = app
         self.run_name = run_name
         self.input_id = input_id
         self.tokens = []
+        self.context_keys_added = []
         self.span_context = None
 
     # Calling set_baggage does not actually add the baggage to the current context, but returns a new one
     # To avoid issues with remembering to add/remove the baggage, we attach it to the runtime context.
     def attach_to_context(self, key: str, value: object):
         self.tokens.append(context_api.attach(set_baggage(key, value)))
+        self.context_keys_added.append(key)
 
     # For use as a context manager.
     def __enter__(self):
@@ -130,7 +138,10 @@ def __enter__(self):
 
         tracer = trace.get_tracer_provider().get_tracer(TRULENS_SERVICE_NAME)
 
-        self.attach_to_context(SpanAttributes.RECORD_ID, otel_record_id)
+        # There might be nested records - in those scenarios use the outer one.
+        if not get_baggage(SpanAttributes.RECORD_ID):
+            self.attach_to_context(SpanAttributes.RECORD_ID, otel_record_id)
+
         self.attach_to_context(SpanAttributes.APP_NAME, self.app.app_name)
         self.attach_to_context(SpanAttributes.APP_VERSION, self.app.app_version)
 
@@ -171,16 +182,12 @@ def __exit__(self, exc_type, exc_value, exc_tb):
             # TODO[SNOW-1854360]: Add in feature function spans.
             self.span_context.__exit__(exc_type, exc_value, exc_tb)
 
-        remove_baggage(SpanAttributes.RECORD_ID)
-        remove_baggage(SpanAttributes.APP_NAME)
-        remove_baggage(SpanAttributes.APP_VERSION)
-
-        # Safe to remove baggage keys even if we did not set them
-        remove_baggage(SpanAttributes.RUN_NAME)
-        remove_baggage(SpanAttributes.INPUT_ID)
-
         logger.debug("Exiting the OTEL app context.")
 
+        # Clearing the context / baggage added.
+        while self.context_keys_added:
+            remove_baggage(self.context_keys_added.pop())
+
         while self.tokens:
             # Clearing the context once we're done with this root span.
             # See https://github.com/open-telemetry/opentelemetry-python/issues/2432#issuecomment-1593458684

src/core/trulens/experimental/otel_tracing/core/session.py

@@ -10,7 +10,7 @@
 from trulens.core.database.connector import DBConnector
 from trulens.core.utils import python as python_utils
 from trulens.core.utils import text as text_utils
-from trulens.experimental.otel_tracing.core.exporter import (
+from trulens.experimental.otel_tracing.core.exporter.connector import (
     TruLensOTELSpanExporter,
 )