Guard LLM-generated sql_query rules

src/databricks/labs/dqx/llm/llm_core.py

@@ -7,12 +7,59 @@
 
 from databricks.labs.dqx.config import LLMModelConfig
 from databricks.labs.dqx.llm.llm_utils import create_optimizer_training_set, create_optimizer_training_set_with_stats
+from databricks.labs.dqx.utils import is_sql_query_safe
 from databricks.labs.dqx.llm.optimizers import BootstrapFewShotOptimizer
 from databricks.labs.dqx.llm.validators import RuleValidator
 
 logger = logging.getLogger(__name__)
 
 
+def _filter_unsafe_sql_rules(rules: object) -> list[dict]:
+    """Filter out any sql_query rules whose query argument fails the SQL safety check.
+
+    Accepts a pre-parsed rules value and drops every rule where *check.function* is
+    *sql_query* and the *query* argument is not a string or does not pass
+    *is_sql_query_safe*. All other rules are returned unchanged. Non-list input
+    returns an empty list.
+
+    Args:
+        rules: Parsed rules value (expected to be a list of rule dicts).
+
+    Returns:
+        List of safe rules with unsafe sql_query rules removed.
+    """
+    if not isinstance(rules, list):
+        return []
+
+    safe_rules: list[dict] = []
+    for rule in rules:
+        if not isinstance(rule, dict):
+            continue
+
+        check = rule.get("check", {})
+        if not isinstance(check, dict):
+            continue
+
+        if check.get("function") == "sql_query":
+            arguments = check.get("arguments", {})
+            if not isinstance(arguments, dict):
+                continue
+
+            query = arguments.get("query")
+            if not isinstance(query, str):
+                logger.warning(f"LLM-generated sql_query rule dropped: non-string query argument {query!r}")
+                continue
+            if not is_sql_query_safe(query):
+                # Strip control characters before logging to prevent log injection (CWE-117)
+                safe_repr = repr(query.replace("\n", " ").replace("\r", " "))
+                logger.warning(f"LLM-generated sql_query rule dropped: unsafe SQL query {safe_repr}")
+                continue
+
+        safe_rules.append(rule)
+
+    return safe_rules
+
+
 class LLMModelConfigurator:
     """
     Configures DSPy language models.
@@ -147,13 +194,15 @@ def forward(
             schema_info=schema_info, business_description=business_description, available_functions=available_functions
         )
 
-        # Validate JSON output
+        # Validate JSON output and filter unsafe sql_query rules
         if result.quality_rules:
             try:
-                json.loads(result.quality_rules)
+                parsed = json.loads(result.quality_rules)
             except json.JSONDecodeError as e:
                 logger.warning(f"Generated invalid JSON: {e}. Returning empty rules.")
                 result.quality_rules = "[]"
+            else:
+                result.quality_rules = json.dumps(_filter_unsafe_sql_rules(parsed))
 
         return result
 
@@ -275,13 +324,15 @@ def forward(
             business_description=business_description or "",
         )
 
-        # Validate JSON output
+        # Validate JSON output and filter unsafe sql_query rules
         if result.quality_rules:
             try:
-                json.loads(result.quality_rules)
+                parsed = json.loads(result.quality_rules)
             except json.JSONDecodeError as e:
                 logger.warning(f"Generated invalid JSON: {e}. Returning empty rules.")
                 result.quality_rules = "[]"
+            else:
+                result.quality_rules = json.dumps(_filter_unsafe_sql_rules(parsed))
 
         return result
 

src/databricks/labs/dqx/llm/resources/training_examples_with_stats.yml

@@ -104,7 +104,7 @@
       mean: null
       min: "active"
       max: "suspended"
-  quality_rules: '[{"criticality":"warn","check":{"function":"is_not_null_and_not_empty","arguments":{"column":"email"}}},{"criticality":"error","check":{"function":"is_email","arguments":{"column":"email"}}}]'
+  quality_rules: '[{"criticality":"warn","check":{"function":"is_not_null_and_not_empty","arguments":{"column":"email"}}},{"criticality":"error","check":{"function":"is_valid_email","arguments":{"column":"email"}}}]'
   reasoning: "Email should be present for all user accounts and follow valid format. Combining completeness and format validation ensures high-quality contact information."
 
 - name: "order_total_within_daily_average_sql_query"

tests/conftest.py

@@ -899,13 +899,25 @@ def _lakebase_create_catalog(
 
 @retried(on=[BadRequest, TooManyRequests, RequestLimitExceeded], timeout=timedelta(minutes=2))
 def _lakebase_delete_catalog(workspace: WorkspaceClient, catalog_name: str) -> None:
-    """Delete database catalog; retries on rate limits."""
+    """Delete the database catalog and ensure its Unity Catalog catalog is removed; retries on limits.
+
+    delete_database_catalog tears down the database federation, but the UC catalog object - which is
+    what counts toward the per-metastore catalog limit (1000) - is removed via catalogs.delete. We
+    call both so the metastore slot is freed regardless of what delete_database_catalog leaves behind.
+    Both are NotFound-safe so a prior/partial deletion is tolerated.
+    """
     try:
         workspace.database.delete_database_catalog(name=catalog_name)
-        logger.info(f"Successfully deleted database catalog: {catalog_name}")
+        logger.info(f"Deleted database catalog: {catalog_name}")
     except NotFound:
         logger.info(f"Database catalog {catalog_name} not found (already deleted)")
 
+    try:
+        workspace.catalogs.delete(name=catalog_name, force=True)
+        logger.info(f"Deleted Unity Catalog catalog backing database catalog: {catalog_name}")
+    except NotFound:
+        logger.info(f"Unity Catalog catalog {catalog_name} not found (already deleted)")
+
 
 @retried(on=[BadRequest, TooManyRequests, RequestLimitExceeded], timeout=timedelta(minutes=2))
 def _lakebase_delete_database_instance(workspace: WorkspaceClient, instance_name: str) -> None:
@@ -932,8 +944,13 @@ def create() -> LakebaseInstance:
         return LakebaseInstance(name=instance_name, catalog_name=catalog_name, database_name=database_name)
 
     def delete(instance: LakebaseInstance) -> None:
-        _lakebase_delete_catalog(ws, instance.catalog_name)
-        _lakebase_delete_database_instance(ws, instance.name)
+        # Always attempt instance deletion even if catalog deletion fails, otherwise a failed
+        # catalog delete would orphan the instance. Catalogs count toward the metastore limit,
+        # so deleting the catalog first is intentional.
+        try:
+            _lakebase_delete_catalog(ws, instance.catalog_name)
+        finally:
+            _lakebase_delete_database_instance(ws, instance.name)
 
     yield from factory("lakebase", create, delete)
 

tests/integration/test_ai_rules_generator.py

@@ -126,11 +126,11 @@ def not_ends_with_suffix(column: str, suffix: str):
 
     expected_checks = EXPECTED_CHECKS + [
         {
-            'check': {
-                'arguments': {'column': 'email', 'suffix': '@gmail.com'},
-                'function': 'not_ends_with_suffix',
+            "check": {
+                "arguments": {"column": "email", "suffix": "@gmail.com"},
+                "function": "not_ends_with_suffix",
             },
-            'criticality': 'error',
+            "criticality": "error",
         }
     ]
     assert actual_checks == expected_checks
@@ -208,6 +208,34 @@ def test_generate_dq_rules_ai_assisted_with_summary_stats_only(ws, spark):
     assert not DQEngineCore.validate_checks(actual_checks).has_errors
 
 
+def test_generate_dq_rules_ai_assisted_email_format_with_summary_stats(ws, spark):
+    """Email column from summary stats should yield a valid is_valid_email check.
+
+    Regression for the stats training example that referenced the non-existent
+    `is_email` function (corrected to `is_valid_email`). Previously the model
+    imitated the bad example and the rule was silently dropped by validation,
+    producing no email-format check at all.
+    """
+    user_input = "Email addresses must be present and follow a valid email format."
+    summary_stats = {
+        "email": {"mean": None, "min": "alice@example.com", "max": "zoe@example.com"},
+        "username": {"mean": None, "min": "alice", "max": "zoe"},
+        "account_status": {"mean": None, "min": "active", "max": "suspended"},
+    }
+
+    generator = DQGenerator(ws, spark)
+    actual_checks = generator.generate_dq_rules_ai_assisted(user_input=user_input, summary_stats=summary_stats)
+
+    assert not DQEngineCore.validate_checks(actual_checks).has_errors
+    email_format_checks = [
+        c
+        for c in actual_checks
+        if c.get("check", {}).get("function") == "is_valid_email"
+        and c.get("check", {}).get("arguments", {}).get("column") == "email"
+    ]
+    assert len(email_format_checks) == 1, f"Expected one is_valid_email check for email, got: {actual_checks}"
+
+
 def test_generate_dq_rules_ai_assisted_sql_query_with_custom_input_placeholder(ws, spark):
     user_input = "Each order's total must not exceed 10 times the average order total for that day. Use sql_query for the check function."
 

tests/integration/test_cleanup.py

@@ -1,7 +1,9 @@
 import logging
 import time
 
+from databricks.sdk import WorkspaceClient
 from databricks.sdk.errors import PermissionDenied
+from databricks.sdk.service.catalog import SchemaInfo
 
 from tests.constants import TEST_CATALOG
 
@@ -11,7 +13,7 @@
 MAX_AGE_MS = 24 * 60 * 60 * 1000  # only drop objects older than 1 day
 
 
-def test_remove_orhpaned_test_schemas(ws):
+def test_remove_orhpaned_test_schemas(ws: WorkspaceClient) -> None:
     """Maintenance sweep: drop every schema in the test catalog whose name starts with
     ``dummy_`` and that is older than one day, cascading the deletion to all contained
     tables and registered models.
