diff --git a/src/google/adk/sessions/migration/_schema_check_utils.py b/src/google/adk/sessions/migration/_schema_check_utils.py index 284f82afe7..249161c84c 100644 --- a/src/google/adk/sessions/migration/_schema_check_utils.py +++ b/src/google/adk/sessions/migration/_schema_check_utils.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. """Database schema version check utility.""" + from __future__ import annotations import logging @@ -32,8 +33,11 @@ def _get_schema_version_impl(inspector, connection) -> str: """Gets DB schema version using inspector and connection.""" if inspector.has_table("adk_internal_metadata"): try: + key_col = inspector.dialect.identifier_preparer.quote("key") result = connection.execute( - text("SELECT value FROM adk_internal_metadata WHERE key = :key"), + text( + f"SELECT value FROM adk_internal_metadata WHERE {key_col} = :key" + ), {"key": SCHEMA_VERSION_KEY}, ).fetchone() if result: @@ -49,6 +53,7 @@ def _get_schema_version_impl(inspector, connection) -> str: e, ) raise + # Metadata table doesn't exist, check for v0 schema. # V0 schema has an 'events' table with an 'actions' column. if inspector.has_table("events"): diff --git a/tests/unittests/sessions/migration/test_database_schema.py b/tests/unittests/sessions/migration/test_database_schema.py index 239da2f1e2..d08bb97ba0 100644 --- a/tests/unittests/sessions/migration/test_database_schema.py +++ b/tests/unittests/sessions/migration/test_database_schema.py @@ -51,12 +51,18 @@ async def test_new_db_uses_latest_schema(tmp_path): lambda sync_conn: inspect(sync_conn).has_table('adk_internal_metadata') ) assert has_metadata_table - schema_version = await conn.run_sync( - lambda sync_conn: sync_conn.execute( - text('SELECT value FROM adk_internal_metadata WHERE key = :key'), - {'key': _schema_check_utils.SCHEMA_VERSION_KEY}, - ).scalar_one_or_none() - ) + + def get_schema_version(sync_conn): + inspector = inspect(sync_conn) + key_col = inspector.dialect.identifier_preparer.quote('key') + return sync_conn.execute( + text( + f'SELECT value FROM adk_internal_metadata WHERE {key_col} = :key' + ), + {'key': _schema_check_utils.SCHEMA_VERSION_KEY}, + ).scalar_one_or_none() + + schema_version = await conn.run_sync(get_schema_version) assert schema_version == _schema_check_utils.LATEST_SCHEMA_VERSION # Verify events table columns for v1