Skip to content

Commit 28743ec

Browse files
mittalpkcopybara-github
mittalpk
authored and
copybara-github
committed
fix: keep span status and events in SqliteSpanExporter
Merge #7486 Opening an existing spans database now adds three nullable columns in place as the next schema migration. Spans written before the upgrade read back with status UNSET and no events. Fixes #7484 PiperOrigin-RevId: 1013877695
1 parent 512d56b commit 28743ec

2 files changed

Lines changed: 298 additions & 2 deletions

File tree

‎src/google/adk/telemetry/sqlite_span_exporter.py‎

Lines changed: 78 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,11 +25,14 @@
2525
from typing import Mapping
2626
from typing import Sequence
2727

28+
from opentelemetry.sdk.trace import Event
2829
from opentelemetry.sdk.trace import ReadableSpan
2930
from opentelemetry.sdk.trace.export import SpanExporter
3031
from opentelemetry.sdk.trace.export import SpanExportResult
3132
from opentelemetry.trace import SpanContext
3233
from opentelemetry.trace import SpanKind
34+
from opentelemetry.trace import Status
35+
from opentelemetry.trace import StatusCode
3336
from opentelemetry.trace import TraceFlags
3437
from opentelemetry.trace import TraceState
3538
from opentelemetry.util.types import AttributeValue
@@ -65,6 +68,11 @@
6568
_MIGRATIONS: tuple[tuple[str, ...], ...] = (
6669
(_CREATE_SPANS_TABLE, _CREATE_SESSION_INDEX, _CREATE_TRACE_INDEX),
6770
("ALTER TABLE spans ADD COLUMN kind TEXT NOT NULL DEFAULT 'INTERNAL'",),
71+
(
72+
"ALTER TABLE spans ADD COLUMN status_code TEXT",
73+
"ALTER TABLE spans ADD COLUMN status_description TEXT",
74+
"ALTER TABLE spans ADD COLUMN events_json TEXT",
75+
),
6876
)
6977

7078
_INSERT_SPAN = """
@@ -78,8 +86,11 @@
7886
session_id,
7987
invocation_id,
8088
attributes_json,
81-
kind
82-
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?);
89+
kind,
90+
status_code,
91+
status_description,
92+
events_json
93+
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);
8394
"""
8495

8596
_DEFAULT_TIMEOUT_SECONDS = 30.0
@@ -165,6 +176,64 @@ def _deserialize_attributes(
165176
return {}
166177
return cast(dict[str, AttributeValue], decoded)
167178

179+
def _serialize_events(self, events: Sequence[Event]) -> str:
180+
try:
181+
return json.dumps(
182+
[
183+
{
184+
"name": event.name,
185+
"timestamp": event.timestamp,
186+
"attributes": dict(event.attributes or {}),
187+
}
188+
for event in events
189+
],
190+
ensure_ascii=False,
191+
default=lambda o: "<not serializable>",
192+
)
193+
except (TypeError, ValueError) as e:
194+
logger.debug("Failed to serialize span events: %r", e)
195+
return "[]"
196+
197+
def _deserialize_events(self, events_json: object) -> list[Event]:
198+
if not isinstance(events_json, (str, bytes, bytearray)):
199+
return []
200+
try:
201+
decoded: object = json.loads(events_json)
202+
except (json.JSONDecodeError, TypeError) as e:
203+
logger.debug("Failed to deserialize span events: %r", e)
204+
return []
205+
if not isinstance(decoded, list):
206+
return []
207+
events: list[Event] = []
208+
for item in decoded:
209+
if not isinstance(item, dict) or not isinstance(item.get("name"), str):
210+
continue
211+
timestamp = item.get("timestamp")
212+
# Event substitutes time_ns() for a missing timestamp; drop the event.
213+
if not isinstance(timestamp, int) or isinstance(timestamp, bool):
214+
continue
215+
attributes = item.get("attributes")
216+
events.append(
217+
Event(
218+
name=item["name"],
219+
attributes=attributes if isinstance(attributes, dict) else None,
220+
timestamp=timestamp,
221+
)
222+
)
223+
return events
224+
225+
def _deserialize_status(
226+
self, status_code: object, status_description: object
227+
) -> Status:
228+
try:
229+
code = StatusCode[str(status_code)]
230+
except KeyError:
231+
return Status(StatusCode.UNSET)
232+
# OpenTelemetry keeps a description only for an error status.
233+
if code is StatusCode.ERROR and isinstance(status_description, str):
234+
return Status(code, status_description)
235+
return Status(code)
236+
168237
def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
169238
try:
170239
with self._lock:
@@ -200,6 +269,9 @@ def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
200269
invocation_id,
201270
self._serialize_attributes(attributes),
202271
span.kind.name,
272+
span.status.status_code.name,
273+
span.status.description,
274+
self._serialize_events(span.events),
203275
))
204276
conn.executemany(_INSERT_SPAN, rows)
205277
conn.commit()
@@ -255,6 +327,10 @@ def _row_to_readable_span(self, row: sqlite3.Row) -> ReadableSpan:
255327
context=context,
256328
parent=parent,
257329
kind=SpanKind[row["kind"]],
330+
status=self._deserialize_status(
331+
row["status_code"], row["status_description"]
332+
),
333+
events=self._deserialize_events(row["events_json"]),
258334
attributes=attributes,
259335
start_time=row["start_time_unix_nano"],
260336
end_time=row["end_time_unix_nano"],

