dbagent-cli 0.11.3__tar.gz → 0.11.4__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/PKG-INFO +1 -1
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/__init__.py +1 -1
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/agent/generator.py +2 -2
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/agent/pipeline.py +1 -1
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/cli.py +3 -3
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/connectors/base.py +2 -2
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent_cli.egg-info/PKG-INFO +1 -1
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/pyproject.toml +1 -1
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_direct_sql.py +1 -1
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_generator.py +1 -1
- dbagent_cli-0.11.4/tests/test_ground_truth_sql_validation.py +150 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_knowledge.py +8 -8
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_prompt_queue.py +2 -2
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_prompt_queue_and_sql.py +7 -7
- dbagent_cli-0.11.3/tests/test_ground_truth_sql_validation.py +0 -150
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/README.md +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/agent/confidence.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/agent/conversation.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/agent/doctor.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/agent/error_classifier.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/agent/recommender.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/agent/validator.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/benchmark.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/config.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/connectors/factory.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/connectors/mongo.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/connectors/relational.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/errors.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/execution/__init__.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/execution/cache.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/execution/changeset.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/execution/cost.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/execution/prompt_queue.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/execution/recovery.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/execution/sandbox.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/execution/streaming.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/llm/anthropic_provider.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/llm/auto_setup.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/llm/base.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/llm/deepseek_provider.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/llm/factory.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/llm/gemini_provider.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/llm/groq_provider.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/llm/mistral_provider.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/llm/mock_provider.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/llm/ollama_provider.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/llm/openai_provider.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/llm/openrouter_provider.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/plugins.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/sample_db.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/schema/cache.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/schema/formatter.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/schema/graph.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/schema/ground_truth.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/schema/models.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/schema/retriever.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/schema/selector.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/security/__init__.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/security/audit.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/security/environment.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/security/masking.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/security/modes.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/security/rbac.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/semantic/__init__.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/semantic/dictionary.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/semantic/knowledge.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/semantic/templates.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/sql/__init__.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/sql/optimizer.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/sql/validator.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/ui/chat_screen.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/ui/console.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/ui/grid_window.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/ui/live_streamer.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/ui/viewer.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/ui/visualizer.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent_cli.egg-info/SOURCES.txt +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent_cli.egg-info/dependency_links.txt +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent_cli.egg-info/entry_points.txt +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent_cli.egg-info/requires.txt +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent_cli.egg-info/top_level.txt +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/setup.cfg +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_benchmarks.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_cache.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_cli.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_connectors.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_demo_and_sample_db.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_doctor_and_viz.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_formats_and_exit_codes.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_graph.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_multi_turn_conversation.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_plugins.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_providers.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_recovery.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_safety.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_sandbox.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_schema.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_security.py +0 -0
- {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_semantic.py +0 -0
|
@@ -60,7 +60,7 @@ Generate ONLY the raw executable SQL query for the given database. No markdown f
|
|
|
60
60
|
- If the column is VARCHAR(1) or CHAR(1) or text, check if values use 'Y'/'N', 'YES'/'NO', '1'/'0', 'ACTIVE'/'INACTIVE', 'TRUE'/'FALSE'.
|
|
61
61
|
- If sample values show 'Y'/'N' (or if it's a single-character VARCHAR/CHAR column), use `active = 'Y'` or `active IN ('Y', '1', 'true', 'TRUE')` for active/enabled status.
|
|
62
62
|
10. For Deletions & Cascading Mapping Cleanup (when deleting a record/user without foreign key constraint violations):
|
|
63
|
-
- For PostgreSQL: Use a single Common Table Expression (CTE) with data-modifying statements deleting from all child, mapping, activity, profile, and preference tables first, and finally deleting from the parent table (e.g. `WITH target AS (SELECT
|
|
63
|
+
- For PostgreSQL: Use a single Common Table Expression (CTE) with data-modifying statements deleting from all child, mapping, activity, profile, and preference tables first, and finally deleting from the parent table (e.g. `WITH target AS (SELECT id FROM users WHERE status = 'inactive'), del_orders AS (DELETE FROM orders WHERE user_id IN (SELECT id FROM target)), del_logs AS (DELETE FROM user_logs WHERE user_id IN (SELECT id FROM target)) DELETE FROM users WHERE id IN (SELECT id FROM target);`).
|
|
64
64
|
- For MySQL / SQLite: Generate clean, ordered DELETE statements deleting referencing child records before the parent record.
|
|
65
65
|
- Always delete from all referencing tables present in the schema to avoid foreign key violations.
|
|
66
66
|
11. Multi-Turn Session Memory & Query Refinement:
|
|
@@ -201,7 +201,7 @@ class ScriptGenerator:
|
|
|
201
201
|
if re.search(r'\b(?:is|are)\b.*\b(?:same|different|difference|related|equal|linked)\b', prompt_lower):
|
|
202
202
|
return "analysis"
|
|
203
203
|
|
|
204
|
-
# Contextual clarification for analysis (e.g. "i mean
|
|
204
|
+
# Contextual clarification for analysis (e.g. "i mean user_profiles", "i mean table_x")
|
|
205
205
|
if re.search(r'\bi mean\b', prompt_lower) and not re.search(r'\b(insert|update|delete|drop)\b', prompt_lower):
|
|
206
206
|
return "analysis"
|
|
207
207
|
|
|
@@ -176,7 +176,7 @@ class QueryPipeline:
|
|
|
176
176
|
all_table_names = self.connector.get_table_names()
|
|
177
177
|
stems = set()
|
|
178
178
|
for t in target_names:
|
|
179
|
-
unprefixed = re.sub(r'^(?:
|
|
179
|
+
unprefixed = re.sub(r'^(?:tbl_|table_|app_|dim_|fact_|rel_|raw_)', '', t.lower())
|
|
180
180
|
unprefixed = re.sub(r'(?:s|es)$', '', unprefixed)
|
|
181
181
|
if len(unprefixed) >= 3:
|
|
182
182
|
stems.add(unprefixed)
|
|
@@ -808,7 +808,7 @@ def chat_command(
|
|
|
808
808
|
|
|
809
809
|
if cmd_name == ":table":
|
|
810
810
|
if not cmd_arg:
|
|
811
|
-
print_warning("Usage: :table <table_name> (e.g. :table
|
|
811
|
+
print_warning("Usage: :table <table_name> (e.g. :table users)")
|
|
812
812
|
continue
|
|
813
813
|
exact, fuzzy = connector.resolve_tables(cmd_arg)
|
|
814
814
|
target = exact[0] if exact else (fuzzy[0] if fuzzy else cmd_arg)
|
|
@@ -823,7 +823,7 @@ def chat_command(
|
|
|
823
823
|
|
|
824
824
|
if cmd_name in [":fk", ":relations", ":mapped"]:
|
|
825
825
|
if not cmd_arg:
|
|
826
|
-
print_warning("Usage: :fk <table_name> (e.g. :fk
|
|
826
|
+
print_warning("Usage: :fk <table_name> (e.g. :fk users)")
|
|
827
827
|
continue
|
|
828
828
|
exact, fuzzy = connector.resolve_tables(cmd_arg)
|
|
829
829
|
target = exact[0] if exact else (fuzzy[0] if fuzzy else cmd_arg)
|
|
@@ -931,7 +931,7 @@ def chat_command(
|
|
|
931
931
|
|
|
932
932
|
if cmd_name in [":learn", ":remember"]:
|
|
933
933
|
if not cmd_arg or "=" not in cmd_arg:
|
|
934
|
-
print_warning("Usage: :learn <topic> = <description> (e.g. :learn
|
|
934
|
+
print_warning("Usage: :learn <topic> = <description> (e.g. :learn user_roles = users map to roles via user_role_map)")
|
|
935
935
|
continue
|
|
936
936
|
parts = cmd_arg.split("=", 1)
|
|
937
937
|
topic = parts[0].strip()
|
|
@@ -87,8 +87,8 @@ class BaseConnector(ABC):
|
|
|
87
87
|
exact.append(t)
|
|
88
88
|
continue
|
|
89
89
|
|
|
90
|
-
# 2. Match without common prefixes (e.g. "
|
|
91
|
-
unprefixed = re.sub(r'^(?:
|
|
90
|
+
# 2. Match without common database table prefixes (e.g. "tbl_orders" -> "orders")
|
|
91
|
+
unprefixed = re.sub(r'^(?:tbl_|table_|app_|dim_|fact_|rel_|raw_)', '', t_lower)
|
|
92
92
|
if unprefixed and len(unprefixed) >= 3 and unprefixed not in stop_words:
|
|
93
93
|
if (unprefixed in tokens or
|
|
94
94
|
f"{unprefixed}s" in tokens or
|
|
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "dbagent-cli"
|
|
7
|
-
version = "0.11.
|
|
7
|
+
version = "0.11.4"
|
|
8
8
|
description = "The Privacy-First AI Database Copilot, Schema Graph Introspector & Post-Commit Recovery Engine (CLI)"
|
|
9
9
|
readme = "README.md"
|
|
10
10
|
requires-python = ">=3.9"
|
|
@@ -11,7 +11,7 @@ from dbagent.llm.mock_provider import MockProvider
|
|
|
11
11
|
def test_is_direct_sql_detection():
|
|
12
12
|
# Direct SQL queries
|
|
13
13
|
assert SQLValidator.is_direct_sql("SELECT * FROM users;") is True
|
|
14
|
-
assert SQLValidator.is_direct_sql("select
|
|
14
|
+
assert SQLValidator.is_direct_sql("select id, username from users where active = 'Y'") is True
|
|
15
15
|
assert SQLValidator.is_direct_sql("SELECT count(*) FROM orders") is True
|
|
16
16
|
assert SQLValidator.is_direct_sql("SELECT 1;") is True
|
|
17
17
|
assert SQLValidator.is_direct_sql("WITH cte AS (SELECT 1) SELECT * FROM cte") is True
|
|
@@ -74,7 +74,7 @@ def test_validator_destructive_delete_without_where():
|
|
|
74
74
|
|
|
75
75
|
|
|
76
76
|
def test_classify_intent_analysis():
|
|
77
|
-
prompt = "analyze the database and give me the linkage architecture
|
|
77
|
+
prompt = "analyze the database and give me the linkage architecture of how users and accounts are related and how they are mapped and how to know which devices in customer_devices are mapped to which user"
|
|
78
78
|
assert ScriptGenerator.classify_intent(prompt) == "analysis"
|
|
79
79
|
|
|
80
80
|
assert ScriptGenerator.classify_intent("how are users and orders related?") == "analysis"
|
|
@@ -0,0 +1,150 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Unit tests for Schema Ground-Truth SQL Validation and Auto-Correction.
|
|
3
|
+
Verifies table, column, PK, and FK validation against physical database schema using generic domain models.
|
|
4
|
+
"""
|
|
5
|
+
|
|
6
|
+
import pytest
|
|
7
|
+
from unittest.mock import MagicMock
|
|
8
|
+
from dbagent.schema.models import DatabaseSchema, TableModel, ColumnModel, ForeignKeyModel
|
|
9
|
+
from dbagent.schema.ground_truth import SchemaGroundTruthValidator
|
|
10
|
+
from dbagent.agent.pipeline import QueryPipeline, PipelineResult
|
|
11
|
+
from dbagent.security.modes import OperationMode
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
@pytest.fixture
|
|
15
|
+
def sample_schema():
|
|
16
|
+
"""Create a sample generic schema with realistic primary keys and foreign keys."""
|
|
17
|
+
customers = TableModel(
|
|
18
|
+
name="customers",
|
|
19
|
+
columns=[
|
|
20
|
+
ColumnModel(name="customer_id", data_type="INTEGER", is_primary_key=True),
|
|
21
|
+
ColumnModel(name="customer_name", data_type="VARCHAR(255)"),
|
|
22
|
+
ColumnModel(name="account_id", data_type="INTEGER"),
|
|
23
|
+
ColumnModel(name="region_id", data_type="INTEGER"),
|
|
24
|
+
ColumnModel(name="is_active", data_type="BOOLEAN"),
|
|
25
|
+
],
|
|
26
|
+
foreign_keys=[
|
|
27
|
+
ForeignKeyModel(
|
|
28
|
+
constrained_columns=["account_id"],
|
|
29
|
+
referred_table="accounts",
|
|
30
|
+
referred_columns=["account_id"],
|
|
31
|
+
),
|
|
32
|
+
ForeignKeyModel(
|
|
33
|
+
constrained_columns=["region_id"],
|
|
34
|
+
referred_table="regions",
|
|
35
|
+
referred_columns=["region_id"],
|
|
36
|
+
),
|
|
37
|
+
],
|
|
38
|
+
primary_key=["customer_id"],
|
|
39
|
+
)
|
|
40
|
+
|
|
41
|
+
accounts = TableModel(
|
|
42
|
+
name="accounts",
|
|
43
|
+
columns=[
|
|
44
|
+
ColumnModel(name="account_id", data_type="INTEGER", is_primary_key=True),
|
|
45
|
+
ColumnModel(name="account_name", data_type="VARCHAR(255)"),
|
|
46
|
+
ColumnModel(name="billing_code", data_type="VARCHAR(50)"),
|
|
47
|
+
],
|
|
48
|
+
foreign_keys=[],
|
|
49
|
+
primary_key=["account_id"],
|
|
50
|
+
)
|
|
51
|
+
|
|
52
|
+
regions = TableModel(
|
|
53
|
+
name="regions",
|
|
54
|
+
columns=[
|
|
55
|
+
ColumnModel(name="region_id", data_type="INTEGER", is_primary_key=True),
|
|
56
|
+
ColumnModel(name="region_name", data_type="VARCHAR(255)"),
|
|
57
|
+
],
|
|
58
|
+
foreign_keys=[],
|
|
59
|
+
primary_key=["region_id"],
|
|
60
|
+
)
|
|
61
|
+
|
|
62
|
+
return DatabaseSchema(
|
|
63
|
+
database_name="sample_db",
|
|
64
|
+
dialect_name="postgresql",
|
|
65
|
+
tables=[customers, accounts, regions],
|
|
66
|
+
)
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
def test_valid_select_query_passes_ground_truth(sample_schema):
|
|
70
|
+
sql = "SELECT c.customer_id, c.customer_name, a.account_name FROM customers c JOIN accounts a ON c.account_id = a.account_id WHERE c.is_active = true;"
|
|
71
|
+
is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
|
|
72
|
+
assert is_valid is True
|
|
73
|
+
assert len(errors) == 0
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
def test_valid_update_query_passes_ground_truth(sample_schema):
|
|
77
|
+
sql = "UPDATE customers SET account_id = 10, region_id = 20 WHERE customer_id = 100;"
|
|
78
|
+
is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
|
|
79
|
+
assert is_valid is True
|
|
80
|
+
assert len(errors) == 0
|
|
81
|
+
|
|
82
|
+
|
|
83
|
+
def test_hallucinated_column_in_update_set_detected(sample_schema):
|
|
84
|
+
sql = "UPDATE customers SET project = 5, site_group = 10 WHERE customer_id = 100;"
|
|
85
|
+
is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
|
|
86
|
+
assert is_valid is False
|
|
87
|
+
assert any("Column 'project' does not exist in table 'customers'" in e for e in errors)
|
|
88
|
+
assert any("Column 'site_group' does not exist in table 'customers'" in e for e in errors)
|
|
89
|
+
assert any("Valid foreign keys" in e for e in errors)
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
def test_hallucinated_column_in_where_clause_detected(sample_schema):
|
|
93
|
+
sql = "UPDATE customers SET account_id = 5 WHERE legacy_user_code = 100;"
|
|
94
|
+
is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
|
|
95
|
+
assert is_valid is False
|
|
96
|
+
assert any("Column 'legacy_user_code' in WHERE clause does not exist in table 'customers'" in e for e in errors)
|
|
97
|
+
assert any("Primary key is 'customer_id'" in e for e in errors)
|
|
98
|
+
|
|
99
|
+
|
|
100
|
+
def test_hallucinated_table_name_detected(sample_schema):
|
|
101
|
+
sql = "SELECT * FROM non_existent_table WHERE id = 1;"
|
|
102
|
+
is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
|
|
103
|
+
assert is_valid is False
|
|
104
|
+
assert any("Table 'non_existent_table' does not exist" in e for e in errors)
|
|
105
|
+
|
|
106
|
+
|
|
107
|
+
def test_hallucinated_qualified_column_detected(sample_schema):
|
|
108
|
+
sql = "SELECT c.customer_id, c.fake_column FROM customers c;"
|
|
109
|
+
is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
|
|
110
|
+
assert is_valid is False
|
|
111
|
+
assert any("Column 'fake_column' does not exist in table 'customers'" in e for e in errors)
|
|
112
|
+
|
|
113
|
+
|
|
114
|
+
def test_multi_statement_validation(sample_schema):
|
|
115
|
+
sql = """
|
|
116
|
+
UPDATE customers SET account_id = 2 WHERE customer_id = 1;
|
|
117
|
+
UPDATE customers SET bad_col = 3 WHERE customer_id = 2;
|
|
118
|
+
"""
|
|
119
|
+
is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
|
|
120
|
+
assert is_valid is False
|
|
121
|
+
assert any("Column 'bad_col' does not exist in table 'customers'" in e for e in errors)
|
|
122
|
+
|
|
123
|
+
|
|
124
|
+
def test_pipeline_auto_corrects_hallucinated_columns(sample_schema):
|
|
125
|
+
mock_connector = MagicMock()
|
|
126
|
+
mock_connector.get_table_names.return_value = ["customers", "accounts", "regions"]
|
|
127
|
+
mock_connector.resolve_tables.return_value = (["customers"], [])
|
|
128
|
+
mock_connector.inspect_schema.return_value = sample_schema
|
|
129
|
+
mock_connector.execute_query.return_value = (["customer_id", "account_id"], [{"customer_id": 1, "account_id": 2}], None)
|
|
130
|
+
|
|
131
|
+
mock_llm = MagicMock()
|
|
132
|
+
mock_llm.generate.side_effect = [
|
|
133
|
+
"UPDATE customers SET account = 2 WHERE cust_id = 1;",
|
|
134
|
+
"UPDATE customers SET account_id = 2 WHERE customer_id = 1;",
|
|
135
|
+
]
|
|
136
|
+
|
|
137
|
+
pipeline = QueryPipeline(
|
|
138
|
+
connector=mock_connector,
|
|
139
|
+
llm=mock_llm,
|
|
140
|
+
mode=OperationMode.CONFIRM,
|
|
141
|
+
auto_execute=True,
|
|
142
|
+
)
|
|
143
|
+
|
|
144
|
+
result = pipeline.run("Change the account for customer 1 to account 2", force=True)
|
|
145
|
+
|
|
146
|
+
assert result.success is True
|
|
147
|
+
assert "account_id = 2" in result.sql
|
|
148
|
+
assert "customer_id = 1" in result.sql
|
|
149
|
+
assert result.retries >= 1
|
|
150
|
+
assert result.ground_truth is not None
|
|
@@ -17,27 +17,27 @@ def temp_kb(tmp_path):
|
|
|
17
17
|
def test_knowledge_base_learn_and_find(temp_kb):
|
|
18
18
|
rule = temp_kb.learn(
|
|
19
19
|
db_name="test_db",
|
|
20
|
-
topic="
|
|
21
|
-
rule_description="
|
|
22
|
-
relevant_tables=["
|
|
20
|
+
topic="User to Role Permission Mapping",
|
|
21
|
+
rule_description="users map to permissions through user_roles -> roles -> role_permissions",
|
|
22
|
+
relevant_tables=["users", "user_roles", "roles", "role_permissions", "permissions"],
|
|
23
23
|
source="user_correction",
|
|
24
24
|
)
|
|
25
25
|
assert rule.rule_id.startswith("rule_")
|
|
26
|
-
assert rule.topic == "
|
|
26
|
+
assert rule.topic == "User to Role Permission Mapping"
|
|
27
27
|
|
|
28
28
|
# List all
|
|
29
29
|
all_rules = temp_kb.list_all("test_db")
|
|
30
30
|
assert len(all_rules) == 1
|
|
31
31
|
|
|
32
32
|
# Find relevant by prompt
|
|
33
|
-
matched = temp_kb.find_relevant("test_db", "how to map
|
|
33
|
+
matched = temp_kb.find_relevant("test_db", "how to map users to permissions")
|
|
34
34
|
assert len(matched) == 1
|
|
35
|
-
assert matched[0].topic == "
|
|
35
|
+
assert matched[0].topic == "User to Role Permission Mapping"
|
|
36
36
|
|
|
37
37
|
# Context formatting
|
|
38
|
-
ctx = temp_kb.format_prompt_context("test_db", "how to map
|
|
38
|
+
ctx = temp_kb.format_prompt_context("test_db", "how to map users to permissions")
|
|
39
39
|
assert "Persistent Database Domain Knowledge" in ctx
|
|
40
|
-
assert "
|
|
40
|
+
assert "User to Role Permission Mapping" in ctx
|
|
41
41
|
|
|
42
42
|
# Delete
|
|
43
43
|
deleted = temp_kb.delete("test_db", rule.rule_id)
|
|
@@ -26,10 +26,10 @@ def test_prompt_queue_fifo():
|
|
|
26
26
|
|
|
27
27
|
def test_split_and_enqueue_chained():
|
|
28
28
|
qm = PromptQueueManager()
|
|
29
|
-
raw = "how
|
|
29
|
+
raw = "how users and orders were linked ; give the flow diagram ; count total users"
|
|
30
30
|
first = qm.split_and_enqueue_chained(raw)
|
|
31
31
|
|
|
32
|
-
assert first == "how
|
|
32
|
+
assert first == "how users and orders were linked"
|
|
33
33
|
assert qm.count() == 2
|
|
34
34
|
assert qm.list_all() == ["give the flow diagram", "count total users"]
|
|
35
35
|
|
|
@@ -11,10 +11,10 @@ def test_queue_does_not_split_multiline_sql_with_semicolons():
|
|
|
11
11
|
queue_mgr = PromptQueueManager()
|
|
12
12
|
|
|
13
13
|
sql_input = """
|
|
14
|
-
DELETE FROM
|
|
15
|
-
DELETE FROM
|
|
14
|
+
DELETE FROM customer_user_map WHERE user_id = 5201;
|
|
15
|
+
DELETE FROM user_access WHERE user_id = 5201;
|
|
16
16
|
DELETE FROM user_activity WHERE user_id = 5201;
|
|
17
|
-
DELETE FROM
|
|
17
|
+
DELETE FROM users WHERE user_id = 5201;
|
|
18
18
|
COMMIT; correct this
|
|
19
19
|
"""
|
|
20
20
|
|
|
@@ -27,12 +27,12 @@ def test_classify_query_with_cte_delete():
|
|
|
27
27
|
"""Verify that data-modifying WITH CTE statements are classified as 'write'."""
|
|
28
28
|
cte_delete = """
|
|
29
29
|
WITH target_user AS (
|
|
30
|
-
SELECT
|
|
30
|
+
SELECT id FROM users WHERE username = 'target_user'
|
|
31
31
|
),
|
|
32
32
|
del_access AS (
|
|
33
|
-
DELETE FROM
|
|
33
|
+
DELETE FROM user_access WHERE user_id IN (SELECT id FROM target_user)
|
|
34
34
|
)
|
|
35
|
-
DELETE FROM
|
|
35
|
+
DELETE FROM users WHERE id IN (SELECT id FROM target_user);
|
|
36
36
|
"""
|
|
37
37
|
assert SQLValidator.classify_query(cte_delete) == "write"
|
|
38
38
|
|
|
@@ -41,7 +41,7 @@ def test_classify_query_with_cte_read():
|
|
|
41
41
|
"""Verify that regular SELECT WITH CTE statements remain classified as 'read'."""
|
|
42
42
|
cte_read = """
|
|
43
43
|
WITH active_users AS (
|
|
44
|
-
SELECT
|
|
44
|
+
SELECT id, username FROM users WHERE active = 'Y'
|
|
45
45
|
)
|
|
46
46
|
SELECT * FROM active_users LIMIT 10;
|
|
47
47
|
"""
|
|
@@ -1,150 +0,0 @@
|
|
|
1
|
-
"""
|
|
2
|
-
Unit tests for Schema Ground-Truth SQL Validation and Auto-Correction.
|
|
3
|
-
Verifies table, column, PK, and FK validation against physical database schema.
|
|
4
|
-
"""
|
|
5
|
-
|
|
6
|
-
import pytest
|
|
7
|
-
from unittest.mock import MagicMock
|
|
8
|
-
from dbagent.schema.models import DatabaseSchema, TableModel, ColumnModel, ForeignKeyModel
|
|
9
|
-
from dbagent.schema.ground_truth import SchemaGroundTruthValidator
|
|
10
|
-
from dbagent.agent.pipeline import QueryPipeline, PipelineResult
|
|
11
|
-
from dbagent.security.modes import OperationMode
|
|
12
|
-
|
|
13
|
-
|
|
14
|
-
@pytest.fixture
|
|
15
|
-
def sample_schema():
|
|
16
|
-
"""Create a sample schema with realistic primary keys and foreign keys."""
|
|
17
|
-
ivis_site = TableModel(
|
|
18
|
-
name="ivis_site",
|
|
19
|
-
columns=[
|
|
20
|
-
ColumnModel(name="pk_site_id", data_type="INTEGER", is_primary_key=True),
|
|
21
|
-
ColumnModel(name="site_name", data_type="VARCHAR(255)"),
|
|
22
|
-
ColumnModel(name="fk_project_id", data_type="INTEGER"),
|
|
23
|
-
ColumnModel(name="fk_site_group_id", data_type="INTEGER"),
|
|
24
|
-
ColumnModel(name="is_active", data_type="BOOLEAN"),
|
|
25
|
-
],
|
|
26
|
-
foreign_keys=[
|
|
27
|
-
ForeignKeyModel(
|
|
28
|
-
constrained_columns=["fk_project_id"],
|
|
29
|
-
referred_table="ivis_project",
|
|
30
|
-
referred_columns=["pk_project_id"],
|
|
31
|
-
),
|
|
32
|
-
ForeignKeyModel(
|
|
33
|
-
constrained_columns=["fk_site_group_id"],
|
|
34
|
-
referred_table="ivis_site_group",
|
|
35
|
-
referred_columns=["pk_site_group_id"],
|
|
36
|
-
),
|
|
37
|
-
],
|
|
38
|
-
primary_key=["pk_site_id"],
|
|
39
|
-
)
|
|
40
|
-
|
|
41
|
-
ivis_project = TableModel(
|
|
42
|
-
name="ivis_project",
|
|
43
|
-
columns=[
|
|
44
|
-
ColumnModel(name="pk_project_id", data_type="INTEGER", is_primary_key=True),
|
|
45
|
-
ColumnModel(name="project_name", data_type="VARCHAR(255)"),
|
|
46
|
-
ColumnModel(name="customer_code", data_type="VARCHAR(50)"),
|
|
47
|
-
],
|
|
48
|
-
foreign_keys=[],
|
|
49
|
-
primary_key=["pk_project_id"],
|
|
50
|
-
)
|
|
51
|
-
|
|
52
|
-
ivis_site_group = TableModel(
|
|
53
|
-
name="ivis_site_group",
|
|
54
|
-
columns=[
|
|
55
|
-
ColumnModel(name="pk_site_group_id", data_type="INTEGER", is_primary_key=True),
|
|
56
|
-
ColumnModel(name="group_name", data_type="VARCHAR(255)"),
|
|
57
|
-
],
|
|
58
|
-
foreign_keys=[],
|
|
59
|
-
primary_key=["pk_site_group_id"],
|
|
60
|
-
)
|
|
61
|
-
|
|
62
|
-
return DatabaseSchema(
|
|
63
|
-
database_name="iportal",
|
|
64
|
-
dialect_name="postgresql",
|
|
65
|
-
tables=[ivis_site, ivis_project, ivis_site_group],
|
|
66
|
-
)
|
|
67
|
-
|
|
68
|
-
|
|
69
|
-
def test_valid_select_query_passes_ground_truth(sample_schema):
|
|
70
|
-
sql = "SELECT s.pk_site_id, s.site_name, p.project_name FROM ivis_site s JOIN ivis_project p ON s.fk_project_id = p.pk_project_id WHERE s.is_active = true;"
|
|
71
|
-
is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
|
|
72
|
-
assert is_valid is True
|
|
73
|
-
assert len(errors) == 0
|
|
74
|
-
|
|
75
|
-
|
|
76
|
-
def test_valid_update_query_passes_ground_truth(sample_schema):
|
|
77
|
-
sql = "UPDATE ivis_site SET fk_project_id = 10, fk_site_group_id = 20 WHERE pk_site_id = 100;"
|
|
78
|
-
is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
|
|
79
|
-
assert is_valid is True
|
|
80
|
-
assert len(errors) == 0
|
|
81
|
-
|
|
82
|
-
|
|
83
|
-
def test_hallucinated_column_in_update_set_detected(sample_schema):
|
|
84
|
-
sql = "UPDATE ivis_site SET project = 5, site_group = 10 WHERE pk_site_id = 100;"
|
|
85
|
-
is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
|
|
86
|
-
assert is_valid is False
|
|
87
|
-
assert any("Column 'project' does not exist in table 'ivis_site'" in e for e in errors)
|
|
88
|
-
assert any("Column 'site_group' does not exist in table 'ivis_site'" in e for e in errors)
|
|
89
|
-
assert any("Valid foreign keys" in e for e in errors)
|
|
90
|
-
|
|
91
|
-
|
|
92
|
-
def test_hallucinated_column_in_where_clause_detected(sample_schema):
|
|
93
|
-
sql = "UPDATE ivis_site SET fk_project_id = 5 WHERE customer_site_id = 100;"
|
|
94
|
-
is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
|
|
95
|
-
assert is_valid is False
|
|
96
|
-
assert any("Column 'customer_site_id' in WHERE clause does not exist in table 'ivis_site'" in e for e in errors)
|
|
97
|
-
assert any("Primary key is 'pk_site_id'" in e for e in errors)
|
|
98
|
-
|
|
99
|
-
|
|
100
|
-
def test_hallucinated_table_name_detected(sample_schema):
|
|
101
|
-
sql = "SELECT * FROM non_existent_table WHERE id = 1;"
|
|
102
|
-
is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
|
|
103
|
-
assert is_valid is False
|
|
104
|
-
assert any("Table 'non_existent_table' does not exist" in e for e in errors)
|
|
105
|
-
|
|
106
|
-
|
|
107
|
-
def test_hallucinated_qualified_column_detected(sample_schema):
|
|
108
|
-
sql = "SELECT s.pk_site_id, s.fake_column FROM ivis_site s;"
|
|
109
|
-
is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
|
|
110
|
-
assert is_valid is False
|
|
111
|
-
assert any("Column 'fake_column' does not exist in table 'ivis_site'" in e for e in errors)
|
|
112
|
-
|
|
113
|
-
|
|
114
|
-
def test_multi_statement_validation(sample_schema):
|
|
115
|
-
sql = """
|
|
116
|
-
UPDATE ivis_site SET fk_project_id = 2 WHERE pk_site_id = 1;
|
|
117
|
-
UPDATE ivis_site SET bad_col = 3 WHERE pk_site_id = 2;
|
|
118
|
-
"""
|
|
119
|
-
is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
|
|
120
|
-
assert is_valid is False
|
|
121
|
-
assert any("Column 'bad_col' does not exist in table 'ivis_site'" in e for e in errors)
|
|
122
|
-
|
|
123
|
-
|
|
124
|
-
def test_pipeline_auto_corrects_hallucinated_columns(sample_schema):
|
|
125
|
-
mock_connector = MagicMock()
|
|
126
|
-
mock_connector.get_table_names.return_value = ["ivis_site", "ivis_project", "ivis_site_group"]
|
|
127
|
-
mock_connector.resolve_tables.return_value = (["ivis_site"], [])
|
|
128
|
-
mock_connector.inspect_schema.return_value = sample_schema
|
|
129
|
-
mock_connector.execute_query.return_value = (["pk_site_id", "fk_project_id"], [{"pk_site_id": 1, "fk_project_id": 2}], None)
|
|
130
|
-
|
|
131
|
-
mock_llm = MagicMock()
|
|
132
|
-
mock_llm.generate.side_effect = [
|
|
133
|
-
"UPDATE ivis_site SET project = 2 WHERE site_id = 1;",
|
|
134
|
-
"UPDATE ivis_site SET fk_project_id = 2 WHERE pk_site_id = 1;",
|
|
135
|
-
]
|
|
136
|
-
|
|
137
|
-
pipeline = QueryPipeline(
|
|
138
|
-
connector=mock_connector,
|
|
139
|
-
llm=mock_llm,
|
|
140
|
-
mode=OperationMode.CONFIRM,
|
|
141
|
-
auto_execute=True,
|
|
142
|
-
)
|
|
143
|
-
|
|
144
|
-
result = pipeline.run("Change the project for site 1 to project 2", force=True)
|
|
145
|
-
|
|
146
|
-
assert result.success is True
|
|
147
|
-
assert "fk_project_id = 2" in result.sql
|
|
148
|
-
assert "pk_site_id = 1" in result.sql
|
|
149
|
-
assert result.retries >= 1
|
|
150
|
-
assert result.ground_truth is not None
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|