Skip to content

Commit cd6dee0

Browse files
DeanChensjcopybara-github
DeanChensj
authored and
copybara-github
committed
fix(sessions): avoid runtime index drop and guard concurrent index creation
Previously, `_ensure_schema_indexes_exist` automatically dropped superseded indexes (such as `idx_events_app_user_session_ts`) during `DatabaseSessionService.prepare_tables()`. On existing deployments with large or actively queried `events` tables, dropping an index at runtime can acquire an `ACCESS EXCLUSIVE` lock in PostgreSQL, break zero-downtime rolling updates, and trigger cross-container TOCTOU race failures when multiple service instances start simultaneously. This change: 1. Removes automatic runtime `DROP INDEX` of superseded indexes so existing deployments retain backward-compatible indexes without exclusive table locks or rolling-update race conditions (operators can drop obsolete indexes out-of-band via `DROP INDEX CONCURRENTLY` if desired). 2. Wraps `index.create(bind=connection, checkfirst=True)` in a savepoint (`connection.begin_nested()`) and catches `(OperationalError, ProgrammingError)` when another container instance concurrently creates the same index between the `checkfirst` check and DDL execution. 3. Documents out-of-band `CREATE INDEX CONCURRENTLY` preparation in `prepare_tables()` and `_ensure_schema_indexes_exist()`. Co-authored-by: Shangjie Chen <[email protected]> PiperOrigin-RevId: 986182210
1 parent b1cba29 commit cd6dee0

2 files changed

Lines changed: 104 additions & 22 deletions

File tree

‎src/google/adk/sessions/database_session_service.py‎

Lines changed: 31 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -32,20 +32,18 @@
3232
from google.adk.platform import uuid as platform_uuid
3333

