From 4069887b1fd5cc582a0b95302bfb7a07c6b9968b Mon Sep 17 00:00:00 2001 From: samecho Date: Mon, 26 Oct 2026 15:53:44 -0400 Subject: [PATCH 1/2] fix(telemetry): preserve span links in SQLite traces Store link contexts and attributes so they survive exporter restart. Upgrade existing span tables without losing rows, and serialize schema initialization across exporter instances. Related: #7484 --- .../adk/telemetry/sqlite_span_exporter.py | 82 ++++++++- .../telemetry/test_sqlite_span_exporter.py | 174 ++++++++++++++++++ 2 files changed, 247 insertions(+), 9 deletions(-) diff --git a/src/google/adk/telemetry/sqlite_span_exporter.py b/src/google/adk/telemetry/sqlite_span_exporter.py index eae4c11bd15..da3d069430b 100644 --- a/src/google/adk/telemetry/sqlite_span_exporter.py +++ b/src/google/adk/telemetry/sqlite_span_exporter.py @@ -29,6 +29,7 @@ from opentelemetry.sdk.trace import ReadableSpan from opentelemetry.sdk.trace.export import SpanExporter from opentelemetry.sdk.trace.export import SpanExportResult +from opentelemetry.trace import Link from opentelemetry.trace import SpanContext from opentelemetry.trace import TraceFlags from opentelemetry.trace import TraceState @@ -46,7 +47,8 @@ end_time_unix_nano INTEGER, session_id TEXT, invocation_id TEXT, - attributes_json TEXT + attributes_json TEXT, + links_json TEXT NOT NULL DEFAULT '[]' ); """ @@ -68,8 +70,9 @@ end_time_unix_nano, session_id, invocation_id, - attributes_json -) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?); + attributes_json, + links_json +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?); """ _DEFAULT_TIMEOUT_SECONDS = 30.0 @@ -78,8 +81,9 @@ class SqliteSpanExporter(SpanExporter): """Exports spans to a local SQLite database. - This is intended for local development (e.g. `adk web`) to allow reloading - traces for older sessions after process restart. + This is intended for local development, allowing traces to be reloaded + after process restart. Existing databases are upgraded in place; spans + written by older versions have no links. """ def __init__(self, *, db_path: str): @@ -101,10 +105,20 @@ def _get_connection(self) -> sqlite3.Connection: def _ensure_schema(self) -> None: with self._lock: conn = self._get_connection() - conn.execute(_CREATE_SPANS_TABLE) - conn.execute(_CREATE_SESSION_INDEX) - conn.execute(_CREATE_TRACE_INDEX) - conn.commit() + with conn: + # Serialize the schema check with other exporter instances. + conn.execute("BEGIN IMMEDIATE") + conn.execute(_CREATE_SPANS_TABLE) + columns = { + row["name"] for row in conn.execute("PRAGMA table_info(spans)") + } + if "links_json" not in columns: + conn.execute( + "ALTER TABLE spans ADD COLUMN links_json TEXT NOT NULL DEFAULT" + " '[]'" + ) + conn.execute(_CREATE_SESSION_INDEX) + conn.execute(_CREATE_TRACE_INDEX) def _serialize_attributes( self, attributes: Mapping[str, AttributeValue] @@ -133,6 +147,54 @@ def _deserialize_attributes( return {} return cast(dict[str, AttributeValue], decoded) + def _serialize_links(self, links: Sequence[Link]) -> str: + return json.dumps( + [ + { + "trace_id": format(link.context.trace_id, "032x"), + "span_id": format(link.context.span_id, "016x"), + "is_remote": link.context.is_remote, + "trace_flags": int(link.context.trace_flags), + "trace_state": link.context.trace_state.to_header(), + "attributes": ( + dict(link.attributes) + if link.attributes is not None + else None + ), + } + for link in links + ], + ensure_ascii=False, + default=lambda o: "", + ) + + def _deserialize_links(self, links_json: object) -> list[Link]: + if not isinstance(links_json, (str, bytes, bytearray)): + return [] + try: + decoded: object = json.loads(links_json) + except (json.JSONDecodeError, TypeError) as e: + logger.debug("Failed to deserialize span links: %r", e) + return [] + if not isinstance(decoded, list): + return [] + links: list[Link] = [] + for item in decoded: + if not isinstance(item, dict): + continue + try: + context = SpanContext( + trace_id=int(item["trace_id"], 16), + span_id=int(item["span_id"], 16), + is_remote=item["is_remote"], + trace_flags=TraceFlags(item["trace_flags"]), + trace_state=TraceState.from_header([item["trace_state"]]), + ) + links.append(Link(context, attributes=item["attributes"])) + except (KeyError, TypeError, ValueError) as e: + logger.debug("Failed to deserialize span link: %r", e) + return links + def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: try: with self._lock: @@ -167,6 +229,7 @@ def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: session_id, invocation_id, self._serialize_attributes(attributes), + self._serialize_links(span.links), )) conn.executemany(_INSERT_SPAN, rows) conn.commit() @@ -222,6 +285,7 @@ def _row_to_readable_span(self, row: sqlite3.Row) -> ReadableSpan: context=context, parent=parent, attributes=attributes, + links=self._deserialize_links(row["links_json"]), start_time=row["start_time_unix_nano"], end_time=row["end_time_unix_nano"], ) diff --git a/tests/unittests/telemetry/test_sqlite_span_exporter.py b/tests/unittests/telemetry/test_sqlite_span_exporter.py index 214371755af..430872c4f35 100644 --- a/tests/unittests/telemetry/test_sqlite_span_exporter.py +++ b/tests/unittests/telemetry/test_sqlite_span_exporter.py @@ -14,15 +14,22 @@ from __future__ import annotations +from concurrent.futures import ThreadPoolExecutor import json from pathlib import Path +import sqlite3 +import threading from google.adk.telemetry.sqlite_span_exporter import SqliteSpanExporter from opentelemetry.sdk.trace import ReadableSpan +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export import SpanExportResult +from opentelemetry.trace import Link from opentelemetry.trace import SpanContext from opentelemetry.trace import TraceFlags from opentelemetry.trace import TraceState +import pytest def _create_span( @@ -34,6 +41,7 @@ def _create_span( attributes: dict | None = None, start_time: int = 1000, end_time: int = 2000, + links: list[Link] | None = None, ) -> ReadableSpan: """Helper to create ReadableSpan instances for testing.""" context = SpanContext( @@ -61,6 +69,7 @@ def _create_span( attributes=attributes or {}, start_time=start_time, end_time=end_time, + links=links or [], ) @@ -460,3 +469,168 @@ def test_get_spans_ordered_by_start_time(tmp_path): assert result[0].context.span_id == 0x100 assert result[1].context.span_id == 0x200 assert result[2].context.span_id == 0x300 + + +@pytest.mark.parametrize("is_remote,trace_flags", [(False, 0), (True, 1)]) +def test_links_survive_exporter_restart(tmp_path, is_remote, trace_flags): + db_path = str(tmp_path / "spans.db") + linked_context = SpanContext( + trace_id=0x123456789ABCDEF0123456789ABCDEF0, + span_id=0x123456789ABCDEF0, + is_remote=is_remote, + trace_flags=TraceFlags(trace_flags), + trace_state=TraceState([("vendor", "value"), ("other", "state")]), + ) + other_context = SpanContext( + trace_id=0xABC, + span_id=0xDEF, + is_remote=False, + trace_flags=TraceFlags(0), + ) + links = [ + Link(linked_context, {"source": "café", "attempts": 2, "ids": [1, 2]}), + Link(other_context), + ] + provider = TracerProvider() + provider.add_span_processor( + SimpleSpanProcessor(SqliteSpanExporter(db_path=db_path)) + ) + with provider.get_tracer(__name__).start_as_current_span( + "linked-operation", + attributes={"gcp.vertex.agent.session_id": "linked-session"}, + links=links, + ): + pass + provider.shutdown() + + reopened = SqliteSpanExporter(db_path=db_path) + try: + spans = reopened.get_all_spans_for_session("linked-session") + finally: + reopened.shutdown() + + assert len(spans) == 1 + assert len(spans[0].links) == 2 + restored, other = spans[0].links + assert restored.context == linked_context + assert restored.attributes["source"] == "café" + assert restored.attributes["attempts"] == 2 + assert list(restored.attributes["ids"]) == [1, 2] + assert other.context == other_context + assert not other.attributes + + +@pytest.mark.parametrize("initializers", [1, 2]) +@pytest.mark.parametrize( + "extra_columns", + [ + "", + ", status_code TEXT, status_description TEXT, events_json TEXT", + ", kind TEXT NOT NULL DEFAULT 'INTERNAL'", + ( + ", status_code TEXT, status_description TEXT, events_json TEXT," + " kind TEXT NOT NULL DEFAULT 'INTERNAL'" + ), + ], +) +def test_legacy_database_retains_spans_and_accepts_links( + tmp_path, initializers, extra_columns +): + """Opening an old database preserves its rows, including concurrent opens.""" + db_path = str(tmp_path / "spans.db") + with sqlite3.connect(db_path) as conn: + conn.execute( + "CREATE TABLE spans (span_id TEXT PRIMARY KEY, trace_id TEXT NOT NULL," + " parent_span_id TEXT, name TEXT NOT NULL, start_time_unix_nano" + " INTEGER, end_time_unix_nano INTEGER, session_id TEXT, invocation_id" + f" TEXT, attributes_json TEXT{extra_columns})" + ) + conn.execute( + "INSERT INTO spans (span_id, trace_id, name, session_id)" + " VALUES ('111', '222', 'legacy', 'session')" + ) + barrier = threading.Barrier(initializers) + + def open_database(_: int) -> None: + barrier.wait(timeout=10) + exporter = SqliteSpanExporter(db_path=db_path) + exporter.shutdown() + + with ThreadPoolExecutor(max_workers=initializers) as pool: + list(pool.map(open_database, range(initializers))) + + link = Link( + SpanContext( + trace_id=0x333, + span_id=0x444, + is_remote=True, + trace_flags=TraceFlags(1), + ) + ) + exporter = SqliteSpanExporter(db_path=db_path) + try: + result = exporter.export([ + _create_span( + name="new", + attributes={"gcp.vertex.agent.session_id": "session"}, + links=[link], + ) + ]) + with sqlite3.connect(db_path) as conn: + conn.execute( + "INSERT INTO spans (span_id, trace_id, name, session_id)" + " VALUES ('333', '222', 'old-writer', 'session')" + ) + spans = exporter.get_all_spans_for_session("session") + finally: + exporter.shutdown() + + assert result == SpanExportResult.SUCCESS + by_name = {span.name: span for span in spans} + assert set(by_name) == {"legacy", "new", "old-writer"} + assert not by_name["old-writer"].links + assert not by_name["legacy"].links + assert len(by_name["new"].links) == 1 + assert by_name["new"].links[0].context == link.context + + +@pytest.mark.parametrize("links", [None, []]) +def test_span_without_links_survives_restart(tmp_path, links): + db_path = str(tmp_path / "spans.db") + exporter = SqliteSpanExporter(db_path=db_path) + result = exporter.export([ + _create_span( + attributes={"gcp.vertex.agent.session_id": "session"}, links=links + ) + ]) + exporter.shutdown() + + reopened = SqliteSpanExporter(db_path=db_path) + try: + spans = reopened.get_all_spans_for_session("session") + finally: + reopened.shutdown() + + assert result == SpanExportResult.SUCCESS + assert len(spans) == 1 + assert not spans[0].links + + +@pytest.mark.parametrize("links_json", ["not json", "{}", "null", "[null, {}]"]) +def test_invalid_links_do_not_prevent_loading_span(tmp_path, links_json): + db_path = str(tmp_path / "spans.db") + exporter = SqliteSpanExporter(db_path=db_path) + try: + with sqlite3.connect(db_path) as conn: + conn.execute( + "INSERT INTO spans (span_id, trace_id, name, session_id, links_json)" + " VALUES ('111', '222', 'readable', 'session', ?)", + (links_json,), + ) + spans = exporter.get_all_spans_for_session("session") + finally: + exporter.shutdown() + + assert len(spans) == 1 + assert spans[0].name == "readable" + assert not spans[0].links From 9ee54e5edcab07d9816563c2e250ed99462356ee Mon Sep 17 00:00:00 2001 From: samecho Date: Tue, 27 Oct 2026 03:22:51 -0400 Subject: [PATCH 2/2] fix(telemetry): keep saved traces readable during writes Read schema metadata before acquiring SQLite's write reservation. Only take BEGIN IMMEDIATE when a schema update is needed, keeping the column recheck under that transaction for concurrent migrations. Related: #7484 --- .../adk/telemetry/sqlite_span_exporter.py | 10 +++ .../telemetry/test_sqlite_span_exporter.py | 79 +++++++++++++++++++ 2 files changed, 89 insertions(+) diff --git a/src/google/adk/telemetry/sqlite_span_exporter.py b/src/google/adk/telemetry/sqlite_span_exporter.py index da3d069430b..392b4c88bbf 100644 --- a/src/google/adk/telemetry/sqlite_span_exporter.py +++ b/src/google/adk/telemetry/sqlite_span_exporter.py @@ -105,6 +105,16 @@ def _get_connection(self) -> sqlite3.Connection: def _ensure_schema(self) -> None: with self._lock: conn = self._get_connection() + # Readers of a current schema must not wait for unrelated writers. + columns = { + row["name"] for row in conn.execute("PRAGMA table_info(spans)") + } + indexes = { + row["name"] for row in conn.execute("PRAGMA index_list(spans)") + } + required_indexes = {"spans_session_id_idx", "spans_trace_id_idx"} + if "links_json" in columns and required_indexes <= indexes: + return with conn: # Serialize the schema check with other exporter instances. conn.execute("BEGIN IMMEDIATE") diff --git a/tests/unittests/telemetry/test_sqlite_span_exporter.py b/tests/unittests/telemetry/test_sqlite_span_exporter.py index 430872c4f35..b9d0373d36b 100644 --- a/tests/unittests/telemetry/test_sqlite_span_exporter.py +++ b/tests/unittests/telemetry/test_sqlite_span_exporter.py @@ -15,6 +15,7 @@ from __future__ import annotations from concurrent.futures import ThreadPoolExecutor +from contextlib import closing import json from pathlib import Path import sqlite3 @@ -634,3 +635,81 @@ def test_invalid_links_do_not_prevent_loading_span(tmp_path, links_json): assert len(spans) == 1 assert spans[0].name == "readable" assert not spans[0].links + + +def test_current_database_readable_during_wal_write(tmp_path: Path) -> None: + """A saved trace remains readable while another connection is writing.""" + db_path = str(tmp_path / "spans.db") + link = Link( + SpanContext( + trace_id=0x123, + span_id=0x456, + is_remote=True, + trace_flags=TraceFlags(1), + ) + ) + exporter = SqliteSpanExporter(db_path=db_path) + try: + assert ( + exporter.export([ + _create_span( + name="saved-span", + attributes={"gcp.vertex.agent.session_id": "archived-session"}, + links=[link], + ) + ]) + == SpanExportResult.SUCCESS + ) + finally: + exporter.shutdown() + + with closing(sqlite3.connect(db_path)) as writer: + assert writer.execute("PRAGMA journal_mode=WAL").fetchone()[0] == "wal" + writer.execute("BEGIN IMMEDIATE") + writer.execute("UPDATE spans SET name = 'uncommitted'") + reader = None + try: + reader = SqliteSpanExporter(db_path=db_path) + spans = reader.get_all_spans_for_session("archived-session") + finally: + if reader is not None: + reader.shutdown() + writer.rollback() + + assert len(spans) == 1 + assert spans[0].name == "saved-span" + assert len(spans[0].links) == 1 + assert spans[0].links[0].context == link.context + + +@pytest.mark.parametrize( + "index_name", ["spans_session_id_idx", "spans_trace_id_idx"] +) +def test_reopening_database_restores_missing_index( + tmp_path: Path, index_name: str +) -> None: + db_path = str(tmp_path / "spans.db") + exporter = SqliteSpanExporter(db_path=db_path) + try: + assert ( + exporter.export([ + _create_span(attributes={"gcp.vertex.agent.session_id": "session"}) + ]) + == SpanExportResult.SUCCESS + ) + finally: + exporter.shutdown() + with closing(sqlite3.connect(db_path)) as conn: + conn.execute(f"DROP INDEX {index_name}") + conn.commit() + + reopened = SqliteSpanExporter(db_path=db_path) + try: + spans = reopened.get_all_spans_for_session("session") + with closing(sqlite3.connect(db_path)) as conn: + indexes = {row[1] for row in conn.execute("PRAGMA index_list(spans)")} + finally: + reopened.shutdown() + + assert len(spans) == 1 + assert index_name in indexes