Skip to content

Commit a84352e

Browse files
committed
feat: make table sample data in SQL prompts configurable
Add a default-on TABLE_SAMPLE_DATA_ENABLED setting. When disabled, skip automatic sample queries and omit sample values from SQL prompt templates. Preserve schema context and explicit query/preview behavior. Add isolated regression tests and document configuration and limitations. Refs #1291
1 parent 4ee923c commit a84352e

5 files changed

Lines changed: 304 additions & 2 deletions

File tree

backend/README.md

Lines changed: 49 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1,49 @@
1-
# FastAPI Project - Backend
1+
# FastAPI Project - Backend
2+
3+
## Table sample data in SQL generation
4+
5+
`TABLE_SAMPLE_DATA_ENABLED` controls whether SQLBot fetches table sample rows
6+
and includes them in SQL-generation prompts. It defaults to `true` to preserve
7+
existing behavior, including the current three-row sample limit.
8+
9+
To disable automatic sampling, set the following environment variable (or add it
10+
to the project-root `.env` used when running the backend from `backend/`):
11+
12+
```dotenv
13+
TABLE_SAMPLE_DATA_ENABLED=false
14+
```
15+
16+
For Docker Compose, pass it explicitly to the SQLBot service:
17+
18+
```yaml
19+
environment:
20+
TABLE_SAMPLE_DATA_ENABLED: "false"
21+
```
22+
23+
A Compose `.env` file alone does not automatically pass all its variables into a
24+
container. For `docker run`, use `-e TABLE_SAMPLE_DATA_ENABLED=false`.
25+
Settings are loaded at process startup: restart a source deployment, or recreate
26+
the container with the new environment configuration.
27+
28+
When disabled, the automatic sample-data helper returns before reading table
29+
metadata or sample rows. The SQL prompt template also ignores pre-existing
30+
`sample_data` values. Schema context, explicit SQL queries, manual data previews
31+
and their existing permission checks remain unchanged. Less sample context may
32+
affect SQL-generation quality.
33+
34+
This is not a global data-loss-prevention switch. User messages, field comments,
35+
SQL examples, query results and analysis prompts can still contain business
36+
data. Existing logs are not deleted or retroactively redacted.
37+
38+
### Focused regression tests
39+
40+
From the repository root, with the backend development dependencies available:
41+
42+
```bash
43+
python -m pytest -q tests/test_table_sample_data.py
44+
```
45+
46+
The tests execute actual source function definitions with dependency doubles to
47+
avoid initializing database drivers, embedding models or X-Pack. They cover
48+
configuration parsing, sampling, prompt rendering and unaffected query/preview
49+
behavior; they are not full application or model-provider integration tests.

backend/apps/datasource/crud/datasource.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -500,6 +500,9 @@ def get_table_sample_data(ds: CoreDatasource, table_name: str, fields: list) ->
500500
def get_tables_sample_data(session: SessionDep, current_user: CurrentUser, ds: CoreDatasource,
501501
table_list: list[str] = None) -> str:
502502
"""Get sample data (3 rows) for all tables to help AI understand the data"""
503+
if not settings.TABLE_SAMPLE_DATA_ENABLED:
504+
return ""
505+
503506
table_objs = get_table_obj_by_ds(session=session, current_user=current_user, ds=ds)
504507
if len(table_objs) == 0:
505508
return ""

backend/apps/template/generate_sql/generator.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,17 @@
22

33
from apps.db.constant import DB
44
from apps.template.template import get_base_template, get_sql_template as get_base_sql_template
5+
from common.core.config import settings
56

67

78
def get_sql_template():
89
template = get_base_template()
9-
return template['template']['sql']
10+
sql_template = template['template']['sql']
11+
if not settings.TABLE_SAMPLE_DATA_ENABLED:
12+
# Do not mutate shared templates or render pre-existing sample values.
13+
sql_template = sql_template.copy()
14+
sql_template['generate_basic_info'] = sql_template['generate_basic_info'].replace('{sample_data}', '')
15+
return sql_template
1016

1117

1218
def get_sql_example_template(db_type: Union[str, DB]):

backend/common/core/config.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,9 @@ def SQLALCHEMY_DATABASE_URI(self) -> PostgresDsn | str:
126126
PG_POOL_RECYCLE: int = 3600
127127
PG_POOL_PRE_PING: bool = True
128128

