Skip to content

Guard LLM-generated sql_query rules#1275

Merged
mwojtyczka merged 7 commits into
databrickslabs:mainfrom
aarushisingh04:fix/guard-llm-sql-rules-generation
Jun 29, 2026
Merged

Guard LLM-generated sql_query rules#1275
mwojtyczka merged 7 commits into
databrickslabs:mainfrom
aarushisingh04:fix/guard-llm-sql-rules-generation

Conversation

@aarushisingh04

@aarushisingh04 aarushisingh04 commented Jun 25, 2026

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.

  • fixed training example for ai-assisted rules generation

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
Comment thread src/databricks/labs/dqx/llm/llm_core.py Outdated
Comment thread src/databricks/labs/dqx/llm/llm_core.py Outdated
Comment thread tests/unit/test_llm_core.py Outdated

@mwojtyczka mwojtyczka left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for the PR. Generally looking good. Left some comments to address.

@mwojtyczka mwojtyczka added the under-review This PR is currently being reviewed by one of DQX maintainers. label Jun 26, 2026
@aarushisingh04

Copy link
Copy Markdown
Contributor Author

@mwojtyczka

appreciate the thorough review! addressed all three points in the latest push:

  • added isinstance(query, str) guard, non-string query is now dropped with a warning instead of crashing
  • refactored _filter_unsafe_sql_rules to accept a pre-parsed list, eliminating the double-parse
  • added missing tests: non-string query, non-dict arguments, DspyRuleGenerationWithSchemaInference path, and empty/None quality_rules at the integration level

@mwojtyczka mwojtyczka left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@mwojtyczka mwojtyczka added Approved to Merge When PR is reviewed and approved. To be merged once all tests pass and removed under-review This PR is currently being reviewed by one of DQX maintainers. labels Jun 29, 2026
@mwojtyczka mwojtyczka changed the title feat(llm): guard LLM-generated sql_query rules with is_sql_query_safe Guard LLM-generated sql_query rules with is_sql_query_safe Jun 29, 2026
@mwojtyczka mwojtyczka changed the title Guard LLM-generated sql_query rules with is_sql_query_safe Guard LLM-generated sql_query rules Jun 29, 2026
@mwojtyczka mwojtyczka merged commit d1d2a9b into databrickslabs:main Jun 29, 2026
35 of 36 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Approved to Merge When PR is reviewed and approved. To be merged once all tests pass

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

2 participants