Skip to content

feat(llm): guard LLM-generated sql_query rules with is_sql_query_safe#1275

Open
aarushisingh04 wants to merge 2 commits into
databrickslabs:mainfrom
aarushisingh04:fix/guard-llm-sql-rules-generation
Open

feat(llm): guard LLM-generated sql_query rules with is_sql_query_safe#1275
aarushisingh04 wants to merge 2 commits into
databrickslabs:mainfrom
aarushisingh04:fix/guard-llm-sql-rules-generation

Conversation

@aarushisingh04

Copy link
Copy Markdown
Contributor

Changes

adds SQL safety validation to both LLM-assisted rule generation paths so that any sql_query rule containing unsafe SQL (DML/DDL) is dropped before being returned to the caller.

  • added _filter_unsafe_sql_rules() in llm/llm_core.py : parses the generated JSON rules array and drops any rule whose check.function is sql_query and whose query argument fails is_sql_query_safe(), logging a warning with the sanitized query string (CWE-117)
  • hooked it into DspyRuleGeneration.forward() and DspyRuleUsingDataStats.forward() via an else clause on the existing JSON-validation block; DspyRuleGenerationWithSchemaInference inherits the fix automatically since it delegates to DspyRuleGeneration

note: max_tokens was already present in LLMModelConfig and forwarded to dspy.LM(), it caps completion length, not prompt input, which is sufficient for the issue's intent of bounding unbounded LLM output.

Linked issues

resolves #1123

Tests

  • manually tested
  • added unit tests
  • added integration tests
  • added end-to-end tests
  • added performance tests

19 unit tests in tests/unit/test_llm_core.py covering: safe SELECT queries pass through, unsafe DROP/DELETE/INSERT/TRUNCATE rules are dropped, non-sql_query rules are unaffected, mixed safe+unsafe batches, null/missing query argument, non-array JSON input, malformed rule entries, invalid JSON passthrough, and warning log emission.

Documentation and Demos

  • added/updated demos
  • added/updated docs
  • added/updated agent skills

@aarushisingh04 aarushisingh04 requested a review from a team as a code owner June 25, 2026 08:28
@aarushisingh04 aarushisingh04 requested review from tombonfert and removed request for a team June 25, 2026 08:28
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[FEATURE]: Guard results of AI-assisted rules generation

1 participant