129+
# Include table sample rows in SQL generation context (disable for sensitive data).
130+
TABLE_SAMPLE_DATA_ENABLED: bool = True
131+
129132
TABLE_EMBEDDING_ENABLED: bool = True
130133
TABLE_EMBEDDING_COUNT: int = 10
131134
DS_EMBEDDING_COUNT: int = 10
@@ -138,6 +141,7 @@ def SQLALCHEMY_DATABASE_URI(self) -> PostgresDsn | str:
138141
'PARSE_REASONING_BLOCK_ENABLED',
139142
'PG_POOL_PRE_PING',
140143
'TABLE_EMBEDDING_ENABLED',
144+
'TABLE_SAMPLE_DATA_ENABLED',
141145
mode='before')
142146
@classmethod
143147
def lowercase_bool(cls, v: Any) -> Any:

tests/test_table_sample_data.py

Lines changed: 241 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,241 @@
1+
"""Isolated regression tests for the table-sample context switch (#1291).
2+
3+
Load function definitions from the actual sources, not copied implementations.
4+
This avoids importing database drivers, embedding models and X-Pack just to test
5+
sampling and prompt construction. Settings uses the real Pydantic loader. These
6+
are unit tests, not full application or model-provider integration tests.
7+
"""
8+
9+
import __future__
10+
import ast
11+
import importlib.util
12+
import json
13+
from copy import deepcopy
14+
from pathlib import Path
15+
from types import SimpleNamespace
16+
from unittest.mock import MagicMock, Mock
17+
18+
import pytest
19+
from pydantic import ValidationError
20+
21+
ROOT = Path(__file__).resolve().parents[1]
22+
BACKEND = ROOT / "backend"
23+
SENTINEL = "SQLBOT_SAMPLE_SENTINEL_1291"
24+
25+
26+
def load_functions(relative_path, names, namespace):
27+
"""Execute unchanged source definitions with explicit dependency doubles."""
28+
path = BACKEND / relative_path
29+
source = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
30+
definitions = [
31+
node for node in source.body
32+
if isinstance(node, ast.FunctionDef) and node.name in names
33+
]
34+
assert {node.name for node in definitions} == set(names)
35+
module = ast.Module(body=definitions, type_ignores=[])
36+
code = compile(
37+
module, str(path), "exec", flags=__future__.annotations.compiler_flag,
38+
dont_inherit=True,
39+
)
40+
exec(code, namespace)
41+
return namespace
42+
43+
44+
@pytest.fixture
45+
def settings_class(monkeypatch, tmp_path):
46+
# Do not load a developer's .env or carry a previous test's configuration.
47+
monkeypatch.chdir(tmp_path)
48+
monkeypatch.delenv("TABLE_SAMPLE_DATA_ENABLED", raising=False)
49+
spec = importlib.util.spec_from_file_location(
50+
"sqlbot_test_settings", BACKEND / "common/core/config.py"
51+
)
52+
module = importlib.util.module_from_spec(spec)
53+
spec.loader.exec_module(module)
54+
return module.Settings
55+
56+
57+
def test_sample_context_is_enabled_by_default(settings_class):
58+
assert settings_class(_env_file=None).TABLE_SAMPLE_DATA_ENABLED is True
59+
60+
61+
@pytest.mark.parametrize("value, expected", [
62+
("true", True), ("false", False), ("TRUE", True), ("FALSE", False),
63+
(" True ", True), (" False ", False), ("1", True), ("0", False),
64+
])
65+
def test_environment_boolean_parsing(settings_class, monkeypatch, value, expected):
66+
monkeypatch.setenv("TABLE_SAMPLE_DATA_ENABLED", value)
67+
assert settings_class(_env_file=None).TABLE_SAMPLE_DATA_ENABLED is expected
68+
69+
70+
def test_invalid_boolean_is_rejected(settings_class, monkeypatch):
71+
monkeypatch.setenv("TABLE_SAMPLE_DATA_ENABLED", "not-a-boolean")
72+
with pytest.raises(ValidationError):
73+
settings_class(_env_file=None)
74+
75+
76+
def test_dotenv_can_disable_sample_context(settings_class, tmp_path):
77+
env_file = tmp_path / "sample.env"
78+
env_file.write_text("TABLE_SAMPLE_DATA_ENABLED=false\n", encoding="utf-8")
79+
assert settings_class(_env_file=env_file).TABLE_SAMPLE_DATA_ENABLED is False
80+
81+
82+
@pytest.fixture
83+
def sampling():
84+
fields = [SimpleNamespace(
85+
field_name="sku", field_type="varchar", custom_comment="Product code",
86+
)]
87+
table = SimpleNamespace(
88+
id=1, table_name="orders", custom_comment="Order records", embedding=None,
89+
)
90+
table_obj = SimpleNamespace(schema="shop", table=table, fields=fields)
91+
namespace = {
92+
"settings": SimpleNamespace(
93+
TABLE_SAMPLE_DATA_ENABLED=True, TABLE_EMBEDDING_ENABLED=False,
94+
),
95+
"get_table_obj_by_ds": Mock(return_value=[table_obj]),
96+
"exec_sql": Mock(return_value={"data": [{"sku": SENTINEL}]}),
97+
"DB": SimpleNamespace(get_db=lambda _: SimpleNamespace(prefix='"', suffix='"')),
98+
"equals_ignore_case": lambda first, second: first.lower() == second.lower(),
99+
}
100+
return load_functions(
101+
"apps/datasource/crud/datasource.py",
102+
{"get_tables_sample_data", "get_table_sample_data", "get_table_schema", "execSql", "preview"},
103+
namespace,
104+
)
105+
106+
107+
def test_disabled_does_not_read_metadata_or_rows(sampling):
108+
sampling["settings"].TABLE_SAMPLE_DATA_ENABLED = False
109+
result = sampling["get_tables_sample_data"](
110+
object(), object(), SimpleNamespace(type="pg"),
111+
)
112+
assert result == ""
113+
sampling["get_table_obj_by_ds"].assert_not_called()
114+
sampling["exec_sql"].assert_not_called()
115+
116+
117+
def test_enabled_preserves_sample_query_and_output(sampling):
118+
session, user = object(), object()
119+
datasource = SimpleNamespace(type="pg")
120+
result = sampling["get_tables_sample_data"](session, user, datasource)
121+
sampling["get_table_obj_by_ds"].assert_called_once_with(
122+
session=session, current_user=user, ds=datasource,
123+
)
124+
sampling["exec_sql"].assert_called_once_with(
125+
ds=datasource, sql='SELECT "sku" FROM "orders" LIMIT 3', origin_column=True,
126+
)
127+
assert result == '# Table: orders\n[\n {\n "sku": "' + SENTINEL + '"\n }\n]'
128+
129+
130+
@pytest.mark.parametrize("table_list", [[], ["another_table"]])
131+
def test_selected_table_filter_is_preserved(sampling, table_list):
132+
assert sampling["get_tables_sample_data"](
133+
object(), object(), SimpleNamespace(type="pg"), table_list,
134+
) == ""
135+
sampling["exec_sql"].assert_not_called()
136+
137+
138+
def test_no_authorized_fields_does_not_sample(sampling):
139+
sampling["get_table_obj_by_ds"].return_value[0].fields = []
140+
assert sampling["get_tables_sample_data"](
141+
object(), object(), SimpleNamespace(type="pg"),
142+
) == ""
143+
sampling["exec_sql"].assert_not_called()
144+
145+
146+
def test_sample_limit_remains_three_rows(sampling):
147+
sampling["exec_sql"].return_value = {"data": [{"sku": n} for n in range(5)]}
148+
result = sampling["get_tables_sample_data"](
149+
object(), object(), SimpleNamespace(type="pg"),
150+
)
151+
assert json.loads(result.split("\n", 1)[1]) == [{"sku": 0}, {"sku": 1}, {"sku": 2}]
152+
153+
154+
def test_disabled_does_not_remove_schema(sampling):
155+
sampling["settings"].TABLE_SAMPLE_DATA_ENABLED = False
156+
schema, tables = sampling["get_table_schema"](
157+
object(), object(), SimpleNamespace(type="pg", table_relation=None),
158+
"Show orders", embedding=False,
159+
)
160+
assert tables == ["orders"]
161+
assert "shop.orders" in schema
162+
assert "sku:varchar, Product code" in schema
163+
sampling["exec_sql"].assert_not_called()
164+
165+
166+
def test_disabled_does_not_block_explicit_sql_execution(sampling):
167+
sampling["settings"].TABLE_SAMPLE_DATA_ENABLED = False
168+
sampling["CoreDatasource"] = SimpleNamespace(id=1)
169+
sampling["select"] = Mock()
170+
session, datasource = Mock(), object()
171+
session.exec.return_value.first.return_value = datasource
172+
result = sampling["execSql"](session, 1, "SELECT 1")
173+
sampling["exec_sql"].assert_called_once_with(datasource, "SELECT 1", True)
174+
assert result is sampling["exec_sql"].return_value
175+
176+
177+
@pytest.fixture
178+
def templates():
179+
template = {"template": {"sql": {
180+
"generate_basic_info": (
181+
"<db-engine>{engine}</db-engine>\n<schema>{schema}</schema>\n"
182+
"<sample-data>{sample_data}</sample-data>"
183+
),
184+
"other_rule": "Preserve other rules",
185+
}}}
186+
namespace = {
187+
"settings": SimpleNamespace(TABLE_SAMPLE_DATA_ENABLED=True),
188+
"get_base_template": Mock(return_value=template),
189+
}
190+
return load_functions(
191+
"apps/template/generate_sql/generator.py", {"get_sql_template"}, namespace,
192+
)
193+
194+
195+
@pytest.mark.parametrize("enabled", [True, False])
196+
def test_template_controls_sample_insertion_without_removing_schema(templates, enabled):
197+
templates["settings"].TABLE_SAMPLE_DATA_ENABLED = enabled
198+
result = templates["get_sql_template"]()
199+
rendered = result["generate_basic_info"].format(
200+
engine="PostgreSQL", schema="orders(sku varchar)", sample_data=SENTINEL,
201+
)
202+
assert (SENTINEL in rendered) is enabled
203+
assert "PostgreSQL" in rendered
204+
assert "orders(sku varchar)" in rendered
205+
assert result["other_rule"] == "Preserve other rules"
206+
207+
208+
def test_disabling_does_not_mutate_shared_templates(templates):
209+
original = deepcopy(templates["get_base_template"].return_value)
210+
templates["settings"].TABLE_SAMPLE_DATA_ENABLED = False
211+
disabled = templates["get_sql_template"]()
212+
assert "{sample_data}" not in disabled["generate_basic_info"]
213+
assert templates["get_base_template"].return_value == original
214+
templates["settings"].TABLE_SAMPLE_DATA_ENABLED = True
215+
assert "{sample_data}" in templates["get_sql_template"]()["generate_basic_info"]
216+
217+
218+
def test_disabled_does_not_block_manual_preview(sampling):
219+
sampling["settings"].TABLE_SAMPLE_DATA_ENABLED = False
220+
sampling.update({
221+
"CoreDatasource": MagicMock(), "CoreField": MagicMock(), "CoreTable": MagicMock(),
222+
"is_normal_user": Mock(return_value=False),
223+
"get_engine_config": Mock(return_value=SimpleNamespace(dbSchema="public")),
224+
})
225+
datasource = SimpleNamespace(type="excel")
226+
field = SimpleNamespace(field_name="sku", checked=True)
227+
table = SimpleNamespace(id=1, table_name="orders")
228+
ds_query, field_query, table_query = Mock(), Mock(), Mock()
229+
ds_query.filter.return_value.first.return_value = datasource
230+
field_query.filter.return_value.order_by.return_value.all.return_value = [field]
231+
table_query.filter.return_value.first.return_value = table
232+
session = Mock()
233+
session.query.side_effect = [ds_query, field_query, table_query]
234+
result = sampling["preview"](session, object(), 1, SimpleNamespace(table=table))
235+
sampling["exec_sql"].assert_called_once()
236+
args = sampling["exec_sql"].call_args.args
237+
assert args[0] is datasource
238+
assert '"public"."orders"' in args[1]
239+
assert "LIMIT 100" in args[1]
240+
assert args[2] is True
241+
assert result is sampling["exec_sql"].return_value

0 commit comments

Comments
 (0)