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.
Files changed (99) hide show
  1. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/PKG-INFO +1 -1
  2. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/__init__.py +1 -1
  3. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/agent/generator.py +2 -2
  4. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/agent/pipeline.py +1 -1
  5. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/cli.py +3 -3
  6. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/connectors/base.py +2 -2
  7. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent_cli.egg-info/PKG-INFO +1 -1
  8. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/pyproject.toml +1 -1
  9. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_direct_sql.py +1 -1
  10. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_generator.py +1 -1
  11. dbagent_cli-0.11.4/tests/test_ground_truth_sql_validation.py +150 -0
  12. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_knowledge.py +8 -8
  13. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_prompt_queue.py +2 -2
  14. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_prompt_queue_and_sql.py +7 -7
  15. dbagent_cli-0.11.3/tests/test_ground_truth_sql_validation.py +0 -150
  16. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/README.md +0 -0
  17. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/agent/confidence.py +0 -0
  18. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/agent/conversation.py +0 -0
  19. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/agent/doctor.py +0 -0
  20. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/agent/error_classifier.py +0 -0
  21. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/agent/recommender.py +0 -0
  22. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/agent/validator.py +0 -0
  23. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/benchmark.py +0 -0
  24. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/config.py +0 -0
  25. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/connectors/factory.py +0 -0
  26. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/connectors/mongo.py +0 -0
  27. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/connectors/relational.py +0 -0
  28. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/errors.py +0 -0
  29. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/execution/__init__.py +0 -0
  30. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/execution/cache.py +0 -0
  31. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/execution/changeset.py +0 -0
  32. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/execution/cost.py +0 -0
  33. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/execution/prompt_queue.py +0 -0
  34. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/execution/recovery.py +0 -0
  35. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/execution/sandbox.py +0 -0
  36. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/execution/streaming.py +0 -0
  37. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/llm/anthropic_provider.py +0 -0
  38. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/llm/auto_setup.py +0 -0
  39. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/llm/base.py +0 -0
  40. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/llm/deepseek_provider.py +0 -0
  41. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/llm/factory.py +0 -0
  42. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/llm/gemini_provider.py +0 -0
  43. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/llm/groq_provider.py +0 -0
  44. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/llm/mistral_provider.py +0 -0
  45. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/llm/mock_provider.py +0 -0
  46. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/llm/ollama_provider.py +0 -0
  47. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/llm/openai_provider.py +0 -0
  48. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/llm/openrouter_provider.py +0 -0
  49. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/plugins.py +0 -0
  50. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/sample_db.py +0 -0
  51. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/schema/cache.py +0 -0
  52. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/schema/formatter.py +0 -0
  53. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/schema/graph.py +0 -0
  54. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/schema/ground_truth.py +0 -0
  55. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/schema/models.py +0 -0
  56. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/schema/retriever.py +0 -0
  57. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/schema/selector.py +0 -0
  58. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/security/__init__.py +0 -0
  59. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/security/audit.py +0 -0
  60. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/security/environment.py +0 -0
  61. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/security/masking.py +0 -0
  62. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/security/modes.py +0 -0
  63. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/security/rbac.py +0 -0
  64. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/semantic/__init__.py +0 -0
  65. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/semantic/dictionary.py +0 -0
  66. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/semantic/knowledge.py +0 -0
  67. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/semantic/templates.py +0 -0
  68. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/sql/__init__.py +0 -0
  69. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/sql/optimizer.py +0 -0
  70. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/sql/validator.py +0 -0
  71. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/ui/chat_screen.py +0 -0
  72. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/ui/console.py +0 -0
  73. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/ui/grid_window.py +0 -0
  74. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/ui/live_streamer.py +0 -0
  75. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/ui/viewer.py +0 -0
  76. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent/ui/visualizer.py +0 -0
  77. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent_cli.egg-info/SOURCES.txt +0 -0
  78. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent_cli.egg-info/dependency_links.txt +0 -0
  79. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent_cli.egg-info/entry_points.txt +0 -0
  80. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent_cli.egg-info/requires.txt +0 -0
  81. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/dbagent_cli.egg-info/top_level.txt +0 -0
  82. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/setup.cfg +0 -0
  83. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_benchmarks.py +0 -0
  84. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_cache.py +0 -0
  85. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_cli.py +0 -0
  86. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_connectors.py +0 -0
  87. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_demo_and_sample_db.py +0 -0
  88. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_doctor_and_viz.py +0 -0
  89. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_formats_and_exit_codes.py +0 -0
  90. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_graph.py +0 -0
  91. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_multi_turn_conversation.py +0 -0
  92. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_plugins.py +0 -0
  93. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_providers.py +0 -0
  94. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_recovery.py +0 -0
  95. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_safety.py +0 -0
  96. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_sandbox.py +0 -0
  97. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_schema.py +0 -0
  98. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_security.py +0 -0
  99. {dbagent_cli-0.11.3 → dbagent_cli-0.11.4}/tests/test_semantic.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: dbagent-cli
