diff --git a/keep/api/bl/correlation_rules_provisioning.py b/keep/api/bl/correlation_rules_provisioning.py new file mode 100644 index 0000000000..9df1321d38 --- /dev/null +++ b/keep/api/bl/correlation_rules_provisioning.py @@ -0,0 +1,210 @@ +import json +import logging +import re + +import celpy + +import keep.api.core.db as db +from keep.api.core.config import config +from keep.api.models.db.rule import CreateIncidentOn, ResolveOn + +logger = logging.getLogger(__name__) + +KEEP_CORRELATION_RULES_ENV_VAR = "KEEP_CORRELATION_RULES" +SYSTEM_ACTOR = "system" + +_REQUIRED_FIELDS = ("ruleName", "celQuery", "sqlQuery", "timeframeInSeconds", "timeUnit") +_CEL_ENV = celpy.Environment() + + +def provision_correlation_rules_from_env(tenant_id: str): + """Provision correlation (rules-engine) rules from ``KEEP_CORRELATION_RULES``. + + Mirrors ``provision_deduplication_rules_from_env``: the env var holds a JSON + array (or a path to a ``.json`` file) of rule specs matching the ``POST /rules`` + schema. On every startup the DB is reconciled to the env: + + - Provisioned rules whose ``ruleName`` is no longer in the env are deleted. + - A provisioned rule with a matching name is updated in place. + - A new provisioned rule is created (``is_provisioned=True``). + - UI-created rules (``is_provisioned=False``) are never touched. + + If the env var is unset, any currently-provisioned rules are deprovisioned. + A spec that fails validation (missing field or an unparseable ``celQuery``) + is logged and skipped so one bad rule does not block the others. + """ + rules_to_provision = get_correlation_rules_to_provision() + + provisioned_rules = [rule for rule in db.get_rules(tenant_id) if rule.is_provisioned] + + if not rules_to_provision: + if provisioned_rules: + logger.info( + "%s unset; deprovisioning %d existing correlation rule(s)", + KEEP_CORRELATION_RULES_ENV_VAR, + len(provisioned_rules), + ) + for rule in provisioned_rules: + db.delete_rule(tenant_id=tenant_id, rule_id=str(rule.id)) + else: + logger.info( + "No correlation rules to provision and none currently provisioned" + ) + return + + provisioned_rules_by_name = {rule.name: rule for rule in provisioned_rules} + + # delete provisioned rules that are no longer in the env + for rule in provisioned_rules: + if str(rule.name) not in rules_to_provision: + logger.info( + "Correlation rule with name '%s' is not in the env, deleting from DB", + rule.name, + ) + db.delete_rule(tenant_id=tenant_id, rule_id=str(rule.id)) + + for rule_name, rule_to_provision in rules_to_provision.items(): + try: + _validate_rule_to_provision(rule_name, rule_to_provision) + except ValueError as exc: + logger.warning( + "Skipping invalid correlation rule '%s': %s", rule_name, exc + ) + continue + + existing_rule = provisioned_rules_by_name.get(rule_name) + if existing_rule is not None: + logger.info( + "Correlation rule with name '%s' already exists, updating in DB", + rule_name, + ) + _update_correlation_rule( + tenant_id, str(existing_rule.id), rule_name, rule_to_provision + ) + continue + + logger.info( + "Correlation rule with name '%s' does not exist, creating in DB", + rule_name, + ) + _create_correlation_rule(tenant_id, rule_name, rule_to_provision) + + +def get_correlation_rules_to_provision() -> dict[str, dict]: + """Read correlation rules from ``KEEP_CORRELATION_RULES`` as a dict keyed by name. + + The env var is either an absolute/relative path to a ``.json`` file or a JSON + string. Its content is a JSON array of rule specs matching the ``POST /rules`` + schema (a single object is also accepted). Specs without a ``ruleName`` are + skipped. + """ + rules_from_env_var = config(key=KEEP_CORRELATION_RULES_ENV_VAR, default=None) + + if not rules_from_env_var: + return None + + if re.compile(r"^(\/|\.\/|\.\.\/).*\.json$").match(rules_from_env_var): + with open(file=rules_from_env_var, mode="r", encoding="utf8") as file: + try: + parsed = json.loads(file.read()) + except json.JSONDecodeError as e: + raise Exception( + f"Error parsing correlation rules from file {rules_from_env_var}: {e}" + ) from e + else: + try: + parsed = json.loads(rules_from_env_var) + except json.JSONDecodeError as e: + raise Exception( + f"Error parsing correlation rules from env var {KEEP_CORRELATION_RULES_ENV_VAR}: {e}" + ) from e + + if isinstance(parsed, dict): + parsed = [parsed] + if not isinstance(parsed, list): + raise Exception( + f"{KEEP_CORRELATION_RULES_ENV_VAR} must be a JSON array of rule specs" + ) + + rules_dict: dict[str, dict] = {} + for rule in parsed: + if not isinstance(rule, dict): + logger.warning("Skipping non-object correlation rule spec: %r", rule) + continue + rule_name = rule.get("ruleName") + if not rule_name: + logger.warning( + "Skipping correlation rule spec without 'ruleName': %r", rule + ) + continue + rules_dict[rule_name] = rule + + return rules_dict or None + + +def _validate_rule_to_provision(rule_name: str, rule_to_provision: dict) -> None: + for field in _REQUIRED_FIELDS: + if not rule_to_provision.get(field): + raise ValueError(f"missing required field '{field}'") + + sql_query = rule_to_provision.get("sqlQuery") + if not isinstance(sql_query, dict) or not sql_query.get("sql"): + raise ValueError("'sqlQuery.sql' is required") + + try: + _CEL_ENV.compile(rule_to_provision["celQuery"]) + except Exception as exc: + raise ValueError(f"unparseable celQuery: {exc}") from exc + + +def _create_correlation_rule( + tenant_id: str, rule_name: str, rule_to_provision: dict +) -> None: + sql_query = rule_to_provision["sqlQuery"] + db.create_rule( + tenant_id=tenant_id, + name=rule_name, + definition={"sql": sql_query.get("sql"), "params": sql_query.get("params")}, + timeframe=rule_to_provision["timeframeInSeconds"], + timeunit=rule_to_provision["timeUnit"], + definition_cel=rule_to_provision["celQuery"], + created_by=SYSTEM_ACTOR, + grouping_criteria=rule_to_provision.get("groupingCriteria", []), + group_description=rule_to_provision.get("groupDescription"), + require_approve=rule_to_provision.get("requireApprove", False), + resolve_on=rule_to_provision.get("resolveOn", ResolveOn.NEVER.value), + create_on=rule_to_provision.get("createOn", CreateIncidentOn.ANY.value), + incident_name_template=rule_to_provision.get("incidentNameTemplate"), + incident_prefix=rule_to_provision.get("incidentPrefix"), + multi_level=rule_to_provision.get("multiLevel", False), + multi_level_property_name=rule_to_provision.get("multiLevelPropertyName"), + threshold=rule_to_provision.get("threshold", 1), + assignee=rule_to_provision.get("assignee"), + is_provisioned=True, + ) + + +def _update_correlation_rule( + tenant_id: str, rule_id: str, rule_name: str, rule_to_provision: dict +) -> None: + sql_query = rule_to_provision["sqlQuery"] + db.update_rule( + tenant_id=tenant_id, + rule_id=rule_id, + name=rule_name, + definition={"sql": sql_query.get("sql"), "params": sql_query.get("params")}, + timeframe=rule_to_provision["timeframeInSeconds"], + timeunit=rule_to_provision["timeUnit"], + definition_cel=rule_to_provision["celQuery"], + updated_by=SYSTEM_ACTOR, + grouping_criteria=rule_to_provision.get("groupingCriteria", []), + require_approve=rule_to_provision.get("requireApprove", False), + resolve_on=rule_to_provision.get("resolveOn", ResolveOn.NEVER.value), + create_on=rule_to_provision.get("createOn", CreateIncidentOn.ANY.value), + incident_name_template=rule_to_provision.get("incidentNameTemplate"), + incident_prefix=rule_to_provision.get("incidentPrefix"), + multi_level=rule_to_provision.get("multiLevel", False), + multi_level_property_name=rule_to_provision.get("multiLevelPropertyName"), + threshold=rule_to_provision.get("threshold", 1), + assignee=rule_to_provision.get("assignee"), + ) diff --git a/keep/api/config.py b/keep/api/config.py index d1312297d4..f4f95ea7a0 100644 --- a/keep/api/config.py +++ b/keep/api/config.py @@ -6,6 +6,9 @@ provision_deduplication_rules_from_env, ) from keep.api.api import AUTH_TYPE +from keep.api.bl.correlation_rules_provisioning import ( + provision_correlation_rules_from_env, +) from keep.api.bl.mapping_rules_provisioning import provision_mapping_rules_from_env from keep.api.core.db_on_start import migrate_db, try_create_single_tenant from keep.api.core.dependencies import SINGLE_TENANT_UUID @@ -40,6 +43,9 @@ def provision_resources(): logger.info("Provisioning mapping rules") provision_mapping_rules_from_env(SINGLE_TENANT_UUID) logger.info("Mapping rules provisioned successfully") + logger.info("Provisioning correlation rules") + provision_correlation_rules_from_env(SINGLE_TENANT_UUID) + logger.info("Correlation rules provisioned successfully") else: logger.info("Provisioning resources is disabled") diff --git a/keep/api/core/db.py b/keep/api/core/db.py index 878d0d877b..7bf18fcf6e 100644 --- a/keep/api/core/db.py +++ b/keep/api/core/db.py @@ -2241,6 +2241,7 @@ def create_rule( multi_level_property_name=None, threshold=1, assignee=None, + is_provisioned=False, ): grouping_criteria = grouping_criteria or [] with Session(engine) as session: @@ -2264,6 +2265,7 @@ def create_rule( multi_level_property_name=multi_level_property_name, threshold=threshold, assignee=assignee, + is_provisioned=is_provisioned, ) session.add(rule) session.commit() diff --git a/keep/api/models/db/migrations/versions/2026-07-12-09-30_b3f1c9a24d70.py b/keep/api/models/db/migrations/versions/2026-07-12-09-30_b3f1c9a24d70.py new file mode 100644 index 0000000000..b479bb859c --- /dev/null +++ b/keep/api/models/db/migrations/versions/2026-07-12-09-30_b3f1c9a24d70.py @@ -0,0 +1,33 @@ +"""feat: add is_provisioned to Rule + +Revision ID: b3f1c9a24d70 +Revises: 67ff7efffed4 +Create Date: 2026-07-12 09:30:00.000000 + +""" + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision = "b3f1c9a24d70" +down_revision = "67ff7efffed4" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + with op.batch_alter_table("rule", schema=None) as batch_op: + batch_op.add_column( + sa.Column( + "is_provisioned", + sa.Boolean(), + nullable=False, + server_default=sa.false(), + ) + ) + + +def downgrade() -> None: + with op.batch_alter_table("rule", schema=None) as batch_op: + batch_op.drop_column("is_provisioned") diff --git a/keep/api/models/db/rule.py b/keep/api/models/db/rule.py index 2669571bf5..f12c672d69 100644 --- a/keep/api/models/db/rule.py +++ b/keep/api/models/db/rule.py @@ -52,6 +52,7 @@ class Rule(SQLModel, table=True): resolve_on: str = ResolveOn.NEVER.value create_on: str = CreateIncidentOn.ANY.value is_deleted: bool = False + is_provisioned: bool = Field(default=False) incident_name_template: str = None incident_prefix: str | None = None multi_level: bool = False diff --git a/tests/test_correlation_rules_provisioning.py b/tests/test_correlation_rules_provisioning.py new file mode 100644 index 0000000000..d1a150845d --- /dev/null +++ b/tests/test_correlation_rules_provisioning.py @@ -0,0 +1,216 @@ +"""Tests for keep.api.bl.correlation_rules_provisioning.provision_correlation_rules_from_env. + +Mirrors tests/test_mapping_rules_provisioning.py by using the real in-memory SQLite +`db_session` fixture from tests/conftest.py rather than patching DB helpers, and +drives the provisioner through the KEEP_CORRELATION_RULES env var (a JSON array of +specs matching the POST /rules schema), the same way deduplication rules are +provisioned from env. +""" + +import json + +from sqlmodel import Session, select + +import keep.api.core.db as db +from keep.api.bl.correlation_rules_provisioning import ( + provision_correlation_rules_from_env, +) +from keep.api.core.dependencies import SINGLE_TENANT_UUID +from keep.api.models.db.rule import Rule + +ENV_VAR = "KEEP_CORRELATION_RULES" + + +def _spec(name, cel_query='severity == "critical"', **overrides): + spec = { + "ruleName": name, + "celQuery": cel_query, + "sqlQuery": {"sql": "severity = :severity", "params": {"severity": "critical"}}, + "timeframeInSeconds": 600, + "timeUnit": "seconds", + "groupingCriteria": [], + "createOn": "any", + "resolveOn": "never", + "threshold": 1, + } + spec.update(overrides) + return spec + + +def _set_env(monkeypatch, *specs): + monkeypatch.setenv(ENV_VAR, json.dumps(list(specs))) + + +def _provisioned_rules(): + return [rule for rule in db.get_rules(SINGLE_TENANT_UUID) if rule.is_provisioned] + + +def _all_rules(): + with Session(db.engine) as session: + return session.exec( + select(Rule).where(Rule.tenant_id == SINGLE_TENANT_UUID) + ).all() + + +def test_creates_new_rule(monkeypatch, db_session): + """Empty DB + one spec in env -> rule created and marked provisioned.""" + _set_env(monkeypatch, _spec("critical-alerts-correlation")) + + provision_correlation_rules_from_env(SINGLE_TENANT_UUID) + + rules = _provisioned_rules() + assert len(rules) == 1 + rule = rules[0] + assert rule.name == "critical-alerts-correlation" + assert rule.is_provisioned is True + assert rule.definition_cel == 'severity == "critical"' + assert rule.definition == { + "sql": "severity = :severity", + "params": {"severity": "critical"}, + } + assert rule.timeframe == 600 + assert rule.created_by == "system" + + +def test_provisions_multiple_rules(monkeypatch, db_session): + _set_env(monkeypatch, _spec("rule-a"), _spec("rule-b")) + + provision_correlation_rules_from_env(SINGLE_TENANT_UUID) + + names = sorted(rule.name for rule in _provisioned_rules()) + assert names == ["rule-a", "rule-b"] + + +def test_is_idempotent(monkeypatch, db_session): + """Running provisioning twice does not create duplicates.""" + _set_env(monkeypatch, _spec("rule-a"), _spec("rule-b")) + + provision_correlation_rules_from_env(SINGLE_TENANT_UUID) + first_ids = sorted(str(rule.id) for rule in _provisioned_rules()) + assert len(first_ids) == 2 + + provision_correlation_rules_from_env(SINGLE_TENANT_UUID) + second_ids = sorted(str(rule.id) for rule in _provisioned_rules()) + + assert first_ids == second_ids + + +def test_updates_existing_provisioned_rule(monkeypatch, db_session): + """A previously-provisioned rule gets its content refreshed from the env, id preserved.""" + _set_env(monkeypatch, _spec("rule-a", timeframeInSeconds=600)) + provision_correlation_rules_from_env(SINGLE_TENANT_UUID) + original = _provisioned_rules()[0] + original_id = original.id + + # Simulate drift (someone edited the rule directly) + with Session(db.engine) as session: + rule = session.exec(select(Rule).where(Rule.id == original_id)).first() + rule.timeframe = 42 + session.add(rule) + session.commit() + + # Re-provision with the env value -> timeframe reset to 600 + _set_env(monkeypatch, _spec("rule-a", timeframeInSeconds=600)) + provision_correlation_rules_from_env(SINGLE_TENANT_UUID) + + refreshed = _provisioned_rules() + assert len(refreshed) == 1 + assert refreshed[0].id == original_id + assert refreshed[0].timeframe == 600 + + +def test_deprovisions_rule_removed_from_env(monkeypatch, db_session): + _set_env(monkeypatch, _spec("rule-a"), _spec("rule-b")) + provision_correlation_rules_from_env(SINGLE_TENANT_UUID) + assert len(_provisioned_rules()) == 2 + + _set_env(monkeypatch, _spec("rule-a")) + provision_correlation_rules_from_env(SINGLE_TENANT_UUID) + + remaining = _provisioned_rules() + assert len(remaining) == 1 + assert remaining[0].name == "rule-a" + + +def test_deprovisions_all_when_env_unset(monkeypatch, db_session): + _set_env(monkeypatch, _spec("rule-a"), _spec("rule-b")) + provision_correlation_rules_from_env(SINGLE_TENANT_UUID) + assert len(_provisioned_rules()) == 2 + + monkeypatch.delenv(ENV_VAR) + provision_correlation_rules_from_env(SINGLE_TENANT_UUID) + + assert len(_provisioned_rules()) == 0 + + +def test_leaves_ui_rules_untouched(monkeypatch, db_session): + """A UI-created rule (is_provisioned=False) is never deleted or modified.""" + ui_rule = db.create_rule( + tenant_id=SINGLE_TENANT_UUID, + name="ui-only-rule", + definition={"sql": "1=1", "params": {}}, + timeframe=300, + timeunit="seconds", + definition_cel='source == "grafana"', + created_by="ui-user@example.com", + ) + ui_rule_id = ui_rule.id + + _set_env(monkeypatch, _spec("provisioned-rule")) + provision_correlation_rules_from_env(SINGLE_TENANT_UUID) + + all_rules = _all_rules() + assert len(all_rules) == 2 + ui = next(rule for rule in all_rules if rule.id == ui_rule_id) + assert ui.is_provisioned is False + assert ui.timeframe == 300 + + +def test_invalid_cel_is_skipped(monkeypatch, db_session): + """A spec with an unparseable celQuery is skipped; valid specs still provision.""" + _set_env( + monkeypatch, + _spec("valid-rule"), + _spec("broken-rule", cel_query="severity == "), + ) + + provision_correlation_rules_from_env(SINGLE_TENANT_UUID) + + rules = _provisioned_rules() + assert len(rules) == 1 + assert rules[0].name == "valid-rule" + + +def test_missing_required_field_is_skipped(monkeypatch, db_session): + """A spec missing a required field is skipped; valid specs still provision.""" + incomplete = _spec("incomplete-rule") + del incomplete["timeframeInSeconds"] + _set_env(monkeypatch, _spec("valid-rule"), incomplete) + + provision_correlation_rules_from_env(SINGLE_TENANT_UUID) + + rules = _provisioned_rules() + assert len(rules) == 1 + assert rules[0].name == "valid-rule" + + +def test_reads_from_json_file_path(monkeypatch, tmp_path, db_session): + """KEEP_CORRELATION_RULES pointing at a .json file is loaded from disk.""" + path = tmp_path / "correlation_rules.json" + path.write_text(json.dumps([_spec("from-file-rule")])) + monkeypatch.setenv(ENV_VAR, str(path)) + + provision_correlation_rules_from_env(SINGLE_TENANT_UUID) + + rules = _provisioned_rules() + assert len(rules) == 1 + assert rules[0].name == "from-file-rule" + + +def test_noop_when_env_unset_and_no_provisioned_rules(monkeypatch, db_session): + monkeypatch.delenv(ENV_VAR, raising=False) + assert len(_all_rules()) == 0 + + provision_correlation_rules_from_env(SINGLE_TENANT_UUID) + + assert len(_all_rules()) == 0