feat: add declarative database tool templates#106
Conversation
- Add safety validation layer for identifiers and column types - Implement SqliteDeclarativeDbStorage storage layer for templates - Integrate declarative db storage with AppContainer - Implement ADK plugin to inject custom schema XML in before_model - Register declarative db tools in registry.py and app agent - Write extensive unit and integration tests in tests/declarative_db/
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a declarative database feature, allowing users to dynamically define custom tables, query templates, and instruction overrides. It includes a new ADK plugin, an SQLite storage layer, validation logic, agent-facing tools, and comprehensive tests. The review feedback highlights several important improvements for the storage layer: preventing metadata desynchronization when recreating existing tables, allowing updates to query templates even when the template limit is reached, safely handling boolean and complex default values in DDL construction, and using async with blocks for database cursors during UPDATE and DELETE operations to avoid resource leaks.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| # Check table counts guardrail | ||
| existing_count_row = await self._fetch_one( | ||
| "SELECT COUNT(*) as count FROM custom_tables WHERE user_id = ?", | ||
| (user_id,), | ||
| ) | ||
| existing_count = existing_count_row["count"] if existing_count_row else 0 | ||
| if existing_count >= 5: | ||
| raise ValueError("Limit of 5 custom tables per user reached") |
There was a problem hiding this comment.
If create_custom_table is called for an existing table, CREATE TABLE IF NOT EXISTS does not modify the physical SQLite table. However, the metadata in custom_table_columns is completely deleted and replaced with the new column definitions. This leads to a critical desynchronization where the metadata columns do not match the physical table columns, causing subsequent queries to fail.
To prevent this, check if the table already exists and raise an error, forcing the user to explicitly delete the table first if they want to modify its schema.
# Check if table already exists to prevent desynchronization
exists_row = await self._fetch_one(
"SELECT 1 FROM custom_tables WHERE user_id = ? AND table_name = ?",
(user_id, table_name),
)
if exists_row:
raise ValueError(
f"Table '{table_name}' already exists. "
"You must delete it first to modify its schema."
)
# Check table counts guardrail
existing_count_row = await self._fetch_one(
"SELECT COUNT(*) as count FROM custom_tables WHERE user_id = ?",
(user_id,),
)
existing_count = existing_count_row["count"] if existing_count_row else 0
if existing_count >= 5:
raise ValueError("Limit of 5 custom tables per user reached")| # Check templates limit (10 per user) | ||
| existing_count_row = await self._fetch_one( | ||
| "SELECT COUNT(*) as count FROM saved_query_templates WHERE user_id = ?", | ||
| (user_id,), | ||
| ) | ||
| existing_count = existing_count_row["count"] if existing_count_row else 0 | ||
| if existing_count >= 10: | ||
| raise ValueError("Limit of 10 saved query templates per user reached") |
There was a problem hiding this comment.
The template limit check (existing_count >= 10) is enforced even when updating/overwriting an existing template. Once a user reaches the limit of 10 templates, they will be blocked from updating any of their existing templates because the count remains at 10.
To fix this, check if the template already exists first, and only enforce the limit for new templates.
# Check if template already exists to allow updates without hitting the limit
exists_row = await self._fetch_one(
"SELECT 1 FROM saved_query_templates WHERE user_id = ? AND template_name = ?",
(user_id, template_name),
)
if not exists_row:
existing_count_row = await self._fetch_one(
"SELECT COUNT(*) as count FROM saved_query_templates WHERE user_id = ?",
(user_id,),
)
existing_count = existing_count_row["count"] if existing_count_row else 0
if existing_count >= 10:
raise ValueError("Limit of 10 saved query templates per user reached")| if default_val is not None: | ||
| # Basic escaping of default string values if string-like | ||
| if isinstance(default_val, str): | ||
| escaped_default = default_val.replace("'", "''") | ||
| part += f" DEFAULT '{escaped_default}'" | ||
| else: | ||
| part += f" DEFAULT {default_val}" |
There was a problem hiding this comment.
When constructing the column DDL, default_val is appended directly to the SQL string if it is not a string. If default_val is a boolean, Python's str(True) produces True, resulting in DEFAULT True. While newer SQLite versions support boolean literals, older versions do not and will fail. Furthermore, if default_val is a complex type like a list or dict, it will cause a SQL syntax error.
To ensure compatibility and safety, explicitly handle booleans (converting them to 1 or 0) and restrict the allowed default value types to scalar types (str, int, float, bool).
| if default_val is not None: | |
| # Basic escaping of default string values if string-like | |
| if isinstance(default_val, str): | |
| escaped_default = default_val.replace("'", "''") | |
| part += f" DEFAULT '{escaped_default}'" | |
| else: | |
| part += f" DEFAULT {default_val}" | |
| if default_val is not None: | |
| if isinstance(default_val, bool): | |
| part += f" DEFAULT {1 if default_val else 0}" | |
| elif isinstance(default_val, str): | |
| escaped_default = default_val.replace("'", "''") | |
| part += f" DEFAULT '{escaped_default}'" | |
| elif isinstance(default_val, (int, float)): | |
| part += f" DEFAULT {default_val}" | |
| else: | |
| raise ValueError(f"Unsupported default value type: {type(default_val).__name__}") |
| # Run write update | ||
| async with self._lock: | ||
| cursor = await self._conn.execute(sql, bindings) | ||
| return cursor.rowcount |
There was a problem hiding this comment.
In execute_query_template, the UPDATE query execution does not use an async with block for the cursor, unlike the SELECT and INSERT paths. This can leave cursors unclosed, potentially leaking database resources.
Use async with self._conn.execute(...) as cursor: to ensure proper cleanup.
| # Run write update | |
| async with self._lock: | |
| cursor = await self._conn.execute(sql, bindings) | |
| return cursor.rowcount | |
| # Run write update | |
| async with self._lock: | |
| async with self._conn.execute(sql, bindings) as cursor: | |
| return cursor.rowcount |
| # Run write delete | ||
| async with self._lock: | ||
| cursor = await self._conn.execute(sql, bindings) | ||
| return cursor.rowcount |
There was a problem hiding this comment.
In execute_query_template, the DELETE query execution does not use an async with block for the cursor, unlike the SELECT and INSERT paths. This can leave cursors unclosed, potentially leaking database resources.
Use async with self._conn.execute(...) as cursor: to ensure proper cleanup.
| # Run write delete | |
| async with self._lock: | |
| cursor = await self._conn.execute(sql, bindings) | |
| return cursor.rowcount | |
| # Run write delete | |
| async with self._lock: | |
| async with self._conn.execute(sql, bindings) as cursor: | |
| return cursor.rowcount |
- Prevent metadata desync when recreating existing tables - Allow query template updates when limit is reached - Format boolean and complex defaults safely in DDL - Prevent UPDATE/DELETE template cursor resource leaks - Achieve 100% test coverage with new unit and integration tests
1 participant
What
Add a safe, sandboxed declarative database capability layer and custom instructions system for Blacki.
Why
This allows the AI agent to define lightweight custom schemas, query templates, and instruction overrides dynamically and securely per user/chat session without generating arbitrary executable python code, modifying core prompts, or risking SQL injection.
How
SqliteDeclarativeDbStorage(SqlStorage)insrc/blacki/declarative_db/storage.pywith multi-user isolation (user-id scoped metadata and md5-hashed physical table prefixing), safe dynamic parameterized query compilation, and transactional schema safety.DeclarativeDbPlugin(BasePlugin)insrc/blacki/declarative_db/plugin.pyto compile active schemas/templates into XML and append them to system instructions inbefore_model_callbackwithout altering conversation turns or system prompt code.tests/declarative_db/test_declarative_db.py.Tests
uv run ruff check,uv run ruff format,uv run mypy ., anduv run pytest(875/875 passing tests)Related Issues
Closes #103