diff --git a/src/google/adk/telemetry/sqlite_span_exporter.py b/src/google/adk/telemetry/sqlite_span_exporter.py index eae4c11bd15..392b4c88bbf 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,30 @@ 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() + # 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") + 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 +157,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 +239,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 +295,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..b9d0373d36b 100644 --- a/tests/unittests/telemetry/test_sqlite_span_exporter.py +++ b/tests/unittests/telemetry/test_sqlite_span_exporter.py @@ -14,15 +14,23 @@ from __future__ import annotations +from concurrent.futures import ThreadPoolExecutor +from contextlib import closing 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 +42,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 +70,7 @@ def _create_span( attributes=attributes or {}, start_time=start_time, end_time=end_time, + links=links or [], ) @@ -460,3 +470,246 @@ 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 + + +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