sift-stack-py 0.3.2__py3-none-any.whl

sift_py/report_templates/_config_test.py

@@ -0,0 +1,34 @@
+import pytest
+from pydantic_core import ValidationError
+
+from sift_py.report_templates.config import ReportTemplateConfig
+
+
+@pytest.fixture
+def report_template_dict() -> dict:
+    return {
+        "name": "report-template",
+        "template_client_key": "template-client-key",
+    }
+
+
+def test_report_template_config(report_template_dict):
+    report_template_config = ReportTemplateConfig(**report_template_dict)
+    assert report_template_config.name == "report-template"
+    assert report_template_config.template_client_key == "template-client-key"
+    assert report_template_config.organization_id == ""
+    assert report_template_config.tags is None
+    assert report_template_config.description is None
+    assert report_template_config.rule_client_keys == []
+    assert report_template_config.archived_date is None
+
+
+def test_report_template_config_validation(report_template_dict):
+    report_template_dict.pop("name")
+    with pytest.raises(ValidationError, match="Field required"):
+        ReportTemplateConfig(**report_template_dict)
+
+    report_template_dict["name"] = "report-template"
+    report_template_dict.pop("template_client_key")
+    with pytest.raises(ValidationError, match="Field required"):
+        ReportTemplateConfig(**report_template_dict)

sift_py/report_templates/_service_test.py

@@ -0,0 +1,94 @@
+from unittest import mock
+
+import pytest
+from sift.report_templates.v1.report_templates_pb2 import ReportTemplate
+
+from sift_py._internal.test_util.channel import MockChannel
+from sift_py.report_templates.config import ReportTemplateConfig
+from sift_py.report_templates.service import ReportTemplateService
+
+
+@pytest.fixture
+def report_template_service():
+    return ReportTemplateService(MockChannel())
+
+
+def test_report_template_service_get_report_template_by_client_key(report_template_service):
+    report_template_client_key = "report-template-client-key"
+
+    with mock.patch.object(
+        ReportTemplateService,
+        "_get_report_template_by_client_key",
+        return_value=ReportTemplate(name="abc"),
+    ) as mock_get_report_template_by_client_key:
+        report_template_service.get_report_template(client_key=report_template_client_key)
+        mock_get_report_template_by_client_key.assert_called_once_with(report_template_client_key)
+
+
+def test_report_template_service_get_report_template_by_id(report_template_service):
+    report_template_id = "report-template-id"
+
+    with mock.patch.object(
+        ReportTemplateService, "_get_report_template_by_id", return_value=ReportTemplate(name="abc")
+    ) as mock_get_report_template_by_id:
+        report_template_service.get_report_template(id=report_template_id)
+        mock_get_report_template_by_id.assert_called_once_with(report_template_id)
+
+
+def test_report_template_service_get_report_template_missing_or_both_client_key_and_id(
+    report_template_service,
+):
+    with pytest.raises(ValueError, match="One of client_key or id must be provided"):
+        report_template_service.get_report_template()
+    with pytest.raises(ValueError, match="One of client_key or id must be provided"):
+        report_template_service.get_report_template(client_key="abc", id="abc")
+
+
+def test_report_template_service_create_report_template(report_template_service):
+    report_template_config = ReportTemplateConfig(
+        name="report-template",
+        template_client_key="template-client-key",
+    )
+
+    with mock.patch.object(
+        ReportTemplateService, "_create_report_template"
+    ) as mock_create_report_template:
+        report_template_service.create_or_update_report_template(report_template_config)
+        mock_create_report_template.assert_called_once_with(report_template_config)
+
+
+def test_report_template_service_update_report_template(report_template_service):
+    report_template_config = ReportTemplateConfig(
+        name="report-template",
+        template_client_key="template-client-key",
+    )
+
+    report_template_config_update = ReportTemplateConfig(
+        name="report-template-updated",
+        template_client_key="template-client-key",
+    )
+
+    with mock.patch.object(
+        ReportTemplateService, "_update_report_template"
+    ) as mock_update_report_template:
+        with mock.patch.object(
+            ReportTemplateService, "_get_report_template_by_client_key"
+        ) as mock_get_report_template_by_client_key:
+            mock_get_report_template_by_client_key.return_value = report_template_config
+            report_template_service.create_or_update_report_template(report_template_config_update)
+            mock_update_report_template.assert_called_once_with(
+                report_template_config_update, report_template_config
+            )
+
+
+def test_report_template_service_missing_template_client_key(report_template_service):
+    report_template_config = ReportTemplateConfig.model_construct(  # Without model validation
+        name="report-template",
+        template_client_key="",
+    )
+
+    with pytest.raises(Exception):
+        report_template_service.create_or_update_report_template(
+            report_template_config,
+            match="Report template report-template requires a template_client_key",
+        )