3434
try:
35-
from sqlalchemy import Column
3635
from sqlalchemy import delete
3736
from sqlalchemy import event
38-
from sqlalchemy import Index
3937
from sqlalchemy import inspect
4038
from sqlalchemy import MetaData
4139
from sqlalchemy import select
42-
from sqlalchemy import String
43-
from sqlalchemy import Table
4440
from sqlalchemy.engine import Connection
4541
from sqlalchemy.engine import make_url
4642
from sqlalchemy.exc import ArgumentError
4743
from sqlalchemy.exc import IntegrityError
4844
from sqlalchemy.exc import InvalidRequestError
45+
from sqlalchemy.exc import OperationalError
46+
from sqlalchemy.exc import ProgrammingError
4947
from sqlalchemy.ext.asyncio import async_sessionmaker
5048
from sqlalchemy.ext.asyncio import AsyncEngine
5149
from sqlalchemy.ext.asyncio import AsyncSession as DatabaseSessionFactory
@@ -237,29 +235,34 @@ def _set_sqlite_pragma(
237235
cursor.close()
238236

239237

240-
_SUPERSEDED_INDEX_NAMES = frozenset({"idx_events_app_user_session_ts"})
241-
242-
243238
def _ensure_schema_indexes_exist(
244239
connection: Connection, metadata: MetaData
245240
) -> None:
246-
"""Ensures indexes declared in metadata exist for existing tables."""
241+
"""Ensures indexes declared in metadata exist for existing tables.
242+
243+
Note:
244+
Superseded indexes (such as ``idx_events_app_user_session_ts``) are
245+
intentionally not dropped at runtime. Dropping an index requires an
246+
``ACCESS EXCLUSIVE`` table lock in PostgreSQL, can race across multiple
247+
service instances starting concurrently, and breaks zero-downtime rolling
248+
updates. Operators of large existing deployments may pre-create new indexes
249+
(e.g. via ``CREATE INDEX CONCURRENTLY``) and drop obsolete indexes
250+
out-of-band during a maintenance window.
251+
"""
247252
logger.debug("Ensuring schema indexes exist for metadata tables.")
248253
for table in metadata.sorted_tables:
249254
for index in sorted(table.indexes, key=lambda item: item.name or ""):
250-
index.create(bind=connection, checkfirst=True)
251-
252-
# Drop obsolete indexes that have been superseded by composite indexes.
253-
inspector = inspect(connection)
254-
if inspector.has_table("events"):
255-
existing_indexes = {idx["name"] for idx in inspector.get_indexes("events")}
256-
for superseded in _SUPERSEDED_INDEX_NAMES:
257-
if superseded in existing_indexes:
258-
logger.info(
259-
"Dropping superseded index %s from events table.", superseded
260-
)
261-
isolated_table = Table("events", MetaData(), Column("id", String))
262-
Index(superseded, isolated_table.c.id).drop(bind=connection)
255+
try:
256+
with connection.begin_nested():
257+
index.create(bind=connection, checkfirst=True)
258+
except (OperationalError, ProgrammingError):
259+
# Another container instance may have created the index concurrently
260+
# between the `checkfirst` inspection and the DDL execution.
261+
existing_indexes = {
262+
idx["name"] for idx in inspect(connection).get_indexes(table.name)
263+
}
264+
if index.name not in existing_indexes:
265+
raise
263266

264267

265268
def _setup_database_schema(connection: Connection, metadata: MetaData) -> None:
@@ -523,6 +526,13 @@ async def prepare_tables(self) -> None:
523526
table-creation cost upfront (e.g. during application startup) instead of
524527
on the first database operation. It is safe to call more than once and
525528
is recommended for latency-sensitive applications.
529+
530+
For existing deployments with large ``events`` tables, operators are
531+
encouraged to pre-create new schema indexes out-of-band (for example,
532+
``CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_events_app_user_session_ts_id
533+
ON events (app_name, user_id, session_id, timestamp DESC, id DESC)``)
534+
before rolling out a new version so that index verification during
535+
``prepare_tables()`` is a fast no-op.
526536
"""
527537
# Early return if tables are already created
528538
if self._tables_created:

‎tests/unittests/sessions/migration/test_database_schema.py‎

Lines changed: 73 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
from google.adk.sessions import database_session_service
1516
from google.adk.sessions.database_session_service import DatabaseSessionService
1617
from google.adk.sessions.migration import _schema_check_utils
1718
from google.adk.sessions.schemas import v0
@@ -23,8 +24,10 @@
2324
from sqlalchemy.dialects import mysql
2425
from sqlalchemy.dialects import postgresql
2526
from sqlalchemy.dialects import sqlite
27+
from sqlalchemy.exc import OperationalError
2628
from sqlalchemy.ext.asyncio import create_async_engine
2729
from sqlalchemy.schema import CreateIndex
30+
from sqlalchemy.schema import Index
2831

2932

3033
async def create_v0_db(db_path):
@@ -302,7 +305,7 @@ async def test_prepare_tables_adds_new_events_index_to_existing_db(tmp_path):
302305
await engine.dispose()
303306

304307
index_names = {idx['name'] for idx in event_indexes}
305-
assert 'idx_events_app_user_session_ts' not in index_names
308+
assert 'idx_events_app_user_session_ts' in index_names
306309
assert 'idx_events_app_user_session_ts_id' in index_names
307310
composite_idx = next(
308311
idx
@@ -479,3 +482,72 @@ def test_get_db_schema_version_from_connection_uses_open_connection(tmp_path):
479482
engine.dispose()
480483

481484
assert version == _schema_check_utils.SCHEMA_VERSION_0_PICKLE
485+
486+
487+
def test_ensure_schema_indexes_exist_tolerates_concurrent_index_creation(
488+
tmp_path, monkeypatch
489+
):
490+
"""If another container creates the index between checkfirst and DDL, it succeeds."""
491+
db_path = tmp_path / 'concurrent_index.db'
492+
engine = create_engine(f'sqlite:///{db_path}')
493+
try:
494+
with engine.begin() as connection:
495+
v1.Base.metadata.create_all(bind=connection)
496+
497+
original_create = Index.create
498+
raised_once = False
499+
500+
def racing_create(self, bind=None, checkfirst=False):
501+
nonlocal raised_once
502+
if not raised_once and self.name == 'idx_events_app_user_session_ts_id':
503+
raised_once = True
504+
assert bind is not None and bind.in_nested_transaction()
505+
return original_create(self, bind=bind, checkfirst=False)
506+
return original_create(self, bind=bind, checkfirst=checkfirst)
507+
508+
monkeypatch.setattr(Index, 'create', racing_create)
509+
510+
with engine.begin() as connection:
511+
database_session_service._ensure_schema_indexes_exist(
512+
connection, v1.Base.metadata
513+
)
514+
assert raised_once
515+
finally:
516+
engine.dispose()
517+
518+
519+
def test_ensure_schema_indexes_exist_reraises_when_index_missing(
520+
tmp_path, monkeypatch
521+
):
522+
"""If index DDL fails and the index is still missing, the error is re-raised."""
523+
db_path = tmp_path / 'failed_index.db'
524+
engine = create_engine(f'sqlite:///{db_path}')
525+
try:
526+
with engine.begin() as connection:
527+
v1.Base.metadata.create_all(bind=connection)
528+
connection.execute(text('DROP INDEX idx_events_app_user_session_ts_id'))
529+
530+
original_create = Index.create
531+
532+
def failing_create(self, bind=None, checkfirst=False):
533+
if self.name == 'idx_events_app_user_session_ts_id':
534+
assert bind is not None and bind.in_nested_transaction()
535+
bind.execute(
536+
text(
537+
'CREATE INDEX idx_events_app_user_session_ts_id '
538+
'ON events (nonexistent_column)'
539+
)
540+
)
541+
return original_create(self, bind=bind, checkfirst=checkfirst)
542+
543+
monkeypatch.setattr(Index, 'create', failing_create)
544+
545+
with (
546+
pytest.raises(OperationalError, match='nonexistent_column'),
547+
engine.begin() as connection,
548+
):
549+
database_session_service._ensure_schema_indexes_exist(
550+
connection, v1.Base.metadata
551+
)
552+
finally:
553+
engine.dispose()

0 commit comments

Comments
 (0)