Skip to content

Commit 4ad3ecc

Browse files
herdiyana256GWeale
authored andcommitted
fix: redact database password from session service errors and logs
Merge #6485 Co-authored-by: George Weale <[email protected]> PiperOrigin-RevId: 960522049
1 parent 9488408 commit 4ad3ecc

7 files changed

Lines changed: 205 additions & 10 deletions

File tree

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -352,16 +352,17 @@ def __init__(
352352
event.listen(db_engine.sync_engine, "connect", _set_sqlite_pragma)
353353

354354
except Exception as e:
355+
redacted_url = _schema_check_utils._redact_db_url(db_url)
355356
if isinstance(e, ArgumentError):
356357
raise ValueError(
357-
f"Invalid database URL format or argument '{db_url}'."
358+
f"Invalid database URL format or argument '{redacted_url}'."
358359
) from e
359360
if isinstance(e, ImportError):
360361
raise ValueError(
361-
f"Database related module not found for URL '{db_url}'."
362+
f"Database related module not found for URL '{redacted_url}'."
362363
) from e
363364
raise ValueError(
364-
f"Failed to create database engine for URL '{db_url}'"
365+
f"Failed to create database engine for URL '{redacted_url}'"
365366
) from e
366367
else:
367368
self._owns_db_engine = False

‎src/google/adk/sessions/migration/_schema_check_utils.py‎

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
from sqlalchemy import create_engine as create_sync_engine
2323
from sqlalchemy import inspect
2424
from sqlalchemy import text
25+
from sqlalchemy.engine import make_url
2526
except ImportError:
2627
pass
2728

@@ -31,6 +32,9 @@
3132

3233
logger = logging.getLogger("google_adk." + __name__)
3334

35+
_UNPARSEABLE_DB_URL = "<unparseable database URL>"
36+
_REDACTED_QUERY_VALUE = "REDACTED"
37+
3438
SCHEMA_VERSION_KEY = "schema_version"
3539
SCHEMA_VERSION_0_PICKLE = "0"
3640
SCHEMA_VERSION_1_JSON = "1"
@@ -125,6 +129,24 @@ def to_sync_url(db_url: str) -> str:
125129
return db_url
126130

127131

132+
def _redact_db_url(db_url: str) -> str:
133+
"""Returns the URL with its credentials masked, for logs and error messages.
134+
135+
A database URL carries the password in the userinfo component, and drivers
136+
also accept secrets as query parameters, so every query value is masked
137+
rather than only the ones with a recognizable name. Redaction happens while
138+
an error is being reported, so it never raises: an unparseable URL yields a
139+
fixed placeholder rather than the original string.
140+
"""
141+
try:
142+
url = make_url(db_url)
143+
if url.query:
144+
url = url.set(query={key: _REDACTED_QUERY_VALUE for key in url.query})
145+
return str(url.render_as_string(hide_password=True))
146+
except Exception: # pylint: disable=broad-except
147+
return _UNPARSEABLE_DB_URL
148+
149+
128150
def get_db_schema_version(db_url: str) -> str:
129151
"""Reads schema version from DB.
130152
@@ -146,7 +168,7 @@ def get_db_schema_version(db_url: str) -> str:
146168
except Exception:
147169
logger.warning(
148170
"Failed to get schema version from database %s.",
149-
db_url,
171+
_redact_db_url(db_url),
150172
)
151173
raise
152174
finally:

‎src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -284,7 +284,10 @@ def migrate(
284284
source_sync_url = _schema_check_utils.to_sync_url(source_db_url)
285285
dest_sync_url = _schema_check_utils.to_sync_url(dest_db_url)
286286

287-
logger.info(f"Connecting to source database: {source_db_url}")
287+
logger.info(
288+
"Connecting to source database: %s",
289+
_schema_check_utils._redact_db_url(source_db_url),
290+
)
288291
if allow_unsafe_unpickling:
289292
logger.warning(
290293
"Unsafe pickle migration mode is enabled. Only use this with a trusted"
@@ -297,7 +300,10 @@ def migrate(
297300
logger.error(f"Failed to connect to source database: {e}")
298301
raise RuntimeError(f"Failed to connect to source database: {e}") from e
299302

300-
logger.info(f"Connecting to destination database: {dest_db_url}")
303+
logger.info(
304+
"Connecting to destination database: %s",
305+
_schema_check_utils._redact_db_url(dest_db_url),
306+
)
301307
try:
302308
dest_engine = create_engine(dest_sync_url)
303309
v1.Base.metadata.create_all(dest_engine)

‎src/google/adk/sessions/migration/migrate_from_sqlalchemy_sqlite.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,10 @@ def migrate(source_db_url: str, dest_db_path: str) -> None:
3838
# them automatically converted to 'sqlite://...' for migration.
3939
source_sync_url = _schema_check_utils.to_sync_url(source_db_url)
4040

41-
logger.info(f"Connecting to source database: {source_db_url}")
41+
logger.info(
42+
"Connecting to source database: %s",
43+
_schema_check_utils._redact_db_url(source_db_url),
44+
)
4245
try:
4346
engine = create_engine(source_sync_url)
4447
v0_schema.Base.metadata.create_all(

‎src/google/adk/sessions/migration/migration_runner.py‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -82,8 +82,9 @@ def upgrade(
8282
current_version = _schema_check_utils.get_db_schema_version(source_db_url)
8383
if current_version == LATEST_VERSION:
8484
logger.info(
85-
f"Database {source_db_url} is already at latest version"
86-
f" {LATEST_VERSION}. No migration needed."
85+
"Database %s is already at latest version %s. No migration needed.",
86+
_schema_check_utils._redact_db_url(source_db_url),
87+
LATEST_VERSION,
8788
)
8889
return
8990

@@ -118,7 +119,10 @@ def upgrade(
118119
logger.debug("Created temp db %s for step %d", out_url, i + 1)
119120

120121
logger.info(
121-
f"Migrating from {in_url} to {out_url} (schema v{end_version})..."
122+
"Migrating from %s to %s (schema v%s)...",
123+
_schema_check_utils._redact_db_url(in_url),
124+
_schema_check_utils._redact_db_url(out_url),
125+
end_version,
122126
)
123127
if migrate_func is migrate_from_sqlalchemy_pickle.migrate:
124128
migrate_func(

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

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,11 @@
1818
import contextlib
1919
from datetime import datetime
2020
from datetime import timezone
21+
import logging
2122
import os
2223
import pickle
2324
import time
25+
from unittest import mock
2426

2527
from fastapi.openapi.models import HTTPBearer
2628
from google.adk.auth.auth_tool import AuthConfig
@@ -29,6 +31,8 @@
2931
from google.adk.events.ui_widget import UiWidget
3032
from google.adk.sessions.migration import _schema_check_utils
3133
from google.adk.sessions.migration import migrate_from_sqlalchemy_pickle as mfsp
34+
from google.adk.sessions.migration import migrate_from_sqlalchemy_sqlite as mfss
35+
from google.adk.sessions.migration import migration_runner
3236
from google.adk.sessions.schemas import v0
3337
from google.adk.sessions.schemas import v1
3438
from google.adk.tools.tool_confirmation import ToolConfirmation
@@ -116,6 +120,118 @@ def test_to_sync_url_empty_string(self):
116120
assert _schema_check_utils.to_sync_url("") == ""
117121

118122

123+
class TestRedactDbUrl:
124+
"""Tests for the _redact_db_url function."""
125+
126+
def test_password_is_masked(self):
127+
redacted = _schema_check_utils._redact_db_url(
128+
"postgresql+asyncpg://user:sup3r-s3cret@host:5432/db"
129+
)
130+
assert redacted == "postgresql+asyncpg://user:***@host:5432/db"
131+
132+
def test_unparseable_url_falls_back_to_placeholder(self):
133+
"""Redaction runs while reporting an error, so it must never raise."""
134+
assert (
135+
_schema_check_utils._redact_db_url("definitely not a url sup3r-s3cret")
136+
== "<unparseable database URL>"
137+
)
138+
139+
def test_query_parameter_values_are_masked(self):
140+
"""Drivers accept secrets as query parameters, so every value is masked."""
141+
redacted = _schema_check_utils._redact_db_url(
142+
"postgresql://user@host:5432/db?password=sup3r-s3cret&sslmode=require"
143+
)
144+
assert redacted == (
145+
"postgresql://user@host:5432/db?password=REDACTED&sslmode=REDACTED"
146+
)
147+
148+
def test_schema_version_failure_warning_hides_password(self, caplog):
149+
db_url = "postgresql+asyncpg://user:sup3r-s3cret@host:5432/db"
150+
151+
with mock.patch.object(
152+
_schema_check_utils,
153+
"create_sync_engine",
154+
side_effect=RuntimeError("boom"),
155+
):
156+
with caplog.at_level(logging.WARNING):
157+
with pytest.raises(RuntimeError):
158+
_schema_check_utils.get_db_schema_version(db_url)
159+
160+
assert "sup3r-s3cret" not in caplog.text
161+
assert "postgresql+asyncpg://user:***@host:5432/db" in caplog.text
162+
163+
164+
_SOURCE_URL = "postgresql+asyncpg://user:sup3r-s3cret@host:5432/src"
165+
_DEST_URL = "postgresql+asyncpg://user:0ther-s3cret@host:5432/dst"
166+
167+
168+
class TestMigrationLogsHidePassword:
169+
"""These entry points log their URLs on every run, not only on failure."""
170+
171+
def test_pickle_migration_connect_logs_are_redacted(self, caplog):
172+
with mock.patch.object(
173+
mfsp,
174+
"create_engine",
175+
side_effect=[mock.MagicMock(), RuntimeError("boom")],
176+
):
177+
with caplog.at_level(logging.INFO):
178+
with pytest.raises(RuntimeError):
179+
mfsp.migrate(_SOURCE_URL, _DEST_URL)
180+
181+
assert "sup3r-s3cret" not in caplog.text
182+
assert "0ther-s3cret" not in caplog.text
183+
assert "postgresql+asyncpg://user:***@host:5432/src" in caplog.text
184+
assert "postgresql+asyncpg://user:***@host:5432/dst" in caplog.text
185+
186+
def test_sqlite_migration_connect_log_is_redacted(self, caplog, tmp_path):
187+
with mock.patch.object(
188+
mfss, "create_engine", side_effect=RuntimeError("boom")
189+
):
190+
with caplog.at_level(logging.INFO):
191+
with pytest.raises(SystemExit):
192+
mfss.migrate(_SOURCE_URL, str(tmp_path / "dest.db"))
193+
194+
assert "sup3r-s3cret" not in caplog.text
195+
assert "postgresql+asyncpg://user:***@host:5432/src" in caplog.text
196+
197+
def test_runner_up_to_date_log_is_redacted(self, caplog):
198+
with mock.patch.object(
199+
_schema_check_utils,
200+
"get_db_schema_version",
201+
return_value=migration_runner.LATEST_VERSION,
202+
):
203+
with caplog.at_level(logging.INFO):
204+
migration_runner.upgrade(_SOURCE_URL, _DEST_URL)
205+
206+
assert "sup3r-s3cret" not in caplog.text
207+
assert "postgresql+asyncpg://user:***@host:5432/src" in caplog.text
208+
209+
def test_runner_migration_step_log_is_redacted(self, caplog):
210+
mock_migrate = mock.Mock()
211+
with mock.patch.object(
212+
_schema_check_utils,
213+
"get_db_schema_version",
214+
return_value=_schema_check_utils.SCHEMA_VERSION_0_PICKLE,
215+
):
216+
with mock.patch.dict(
217+
migration_runner.MIGRATIONS,
218+
{
219+
_schema_check_utils.SCHEMA_VERSION_0_PICKLE: (
220+
_schema_check_utils.SCHEMA_VERSION_1_JSON,
221+
mock_migrate,
222+
)
223+
},
224+
):
225+
with caplog.at_level(logging.INFO):
226+
migration_runner.upgrade(_SOURCE_URL, _DEST_URL)
227+
228+
mock_migrate.assert_called_once_with(_SOURCE_URL, _DEST_URL)
229+
assert "sup3r-s3cret" not in caplog.text
230+
assert "0ther-s3cret" not in caplog.text
231+
assert "postgresql+asyncpg://user:***@host:5432/src" in caplog.text
232+
assert "postgresql+asyncpg://user:***@host:5432/dst" in caplog.text
233+
234+
119235
def test_migrate_from_sqlalchemy_pickle(tmp_path):
120236
"""Tests for migrate_from_sqlalchemy_pickle."""
121237
source_db_path = tmp_path / "source_pickle.db"

‎tests/unittests/sessions/test_session_service.py‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
from sqlalchemy import select
4646
from sqlalchemy import text
4747
from sqlalchemy import update
48+
from sqlalchemy.exc import ArgumentError
4849
from sqlalchemy.ext.asyncio import create_async_engine
4950
from sqlalchemy.pool import StaticPool
5051

@@ -2236,6 +2237,48 @@ async def test_database_session_service_requires_one_argument():
22362237
)
22372238

22382239

2240+
@pytest.mark.parametrize(
2241+
'raised_error',
2242+
[
2243+
RuntimeError('boom'),
2244+
ArgumentError('bad argument'),
2245+
ImportError('no driver'),
2246+
],
2247+
)
2248+
def test_database_session_service_engine_error_hides_password(raised_error):
2249+
"""Engine creation errors must not put the DB password in the message."""
2250+
password = 'sup3r-s3cret'
2251+
db_url = f'postgresql+asyncpg://user:{password}@localhost:5432/db'
2252+
2253+
with mock.patch.object(
2254+
database_session_service,
2255+
'create_async_engine',
2256+
side_effect=raised_error,
2257+
):
2258+
with pytest.raises(ValueError) as exc_info:
2259+
DatabaseSessionService(db_url)
2260+
2261+
message = str(exc_info.value)
2262+
assert password not in message
2263+
# The redacted URL is still there, so the error stays diagnosable.
2264+
assert 'postgresql+asyncpg://user:***@localhost:5432/db' in message
2265+
2266+
2267+
def test_database_session_service_malformed_url_reports_usable_error():
2268+
"""A URL too malformed to parse still yields a usable, leak-free error."""
2269+
# make_url() itself rejects this, so redaction cannot parse it either and
2270+
# must fall back to a placeholder rather than echoing the raw string.
2271+
db_url = 'definitely not a url sup3r-s3cret'
2272+
2273+
with pytest.raises(ValueError) as exc_info:
2274+
DatabaseSessionService(db_url)
2275+
2276+
message = str(exc_info.value)
2277+
assert 'sup3r-s3cret' not in message
2278+
assert 'Invalid database URL format or argument' in message
2279+
assert isinstance(exc_info.value.__cause__, ArgumentError)
2280+
2281+
22392282
@pytest.mark.asyncio
22402283
async def test_database_session_service_sqlite_file_timestamp_read_after_reopen(
22412284
tmp_path,

0 commit comments

Comments
 (0)