sift_py/report_templates/config.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Any, List, Optional
+
+from pydantic import BaseModel, ConfigDict
+
+from sift_py._internal.convert.json import AsJson
+
+
+class ReportTemplateConfig(BaseModel, AsJson):
+    """
+    Configuration for a report template.
+
+    - `name`: Name of the report template.
+    - `template_client_key`: Unique client key to identify the report template.
+    - `organization_id`: Organization ID that the report template belongs to.
+    - `tags`: Tags to associate with the report template.
+    - `description`: Description of the report template.
+    - `rule_client_keys`: List of rule client keys associated with the report template.
+    - `archived_date`: Date when the report template was archived. Setting this field
+        will archive the report template, and unsetting it will unarchive the report template.
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    name: str
+    template_client_key: str
+    organization_id: str = ""
+    tags: Optional[List[str]] = None
+    description: Optional[str] = None
+    rule_client_keys: List[str] = []
+    archived_date: Optional[datetime] = None
+
+    def as_json(self) -> Any:
+        return self.model_dump_json()

sift_py/report_templates/service.py

@@ -0,0 +1,171 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import List, Optional, cast
+
+from google.protobuf.field_mask_pb2 import FieldMask
+from sift.report_templates.v1.report_templates_pb2 import (
+    CreateReportTemplateRequest,
+    CreateReportTemplateRequestClientKeys,
+    GetReportTemplateRequest,
+    GetReportTemplateResponse,
+    ReportTemplate,
+    ReportTemplateRule,
+    ReportTemplateTag,
+    UpdateReportTemplateRequest,
+)
+from sift.report_templates.v1.report_templates_pb2_grpc import ReportTemplateServiceStub
+
+from sift_py._internal.time import to_timestamp_pb
+from sift_py.grpc.transport import SiftChannel
+from sift_py.report_templates.config import ReportTemplateConfig
+from sift_py.yaml.report_templates import load_report_templates
+
+
+class ReportTemplateService:
+    """
+    A service for managing report templates. Allows for creating, updating, and retrieving report
+    templates.
+    """
+
+    _report_template_service_stub: ReportTemplateServiceStub
+
+    def __init__(self, channel: SiftChannel):
+        self._report_template_service_stub = ReportTemplateServiceStub(channel)
+
+    def create_or_update_report_template(self, config: ReportTemplateConfig):
+        """
+        Create or update a report template via a ReportTemplateConfig. The config must contain a
+        template_client_key, otherwise an exception will be raised. If a report template with the
+        same client key exists, it will be updated. Otherwise, a new report template will be created.
+        See `sift_py.report_templates.config.ReportTemplateConfig` for more information on available
+        fields to configure.
+        """
+        if not config.template_client_key:
+            raise Exception(f"Report template {config.name} requires a template_client_key")
+
+        report_template = self._get_report_template_by_client_key(config.template_client_key)
+        if report_template:
+            self._update_report_template(config, report_template)
+        else:
+            self._create_report_template(config)
+
+    def get_report_template(
+        self, client_key: Optional[str] = None, id: Optional[str] = None
+    ) -> Optional[ReportTemplateConfig]:
+        """
+        Retrieve a report template by client key or id. Only one of client_key or id should be
+        provided, otherwise an exception will be raised. If a report template is found, it will be
+        returned as a ReportTemplateConfig object.
+        """
+        report_template = None
+        if (not client_key and not id) or (client_key and id):
+            raise ValueError("One of client_key or id must be provided")
+
+        if id:
+            report_template = self._get_report_template_by_id(id)
+        elif client_key:
+            report_template = self._get_report_template_by_client_key(client_key)
+
+        return (
+            ReportTemplateConfig(
+                name=report_template.name,
+                template_client_key=report_template.client_key,
+                organization_id=report_template.organization_id,
+                tags=[tag.tag_name for tag in report_template.tags],
+                description=report_template.description,
+                rule_client_keys=[rule.client_key for rule in report_template.rules],
+                archived_date=report_template.archived_date.ToDatetime()
+                if report_template.archived_date
+                else None,
+            )
+            if report_template
+            else None
+        )
+
+    def load_report_templates_from_yaml(self, paths: List[Path]) -> List[ReportTemplateConfig]:
+        """
+        Load report templates from YAML definitions. The YAML defined report template must have
+        a client key. If the report template with the given client key exists, it will be updated,
+        otherwise a new report template will be created.
+        See `sift_py.yaml.report_templates.load_report_templates` for more information on the YAML
+        spec for report templates.
+        """
+        report_templates = load_report_templates(paths)
+        [
+            self.create_or_update_report_template(report_template)
+            for report_template in report_templates
+        ]
+        return report_templates
+
+    def _get_report_template_by_id(self, report_template_id: str) -> Optional[ReportTemplate]:
+        req = GetReportTemplateRequest(report_template_id=report_template_id)
+        try:
+            res = cast(
+                GetReportTemplateResponse, self._report_template_service_stub.GetReportTemplate(req)
+            )
+            return cast(ReportTemplate, res.report_template) or None
+        except:
+            return None
+
+    def _get_report_template_by_client_key(self, client_key: str) -> Optional[ReportTemplate]:
+        req = GetReportTemplateRequest(client_key=client_key)
+        try:
+            res = cast(
+                GetReportTemplateResponse, self._report_template_service_stub.GetReportTemplate(req)
+            )
+            return res.report_template or None
+        except:
+            return None
+
+    def _create_report_template(self, config: ReportTemplateConfig):
+        client_keys_req = CreateReportTemplateRequestClientKeys(
+            rule_client_keys=config.rule_client_keys
+        )
+        req = CreateReportTemplateRequest(
+            name=config.name,
+            client_key=config.template_client_key,
+            description=config.description,
+            tag_names=config.tags,
+            organization_id=config.organization_id,
+            rule_client_keys=client_keys_req,
+        )
+        self._report_template_service_stub.CreateReportTemplate(req)
+
+    def _update_report_template(
+        self, config: ReportTemplateConfig, report_template: ReportTemplate
+    ):
+        """
+        Uses the report template id, organization id, and client key from the existing report
+        template. Updates the name, description, tags, rules, and archived date from the config,
+        if they are provided. Always passes all fields to the field mask for updating.
+        """
+        tags = []
+        if config.tags:
+            tags = [ReportTemplateTag(tag_name=tag) for tag in config.tags]
+
+        archived_date = None
+        if config.archived_date:
+            archived_date = to_timestamp_pb(config.archived_date)
+
+        rules = [
+            ReportTemplateRule(client_key=client_key) for client_key in config.rule_client_keys
+        ]
+
+        updated_report_template = ReportTemplate(
+            report_template_id=report_template.report_template_id,
+            organization_id=report_template.organization_id,
+            client_key=report_template.client_key,
+            name=config.name,
+            description=config.description,
+            tags=tags,
+            rules=rules,
+            archived_date=archived_date,
+        )
+
+        field_mask = FieldMask(paths=["name", "description", "tags", "rules", "archived_date"])
+        self._report_template_service_stub.UpdateReportTemplate(
+            UpdateReportTemplateRequest(
+                report_template=updated_report_template, update_mask=field_mask
+            )
+        )

sift_py/rest.py

@@ -0,0 +1,29 @@
+from typing import TypedDict
+
+from typing_extensions import NotRequired
+
+from sift_py.grpc.transport import _clean_uri
+
+
+class SiftRestConfig(TypedDict):
+    """
+    Config class used to to interact with services that use Sift's REST API.`.
+    - `uri`: The URI of Sift's REST API. The scheme portion of the URI i.e. `https://` should be ommitted.
+    - `apikey`: User-generated API key generated via the Sift application.
+    - `use_ssl`: INTERNAL USE. Meant to be used for local development.
+    """
+
+    uri: str
+    apikey: str
+    use_ssl: NotRequired[bool]
+
+
+def compute_uri(restconf: SiftRestConfig) -> str:
+    uri = restconf["uri"]
+    use_ssl = restconf.get("use_ssl", True)
+    clean_uri = _clean_uri(uri, use_ssl)
+
+    if use_ssl:
+        return f"https://{clean_uri}"
+
+    return f"http://{clean_uri}"

sift_py/rule/_config_test.py

@@ -0,0 +1,109 @@
+from sift_py.ingestion.channel import ChannelConfig, ChannelDataType
+
+from .config import (
+    RuleActionCreateDataReviewAnnotation,
+    RuleActionCreatePhaseAnnotation,
+    RuleConfig,
+)
+
+
+def test_rule_config_json():
+    voltage_rule_expression = "$1 > 10"
+    voltage_rule_config = RuleConfig(
+        name="High Voltage",
+        description="Rock & Roll",
+        expression=voltage_rule_expression,
+        action=RuleActionCreatePhaseAnnotation(),
+        channel_references=[
+            {
+                "channel_reference": "$1",
+                "channel_config": ChannelConfig(
+                    name="voltage",
+                    data_type=ChannelDataType.DOUBLE,
+                ),
+            }
+        ],
+    )
+    assert voltage_rule_config.expression == voltage_rule_expression
+
+    overheating_rule_expression = '$1 == "Accelerating" && $2 > $3'
+    overheating_rule_config = RuleConfig(
+        name="overheating",
+        description="checks if vehicle overheats while accelerating",
+        expression=overheating_rule_expression,
+        action=RuleActionCreateDataReviewAnnotation(
+            tags=["foo", "bar"],
+            assignee="foobar@baz.com",
+        ),
+        channel_references=[
+            {
+                "channel_reference": "$1",
+                "channel_config": ChannelConfig(
+                    name="vehicle_state",
+                    component="motor",
+                    data_type=ChannelDataType.INT_32,
+                ),
+            },
+            {
+                "channel_reference": "$2",
+                "channel_config": ChannelConfig(
+                    name="temperature",
+                    component="motor",
+                    data_type=ChannelDataType.INT_32,
+                ),
+            },
+        ],
+        sub_expressions={
+            "$3": 80,
+        },
+    )
+    assert overheating_rule_config.expression == '$1 == "Accelerating" && $2 > 80'
+
+    contains_rule_expression = "contains($1, $2)"
+    contains_rule_config = RuleConfig(
+        name="contains",
+        description="checks if vehicle overheats while accelerating",
+        expression=contains_rule_expression,
+        action=RuleActionCreateDataReviewAnnotation(
+            tags=["foo", "bar"],
+            assignee="foobar@baz.com",
+        ),
+        channel_references=[
+            {
+                "channel_reference": "$1",
+                "channel_config": ChannelConfig(
+                    name="log",
+                    data_type=ChannelDataType.INT_32,
+                ),
+            },
+        ],
+        sub_expressions={
+            "$2": "Error",
+        },
+    )
+    assert contains_rule_config.expression == 'contains($1, "Error")'
+
+
+def test_rule_named_expressions():
+    kinetic_energy_gt_expression = "0.5 * $mass * $1 * $1 > $threshold"
+
+    rule_on_kinetic_energy = RuleConfig(
+        name="rule_onkinetic_energy",
+        description="checks high periods of energy output",
+        expression=kinetic_energy_gt_expression,
+        action=RuleActionCreatePhaseAnnotation(),
+        channel_references=[
+            {
+                "channel_reference": "$1",
+                "channel_config": ChannelConfig(
+                    name="velocity",
+                    data_type=ChannelDataType.INT_32,
+                ),
+            },
+        ],
+        sub_expressions={
+            "$mass": 10,
+            "$threshold": 35,
+        },
+    )
+    assert rule_on_kinetic_energy.expression == "0.5 * 10 * $1 * $1 > 35"

sift_py/rule/_service_test.py

@@ -0,0 +1,168 @@
+from unittest import mock
+
+import pytest
+from sift.assets.v1.assets_pb2 import Asset
+from sift.rules.v1.rules_pb2 import Rule
+
+from sift_py._internal.test_util.channel import MockChannel
+from sift_py.rule.config import RuleActionCreateDataReviewAnnotation, RuleConfig
+from sift_py.rule.service import RuleService
+
+
+@pytest.fixture()
+def rule_service():
+    return RuleService(MockChannel())
+
+
+def test_rule_service_get_rule_from_client_key(rule_service):
+    rule_client_key = "rule-client-key"
+
+    with mock.patch.object(
+        RuleService,
+        "_get_rule_from_client_key",
+        return_value=Rule(name="abc"),
+    ) as mock_get_rule_from_client_key:
+        rule_service.get_rule(rule_client_key)
+        mock_get_rule_from_client_key.assert_called_once_with(rule_client_key)
+
+
+def test_rule_service_get_rule_from_rule_id(rule_service):
+    rule_id = "rule-id"
+
+    with mock.patch.object(
+        RuleService, "_get_rule_from_rule_id", return_value=Rule(name="abc")
+    ) as mock_get_rule_from_rule_id:
+        rule_service.get_rule(rule_id)
+        mock_get_rule_from_rule_id.assert_called_once_with(rule_id)
+
+
+def test_rule_service_get_rule_missing_client_key_and_id(rule_service):
+    with mock.patch.object(
+        RuleService,
+        "_get_rule_from_client_key",
+        return_value=None,
+    ) as mock_get_rule_from_client_key:
+        with mock.patch.object(
+            RuleService, "_get_rule_from_rule_id", return_value=None
+        ) as mock_get_rule_from_rule_id:
+            rule = rule_service.get_rule("")
+            mock_get_rule_from_client_key.asset_called_once()
+            mock_get_rule_from_rule_id.asset_called_once()
+            assert rule is None
+
+
+def test_rule_service_create_rule(rule_service):
+    rule = RuleConfig(
+        name="rule",
+        rule_client_key="rule-client-key",
+        channel_references=[],
+    )
+
+    with mock.patch.object(RuleService, "_create_rule") as mock_create_rule:
+        rule_service.create_or_update_rule(rule)
+        mock_create_rule.assert_called_once_with(rule)
+
+
+def test_rule_service_update_rule(rule_service):
+    rule = RuleConfig(
+        name="rule",
+        rule_client_key="rule-client-key",
+        channel_references=[],
+    )
+
+    with mock.patch.object(RuleService, "_update_rule") as mock_update_rule:
+        with mock.patch.object(
+            RuleService, "_get_rule_from_client_key", return_value=Rule(name=rule.name)
+        ) as mock_get_rule_from_client_key:
+            rule_service.create_or_update_rule(rule)
+            mock_get_rule_from_client_key.assert_called_once_with(rule.rule_client_key)
+            mock_update_rule.assert_called_once()
+
+
+def test_rule_service_load_rules_from_yaml(rule_service):
+    rule_yaml = {
+        "name": "rule",
+        "rule_client_key": "rule-client-key",
+        "channel_references": [{"$1": "channel"}],
+        "description": "description",
+        "expression": "$1 > 0",
+        "assignee": "assignee@abc.com",
+        "type": "review",
+        "asset_names": ["asset"],
+    }
+    with mock.patch.object(RuleService, "create_or_update_rule"):
+        with mock.patch(
+            "sift_py.rule.service.load_rule_modules",
+            return_value=[rule_yaml],
+        ):
+            rule_configs = rule_service.load_rules_from_yaml(["path/to/rules.yml"])
+            assert len(rule_configs) == 1
+
+            rule_config = rule_configs[0]
+            assert rule_config.name == rule_yaml["name"]
+            assert rule_config.rule_client_key == rule_yaml["rule_client_key"]
+            assert rule_config.channel_references == [
+                {
+                    "channel_reference": "$1",
+                    "channel_identifier": "channel",
+                }
+            ]
+            assert rule_config.description == rule_yaml["description"]
+            assert rule_config.expression == rule_yaml["expression"]
+            assert rule_config.action.assignee == rule_yaml["assignee"]
+            assert rule_config.asset_names == rule_yaml["asset_names"]
+            assert isinstance(rule_config.action, RuleActionCreateDataReviewAnnotation)
+
+
+def test_rule_service_attach_asset():
+    rule_config = RuleConfig(
+        name="rule",
+        rule_client_key="rule-client-key",
+        expression="1",
+        channel_references=[],
+        action=RuleActionCreateDataReviewAnnotation(),
+        asset_names=["abc"],
+    )
+    asset = Asset(name="asset", asset_id="asset-id", organization_id="org-id")
+    with mock.patch("sift_py.rule.service.RuleServiceStub", return_value=mock.MagicMock()):
+        rule_service = RuleService(MockChannel())
+        with mock.patch.object(RuleService, "_get_assets", return_value=[asset]):
+            returned_config = rule_service.attach_asset(rule_config, ["asset"])
+            assert "abc" in returned_config.asset_names
+            assert "asset" in returned_config.asset_names
+
+
+def test_rule_service_detach_asset():
+    rule_config = RuleConfig(
+        name="rule",
+        rule_client_key="rule-client-key",
+        expression="1",
+        channel_references=[],
+        action=RuleActionCreateDataReviewAnnotation(),
+        asset_names=["abc", "def"],
+    )
+    asset_abc = Asset(name="abc", asset_id="abc-id", organization_id="org-id")
+    asset_def = Asset(name="def", asset_id="def-id", organization_id="org-id")
+    with mock.patch("sift_py.rule.service.RuleServiceStub", return_value=mock.MagicMock()):
+        rule_service = RuleService(MockChannel())
+        with mock.patch.object(RuleService, "_get_assets", return_value=[asset_abc, asset_def]):
+            returned_config = rule_service.detach_asset(rule_config, ["abc"])
+            assert "abc" not in returned_config.asset_names
+            assert "def" in returned_config.asset_names
+
+
+def test_rule_service_detach_asset_empty_list():
+    rule_config = RuleConfig(
+        name="rule",
+        rule_client_key="rule-client-key",
+        expression="1",
+        channel_references=[],
+        action=RuleActionCreateDataReviewAnnotation(),
+        asset_names=["asset"],
+    )
+    asset = Asset(name="asset", asset_id="asset-id", organization_id="org-id")
+    with mock.patch("sift_py.rule.service.RuleServiceStub", return_value=mock.MagicMock()):
+        rule_service = RuleService(MockChannel())
+        with mock.patch.object(RuleService, "_get_assets", return_value=[asset]):
+            with pytest.raises(ValueError, match="must be associated with at least one asset"):
+                rule_service.detach_asset(rule_config, ["asset"])