‎tests/unittests/telemetry/test_sqlite_span_exporter.py‎

Lines changed: 220 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,15 @@
2020

2121
from google.adk.telemetry.sqlite_span_exporter import _apply_migrations
2222
from google.adk.telemetry.sqlite_span_exporter import SqliteSpanExporter
23+
from opentelemetry.sdk.trace import Event
2324
from opentelemetry.sdk.trace import ReadableSpan
25+
from opentelemetry.sdk.trace import TracerProvider
26+
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
2427
from opentelemetry.sdk.trace.export import SpanExportResult
2528
from opentelemetry.trace import SpanContext
2629
from opentelemetry.trace import SpanKind
30+
from opentelemetry.trace import Status
31+
from opentelemetry.trace import StatusCode
2732
from opentelemetry.trace import TraceFlags
2833
from opentelemetry.trace import TraceState
2934
import pytest
@@ -645,3 +650,218 @@ def test_get_spans_ordered_by_start_time(tmp_path):
645650
assert result[0].context.span_id == 0x100
646651
assert result[1].context.span_id == 0x200
647652
assert result[2].context.span_id == 0x300
653+
654+
655+
def test_round_trip_preserves_error_status_and_events(tmp_path):
656+
exporter = SqliteSpanExporter(db_path=str(tmp_path / "test.db"))
657+
context = SpanContext(
658+
trace_id=0xDEF45,
659+
span_id=0xABC12,
660+
is_remote=False,
661+
trace_flags=TraceFlags(TraceFlags.SAMPLED),
662+
trace_state=TraceState(),
663+
)
664+
span = ReadableSpan(
665+
name="failing_tool",
666+
context=context,
667+
attributes={"gcp.vertex.agent.session_id": "s1"},
668+
start_time=1000,
669+
end_time=2000,
670+
status=Status(StatusCode.ERROR, "ValueError: boom"),
671+
events=[
672+
Event(
673+
name="exception",
674+
attributes={
675+
"exception.type": "ValueError",
676+
"exception.message": "boom",
677+
},
678+
timestamp=1500,
679+
)
680+
],
681+
)
682+
683+
assert exporter.export([span]) == SpanExportResult.SUCCESS
684+
(restored,) = exporter.get_all_spans_for_session("s1")
685+
686+
assert restored.status.status_code is StatusCode.ERROR
687+
assert restored.status.description == "ValueError: boom"
688+
assert len(restored.events) == 1
689+
event = restored.events[0]
690+
assert event.name == "exception"
691+
assert event.timestamp == 1500
692+
assert dict(event.attributes) == {
693+
"exception.type": "ValueError",
694+
"exception.message": "boom",
695+
}
696+
697+
698+
def test_exception_raised_inside_a_real_span_is_read_back(tmp_path):
699+
"""The case from the issue: a run that fails must not read back as clean."""
700+
exporter = SqliteSpanExporter(db_path=str(tmp_path / "test.db"))
701+
provider = TracerProvider()
702+
provider.add_span_processor(SimpleSpanProcessor(exporter))
703+
tracer = provider.get_tracer(__name__)
704+
705+
try:
706+
with tracer.start_as_current_span(
707+
"run", attributes={"gcp.vertex.agent.session_id": "s1"}
708+
):
709+
raise ValueError("boom")
710+
except ValueError:
711+
pass
712+
713+
(restored,) = exporter.get_all_spans_for_session("s1")
714+
715+
assert restored.status.status_code is StatusCode.ERROR
716+
assert "boom" in (restored.status.description or "")
717+
assert [event.name for event in restored.events] == ["exception"]
718+
assert restored.events[0].attributes["exception.type"] == "ValueError"
719+
720+
721+
def test_span_without_status_or_events_reads_back_unset_and_empty(tmp_path):
722+
exporter = SqliteSpanExporter(db_path=str(tmp_path / "test.db"))
723+
exporter.export(
724+
[_create_span(attributes={"gcp.vertex.agent.session_id": "s1"})]
725+
)
726+
727+
(restored,) = exporter.get_all_spans_for_session("s1")
728+
729+
assert restored.status.status_code is StatusCode.UNSET
730+
assert list(restored.events) == []
731+
732+
733+
@pytest.mark.parametrize("user_version", [0, 2])
734+
def test_database_from_older_version_gains_the_new_columns(
735+
tmp_path: Path, user_version: int
736+
) -> None:
737+
"""A file written before status and events were stored keeps working."""
738+
db_path = tmp_path / "old.db"
739+
conn = sqlite3.connect(db_path)
740+
conn.execute("""
741+
CREATE TABLE spans (
742+
span_id TEXT PRIMARY KEY,
743+
trace_id TEXT NOT NULL,
744+
parent_span_id TEXT,
745+
name TEXT NOT NULL,
746+
start_time_unix_nano INTEGER,
747+
end_time_unix_nano INTEGER,
748+
session_id TEXT,
749+
invocation_id TEXT,
750+
attributes_json TEXT
751+
)
752+
""")
753+
if user_version == 2:
754+
conn.execute(
755+
"ALTER TABLE spans ADD COLUMN kind TEXT NOT NULL DEFAULT 'INTERNAL'"
756+
)
757+
conn.execute("PRAGMA user_version = 2")
758+
conn.execute(
759+
"INSERT INTO spans (span_id, trace_id, name, start_time_unix_nano,"
760+
" end_time_unix_nano, session_id, attributes_json)"
761+
" VALUES (?, ?, ?, ?, ?, ?, ?)",
762+
("00000000000abc12", "0" * 27 + "def45", "old", 1, 2, "s1", "{}"),
763+
)
764+
conn.commit()
765+
conn.close()
766+
767+
exporter = SqliteSpanExporter(db_path=str(db_path))
768+
try:
769+
assert (
770+
exporter.export([
771+
_create_span(
772+
span_id=0x999,
773+
name="new",
774+
attributes={"gcp.vertex.agent.session_id": "s1"},
775+
start_time=5,
776+
)
777+
])
778+
== SpanExportResult.SUCCESS
779+
)
780+
781+
old, new = exporter.get_all_spans_for_session("s1")
782+
assert (old.name, new.name) == ("old", "new")
783+
assert old.kind is SpanKind.INTERNAL
784+
assert old.status.status_code is StatusCode.UNSET
785+
assert list(old.events) == []
786+
finally:
787+
exporter.shutdown()
788+
789+
with sqlite3.connect(db_path) as conn:
790+
assert conn.execute("PRAGMA user_version").fetchone()[0] == 3
791+
assert _column_names(conn, "spans")[-4:] == [
792+
"kind",
793+
"status_code",
794+
"status_description",
795+
"events_json",
796+
]
797+
798+
799+
def test_corrupt_stored_status_and_events_read_back_as_defaults(tmp_path):
800+
exporter = SqliteSpanExporter(db_path=str(tmp_path / "test.db"))
801+
conn = exporter._get_connection()
802+
rows = [
803+
("0000000000000001", "BOGUS", "not json"),
804+
("0000000000000002", None, "{}"),
805+
(
806+
"0000000000000003",
807+
"OK",
808+
(
809+
'[null, {"name": 1}, {"name": "e", "attributes": "x",'
810+
' "timestamp": 5}, {"name": "no_ts"}, {"name": "null_ts",'
811+
' "timestamp": null}, {"name": "str_ts", "timestamp": "5"},'
812+
' {"name": "bool_ts", "timestamp": true}]'
813+
),
814+
),
815+
]
816+
for span_id, status_code, events_json in rows:
817+
conn.execute(
818+
"INSERT INTO spans (span_id, trace_id, name, session_id,"
819+
" start_time_unix_nano, status_code, status_description, events_json)"
820+
" VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
821+
(
822+
span_id,
823+
"0" * 27 + "def45",
824+
"corrupt",
825+
"s1",
826+
int(span_id, 16),
827+
status_code,
828+
"ignored",
829+
events_json,
830+
),
831+
)
832+
conn.commit()
833+
834+
first, second, third = exporter.get_all_spans_for_session("s1")
835+
836+
assert first.status.status_code is StatusCode.UNSET
837+
assert list(first.events) == []
838+
assert second.status.status_code is StatusCode.UNSET
839+
assert list(second.events) == []
840+
assert third.status.status_code is StatusCode.OK
841+
assert third.status.description is None
842+
assert [event.name for event in third.events] == ["e"]
843+
assert third.events[0].attributes is None
844+
assert third.events[0].timestamp == 5
845+
846+
847+
def test_unserializable_event_attributes_do_not_fail_export(tmp_path):
848+
exporter = SqliteSpanExporter(db_path=str(tmp_path / "test.db"))
849+
span = ReadableSpan(
850+
name="op",
851+
context=SpanContext(
852+
trace_id=0xDEF45,
853+
span_id=0xABC12,
854+
is_remote=False,
855+
trace_flags=TraceFlags(TraceFlags.SAMPLED),
856+
trace_state=TraceState(),
857+
),
858+
attributes={"gcp.vertex.agent.session_id": "s1"},
859+
start_time=1000,
860+
end_time=2000,
861+
events=[Event(name="bad", attributes={"k": {(1, 2): "tuple key"}})],
862+
)
863+
864+
assert exporter.export([span]) == SpanExportResult.SUCCESS
865+
(restored,) = exporter.get_all_spans_for_session("s1")
866+
assert restored.name == "op"
867+
assert list(restored.events) == []

0 commit comments

Comments
 (0)