@@ -32,6 +34,9 @@ def test_remove_orhpaned_test_schemas(ws):
 
     skipped: list[str] = []
     for schema in schemas:
+        # Listed schemas always carry these fields; the guard also narrows them for the SDK calls below.
+        if schema.full_name is None or schema.catalog_name is None or schema.name is None:
+            continue
         try:
             # Registered models are not removed by a forced schema delete, so drop them first.
             _drop_registered_models(ws, schema.catalog_name, schema.name)
@@ -42,15 +47,17 @@ def test_remove_orhpaned_test_schemas(ws):
             skipped.append(schema.full_name)
             logger.warning(f"Skipping schema {schema.full_name!r}: no permission to drop it ({exc})")
 
-    remaining = {schema.full_name for schema in _list_stale_dummy_schemas(ws, cutoff_ms)}
+    remaining = {
+        schema.full_name for schema in _list_stale_dummy_schemas(ws, cutoff_ms) if schema.full_name is not None
+    }
     not_permission_skipped = remaining - set(skipped)
     assert not not_permission_skipped, (
         f"Stale schemas with prefix {DUMMY_SCHEMA_PREFIX!r} were not dropped "
         f"(and not permission-skipped): {sorted(not_permission_skipped)}"
     )
 
 
-def _list_stale_dummy_schemas(ws, cutoff_ms: int):
+def _list_stale_dummy_schemas(ws: WorkspaceClient, cutoff_ms: int) -> list[SchemaInfo]:
     """Return ``dummy_``-prefixed schemas in the test catalog older than the cutoff."""
     return [
         schema
@@ -61,7 +68,7 @@ def _list_stale_dummy_schemas(ws, cutoff_ms: int):
     ]
 
 
-def _drop_registered_models(ws, catalog_name: str, schema_name: str) -> None:
+def _drop_registered_models(ws: WorkspaceClient, catalog_name: str, schema_name: str) -> None:
     """Delete every registered (Unity Catalog) model in the given schema.
 
     A forced schema delete drops tables and functions but leaves registered models in
@@ -71,7 +78,11 @@ def _drop_registered_models(ws, catalog_name: str, schema_name: str) -> None:
     each version must be deleted first.
     """
     for model in ws.registered_models.list(catalog_name=catalog_name, schema_name=schema_name):
+        if model.full_name is None:
+            continue
         for version in ws.model_versions.list(full_name=model.full_name):
+            if version.version is None:
+                continue
             ws.model_versions.delete(full_name=model.full_name, version=version.version)
             logger.info(f"Deleted model version {model.full_name!r} v{version.version}")
         ws.registered_models.delete(full_name=model.full_name)