3
- Version: 0.11.3
3
+ Version: 0.11.4
4
4
  Summary: The Privacy-First AI Database Copilot, Schema Graph Introspector & Post-Commit Recovery Engine (CLI)
5
5
  Author: Santhosh Gupta
6
6
  License: MIT
@@ -2,5 +2,5 @@
2
2
  DB-Agent: Universal Database Introspection and Script Generation AI Agent (CLI).
3
3
  """
4
4
 
5
- __version__ = "0.11.3"
5
+ __version__ = "0.11.4"
6
6
  __author__ = "Santhosh Gupta"
@@ -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 pk_user_id FROM ivis_user WHERE ...), del_act AS (DELETE FROM user_activity WHERE user_id IN (SELECT pk_user_id FROM target)), del_map AS (DELETE FROM ivis_user_access WHERE fk_user_id IN (SELECT pk_user_id FROM target)) DELETE FROM ivis_user WHERE pk_user_id IN (SELECT pk_user_id FROM target);`).
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 ivis_scope", "i mean table_x")
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'^(?:ivis_|tbl_|table_|app_)', '', t.lower())
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 ivis_users)")
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 ivis_users)")
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 site_user_map = ivis_site maps to ivis_user via site_group -> group)")
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. "ivis_scope" -> "scope", "scopes")
91
- unprefixed = re.sub(r'^(?:ivis_|tbl_|table_|app_)', '', t_lower)
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
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: dbagent-cli
3
- Version: 0.11.3
3
+ Version: 0.11.4
4
4
  Summary: The Privacy-First AI Database Copilot, Schema Graph Introspector & Post-Commit Recovery Engine (CLI)
5
5
  Author: Santhosh Gupta
6
6
  License: MIT
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "dbagent-cli"
7
- version = "0.11.3"
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 pk_user_id, username from ivis_user where active = 'Y'") is True
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 that how ivis_user and ivis_tenant was related and how they were mapped and how to know that which sites in ivis_site were mapped to which user"
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="Site to User Group Mapping",
21
- rule_description="ivis_site maps to ivis_user through ivis_site_group -> ivis_group -> ivis_user_group_map",
22
- relevant_tables=["ivis_site", "ivis_site_group", "ivis_group", "ivis_user_group_map", "ivis_user"],
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 == "Site to User Group Mapping"
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 sites to user")
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 == "Site to User Group Mapping"
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 sites to user")
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 "Site to User Group Mapping" in ctx
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 ivis_user and ivis_site were linked ; give the flow diagram ; count total users"
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 ivis_user and ivis_site were linked"
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 ivis_customer_user_map WHERE fk_user_id = 5201;
15
- DELETE FROM ivis_user_access WHERE fk_user_id = 5201;
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 ivis_user WHERE pk_user_id = 5201;
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 pk_user_id FROM ivis_user WHERE username = 'target_user'
30
+ SELECT id FROM users WHERE username = 'target_user'
31
31
  ),
32
32
  del_access AS (
33
- DELETE FROM ivis_user_access WHERE fk_user_id IN (SELECT pk_user_id FROM target_user)
33
+ DELETE FROM user_access WHERE user_id IN (SELECT id FROM target_user)
34
34
  )
35
- DELETE FROM ivis_user WHERE pk_user_id IN (SELECT pk_user_id FROM target_user);
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 pk_user_id, username FROM ivis_user WHERE active = 'Y'
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